diff --git a/benchmarks/benchmark_compressor_state.py b/benchmarks/benchmark_compressor_state.py new file mode 100644 index 000000000..1a3cd16cd --- /dev/null +++ b/benchmarks/benchmark_compressor_state.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stateful append timing includes validation, staging and synchronous commit.""" + +import argparse +import json +from pathlib import Path +import statistics +import sys +import time + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from rl_engine.backends.rocm.attention.compressor_state import AiterCompressorStateBackend +from rl_engine.contracts.operators.compressor_state import CompressorStateContract +from rl_engine.models.deepseek_v4.compressor_state import CompressorStateCache +from tests.models.deepseek_v4.compressor_state_helpers import tensors + + +def case(backend, ratio, graph, rows, repeats): + c = CompressorStateContract( + ratio=ratio, + graph=graph, + max_requests=1, + max_tokens_per_step=rows, + max_sequence_tokens=max(ratio, rows * (5 * repeats + 3)), + ) + x, a, n, co, si = tensors(c, device="cuda") + arena = CompressorStateCache( + c, + ape=a, + norm=n, + cos=co, + sin=si, + version="benchmark", + writer_rank=0, + expected_arithmetic=c.arithmetic_id, + backend=backend, + ) + h = arena.create_request(0, version="benchmark", writer_rank=0) + start = 0 + for _ in range(3): + arena.append(h, x[start : start + rows], start_position=start) + start += rows + samples = [] + for _ in range(5): + torch.cuda.synchronize() + begin = time.perf_counter() + for _ in range(repeats): + arena.append(h, x[start : start + rows], start_position=start) + start += rows + torch.cuda.synchronize() + samples.append((time.perf_counter() - begin) * 1e6 / repeats) + snapshot = arena.snapshot(h) + arena.release(h) + return { + "ratio": ratio, + "graph": graph, + "rows_per_append": rows, + "append_us": statistics.median(samples), + "processed_tokens": start, + "complete_cache_rows": len(snapshot["rows"]), + "resident_buffer_bytes": c.resident_bytes, + "arithmetic_id": c.arithmetic_id, + "implementation": arena.implementation_fingerprint, + } + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--repeats", type=int, default=30) + p.add_argument("--output", type=Path, required=True) + args = p.parse_args() + backend = AiterCompressorStateBackend() + report = { + "provider": backend.provenance, + "metric": "median of five synchronized loops; validation and device completion included", + "scope": "single writer; not kernel-only latency or full model throughput", + "cases": [ + case(backend, ratio, graph, rows, args.repeats) + for ratio in (4, 128) + for graph in (False, True) + for rows in (1, 256) + ], + } + args.output.write_text(json.dumps(report, indent=2)) + print(json.dumps(report, indent=2)) diff --git a/benchmarks/benchmark_kv_compressor_c4.py b/benchmarks/benchmark_kv_compressor_c4.py new file mode 100644 index 000000000..b191fbc44 --- /dev/null +++ b/benchmarks/benchmark_kv_compressor_c4.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MI300X benchmark for the fixed aiter C4 overlap forward boundary.""" + +import argparse +import json +from pathlib import Path +import sys +import time + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor +from tests.models.deepseek_v4.c4_helpers import fixture + + +def timed(fn, repeats): + for _ in range(5): + fn() + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(repeats): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) * 1e6 / repeats + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--repeats", type=int, default=30) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + backend = AiterC4Backend() + operator = C4Compressor(C4Contract(), backend=backend) + cases = [] + for groups in (1, 8, 64, 128): + values, ids, ape, norm, cos, sin = fixture(groups, device="cuda") + bound = operator.bind(values, ids, version="v1") + + def forward(bound=bound, ape=ape, norm=norm, cos=cos, sin=sin): + operator(bound, ape, norm, cos, sin, version="v1") + + cases.append( + { + "groups": groups, + "tokens": groups * 4, + "input_bytes": values.numel() * values.element_size(), + "output_bytes": groups * 512 * 2, + "bind_us": timed( + lambda values=values, ids=ids: operator.bind(values, ids, version="v1"), + args.repeats, + ), + "forward_us": timed(forward, args.repeats), + "provider_forward_us": timed( + lambda bound=bound, ape=ape, norm=norm, cos=cos, sin=sin: backend.forward( + bound, ape, norm, cos, sin, operator.contract + ), + args.repeats, + ), + } + ) + result = { + "gpu": torch.cuda.get_device_name(), + "arch": torch.cuda.get_device_properties(0).gcnArchName, + "torch": torch.__version__, + "hip": torch.version.hip, + "triton": __import__("triton").__version__, + "backend": backend.provenance, + "contract": operator.contract.fingerprint, + "implementation": operator.implementation_fingerprint, + "metric": ( + "synchronized average over warmup+repeated loops; Python validation and " + "paged-cache gather included" + ), + "scope": ( + "complete C4 forward; aiter overlap pool and norm/RoPE; no live cache " + "commit, no full model" + ), + "cases": cases, + } + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/models/deepseek_v4/compressor-state.md b/docs/models/deepseek_v4/compressor-state.md new file mode 100644 index 000000000..e25759b2a --- /dev/null +++ b/docs/models/deepseek_v4/compressor-state.md @@ -0,0 +1,99 @@ +# Task 21 — request-owned compressor state and single cache commit + +This change is stacked on task 19 / PR #618 at `8793052`; its target remains +`dsv4-rocm`. It adds lifecycle/control code, not numerical GPU kernels. + +## Reused providers and explicit arithmetic + +The backend directly calls aiter `flydsl_fused_compress_attn` to read previous +state and compress/scatter complete rows, then vLLM `save_partial_states` to +update FP32 `[kv, score+ape]` state. The order is fixed: **compress before +updating history**. The installed aiter ABI needs contiguous split readers; +preallocated copies adapt vLLM's packed state without changing its values. + +The admitted main-KV profiles are: + +- `aiter.main_c4.online_fp32.wave64.vllm_partial_f32.v1` +- `aiter.main_c128.online_fp32.wave64.vllm_partial_f32.v1` + +Both use a fixed single-wave online-FP32 recipe. **C128 is not task 18's +eight-wave HCA profile.** The required `expected_arithmetic` constructor +argument rejects a different recipe rather than silently substituting it. +Callers must select this same profile for the paths they compare; shape +compatibility alone is not a numerical compatibility claim. Training/rollout +integration with the eight-wave HCA path is not covered by this profile. + +## API and ownership + +Construct `CompressorStateCache` with a `CompressorStateContract`, FP32 APE/RMS +weights, BF16 position tables, a nonempty model/position version and a writer +rank. The cache owns snapshots of the parameters and a fixed-capacity workspace. +No model downloads or provider replacement occurs during construction. + +1. `create_request(request_id, version=..., writer_rank=...)` returns a handle + bound to this arena, a generation, version and private page slot. +2. `prepare(handle, values, start_position=...)` validates detached BF16 input + and owns its staging copy. Only contiguous new logical tokens are accepted. +3. `commit(lease)` consumes that **exact lease once**, performs ordered native + writes and waits for device completion before publishing the token cursor. +4. `append(...)` combines prepare and commit under the same host lock. +5. `cancel(lease)` aborts a prepared, not-yet-executed transaction. +6. `snapshot(handle)` returns caller-owned complete rows and the FP32 ring, + plus token/group/version metadata. It does not expose mutable arena storage. +7. `release(handle)` clears its state/pages before slot reuse. The old handle + cannot access a later generation with the same logical request ID. +8. `replace_model(...)` requires all requests and leases to be released, a + different version, and fully validated replacement parameters. + +There is one staging lease per arena. Thread races, stale/foreign handles, +wrong writers, duplicate commits, noncontiguous tokens and capacity violations +are rejected. Errors after mutable GPU work poison the arena; callers must +discard it instead of retrying possibly partial writes. + +Only the final ring-capacity suffix of an oversized input chunk is eligible +to write partial state. Earlier writers receive vLLM's invalid slot sentinel, +so one native launch never has competing writers to the same ring slot. + +The request's complete-row interval is determined by the logical cursor: +`[old_tokens // ratio, new_tokens // ratio)`. Incomplete groups are never +committed. The ring also retains the predecessor needed for C4 overlap. + +## Graph and memory + +Graph mode captures the same upstream operations with fixed input, metadata, +state-reader and cache addresses. Warmup/capture uses only invalid plans and +write slots and is checked not to mutate state/cache. Per-request lifecycle, +version validation, metadata preparation and exactly-once checks stay **outside** +capture. Direct external capture of lifecycle methods is rejected. + +Commits synchronously finish their device event. This intentionally conservative +policy prevents early staging/page reuse between threads or streams; asynchronous +communication/compute overlap is a later task. Input producers must order their +writes before `prepare`; private arena storage must not be modified externally. + +The byte budget bounds explicitly owned resident tensors, not PyTorch/FlyDSL +compiler caches, Graph memory pools, admission temporaries or caller snapshots. +Provider-addressing bounds are checked before allocation. + +## Integration and limits + +The vLLM and Megatron opt-in factories share this exact lifecycle implementation. +This is not complete framework model registration. Cache writes are deliberately +non-differentiable; detached inputs are required. Native CP transport must occur +before ingestion; a writer-rank check is not a distributed lock or collective. + +No multi-GPU/RCCL, multi-node, full model or eight-wave C128 training/rollout +equivalence is claimed. Indexer MXFP4 state and cache belong to task 22. + +Run the targeted suite and benchmark in the member container, keeping results +in the mounted output directory: + +```bash +python3 -m pytest tests/contracts/operators tests/models/deepseek_v4 -q \ + --basetemp=/outputs/state-tests --junitxml=/outputs/state-tests.xml +python3 benchmarks/benchmark_compressor_state.py \ + --repeats 30 --output /outputs/state-benchmark.json +``` + +The existing task-19 regression tests require the pinned mLite source through +`RLK_MLITE_DSV4_BACKWARD_SOURCE`; the state backend itself has no mLite dependency. diff --git a/docs/models/deepseek_v4/kv-compressor-c4.md b/docs/models/deepseek_v4/kv-compressor-c4.md new file mode 100644 index 000000000..df64de262 --- /dev/null +++ b/docs/models/deepseek_v4/kv-compressor-c4.md @@ -0,0 +1,81 @@ +# DSV4 item 19: C4 KV compressor + +This item provides the main-KV C4 compressor with the fixed eight-token overlap +window. It is separate from: + +- item 18 C128 main-KV compression; +- item 20 the independent 128-dimensional Indexer C4 compressor; +- item 21 state/cache lifecycle and exactly-once commit; +- item 27 prefill/decode/recompute/Graph runtime integration. + +## Contract + +```text +input: BF16 [tokens, 2048] + [content1024, score1024] where each 1024-wide stream has two 512 halves + int64 logical token IDs [request_id, global_token_position] +output: BF16 [complete_groups, 512] + logical IDs [request_id, global_group_id] +``` + +The compressor consumes four current tokens per output row. For group `g`, the +softmax window is: + +```text +previous group content/gate first half (g-1), then +current group content/gate second half (g) +``` + +The first group has four masked predecessor rows. Incomplete groups are never +padded into a submitted output. Global position is `group_id * 4`; local rank +order is not substituted for it. The optional `emit_group_ids` binding argument +allows a microbatch to include one completed predecessor group as context while +emitting only its requested groups. + +CP callers must provide complete current groups and their overlap predecessor +before binding. The rank owning the current group’s fourth token emits that row. +No collective is performed by this item. + +## Reuse and paged output + +The ROCm backend directly reuses aiter’s +`flydsl_fused_compress_attn` with `overlap=True`, `ratio=4`, `head_dim=512`, +BF16 output, and the fixed one-wave (`k_split_num_waves=1`) arithmetic path. +The private output uses the provider’s real compressed-cache layout: +`k_per_block=64`, a per-request block table, and cache rows selected by +`group_id % 64`. This is important for correctness; treating every row as a +one-row block is not an accepted shortcut. + +The overlap input is materialized as eight direct input rows per plan row. This +reuses aiter’s optimized online-softmax/norm/RoPE kernel while leaving live +request state mutation and cache commit to item 21. The first group’s four +predecessor scores use a finite BF16 floor below the admitted score range, +avoiding the input-phase `-inf - (-inf)` case; the independent oracle retains +the exact masked semantics. + +Backward reuses the pinned mLite/Megatron `compressed_compact_graph` ratio-4 +functional graph. One output’s VJP is evaluated against its previous/current +context. Overlapping row contributions are applied in a fixed Python logical +order, not with unordered floating-point atomics. APE and RMS rows are +registered with RLK’s existing `CanonicalBackwardSession`. + +Provider source hashes are checked before admission; missing or different +aiter/mLite sources fail closed. No PyTorch numerical fallback is exposed. + +## Verification + +The MI300X tests cover: + +- complete and partial groups, first-group masking, overlap previous/current + halves, global positions, request permutation and padding; +- CP final-token ownership, context predecessor admission and block-table + mapping; +- softmax/content scale extremes, fixed overlap plan, `k_per_block=64`, + output/cache isolation, streams and fresh-process/JIT fingerprints; +- independent FP64 forward arithmetic; +- full/microbatch/reversed-order backward, overlap gradient accumulation and + independent autograd/finite-difference checks; +- malformed ABI/metadata, budget, mutation, version and Graph rejection. + +This item does not claim RCCL/multi-GPU CP, live cache commit, Graph replay, +prefill/decode framework integration or full-model acceptance. diff --git a/reports/compressor-state-mi300x-20261029.md b/reports/compressor-state-mi300x-20261029.md new file mode 100644 index 000000000..c454af796 --- /dev/null +++ b/reports/compressor-state-mi300x-20261029.md @@ -0,0 +1,70 @@ +# Task 21 — MI300X state/cache acceptance, 2026-10-29 + +Parent: `879305251c37e1b8bb0cb44409f1762752d03561` (task 19 / PR #618). +Target: `dsv4-rocm`. Source lives in the member workspace; logs and evidence +are under the mounted `/outputs/item21-20261029/`. + +## Environment and reuse + +MI300X `gfx942`; Torch `2.11.0+gitd0c8b1f`; HIP `7.2.53211`; Triton 3.6.0; +FlyDSL 0.2.0; aiter 0.1.16.post3; vLLM 0.26.0+rocm723. + +- Existing aiter compression/RMSNorm/RoPE/cache scatter. +- Existing vLLM FP32 packed partial-state writer. +- Preallocated layout-only reader copies and Torch Graph capture/replay. +- No new numerical GPU kernel. + +The explicit single-wave C128 profile is distinct from task 18's eight-wave +HCA profile. An arithmetic handshake rejects mixing them. The tests below do +not claim equivalence to a different reduction recipe. + +## Test results + +- Full targeted contract + DeepSeek-V4 suite: **316 passed twice**, zero + failures/errors/skips; 54.58 s with coverage, 39.50 s independent repeat. +- 39 new CPU tests and **50 real GPU tests**; 85 C4 prerequisite regressions + and 142 existing operator-contract regressions. +- Each full run reported 16 warnings. +- Genuine FP64 pool/norm/rotation numerical reference; exact FP32 ring bytes. +- Exact cache/state bytes across full input, uneven chunks, single-token decode, + real Graph replay, multiple seeds, interleaved requests and page boundaries. +- Request generation/reuse, writer ownership, version replacement, duplicate + lease rejection, incomplete groups, original input alias reuse and output + snapshot isolation. +- Thread races, cross-stream staging handoff, no writes during Graph warmup, + stable Graph addresses and 40 release/recreate cycles without live-allocation + growth. +- Fault injection after mutable work, release and model replacement verifies + fail-closed poisoned arenas rather than unsafe retry. +- Ruff lint and format passed for all ten new Python files. + +Python statement coverage: **316/320 (98.75%)**; branch coverage **27/28 +(96.43%)**. Native GPU instruction coverage is not represented by this report. + +## Stateful append latency + +Microseconds; median of five synchronized loops of 30 appends after warmup. +Includes host validation, staging, native execution and device completion. +The table is not raw kernel latency, full-model throughput or an overlap claim. + +| Ratio | Graph | Tokens/append | Append (us) | +|---:|:---:|---:|---:| +| 4 | off | 1 | 369.13 | +| 4 | off | 256 | 449.62 | +| 4 | on | 1 | 244.97 | +| 4 | on | 256 | 351.54 | +| 128 | off | 1 | 354.77 | +| 128 | off | 256 | 426.40 | +| 128 | on | 1 | 251.08 | +| 128 | on | 256 | 354.24 | + +The prefill cases process 39,168 consecutive tokens in each arena, validating +capacity/page traversal beyond the shorter numerical test fixtures. The report +does not claim arbitrary shapes or precision profiles are validated. + +## Remaining scope + +No actual distributed CP/RCCL transport, multi-node, full framework registration, +full model, cache-write autograd or compatibility with the eight-wave C128 +profile is claimed. Those require their own integration/acceptance. This PR +implements the explicit state/cache boundary described in its contract. diff --git a/reports/kv-compressor-c4-mi300x-20261028.md b/reports/kv-compressor-c4-mi300x-20261028.md new file mode 100644 index 000000000..a8d4d5b7f --- /dev/null +++ b/reports/kv-compressor-c4-mi300x-20261028.md @@ -0,0 +1,47 @@ +# Item 19 C4 KV compressor — MI300X validation + +Branch: `feat/dsv4-kv-c4-bi` +Target: `dsv4-rocm` +Hardware: AMD Instinct MI300X VF (`gfx942`) +Environment: ROCm/HIP 7.14.60850, PyTorch 2.12.0+rocm7.14.0a20260608, +Triton 3.7.0, FlyDSL 0.3.4.1, pinned aiter commit +`326cffa38c9eda16bef19375ba55334e4d3acd59`, pinned mLite commit +`327a238239448de6771beea185dc6902a8de3eaa`. + +## Test evidence + +- CPU-safe C4 tests: **53 passed**. +- Real MI300X C4 tests: **26 passed**, 0 skipped, 1 existing warning, 19.49s. +- Full targeted contract + DeepSeek-V4 suite: **221 passed**, 0 skipped, 1 + existing warning, 20.33s without coverage and 26.54s with coverage. +- Scoped Python wrapper coverage: **336/371 statements (90%)** and **80/86 + branches**. Coverage does not measure FlyDSL/native GPU instructions; the + unexecuted wrapper lines are provider-constructor and direct native VJP + plumbing branches exercised by the real GPU tests. +- Ruff format/lint passed for all ten Python files. +- Real GPU coverage includes aiter overlap compression, true `k_per_block=64` + paged output mapping, mLite ratio-4 functional VJP, fixed-order overlap + gradients, streams, fresh processes/JIT isolation and CP ownership. +- The forward oracle is independent FP64 softmax/pool/RMS/rotation arithmetic. +- The backward checks compare fixed-order microbatch/reverse-order gradients to + independent autograd and finite-difference results. + +This item does not claim RCCL/multi-GPU CP, Graph replay, live cache commit, +prefill/decode runtime or full-model acceptance. + +The benchmark records bind, provider forward and full operator timing with +Python validation, paged-cache gather and synchronization included. It is not a +kernel-only or throughput claim. + +## Synchronized MI300X benchmark (30 repetitions) + +| Complete groups | Tokens | Bind (us) | Provider (us) | Full forward (us) | +|---:|---:|---:|---:|---:| +| 1 | 4 | 429.67 | 95.44 | 440.68 | +| 8 | 32 | 944.07 | 221.16 | 781.28 | +| 64 | 256 | 2,805.11 | 533.89 | 1,096.21 | +| 128 | 512 | 5,045.03 | 380.60 | 912.69 | + +These are synchronized average wall times with warmup, Python validation, +paged-cache gather and dispatch included. They are not kernel-only latency or +a speedup claim. The raw benchmark records the provider source fingerprints. diff --git a/reports/kv-compressor-c4-new-host-20261029.md b/reports/kv-compressor-c4-new-host-20261029.md new file mode 100644 index 000000000..82f5c8a67 --- /dev/null +++ b/reports/kv-compressor-c4-new-host-20261029.md @@ -0,0 +1,55 @@ +# C4 follow-up: new MI300X environment, 2026-10-29 + +This is a corrective follow-up to task 19 / PR #618, not a new task PR. + +## Corrections + +- The first group's four missing predecessors now use the existing aiter + state-phase negative-position mask. The old finite `-2**30` gate floor was + not a valid universal mask over the admitted gate domain. No handwritten + numerical kernel or approximate masking workaround is used. +- Validate the predecessor context as well as the emitted current-group rows. + An `emit_group_ids` filter cannot hide non-finite predecessor data. +- The independent oracle now preserves FP64 throughout softmax, pooling, + normalization and rotation. The previous helper converted its inputs to + FP32 internally, so the earlier report's description of that comparison as + fully FP64 was incorrect. New regression tests verify FP64 perturbations + survive and gradients remain FP64. +- Backward passes host-known unique context indices instead of iterating GPU + index tensors and synchronizing each index. + +## Dependency profile + +The new image's existing aiter 0.1.16.post3 / FlyDSL 0.2.0 provider is admitted +by its exact source hash, alongside the previously pinned provider ABI. +The older ABI has a fixed single-wave path without a batch-size dispatch. +No numerical fallback is selected on a library error. + +- PyTorch: `2.11.0+gitd0c8b1f` +- HIP: `7.2.53211` +- Triton: `3.6.0` +- GPU architecture: `gfx942` +- aiter C4 source SHA256: + `0c2019ca3a4a0d6dc7d5fdd20ebcce658ff5f504fdc880dffd61c9c398029451` +- mLite functional backward source SHA256: + `6a1e78ca89b83f56391b5e5ac409e765338f1e3a625251d0d4c877c1e80b4a68` + +## Tests + +- C4 CPU/GPU tests: **85 passed**, one warning, 23.66 seconds. +- Full targeted contract + DeepSeek-V4 suite: **227 passed**, zero skips, + one warning, 34.38 seconds with coverage. +- Includes four extreme first-group score biases, hidden non-finite context, + real FP64 oracle checks, real aiter launches, mLite VJP, streams and fresh + subprocess runs. +- Changed-file Ruff lint/format checks pass; the 30-repetition benchmark was + re-run on the new environment. +- Python wrapper coverage remains approximately 90%. Native instructions and + autograd-internal tracing gaps are not evidence of native coverage; earlier + statements describing all uncovered lines as constructor failures were too + broad. + +Logs, JUnit XML, coverage and benchmark JSON are under the member's mounted +`/outputs/c4-new-host-20261029/`, not the shared Docker writable layer. +The task-19 limits remain: no RCCL/multi-GPU CP, Graph replay, live cache +commit, full-model or end-to-end framework acceptance. diff --git a/rl_engine/backends/rocm/attention/compressor_state.py b/rl_engine/backends/rocm/attention/compressor_state.py new file mode 100644 index 000000000..dfe6b96e0 --- /dev/null +++ b/rl_engine/backends/rocm/attention/compressor_state.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Existing aiter compression/cache scatter and vLLM FP32 partial-state writes.""" + +import importlib +import importlib.metadata + +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import ( + AITer_C4_LEGACY_SHA256, + AITer_C4_SHA256, + C4Error, + require, + source_hash, +) + + +class AiterCompressorStateBackend: + device_type = "cuda" + supports_graph = True + + def __init__(self): + require(torch.version.hip is not None and torch.cuda.is_available(), "ROCm GPU required") + self.device = torch.device("cuda", torch.cuda.current_device()) + arch = torch.cuda.get_device_properties(self.device).gcnArchName.split(":")[0] + require(arch == "gfx942", "unvalidated compressor-state architecture") + try: + native = importlib.import_module("aiter.ops.flydsl.kernels.fused_compress_attn") + writer = importlib.import_module( + "vllm.models.deepseek_v4.common.ops.save_partial_states" + ) + self.compress = native.flydsl_fused_compress_attn + self.save_partial = writer.save_partial_states + except (ImportError, AttributeError) as exc: + raise C4Error("required aiter/vLLM state ABI unavailable") from exc + h = source_hash(native.__file__) + require(h in (AITer_C4_LEGACY_SHA256, AITer_C4_SHA256), "unvalidated compressor source") + self.options = ( + {"quant_mode": "none", "k_split_num_waves": 1} if h == AITer_C4_SHA256 else {} + ) + self.provenance = { + "arch": arch, + "torch": torch.__version__, + "hip": torch.version.hip, + "triton": importlib.metadata.version("triton"), + "flydsl": importlib.metadata.version("flydsl"), + "compressor": h, + "state_writer": source_hash(writer.__file__), + "tensor_shim": source_hash( + importlib.import_module("aiter.ops.flydsl.kernels.tensor_shim").__file__ + ), + "bridge": source_hash(__file__), + "recipe": "fixed single-wave online FP32", + } + + def execute(self, arena): + with torch.cuda.device(self.device): + self._execute(arena) + + def _execute(self, arena): + c = arena.contract + # Layout adaptation only. Both destination buffers are preallocated, + # so eager and captured execution use the same reader representation. + arena._read_kv.copy_(arena._state[..., : c.width]) + arena._read_score.copy_(arena._state[..., c.width :]) + self.compress( + kv_in=arena._input[:, : c.width], + score_in=arena._input[:, c.width :], + kv_state=arena._read_kv, + score_state=arena._read_score, + plan_gpu=arena._plan, + state_slot_mapping=arena._state_slots, + ape=arena._ape, + rms_weight=arena._norm, + rms_eps=arena.rms_eps, + cos_cache=arena._cos, + sin_cache=arena._sin, + kv_cache=arena._cache, + block_tables=arena._blocks, + k_per_block=64, + overlap=c.ratio == 4, + ratio=c.ratio, + head_dim=512, + rope_head_dim=64, + quant=False, + **self.options, + ) + # Ordering is essential: compression consumes PREVIOUS partial state. + # Duplicate ring writers are removed by prepare() before this launch. + self.save_partial( + kv=arena._input[:, : c.width], + score=arena._input[:, c.width :], + ape=arena._ape, + positions=arena._positions, + state_cache=arena._state, + slot_mapping=arena._write_slots, + block_size=c.state_rows, + state_width=c.width, + compress_ratio=c.ratio, + ) diff --git a/rl_engine/backends/rocm/attention/kv_compressor_c4.py b/rl_engine/backends/rocm/attention/kv_compressor_c4.py new file mode 100644 index 000000000..fa00be5b0 --- /dev/null +++ b/rl_engine/backends/rocm/attention/kv_compressor_c4.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reuse aiter C4 overlap compression and the pinned mLite functional VJP.""" + +import importlib +import importlib.metadata +import importlib.util +import os + +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import ( + AITer_C4_SHA256, + AITer_C4_LEGACY_SHA256, + C4_BACKEND_ID, + C4Error, + MLITE_BACKWARD_SHA256, + require, + source_hash, +) +from rl_engine.ops.autograd.backward_runtime import record_backward + + +class AiterC4Backend: + backend_id = C4_BACKEND_ID + device_type = "cuda" + + def __init__(self, *, megatron_backward_path=None): + require(torch.version.hip is not None and torch.cuda.is_available(), "C4 requires ROCm") + self.device = torch.device("cuda", torch.cuda.current_device()) + arch = torch.cuda.get_device_properties(self.device).gcnArchName.split(":")[0] + require(arch == "gfx942", f"unvalidated C4 architecture: {arch}") + names = ("aiter.ops.flydsl.kernels.fused_compress_attn",) + try: + self.aiter = importlib.import_module(names[0]) + require( + callable(self.aiter.flydsl_fused_compress_attn), + "aiter C4 entrypoint is not callable", + ) + except (ImportError, AttributeError) as exc: + raise C4Error("pinned aiter C4 ABI unavailable") from exc + provider_hash = source_hash(getattr(self.aiter, "__file__", None)) + require( + provider_hash in (AITer_C4_SHA256, AITer_C4_LEGACY_SHA256), + "unvalidated aiter C4 source revision", + ) + # Two explicitly identified upstream ABIs, not a silent numerical fallback. + self.launch_options = ( + {"quant_mode": "none", "k_split_num_waves": 1} + if provider_hash == AITer_C4_SHA256 + else {} + ) + path = megatron_backward_path or os.environ.get("RLK_MLITE_DSV4_BACKWARD_SOURCE") + if path is None: + try: + module = importlib.import_module( + "megatron.lite.model.deepseek_v4.vllm.primitive.attention.backward" + ) + path = module.__file__ + except (ImportError, AttributeError) as exc: + raise C4Error( + "provide the pinned mLite backward.py source or install mLite" + ) from exc + require(source_hash(path) == MLITE_BACKWARD_SHA256, "unvalidated mLite backward source") + spec = importlib.util.spec_from_file_location("_rlk_pinned_mlite_c4_backward", path) + require(spec is not None and spec.loader is not None, "mLite source loader unavailable") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + self.functional = module.compressed_compact_graph + self.provenance = { + "backend": self.backend_id, + "arch": arch, + "torch": torch.__version__, + "hip": torch.version.hip, + "triton": importlib.import_module("triton").__version__, + "flydsl": importlib.metadata.version("flydsl"), + "aiter_c4": provider_hash, + "launch_options": self.launch_options, + "padding": "native state-phase negative-position masking", + "mlite_backward": MLITE_BACKWARD_SHA256, + "adapter": source_hash(__file__), + } + + def forward(self, bound, ape, norm, cos, sin, contract): + groups = len(bound.logical_keys) + if groups == 0: + return torch.empty((0, 512), dtype=torch.bfloat16, device=bound.values.device) + blocks = max((block for block, _slot in bound.cache_slots_host), default=-1) + 1 + cache = torch.empty((blocks, 64, 512), dtype=torch.bfloat16, device=bound.values.device) + provider_values = bound.context_values.reshape(-1, bound.context_values.shape[-1]) + state_kv = bound.kv_state + state_score = bound.score_state + state_slots = bound.state_slots + plan = bound.plan.contiguous() + block_table = bound.block_table + self.aiter.flydsl_fused_compress_attn( + kv_in=provider_values[:, :1024], + score_in=provider_values[:, 1024:], + kv_state=state_kv, + score_state=state_score, + plan_gpu=plan, + state_slot_mapping=state_slots, + ape=ape, + rms_weight=norm, + rms_eps=contract.rms_eps, + cos_cache=cos, + sin_cache=sin, + kv_cache=cache, + block_tables=block_table, + k_per_block=64, + overlap=True, + ratio=4, + head_dim=512, + rope_head_dim=64, + quant=False, + use_ue8m0=True, + preshuffle=True, + **self.launch_options, + ) + return torch.stack( + [cache[block, slot] for block, slot in bound.cache_slots_host], + dim=0, + ).contiguous() + + def vjp(self, bound, ape, norm, cos, sin, grad, contract): + contributions, ape_rows, norm_rows = [], [], [] + cache = torch.cat((cos, sin), dim=-1) + with torch.inference_mode(False), torch.enable_grad(): + for index, (_request, group) in enumerate(bound.logical_keys): + context = bound.context_values[index].float().detach().requires_grad_() + a = ape.detach().requires_grad_() + n = norm.detach().requires_grad_() + group_ids = torch.tensor( + [group - 1, group] if group > 0 else [0], + dtype=torch.int64, + device=context.device, + ) + source = context if group > 0 else context[4:] + y = self.functional( + source, + a, + n, + group_ids, + cache, + ratio=4, + head_dim=512, + rope_dim=64, + eps=contract.rms_eps, + ) + if group > 0: + grad_out = torch.cat((torch.zeros_like(y[:1]), grad[index : index + 1].float())) + dx, da, dn = torch.autograd.grad(y, (context, a, n), grad_outputs=grad_out) + else: + dx, da, dn = torch.autograd.grad( + y, (source, a, n), grad_outputs=grad[index : index + 1].float() + ) + dx = torch.cat((torch.zeros_like(context[:4]), dx)) + contributions.append((bound.context_indices_host[index], dx)) + ape_rows.append(da) + norm_rows.append(dn) + record_backward( + "kv_compressor_c4", + kernel_id="megatron.compressed_compact_graph.ratio4+rlk.fixed_order_rows", + impl="upstream_functional_vjp", + family="c4", + ) + return contributions, torch.stack(ape_rows), torch.stack(norm_rows) diff --git a/rl_engine/contracts/operators/compressor_state.py b/rl_engine/contracts/operators/compressor_state.py new file mode 100644 index 000000000..4defad9b7 --- /dev/null +++ b/rl_engine/contracts/operators/compressor_state.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Versioned C4/C128 state contract; distinct from the eight-wave HCA profile.""" + +from dataclasses import asdict, dataclass +import hashlib +import json + +from rl_engine.contracts.operators.kv_compressor_c4 import C4Error, require + +CompressorStateError = C4Error + + +@dataclass(frozen=True) +class CompressorStateContract: + ratio: int = 4 + max_requests: int = 4 + max_tokens_per_step: int = 256 + max_sequence_tokens: int = 4096 + max_arena_bytes: int = 128 * 1024 * 1024 + world_size: int = 1 + graph: bool = False + + def __post_init__(self): + for name in ( + "ratio", + "max_requests", + "max_tokens_per_step", + "max_sequence_tokens", + "max_arena_bytes", + "world_size", + ): + require( + type(getattr(self, name)) is int and getattr(self, name) > 0, + f"{name} must be a positive integer", + ) + require(self.ratio in (4, 128), "only main-KV C4/C128 are admitted") + require(type(self.graph) is bool, "graph must be a bool") + require(self.max_sequence_tokens >= self.ratio, "capacity must fit one complete group") + require( + self.max_sequence_tokens < 2**31 and self.max_tokens_per_step < 2**31, + "provider int32 capacity exceeded", + ) + require(self.resident_bytes < 2**31, "provider buffer addressing envelope exceeded") + require(self.resident_bytes <= self.max_arena_bytes, "arena byte budget exceeded") + + @property + def width(self): + return 1024 if self.ratio == 4 else 512 + + @property + def state_rows(self): + return 8 if self.ratio == 4 else 128 + + @property + def plan_rows(self): + return (self.max_tokens_per_step + self.ratio - 1) // self.ratio + + @property + def pages_per_request(self): + return (self.max_sequence_tokens // self.ratio + 63) // 64 + + @property + def resident_bytes(self): + # Packed state + persistent contiguous reader copies, cache, staging, + # positions/write metadata, plan, tables, and model/position constants. + return ( + self.max_requests * self.state_rows * self.width * 4 * 4 + + self.max_requests * self.pages_per_request * 64 * 512 * 2 + + self.max_tokens_per_step * (self.width * 4 + 16) + + self.plan_rows * 16 + + self.max_requests * (self.pages_per_request + 1) * 4 + + self.ratio * self.width * 4 + + 512 * 4 + + self.max_sequence_tokens * 64 * 2 + ) + + @property + def arithmetic_id(self): + # This explicit single-wave profile must not be mixed with task 18's + # eight-wave C128 HCA recipe merely because shapes agree. + return f"aiter.main_c{self.ratio}.online_fp32.wave64.vllm_partial_f32.v1" + + @property + def fingerprint(self): + return hashlib.sha256(json.dumps(asdict(self), sort_keys=True).encode()).hexdigest() diff --git a/rl_engine/contracts/operators/kv_compressor_c4.py b/rl_engine/contracts/operators/kv_compressor_c4.py new file mode 100644 index 000000000..5db416de2 --- /dev/null +++ b/rl_engine/contracts/operators/kv_compressor_c4.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Item 19: overlap-window C4 main-KV compressor contract.""" + +from dataclasses import asdict, dataclass +import hashlib +import json +import math +from pathlib import Path + +from rl_engine.contracts.operators.attention import AttentionContractError + +C4_BACKEND_ID = "aiter.csa.c4.gfx942.bf16.overlap.fixed1.v1" +AITer_C4_SHA256 = "53cdd2801b6b2f5866b02d18bff1b2d56a19289a8048cdc14fe8bb687e572a14" +AITer_C4_LEGACY_SHA256 = "0c2019ca3a4a0d6dc7d5fdd20ebcce658ff5f504fdc880dffd61c9c398029451" +MLITE_BACKWARD_SHA256 = "6a1e78ca89b83f56391b5e5ac409e765338f1e3a625251d0d4c877c1e80b4a68" + + +class C4Error(AttentionContractError): + pass + + +def require(ok, message): + if not ok: + raise C4Error(message) + + +def source_hash(path): + require(isinstance(path, (str, Path)) and bool(path), "provider source is unavailable") + try: + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + except (OSError, ValueError) as exc: + raise C4Error("provider source is unreadable") from exc + + +@dataclass(frozen=True) +class C4Contract: + head_dim: int = 512 + ratio: int = 4 + overlap: int = 8 + rope_dim: int = 64 + rms_eps: float = 1e-6 + max_input_rows: int = 524288 + max_groups: int = 8192 + max_block_table_bytes: int = 16 * 1024 * 1024 + max_backward_groups: int = 128 + backend_id: str = C4_BACKEND_ID + output_dtype: str = "bf16" + graph_capture: bool = False + + def __post_init__(self): + for name in ( + "head_dim", + "ratio", + "overlap", + "rope_dim", + "max_input_rows", + "max_groups", + "max_block_table_bytes", + "max_backward_groups", + ): + require( + type(getattr(self, name)) is int and getattr(self, name) > 0, + f"{name} must be a positive integer", + ) + require( + (self.head_dim, self.ratio, self.overlap, self.rope_dim) == (512, 4, 8, 64), + "validated C4 main-KV profile is D=512, ratio=4, overlap=8, RoPE=64", + ) + require( + type(self.rms_eps) in (int, float) + and math.isfinite(self.rms_eps) + and 1e-12 <= self.rms_eps <= 1, + "invalid RMS epsilon", + ) + require( + self.max_input_rows < 2**31 and self.max_groups < 2**24, + "int32 provider indexing capacity exceeded", + ) + require(self.backend_id == C4_BACKEND_ID, "unknown C4 backend; no fallback") + require(self.output_dtype == "bf16", "C4 emits BF16 independent rows") + require(self.graph_capture is False, "Graph capture belongs to runtime integration") + + @property + def input_width(self): + return 4 * self.head_dim + + @property + def fingerprint(self): + return hashlib.sha256(json.dumps(asdict(self), sort_keys=True).encode()).hexdigest() + + @property + def arithmetic(self): + return { + "pool": ( + "aiter FP32 online softmax;K=8;overlap previous first4/current last4;fixed NW=1" + ), + "norm_rope": "FP32 RMSNorm;BF16 cos/sin;GPT-J tail;BF16 output", + "position": "global group ID * 4; first group previous window is -inf padded", + "backward": "mLite ratio4 compact graph;fixed logical row accumulation", + } diff --git a/rl_engine/integrations/common/compressor_state.py b/rl_engine/integrations/common/compressor_state.py new file mode 100644 index 000000000..816eb68cb --- /dev/null +++ b/rl_engine/integrations/common/compressor_state.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit arithmetic handshake; no implicit C128 HCA-to-legacy substitution.""" + +from rl_engine.models.deepseek_v4.compressor_state import CompressorStateCache + + +def bind_compressor_state(contract, *, expected_arithmetic, **kwargs): + return CompressorStateCache(contract, expected_arithmetic=expected_arithmetic, **kwargs) diff --git a/rl_engine/integrations/common/kv_compressor_c4.py b/rl_engine/integrations/common/kv_compressor_c4.py new file mode 100644 index 000000000..e3710bcb1 --- /dev/null +++ b/rl_engine/integrations/common/kv_compressor_c4.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in C4 operator boundary; item 21 owns state/cache lifetime.""" + +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor + + +def bind_c4_compressor(*, contract, backend=None, expected_implementation=None): + return C4Compressor(contract, backend=backend, expected_implementation=expected_implementation) diff --git a/rl_engine/integrations/engines/rollout/vllm/compressor_state.py b/rl_engine/integrations/engines/rollout/vllm/compressor_state.py new file mode 100644 index 000000000..75105685d --- /dev/null +++ b/rl_engine/integrations/engines/rollout/vllm/compressor_state.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Lifecycle boundary around upstream state/cache kernels, not full model registration.""" + +from rl_engine.integrations.common.compressor_state import bind_compressor_state + +__all__ = ["bind_compressor_state"] diff --git a/rl_engine/integrations/engines/rollout/vllm/kv_compressor_c4.py b/rl_engine/integrations/engines/rollout/vllm/kv_compressor_c4.py new file mode 100644 index 000000000..a0a26b99a --- /dev/null +++ b/rl_engine/integrations/engines/rollout/vllm/kv_compressor_c4.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""vLLM rollout adapter for complete C4 groups; no live cache commit.""" + +from rl_engine.integrations.common.kv_compressor_c4 import bind_c4_compressor + +__all__ = ["bind_c4_compressor"] diff --git a/rl_engine/integrations/engines/train/megatron/compressor_state.py b/rl_engine/integrations/engines/train/megatron/compressor_state.py new file mode 100644 index 000000000..70c3eec25 --- /dev/null +++ b/rl_engine/integrations/engines/train/megatron/compressor_state.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cache state is non-differentiable; upstream transport must precede ingestion.""" + +from rl_engine.integrations.common.compressor_state import bind_compressor_state + +__all__ = ["bind_compressor_state"] diff --git a/rl_engine/integrations/engines/train/megatron/kv_compressor_c4.py b/rl_engine/integrations/engines/train/megatron/kv_compressor_c4.py new file mode 100644 index 000000000..79d4aa306 --- /dev/null +++ b/rl_engine/integrations/engines/train/megatron/kv_compressor_c4.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Megatron adapter; reuse native transport and RLK canonical backward lifetime.""" + +from rl_engine.integrations.common.kv_compressor_c4 import bind_c4_compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session + +__all__ = ["bind_c4_compressor", "canonical_backward_session"] diff --git a/rl_engine/models/deepseek_v4/compressor_state.py b/rl_engine/models/deepseek_v4/compressor_state.py new file mode 100644 index 000000000..483d2e427 --- /dev/null +++ b/rl_engine/models/deepseek_v4/compressor_state.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-checked request leases around existing compressor/state GPU operators.""" + +from dataclasses import dataclass, field +import hashlib +import json +import math +from threading import RLock + +import torch + +from rl_engine.contracts.operators.compressor_state import CompressorStateContract +from rl_engine.contracts.operators.kv_compressor_c4 import require, source_hash + + +@dataclass(frozen=True) +class RequestHandle: + request_id: int + generation: int + version: str + writer_rank: int + slot: int + _arena: object = field(repr=False, compare=True) + + +@dataclass(frozen=True) +class CommitLease: + handle: RequestHandle + start: int + stop: int + first_group: int + end_group: int + + +@dataclass(frozen=True) +class CommitReceipt: + request_id: int + generation: int + version: str + first_group: int + end_group: int + tokens: int + arithmetic_id: str + implementation_fingerprint: str + + +@dataclass +class _Request: + handle: RequestHandle + tokens: int = 0 + commits: int = 0 + + +class CompressorStateCache: + """One staging lease at a time; completed commits are synchronous. + + This conservative ownership policy avoids early buffer reuse across host + threads and GPU streams. Communication overlap is a separate task. Input, + parameter, state and output ownership is explicit; cache writes have no VJP. + """ + + def __init__( + self, + contract, + *, + ape, + norm, + cos, + sin, + version, + writer_rank, + expected_arithmetic, + rms_eps=1e-6, + backend=None, + expected_implementation=None, + ): + require(type(contract) is CompressorStateContract, "typed state contract required") + require( + expected_arithmetic == contract.arithmetic_id, + "compressor arithmetic mismatch; C128 HCA eight-wave is a different profile", + ) + require( + isinstance(version, str) and bool(version), "explicit model/position version required" + ) + require( + type(writer_rank) is int and 0 <= writer_rank < contract.world_size, + "invalid writer rank", + ) + require( + type(rms_eps) in (int, float) and math.isfinite(rms_eps) and 1e-12 <= rms_eps <= 1, + "invalid RMS epsilon", + ) + if backend is None: + from rl_engine.backends.rocm.attention.compressor_state import ( + AiterCompressorStateBackend, + ) + + backend = AiterCompressorStateBackend() + self.contract, self.backend, self.rms_eps = contract, backend, rms_eps + require( + isinstance(ape, torch.Tensor) and ape.device.type == backend.device_type, + "parameter/backend device mismatch", + ) + self.device = ape.device + if hasattr(backend, "device"): + require(self.device == backend.device, "wrong provider device index") + self._eager() + self._lock, self._identity = RLock(), object() + self._version, self.writer_rank = version, writer_rank + self._requests, self._slots = {}, {} + self._generation, self._pending, self._poisoned = 0, None, False + self._graph = None + c = contract + with torch.inference_mode(False), torch.no_grad(): + self._ape, self._norm, self._cos, self._sin = self._parameter_copies( + ape, norm, cos, sin + ) + self._state = torch.zeros( + (c.max_requests, c.state_rows, 2 * c.width), device=self.device, dtype=torch.float32 + ) + self._read_kv = torch.empty( + (c.max_requests, c.state_rows, c.width), device=self.device, dtype=torch.float32 + ) + self._read_score = torch.empty_like(self._read_kv) + self._cache = torch.zeros( + (c.max_requests * c.pages_per_request, 64, 512), + device=self.device, + dtype=torch.bfloat16, + ) + self._input = torch.zeros( + (c.max_tokens_per_step, 2 * c.width), device=self.device, dtype=torch.bfloat16 + ) + self._positions = torch.zeros( + c.max_tokens_per_step, device=self.device, dtype=torch.int64 + ) + self._write_slots = torch.full_like(self._positions, -1) + self._plan = torch.zeros((c.plan_rows, 4), device=self.device, dtype=torch.int32) + self._plan[:, 2] = -1 + self._state_slots = torch.arange(c.max_requests, device=self.device, dtype=torch.int32) + self._blocks = torch.arange( + c.max_requests * c.pages_per_request, device=self.device, dtype=torch.int32 + ).reshape(c.max_requests, -1) + identity = { + "contract": c.fingerprint, + "arithmetic": c.arithmetic_id, + "eps": rms_eps, + "provider": getattr(backend, "provenance", {"test_backend": type(backend).__name__}), + "model_source": source_hash(__file__), + } + self.implementation_fingerprint = hashlib.sha256( + json.dumps(identity, sort_keys=True).encode() + ).hexdigest() + if expected_implementation is not None: + require( + expected_implementation == self.implementation_fingerprint, + "state implementation fingerprint mismatch", + ) + if c.graph: + require( + self.device.type == "cuda" and backend.supports_graph, + "Graph is unsupported by this backend", + ) + self._capture() + + def _eager(self): + if self.device.type == "cuda": + with torch.cuda.device(self.device): + require( + not torch.cuda.is_current_stream_capturing(), + "host lifecycle methods cannot run inside an external Graph capture", + ) + + def _ensure(self): + self._eager() + require( + not self._poisoned, + "state arena is poisoned; discard it rather than retrying a partial write", + ) + + def _parameter_copies(self, ape, norm, cos, sin): + c = self.contract + result = [] + for tensor, shape, dtype, name in ( + (ape, (c.ratio, c.width), torch.float32, "APE"), + (norm, (512,), torch.float32, "RMS"), + (cos, (c.max_sequence_tokens, 32), torch.bfloat16, "cos"), + (sin, (c.max_sequence_tokens, 32), torch.bfloat16, "sin"), + ): + require( + isinstance(tensor, torch.Tensor) + and tuple(tensor.shape) == shape + and tensor.dtype == dtype + and tensor.device == self.device + and tensor.is_contiguous(), + f"invalid {name} shape/dtype/device/layout", + ) + snapshot = tensor.detach().clone() + limit = 1 if name in ("cos", "sin") else 2.0**20 + require( + bool((torch.isfinite(snapshot) & (snapshot.abs() <= limit)).all()), + f"{name} outside finite admitted domain", + ) + result.append(snapshot) + return tuple(result) + + def _complete_device(self): + if self.device.type == "cuda": + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(self.device)) + event.synchronize() + + def _capture(self): + # Warmup/capture only PAD plans and invalid state writers. Native + # kernels compile at full static capacity without touching real state. + stream = torch.cuda.Stream(device=self.device) + stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(stream), torch.no_grad(): + for _ in range(3): + self.backend.execute(self) + torch.cuda.current_stream(self.device).wait_stream(stream) + self._graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self._graph), torch.no_grad(): + self.backend.execute(self) + self._complete_device() + require( + bool((self._state == 0).all()) and bool((self._cache == 0).all()), + "provider warmup/capture wrote real state", + ) + + def _request(self, handle): + require( + isinstance(handle, RequestHandle) and handle._arena is self._identity, + "foreign request handle", + ) + entry = self._requests.get(handle.request_id) + require( + entry is not None and entry.handle is handle and handle.version == self._version, + "stale request generation or model version", + ) + return entry + + def _clear_slot(self, slot): + try: + c = self.contract + page = slot * c.pages_per_request + self._state[slot].zero_() + self._cache[page : page + c.pages_per_request].zero_() + self._complete_device() + except BaseException: + self._poisoned = True + raise + + def create_request(self, request_id, *, version, writer_rank): + with self._lock, torch.inference_mode(False), torch.no_grad(): + self._ensure() + require( + type(request_id) is int and request_id >= 0, + "request_id must be a nonnegative integer", + ) + require(version == self._version, "model/position version mismatch") + require( + type(writer_rank) is int and 0 <= writer_rank < self.contract.world_size, + "invalid request writer rank", + ) + require(request_id not in self._requests, "request already active") + available = [i for i in range(self.contract.max_requests) if i not in self._slots] + require(bool(available), "request capacity exhausted") + slot = available[0] + self._generation += 1 + h = RequestHandle( + request_id, self._generation, version, writer_rank, slot, self._identity + ) + self._clear_slot(slot) + self._requests[request_id] = _Request(h) + self._slots[slot] = request_id + return h + + def prepare(self, handle, values, *, start_position): + with self._lock, torch.inference_mode(False), torch.no_grad(): + self._ensure() + entry = self._request(handle) + require(handle.writer_rank == self.writer_rank, "only the declared writer may commit") + require(self._pending is None, "another staging lease is outstanding") + require( + type(start_position) is int and start_position == entry.tokens, + "noncontiguous start or already committed token range", + ) + c = self.contract + require( + isinstance(values, torch.Tensor) + and values.dtype == torch.bfloat16 + and values.device == self.device + and values.ndim == 2 + and values.shape[1] == 2 * c.width + and values.is_contiguous() + and not values.requires_grad, + "cache inputs must be detached contiguous BF16 rows", + ) + count = values.shape[0] + require(count <= c.max_tokens_per_step, "step input capacity exceeded") + stop = start_position + count + require(stop <= c.max_sequence_tokens, "sequence capacity exceeded") + # Own data before inspecting content; original aliases can be reused + # after prepare returns. Concurrent writes during copying are forbidden. + copy = values.detach().clone() + require( + bool((torch.isfinite(copy.float()) & (copy.float().abs() <= 2.0**30)).all()), + "input outside finite FP32 accumulation domain", + ) + plan = [ + [ + p - start_position, + handle.slot, + p, + max(0, c.state_rows - (p - start_position + 1)), + ] + for p in range(start_position, stop) + if (p + 1) % c.ratio == 0 + ] + self._input.zero_() + self._input[:count].copy_(copy) + self._positions.zero_() + self._write_slots.fill_(-1) + if count: + positions = torch.arange( + start_position, stop, dtype=torch.int64, device=self.device + ) + self._positions[:count].copy_(positions) + first_writer = max(0, count - c.state_rows) + self._write_slots[first_writer:count].copy_( + handle.slot * c.state_rows + positions[first_writer:].remainder(c.state_rows) + ) + self._plan.zero_() + self._plan[:, 2] = -1 + if plan: + self._plan[: len(plan)].copy_( + torch.tensor(plan, dtype=torch.int32, device=self.device) + ) + # A different caller thread/stream may own commit(). Finish the + # staging writes before handing off this exclusive lease. + self._complete_device() + lease = CommitLease( + handle, start_position, stop, start_position // c.ratio, stop // c.ratio + ) + self._pending = lease + return lease + + def cancel(self, lease): + with self._lock: + self._ensure() + require( + isinstance(lease, CommitLease) and lease is self._pending, + "lease is stale, foreign or already consumed", + ) + self._pending = None + + def commit(self, lease): + with self._lock, torch.inference_mode(False), torch.no_grad(): + self._ensure() + require( + isinstance(lease, CommitLease) and lease is self._pending, + "lease is stale, foreign or already consumed", + ) + entry = self._request(lease.handle) + require(entry.tokens == lease.start, "token cursor changed after preparation") + try: + if lease.stop != lease.start: + if self._graph is None: + self.backend.execute(self) + else: + self._graph.replay() + self._complete_device() + except BaseException: + self._poisoned = True + raise + # Publish the cursor only after the ordered device writes complete. + entry.tokens = lease.stop + entry.commits += 1 + self._pending = None + return CommitReceipt( + lease.handle.request_id, + lease.handle.generation, + lease.handle.version, + lease.first_group, + lease.end_group, + lease.stop, + self.contract.arithmetic_id, + self.implementation_fingerprint, + ) + + def append(self, handle, values, *, start_position): + # Keep preparation and commit atomic relative to other host callers. + with self._lock: + return self.commit(self.prepare(handle, values, start_position=start_position)) + + def snapshot(self, handle): + with self._lock, torch.no_grad(): + self._ensure() + entry = self._request(handle) + c = self.contract + page = handle.slot * c.pages_per_request + rows = ( + self._cache[page : page + c.pages_per_request] + .flatten(0, 1)[: entry.tokens // c.ratio] + .clone() + ) + state = self._state[handle.slot].clone() + self._complete_device() + return { + "rows": rows, + "partial_ring": state, + "tokens": entry.tokens, + "incomplete_tokens": entry.tokens % c.ratio, + "commits": entry.commits, + "generation": handle.generation, + "version": handle.version, + "arithmetic_id": c.arithmetic_id, + } + + def release(self, handle): + with self._lock, torch.no_grad(): + self._ensure() + self._request(handle) + require( + self._pending is None or self._pending.handle is not handle, + "cancel the outstanding lease before releasing its request", + ) + self._clear_slot(handle.slot) + del self._requests[handle.request_id] + del self._slots[handle.slot] + + def replace_model(self, *, version, ape, norm, cos, sin): + with self._lock, torch.inference_mode(False), torch.no_grad(): + self._ensure() + require( + not self._requests and self._pending is None, + "release all requests before model replacement", + ) + require( + isinstance(version, str) and bool(version) and version != self._version, + "new model/position version required", + ) + copies = self._parameter_copies(ape, norm, cos, sin) + try: + for target, value in zip( + (self._ape, self._norm, self._cos, self._sin), copies, strict=True + ): + target.copy_(value) + self._complete_device() + except BaseException: + self._poisoned = True + raise + self._version = version diff --git a/rl_engine/models/deepseek_v4/kv_compressor_c4.py b/rl_engine/models/deepseek_v4/kv_compressor_c4.py new file mode 100644 index 000000000..b1d3bc396 --- /dev/null +++ b/rl_engine/models/deepseek_v4/kv_compressor_c4.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Complete-group C4 admission, fixed overlap state and canonical VJP wiring.""" + +from dataclasses import dataclass, field +import hashlib +import json + +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract, require, source_hash +from rl_engine.ops.autograd.canonical_backward import active_session + + +def eager(device): + if device.type == "cuda": + with torch.cuda.device(device): + require(not torch.cuda.is_current_stream_capturing(), "Graph capture is not validated") + + +def stamp(tensor): + return ( + tensor.data_ptr(), + None if tensor.is_inference() else tensor._version, + tensor.shape, + tensor.stride(), + tensor.dtype, + tensor.device, + ) + + +@dataclass(frozen=True) +class C4Groups: + source: torch.Tensor = field(repr=False) + values: torch.Tensor = field(repr=False) + context_values: torch.Tensor = field(repr=False) + context_indices: torch.Tensor = field(repr=False) + selector: torch.Tensor = field(repr=False) + group_ids: torch.Tensor = field(repr=False) + plan: torch.Tensor = field(repr=False) + block_table: torch.Tensor = field(repr=False) + kv_state: torch.Tensor = field(repr=False) + score_state: torch.Tensor = field(repr=False) + state_slots: torch.Tensor = field(repr=False) + logical_keys: tuple[tuple[int, int], ...] + request_spans: tuple[tuple[int, int], ...] + incomplete_groups: tuple[tuple[int, int], ...] + unowned_groups: tuple[tuple[int, int], ...] + context_indices_host: tuple[tuple[int, ...], ...] + cache_slots_host: tuple[tuple[int, int], ...] + version: str + owner_rank: int + cp_size: int + contract: C4Contract + _stamps: tuple = field(repr=False) + _source_stamp: tuple = field(repr=False) + + @property + def tensors(self): + return ( + self.values, + self.context_values, + self.context_indices, + self.selector, + self.group_ids, + self.plan, + self.block_table, + self.kv_state, + self.score_state, + self.state_slots, + ) + + def validate(self): + require(tuple(stamp(t) for t in self.tensors) == self._stamps, "bound C4 storage changed") + if self.source.requires_grad: + require(stamp(self.source) == self._source_stamp, "training C4 source changed") + + +@dataclass(frozen=True) +class C4Output: + values: torch.Tensor + group_ids: torch.Tensor + owner_rank: int + version: str + contract_fingerprint: str + implementation_fingerprint: str + + +def bind_complete_groups( + values, + token_ids, + *, + version, + contract, + owners=None, + cp_size=1, + rank=0, + emit_group_ids=None, +): + require(isinstance(contract, C4Contract), "typed C4 contract required") + require( + isinstance(version, str) and bool(version), + "explicit model/position/request version required", + ) + require( + isinstance(values, torch.Tensor) + and values.dtype == torch.bfloat16 + and values.ndim == 2 + and values.shape[1] == contract.input_width + and values.is_contiguous(), + "projected [tokens, 4*head_dim] contiguous BF16 input required", + ) + eager(values.device) + require(values.shape[0] <= contract.max_input_rows, "input capacity exceeded") + require( + isinstance(token_ids, torch.Tensor) + and token_ids.dtype == torch.int64 + and tuple(token_ids.shape) == (values.shape[0], 2) + and token_ids.device == values.device, + "token IDs must be int64 [tokens,2] on the input device", + ) + require( + type(cp_size) is int and cp_size > 0 and type(rank) is int and 0 <= rank < cp_size, + "invalid CP size/rank", + ) + host_ids = tuple(tuple(row) for row in token_ids.cpu().tolist()) + if owners is None: + require(cp_size == 1, "CP token ownership metadata required") + host_owners = (0,) * len(host_ids) + else: + require( + isinstance(owners, torch.Tensor) + and owners.dtype == torch.int64 + and tuple(owners.shape) == (values.shape[0],) + and owners.device == values.device, + "owners must be int64 [tokens] on the input device", + ) + host_owners = tuple(owners.cpu().tolist()) + emit = None if emit_group_ids is None else set(tuple(key) for key in emit_group_ids) + if emit is not None: + require( + all(len(key) == 2 and key[0] >= 0 and key[1] >= 0 for key in emit), + "emit_group_ids must contain nonnegative [request,group] keys", + ) + groups, seen = {}, set() + for physical, (key, owner) in enumerate(zip(host_ids, host_owners, strict=True)): + if key == (-1, -1): + continue + request, position = key + require(request >= 0 and 0 <= position < 2**31, "invalid logical token ID") + require(key not in seen, "duplicate logical token ID") + require(0 <= owner < cp_size, "token owner outside CP group") + seen.add(key) + group, offset = divmod(position, 4) + groups.setdefault((request, group), {})[offset] = (physical, owner) + + complete = {key: entries for key, entries in groups.items() if len(entries) == 4} + incomplete = [key for key, entries in groups.items() if len(entries) != 4] + selected, unowned = [], [] + for key, entries in sorted(complete.items()): + request, group = key + if emit is not None and key not in emit: + continue + if group > 0 and (request, group - 1) not in complete: + incomplete.append(key) + continue + if entries[3][1] != rank: + unowned.append(key) + continue + selected.append(key) + selected = sorted(selected) + require(len(selected) <= contract.max_groups, "complete-group capacity exceeded") + request_order = {request: i for i, request in enumerate(sorted({key[0] for key in selected}))} + max_block_column = max((key[1] // 64 + 1 for key in selected), default=0) + require( + len(request_order) * max_block_column * 4 <= contract.max_block_table_bytes, + "C4 block-table byte budget exceeded", + ) + block_order = { + key: index + for index, key in enumerate(sorted({(request, group // 64) for request, group in selected})) + } + + current_indices = [] + context_indices_host = [] + for request, group in selected: + current = [complete[(request, group)][offset][0] for offset in range(4)] + previous = ( + [complete[(request, group - 1)][offset][0] for offset in range(4)] + if group > 0 + else [-1] * 4 + ) + current_indices.extend(current) + context_indices_host.append(tuple(previous + current)) + cache_slots_host = tuple( + (block_order[(request, group // 64)], group % 64) for request, group in selected + ) + with torch.inference_mode(False), torch.no_grad(): + selector = torch.tensor(current_indices, dtype=torch.int64, device=values.device) + packed = values.index_select(0, selector).detach() + context_indices = torch.tensor( + context_indices_host, dtype=torch.int64, device=values.device + ).reshape(-1, 8) + safe_context = context_indices.clamp_min(0) + context = ( + values.index_select(0, safe_context.reshape(-1)) + .reshape(-1, 8, contract.input_width) + .detach() + ) + if context.numel(): + context[:, :4] = torch.where( + (context_indices[:, :4] >= 0).unsqueeze(-1), + context[:, :4], + torch.zeros_like(context[:, :4]), + ) + require( + bool((torch.isfinite(context.float()) & (context.float().abs() <= 2.0**30)).all()), + "selected input outside finite C4 FP32 accumulation domain", + ) + group_ids = torch.tensor(selected, dtype=torch.int64, device=values.device).reshape(-1, 2) + plan = torch.tensor( + [ + [i * 8 + 7, request_order[request], group * 4 + 3, 4 if group == 0 else 0] + for i, (request, group) in enumerate(selected) + ], + dtype=torch.int32, + device=values.device, + ).reshape(-1, 4) + table = torch.full( + (len(request_order), max_block_column), + -1, + dtype=torch.int32, + device=values.device, + ) + for (request, group), (block, _slot) in zip(selected, cache_slots_host, strict=True): + table[request_order[request], group // 64] = block + slots = torch.arange(len(request_order), dtype=torch.int32, device=values.device) + state_kv = torch.zeros( + (len(request_order), 8, 2 * contract.head_dim), + dtype=torch.float32, + device=values.device, + ) + state_score = torch.zeros_like(state_kv) + spans = [] + cursor = 0 + for request in sorted({key[0] for key in selected}): + length = sum(key[0] == request for key in selected) + spans.append((cursor, cursor + length)) + cursor += length + tensors = ( + packed, + context, + context_indices, + selector, + group_ids, + plan, + table, + state_kv, + state_score, + slots, + ) + return C4Groups( + values, + *tensors, + tuple(selected), + tuple(spans), + tuple(sorted(set(incomplete))), + tuple(unowned), + tuple(context_indices_host), + cache_slots_host, + version, + rank, + cp_size, + contract, + tuple(stamp(tensor) for tensor in tensors), + stamp(values), + ) + + +def canonical_sum(rows): + result = torch.zeros(rows.shape[1:], dtype=rows.dtype, device=rows.device) + for row in rows.unbind(0): + result = torch.add(result, row) + return result + + +def register_parameters(session, prefix, parameters, bound, implementation): + require(session is not None, "trainable C4 parameters require canonical_backward_session") + require(not session.backward_started, "register all C4 calls before backward") + require(isinstance(prefix, str) and bool(prefix), "stable layer parameter_id required") + bindings = getattr(session, "_c4_bindings", {}) + signature = (implementation, bound.version) + for suffix, parameter in parameters.items(): + if not parameter.requires_grad or not bound.logical_keys: + continue + entry = bindings.get(prefix + "." + suffix) + if entry is not None: + require( + entry[0] is parameter and entry[1] == signature, + "canonical parameter identity/version mismatch", + ) + require( + not entry[2].intersection(bound.logical_keys), + "duplicate C4 group across microbatches", + ) + slots = {} + for suffix, parameter in parameters.items(): + if parameter.requires_grad and bound.logical_keys: + key = prefix + "." + suffix + entry = bindings.setdefault(key, (parameter, signature, set())) + entry[2].update(bound.logical_keys) + slots[suffix] = (key, session.register(key, bound.group_ids)) + session._c4_bindings = bindings + return slots + + +class _C4Autograd(torch.autograd.Function): + @staticmethod + def forward(ctx, source, ape, norm, cos, sin, operator, bound, parameter_id): + ctx.operator, ctx.bound = operator, bound + ctx.session = active_session() + ctx.slots = {} + if (ape.requires_grad or norm.requires_grad) and bound.logical_keys: + ctx.slots = register_parameters( + ctx.session, + parameter_id, + {"ape": ape, "norm": norm}, + bound, + operator.implementation_fingerprint, + ) + ctx.save_for_backward(source, ape, norm, cos, sin) + return operator.backend.forward( + bound, ape.detach(), norm.detach(), cos, sin, operator.contract + ) + + @staticmethod + def backward(ctx, grad): + require(not torch.is_grad_enabled(), "higher-order C4 gradients are not validated") + source, ape, norm, cos, sin = ctx.saved_tensors + op, bound = ctx.operator, ctx.bound + eager(bound.values.device) + bound.validate() + if not bound.logical_keys: + return ( + torch.zeros_like(source), + torch.zeros_like(ape), + torch.zeros_like(norm), + None, + None, + None, + None, + None, + ) + contributions, da, dn = op.backend.vjp(bound, ape, norm, cos, sin, grad, op.contract) + dsource = torch.zeros_like(source) + for indices, rows in contributions: + for physical, row in zip(indices, rows, strict=True): + if physical >= 0: + dsource[physical] = dsource[physical] + row.to(source.dtype) + parameter_grads = [] + for suffix, rows, parameter in (("ape", da, ape), ("norm", dn, norm)): + if suffix not in ctx.slots: + parameter_grads.append(None) + continue + key, slot = ctx.slots[suffix] + reduced = ctx.session.submit_rows(key, slot, rows, canonical_sum) + parameter_grads.append(None if reduced is None else reduced.reshape(parameter.shape)) + return dsource, *parameter_grads, None, None, None, None, None + + +class C4Compressor: + def __init__(self, contract, *, backend=None, expected_implementation=None): + require(isinstance(contract, C4Contract), "typed C4 contract required") + if backend is None: + from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend + + backend = AiterC4Backend() + require(backend.backend_id == contract.backend_id, "C4 backend identity mismatch") + self.contract, self.backend = contract, backend + identity = { + "contract": contract.fingerprint, + "arithmetic": contract.arithmetic, + "backend": getattr(backend, "provenance", {"test_provider": type(backend).__name__}), + "bridge": source_hash(__file__), + "canonical": source_hash(active_session.__code__.co_filename), + } + self.implementation_fingerprint = hashlib.sha256( + json.dumps(identity, sort_keys=True).encode() + ).hexdigest() + if expected_implementation is not None: + require( + expected_implementation == self.implementation_fingerprint, + "C4 implementation fingerprint mismatch", + ) + + def bind( + self, + values, + token_ids, + *, + version, + owners=None, + cp_size=1, + rank=0, + emit_group_ids=None, + ): + require( + isinstance(values, torch.Tensor) and values.device.type == self.backend.device_type, + "C4 input/backend device mismatch", + ) + if hasattr(self.backend, "device"): + require(values.device == self.backend.device, "C4 device index mismatch") + return bind_complete_groups( + values, + token_ids, + version=version, + contract=self.contract, + owners=owners, + cp_size=cp_size, + rank=rank, + emit_group_ids=emit_group_ids, + ) + + def __call__(self, bound, ape, norm, cos, sin, *, version, parameter_id=""): + require( + isinstance(bound, C4Groups) and bound.contract == self.contract, + "C4 bound contract mismatch", + ) + require(bound.version == version, "C4 model/position/request version mismatch") + require(bound.values.device.type == self.backend.device_type, "C4 bound backend mismatch") + if hasattr(self.backend, "device"): + require(bound.values.device == self.backend.device, "C4 bound device index mismatch") + eager(bound.values.device) + bound.validate() + require( + isinstance(ape, torch.Tensor) + and tuple(ape.shape) == (4, 1024) + and ape.dtype == torch.float32 + and ape.is_contiguous() + and ape.device == bound.values.device, + "APE must be contiguous FP32 [4,1024]", + ) + require( + isinstance(norm, torch.Tensor) + and tuple(norm.shape) == (512,) + and norm.dtype == torch.float32 + and norm.is_contiguous() + and norm.device == bound.values.device, + "RMS weight must be contiguous FP32 [512]", + ) + for value, name in ((ape, "APE"), (norm, "RMS weight")): + require( + bool((torch.isfinite(value) & (value.abs() <= 2.0**20)).all()), + f"{name} outside finite C4 parameter domain", + ) + require( + isinstance(cos, torch.Tensor) + and isinstance(sin, torch.Tensor) + and cos.shape == sin.shape + and cos.ndim == 2 + and cos.shape[1] == 32 + and cos.dtype == torch.bfloat16 + and sin.dtype == torch.bfloat16 + and cos.is_contiguous() + and sin.is_contiguous() + and cos.device == bound.values.device + and sin.device == bound.values.device + and not cos.requires_grad + and not sin.requires_grad, + "RoPE tables must be matching contiguous BF16 [positions,32]", + ) + if bound.logical_keys: + max_position = max(group * 4 for _request, group in bound.logical_keys) + require(max_position < cos.shape[0], "RoPE table misses a global C4 position") + for table in (cos, sin): + # Do not create a transient GPU gather immediately before the + # FlyDSL launch; the provider applies this global offset in-kernel. + require( + bool((torch.isfinite(table) & (table.abs() <= 1)).all()), "invalid RoPE values" + ) + training = torch.is_grad_enabled() and any( + tensor.requires_grad for tensor in (bound.source, ape, norm) + ) + if training: + require( + len(bound.logical_keys) <= self.contract.max_backward_groups, + "C4 backward group budget exceeded", + ) + output = _C4Autograd.apply(bound.source, ape, norm, cos, sin, self, bound, parameter_id) + else: + with torch.no_grad(): + output = self.backend.forward(bound, ape, norm, cos, sin, self.contract) + require( + output.shape == (len(bound.logical_keys), 512) + and output.dtype == torch.bfloat16 + and output.device == bound.values.device, + "C4 provider returned invalid output", + ) + return C4Output( + output, + bound.group_ids, + bound.owner_rank, + version, + self.contract.fingerprint, + self.implementation_fingerprint, + ) diff --git a/tests/models/deepseek_v4/c4_helpers.py b/tests/models/deepseek_v4/c4_helpers.py new file mode 100644 index 000000000..1e7922153 --- /dev/null +++ b/tests/models/deepseek_v4/c4_helpers.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Independent ratio-4 overlap oracle, not a production fallback.""" + +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import C4_BACKEND_ID + + +def context_reference(context, ape, norm, group, cos, sin, eps): + if context.dtype != torch.float64: + context = context.float() + # Each source row is [content1024, score1024]. The overlap window uses + # previous group's first half and current group's second half. + previous = context[:4, :1024] + previous_gate = context[:4, 1024:] + current = context[4:, :1024] + current_gate = context[4:, 1024:] + if group == 0: + content = torch.cat((torch.zeros_like(previous[:, :512]), current[:, 512:]), dim=0) + gate = torch.cat( + ( + torch.full_like(previous_gate[:, :512], float("-inf")), + current_gate[:, 512:], + ), + dim=0, + ) + else: + content = torch.cat((previous[:, :512], current[:, 512:]), dim=0) + gate = torch.cat((previous_gate[:, :512], current_gate[:, 512:]), dim=0) + ape_window = torch.cat((ape[:, :512], ape[:, 512:]), dim=0) + gate = gate + ape_window + weights = torch.softmax(gate, dim=0) + pooled = (content * weights).sum(0) + pooled = pooled * torch.rsqrt(pooled.square().mean(-1, keepdim=True) + eps) + pooled = pooled * norm.to(context.dtype) + selected_cos = cos[group * 4].to(context.dtype) + selected_sin = sin[group * 4].to(context.dtype) + tail = pooled[-64:].reshape(32, 2) + rotated = torch.stack( + ( + tail[:, 0] * selected_cos - tail[:, 1] * selected_sin, + tail[:, 0] * selected_sin + tail[:, 1] * selected_cos, + ), + dim=-1, + ).flatten() + return torch.cat((pooled[:-64], rotated), dim=0) + + +def reference(values, ape, norm, groups, cos, sin, eps): + values = values.reshape(-1, 4, 2048) + outputs = [] + for group in range(values.shape[0]): + if group == 0: + context = torch.cat((torch.zeros_like(values[group]), values[group]), dim=0) + else: + context = torch.cat((values[group - 1], values[group]), dim=0) + outputs.append(context_reference(context, ape, norm, int(groups[group]), cos, sin, eps)) + return ( + torch.stack(outputs).to(values.dtype) + if outputs + else torch.empty((0, 512), dtype=values.dtype, device=values.device) + ) + + +def fixture(groups=4, *, request=7, start=0, seed=19, device="cpu", requires_grad=False): + generator = torch.Generator().manual_seed(seed) + rows = groups * 4 + values = torch.randn((rows, 2048), generator=generator).bfloat16().to(device) + values.requires_grad_(requires_grad) + positions = torch.arange(start, start + rows, dtype=torch.int64, device=device) + ids = torch.stack((torch.full_like(positions, request), positions), dim=-1) + ape = ( + (torch.randn((4, 1024), generator=generator) * 0.1).to(device).requires_grad_(requires_grad) + ) + norm = ( + (torch.randn((512,), generator=generator) * 0.1 + 1) + .to(device) + .requires_grad_(requires_grad) + ) + length = max(4, start + rows) + angles = torch.arange(length, dtype=torch.float32).unsqueeze(1) * ( + 10000.0 ** (-torch.arange(0, 64, 2).float() / 64) + ) + cos = angles.cos().bfloat16().to(device) + sin = angles.sin().bfloat16().to(device) + return values, ids, ape, norm, cos, sin + + +def byte_equal(actual, expected): + assert actual.shape == expected.shape and actual.dtype == expected.dtype + assert torch.equal( + actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8) + ) + + +class ReferenceC4Backend: + backend_id = C4_BACKEND_ID + device_type = "cpu" + + def __init__(self): + self.calls = 0 + + def forward(self, bound, ape, norm, cos, sin, contract): + self.calls += 1 + if not bound.logical_keys: + return torch.empty((0, 512), dtype=torch.bfloat16) + return torch.stack( + [ + context_reference( + bound.context_values[index].float(), + ape, + norm, + group, + cos, + sin, + contract.rms_eps, + ) + for index, (_request, group) in enumerate(bound.logical_keys) + ], + dim=0, + ).bfloat16() + + def vjp(self, bound, ape, norm, cos, sin, grad, contract): + contributions, ape_rows, norm_rows = [], [], [] + with torch.enable_grad(): + for index, (_request, group) in enumerate(bound.logical_keys): + context = bound.context_values[index].float().detach().requires_grad_() + a, n = ape.detach().requires_grad_(), norm.detach().requires_grad_() + y = context_reference(context, a, n, group, cos, sin, contract.rms_eps) + dx, da, dn = torch.autograd.grad( + y, (context, a, n), grad_outputs=grad[index].float() + ) + contributions.append((bound.context_indices[index], dx)) + ape_rows.append(da) + norm_rows.append(dn) + return contributions, torch.stack(ape_rows), torch.stack(norm_rows) diff --git a/tests/models/deepseek_v4/c4_process_probe.py b/tests/models/deepseek_v4/c4_process_probe.py new file mode 100644 index 000000000..a581f1c9d --- /dev/null +++ b/tests/models/deepseek_v4/c4_process_probe.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fresh-process C4 output/gradient fingerprint.""" + +import hashlib +import json +from pathlib import Path +import sys + +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from tests.models.deepseek_v4.c4_helpers import fixture + + +def digest(value): + return hashlib.sha256( + value.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + ).hexdigest() + + +if __name__ == "__main__": + values, ids, ape, norm, cos, sin = fixture(8, device="cuda", requires_grad=True) + operator = C4Compressor(C4Contract()) + with canonical_backward_session() as session: + output = operator( + operator.bind(values, ids, version="v1"), + ape, + norm, + cos, + sin, + version="v1", + parameter_id="probe", + ) + output.values.sum().backward() + session.validate_complete() + Path(sys.argv[1]).write_text( + json.dumps( + { + "output": digest(output.values), + "dvalues": digest(values.grad), + "dape": digest(ape.grad), + "dnorm": digest(norm.grad), + "implementation": operator.implementation_fingerprint, + }, + indent=2, + ) + ) diff --git a/tests/models/deepseek_v4/compressor_state_helpers.py b/tests/models/deepseek_v4/compressor_state_helpers.py new file mode 100644 index 000000000..9791c54fe --- /dev/null +++ b/tests/models/deepseek_v4/compressor_state_helpers.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Independent state/cache reference; used only through explicit test injection.""" + +import torch + +from rl_engine.contracts.operators.compressor_state import CompressorStateContract +from rl_engine.models.deepseek_v4.compressor_state import CompressorStateCache + + +def tensors(contract, *, seed=21, device="cpu"): + c = contract + generator = torch.Generator().manual_seed(seed) + x = torch.randn((c.max_sequence_tokens, c.width * 2), generator=generator).bfloat16().to(device) + ape = (torch.randn((c.ratio, c.width), generator=generator) * 0.1).to(device) + norm = (torch.randn(512, generator=generator) * 0.1 + 1).to(device) + angles = torch.arange(c.max_sequence_tokens).double()[:, None] * ( + 10000.0 ** (-torch.arange(0, 64, 2).double() / 64) + ) + return x, ape, norm, angles.cos().bfloat16().to(device), angles.sin().bfloat16().to(device) + + +def make_cache(ratio=4, *, backend=None, device="cpu", **kwargs): + c = CompressorStateContract(ratio=ratio, **kwargs) + x, ape, norm, cos, sin = tensors(c, device=device) + engine = CompressorStateCache( + c, + ape=ape, + norm=norm, + cos=cos, + sin=sin, + version="v1", + writer_rank=0, + expected_arithmetic=c.arithmetic_id, + backend=backend or ReferenceStateBackend(), + ) + return engine, x, (ape, norm, cos, sin) + + +class ReferenceStateBackend: + device_type = "cpu" + supports_graph = False + + def __init__(self): + self.calls = 0 + + def execute(self, arena): + self.calls += 1 + c = arena.contract + for ragged, slot, position, historical in arena._plan.tolist(): + if position < 0: + continue + contents, scores = [], [] + for k in range(c.state_rows): + absolute = position - c.state_rows + 1 + k + col = 512 if c.ratio == 4 and k >= 4 else 0 + if absolute < 0: + contents.append(torch.zeros(512)) + scores.append(torch.full((512,), -torch.inf)) + elif k < historical: + row = arena._state[slot, absolute % c.state_rows] + contents.append(row[col : col + 512]) + scores.append(row[c.width + col : c.width + col + 512]) + else: + row = arena._input[ragged - (c.state_rows - 1 - k)].float() + contents.append(row[col : col + 512]) + scores.append( + row[c.width + col : c.width + col + 512] + + arena._ape[k % c.ratio, col : col + 512] + ) + score = torch.stack(scores).double() + weights = (score - score.amax(0, keepdim=True)).exp() + weights /= weights.sum(0, keepdim=True) + pooled = (torch.stack(contents).double() * weights).sum(0) + out = ( + pooled * torch.rsqrt(pooled.square().mean() + arena.rms_eps) * arena._norm.double() + ) + group = position // c.ratio + cos, sin = arena._cos[group * c.ratio].double(), arena._sin[group * c.ratio].double() + even, odd = out[-64::2], out[-63::2] + rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), -1).flatten() + out = torch.cat((out[:-64], rotated)) + block = arena._blocks[slot, group // 64] + arena._cache[block, group % 64].copy_(out) + for row, physical in enumerate(arena._write_slots.tolist()): + if physical < 0: + continue + slot, offset = divmod(physical, c.state_rows) + source = arena._input[row].float() + position = int(arena._positions[row]) + arena._state[slot, offset, : c.width] = source[: c.width] + arena._state[slot, offset, c.width :] = ( + source[c.width :] + arena._ape[position % c.ratio] + ) + + +def byte_equal(a, b): + assert a.shape == b.shape and a.dtype == b.dtype + assert torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) diff --git a/tests/models/deepseek_v4/test_compressor_state.py b/tests/models/deepseek_v4/test_compressor_state.py new file mode 100644 index 000000000..2f08f3b11 --- /dev/null +++ b/tests/models/deepseek_v4/test_compressor_state.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import replace +from concurrent.futures import ThreadPoolExecutor + +import pytest +import torch + +from rl_engine.contracts.operators.compressor_state import ( + CompressorStateContract, + CompressorStateError, +) +from rl_engine.models.deepseek_v4.compressor_state import CompressorStateCache +from tests.models.deepseek_v4.compressor_state_helpers import ( + make_cache, + tensors, + ReferenceStateBackend, + byte_equal, +) + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize("ratio", [4, 128]) +def test_request_cursor_incomplete_groups_and_duplicate_commit(ratio): + arena, x, _ = make_cache(ratio, max_sequence_tokens=512) + h = arena.create_request(7, version="v1", writer_rank=0) + lease = arena.prepare(h, x[: ratio - 1], start_position=0) + assert arena.snapshot(h)["tokens"] == 0 + arena.commit(lease) + snapshot = arena.snapshot(h) + assert snapshot["rows"].shape == (0, 512) + assert snapshot["incomplete_tokens"] == ratio - 1 + with pytest.raises(CompressorStateError, match="consumed"): + arena.commit(lease) + byte_equal(arena.snapshot(h)["partial_ring"], snapshot["partial_ring"]) + receipt = arena.append(h, x[ratio - 1 : ratio], start_position=ratio - 1) + assert (receipt.first_group, receipt.end_group) == (0, 1) + assert arena.snapshot(h)["rows"].shape == (1, 512) + with pytest.raises(CompressorStateError, match="committed"): + arena.append(h, x[:1], start_position=0) + + +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("chunk", [1, 3, 7, 127, 256]) +def test_chunked_and_full_state_and_rows_match(ratio, chunk): + n = 269 + baseline, x, _ = make_cache(ratio, max_tokens_per_step=n, max_sequence_tokens=n) + h0 = baseline.create_request(1, version="v1", writer_rank=0) + baseline.append(h0, x, start_position=0) + expected = baseline.snapshot(h0) + arena, x, _ = make_cache(ratio, max_tokens_per_step=n, max_sequence_tokens=n) + h = arena.create_request(1, version="v1", writer_rank=0) + for start in range(0, n, chunk): + arena.append(h, x[start : start + chunk], start_position=start) + out = arena.snapshot(h) + byte_equal(out["rows"], expected["rows"]) + byte_equal(out["partial_ring"], expected["partial_ring"]) + assert out["tokens"] == n + + +def test_cancel_and_prepared_original_alias_reuse(): + arena, x, _ = make_cache(max_sequence_tokens=256) + h = arena.create_request(7, version="v1", writer_rank=0) + lease = arena.prepare(h, x[:4], start_position=0) + with pytest.raises(CompressorStateError, match="outstanding"): + arena.prepare(h, x[:4], start_position=0) + with pytest.raises(CompressorStateError, match="outstanding"): + arena.release(h) + arena.cancel(lease) + assert arena.snapshot(h)["tokens"] == 0 + with pytest.raises(CompressorStateError, match="stale"): + arena.commit(lease) + lease = arena.prepare(h, x[:4], start_position=0) + expected = x[:4].float().clone() + x[:4].data.fill_(torch.nan) + arena.commit(lease) + state = arena.snapshot(h)["partial_ring"] + byte_equal(state[:4, :1024], expected[:, :1024]) + for method in (arena.cancel, arena.commit): + with pytest.raises(CompressorStateError, match="stale"): + method(None) + + +def test_release_reuse_stale_foreign_and_model_barriers(): + arena, x, parameters = make_cache(max_requests=1, max_sequence_tokens=256) + h = arena.create_request(7, version="v1", writer_rank=0) + arena.append(h, x[:8], start_position=0) + with pytest.raises(CompressorStateError, match="release all"): + arena.replace_model( + version="v2", **dict(zip(("ape", "norm", "cos", "sin"), parameters, strict=True)) + ) + with pytest.raises(CompressorStateError, match="capacity"): + arena.create_request(8, version="v1", writer_rank=0) + arena.release(h) + with pytest.raises(CompressorStateError, match="stale"): + arena.release(h) + h2 = arena.create_request(7, version="v1", writer_rank=0) + assert h2.slot == h.slot and h2.generation > h.generation + assert bool((arena.snapshot(h2)["partial_ring"] == 0).all()) + with pytest.raises(CompressorStateError, match="stale"): + arena.snapshot(h) + other, _, _ = make_cache(max_sequence_tokens=256) + with pytest.raises(CompressorStateError, match="foreign"): + other.snapshot(h2) + arena.release(h2) + arena.replace_model( + version="v2", **dict(zip(("ape", "norm", "cos", "sin"), parameters, strict=True)) + ) + with pytest.raises(CompressorStateError, match="version"): + arena.create_request(7, version="v1", writer_rank=0) + new = arena.create_request(7, version="v2", writer_rank=0) + arena.append(new, x[:8], start_position=0) + assert arena.snapshot(new)["version"] == "v2" + + +def test_unique_writer_and_request_validation(): + arena, x, _ = make_cache(world_size=2, max_sequence_tokens=256) + other = arena.create_request(9, version="v1", writer_rank=1) + with pytest.raises(CompressorStateError, match="writer"): + arena.prepare(other, x[:4], start_position=0) + for rid, version, rank in [(True, "v1", 0), (-1, "v1", 0), (10, "v2", 0), (10, "v1", 2)]: + with pytest.raises(CompressorStateError): + arena.create_request(rid, version=version, writer_rank=rank) + with pytest.raises(CompressorStateError, match="already active"): + arena.create_request(9, version="v1", writer_rank=1) + + +@pytest.mark.parametrize( + "change", + [ + {"ratio": 8}, + {"ratio": True}, + {"max_requests": 0}, + {"world_size": 0}, + {"max_sequence_tokens": 1}, + {"max_sequence_tokens": 2**31}, + {"max_arena_bytes": 1}, + {"graph": 1}, + {"max_requests": 2**31, "max_arena_bytes": 2**100}, + ], +) +def test_invalid_capacity_and_contract(change): + with pytest.raises(CompressorStateError): + replace(CompressorStateContract(), **change) + + +@pytest.mark.parametrize( + "bad", ["dtype", "width", "stride", "nan", "inf", "grad", "capacity", "start"] +) +def test_bad_input_before_state_mutation(bad): + arena, x, _ = make_cache(max_sequence_tokens=256, max_tokens_per_step=8) + h = arena.create_request(1, version="v1", writer_rank=0) + source, start = x[:4].clone(), 0 + if bad == "dtype": + source = source.float() + elif bad == "width": + source = source[:, :512].contiguous() + elif bad == "stride": + source = source.T.contiguous().T + elif bad == "nan": + source[0, 0] = torch.nan + elif bad == "inf": + source[0, 0] = torch.inf + elif bad == "grad": + source.requires_grad_() + elif bad == "capacity": + source = x[:9] + else: + start = 1 + with pytest.raises(CompressorStateError): + arena.prepare(h, source, start_position=start) + assert arena.backend.calls == 0 + assert arena.snapshot(h)["tokens"] == 0 + + +def test_sequence_capacity_and_poisoned_arena_never_retries_partial_write(): + arena, x, _ = make_cache(max_sequence_tokens=8) + h = arena.create_request(1, version="v1", writer_rank=0) + arena.append(h, x, start_position=0) + with pytest.raises(CompressorStateError, match="sequence capacity"): + arena.prepare(h, x[:1], start_position=8) + + class Fault(ReferenceStateBackend): + def execute(self, arena): + super().execute(arena) + raise RuntimeError("simulated failure after native writes") + + arena, x, _ = make_cache(backend=Fault(), max_sequence_tokens=128) + h = arena.create_request(1, version="v1", writer_rank=0) + lease = arena.prepare(h, x[:4], start_position=0) + with pytest.raises(RuntimeError, match="simulated"): + arena.commit(lease) + for method, args in [(arena.commit, (lease,)), (arena.snapshot, (h,)), (arena.release, (h,))]: + with pytest.raises(CompressorStateError, match="poisoned"): + method(*args) + + +def test_concurrent_commits_consume_one_lease_once(): + arena, x, _ = make_cache(max_sequence_tokens=128) + h = arena.create_request(1, version="v1", writer_rank=0) + lease = arena.prepare(h, x[:8], start_position=0) + with ThreadPoolExecutor(2) as pool: + futures = [pool.submit(arena.commit, lease) for _ in range(2)] + success = 0 + for f in futures: + try: + f.result() + success += 1 + except CompressorStateError: + pass + assert success == 1 and arena.backend.calls == 1 + assert arena.snapshot(h)["commits"] == 1 + + +def test_explicit_profile_handshake_and_bad_model_inputs(): + c = CompressorStateContract(ratio=128, max_sequence_tokens=128) + _, a, n, co, si = tensors(c) + common = dict( + ape=a, norm=n, cos=co, sin=si, version="v1", writer_rank=0, backend=ReferenceStateBackend() + ) + with pytest.raises(CompressorStateError, match="arithmetic mismatch"): + CompressorStateCache(c, expected_arithmetic="aiter.hca.eight_waves", **common) + for bad in [ + {"ape": a.bfloat16()}, + {"norm": n[:127]}, + {"cos": co.float()}, + {"sin": torch.full_like(si, torch.nan)}, + {"rms_eps": 0}, + ]: + with pytest.raises(CompressorStateError): + CompressorStateCache(c, expected_arithmetic=c.arithmetic_id, **(common | bad)) + with pytest.raises(CompressorStateError, match="fingerprint"): + CompressorStateCache( + c, expected_arithmetic=c.arithmetic_id, expected_implementation="wrong", **common + ) + with pytest.raises(CompressorStateError, match="Graph"): + CompressorStateCache(replace(c, graph=True), expected_arithmetic=c.arithmetic_id, **common) + + +@pytest.mark.parametrize("operation", ["create", "release", "model"]) +def test_lifecycle_device_failure_also_poisons_arena(monkeypatch, operation): + arena, x, parameters = make_cache(max_sequence_tokens=128) + if operation == "release": + h = arena.create_request(7, version="v1", writer_rank=0) + arena.append(h, x[:4], start_position=0) + + def fault(): + raise RuntimeError("completion failed") + + monkeypatch.setattr(arena, "_complete_device", fault) + with pytest.raises(RuntimeError, match="completion"): + if operation == "create": + arena.create_request(7, version="v1", writer_rank=0) + elif operation == "release": + arena.release(h) + else: + arena.replace_model( + version="v2", **dict(zip(("ape", "norm", "cos", "sin"), parameters, strict=True)) + ) + with pytest.raises(CompressorStateError, match="poisoned"): + arena.create_request(8, version="v1", writer_rank=0) + + +def test_framework_factory_reuse_and_immutable_snapshot_outputs(): + from rl_engine.integrations.engines.rollout.vllm.compressor_state import ( + bind_compressor_state as rollout, + ) + from rl_engine.integrations.engines.train.megatron.compressor_state import ( + bind_compressor_state as train, + ) + + assert rollout is train + c = CompressorStateContract(max_sequence_tokens=128) + x, a, n, co, si = tensors(c) + arena = rollout( + c, + expected_arithmetic=c.arithmetic_id, + ape=a, + norm=n, + cos=co, + sin=si, + version="v1", + writer_rank=0, + backend=ReferenceStateBackend(), + ) + h = arena.create_request(7, version="v1", writer_rank=0) + arena.append(h, x[:8], start_position=0) + snapshot = arena.snapshot(h) + snapshot["rows"].fill_(torch.nan) + snapshot["partial_ring"].fill_(torch.nan) + actual = arena.snapshot(h) + assert bool(torch.isfinite(actual["rows"]).all()) + assert bool(torch.isfinite(actual["partial_ring"]).all()) diff --git a/tests/models/deepseek_v4/test_compressor_state_rocm.py b/tests/models/deepseek_v4/test_compressor_state_rocm.py new file mode 100644 index 000000000..1118b0619 --- /dev/null +++ b/tests/models/deepseek_v4/test_compressor_state_rocm.py @@ -0,0 +1,367 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native aiter/vLLM lifecycle, partial-state, single-commit and Graph tests.""" + +from concurrent.futures import ThreadPoolExecutor +import gc +import json + +import pytest +import torch + +from rl_engine.backends.rocm.attention.compressor_state import AiterCompressorStateBackend +from rl_engine.contracts.operators.compressor_state import ( + CompressorStateContract, + CompressorStateError, +) +from rl_engine.models.deepseek_v4.compressor_state import CompressorStateCache +from tests.models.deepseek_v4.compressor_state_helpers import ( + tensors, + byte_equal, + ReferenceStateBackend, +) + +pytestmark = [ + pytest.mark.ws1, + pytest.mark.skipif( + torch.version.hip is None or not torch.cuda.is_available(), reason="ROCm GPU" + ), +] + + +@pytest.fixture(scope="module") +def backend(): + return AiterCompressorStateBackend() + + +def make(backend, ratio, *, graph=False, seed=21, max_tokens=271, **kwargs): + c = CompressorStateContract( + ratio=ratio, + graph=graph, + max_sequence_tokens=max_tokens, + max_tokens_per_step=max_tokens, + **kwargs, + ) + x, a, n, co, si = tensors(c, seed=seed, device="cuda") + engine = CompressorStateCache( + c, + ape=a, + norm=n, + cos=co, + sin=si, + version="v1", + writer_rank=0, + expected_arithmetic=c.arithmetic_id, + backend=backend, + ) + return engine, x, (a, n, co, si) + + +def run(engine, x, chunk): + h = engine.create_request(7, version="v1", writer_rank=0) + rows_written = set() + for first in range(0, len(x), chunk): + receipt = engine.append(h, x[first : first + chunk], start_position=first) + for g in range(receipt.first_group, receipt.end_group): + assert g not in rows_written + rows_written.add(g) + assert rows_written == set(range(len(x) // engine.contract.ratio)) + return engine.snapshot(h), h + + +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("graph", [False, True]) +@pytest.mark.parametrize("chunk", [1, 3, 17, 129, 271]) +def test_complete_state_and_cache_bytes_across_prefill_decode_chunks_graph( + backend, ratio, graph, chunk +): + baseline, x, _ = make(backend, ratio) + expected, _ = run(baseline, x, len(x)) + engine, x, _ = make(backend, ratio, graph=graph) + actual, _ = run(engine, x, chunk) + byte_equal(actual["rows"], expected["rows"]) + byte_equal(actual["partial_ring"], expected["partial_ring"]) + assert actual["rows"].shape == (271 // ratio, 512) + assert actual["incomplete_tokens"] == 271 % ratio + if graph: + assert engine._graph is not None + + +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("seed", [0, 5, 21, 73]) +def test_numerical_cache_vs_independent_fp64_and_exact_ring_state(backend, ratio, seed): + native, x, (a, n, co, si) = make(backend, ratio, seed=seed, max_tokens=265) + actual, _ = run(native, x, 7) + reference = CompressorStateCache( + native.contract, + ape=a.cpu(), + norm=n.cpu(), + cos=co.cpu(), + sin=si.cpu(), + version="v1", + writer_rank=0, + expected_arithmetic=native.contract.arithmetic_id, + backend=ReferenceStateBackend(), + ) + expected, _ = run(reference, x.cpu(), 265) + byte_equal(actual["partial_ring"].cpu(), expected["partial_ring"]) + # The reference rounds its FP64 result to the declared BF16 output. A + # separate byte-exact test above covers BI rather than numerical tolerance. + torch.testing.assert_close( + actual["rows"].float().cpu(), expected["rows"].float(), rtol=0.009, atol=1e-4 + ) + + +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("graph", [False, True]) +def test_interleaved_requests_duplicate_lease_release_page_reuse_and_new_version( + backend, ratio, graph +): + engine, x, parameters = make(backend, ratio, graph=graph, max_requests=2) + h1 = engine.create_request(10, version="v1", writer_rank=0) + h2 = engine.create_request(11, version="v1", writer_rank=0) + for first in range(0, len(x), 17): + for handle, values in [(h1, x), (h2, -x)]: + lease = engine.prepare(handle, values[first : first + 17], start_position=first) + engine.commit(lease) + previous = engine.snapshot(handle) + with pytest.raises(CompressorStateError, match="consumed"): + engine.commit(lease) + byte_equal(engine.snapshot(handle)["rows"], previous["rows"]) + saved = engine.snapshot(h2) + engine.release(h1) + new = engine.create_request(10, version="v1", writer_rank=0) + assert new.slot == h1.slot and new.generation > h1.generation + assert engine.snapshot(new)["rows"].shape == (0, 512) + assert bool((engine.snapshot(new)["partial_ring"] == 0).all()) + with pytest.raises(CompressorStateError, match="stale"): + engine.prepare(h1, x[:1], start_position=0) + engine.append(new, x, start_position=0) + byte_equal(engine.snapshot(h2)["rows"], saved["rows"]) + byte_equal(engine.snapshot(h2)["partial_ring"], saved["partial_ring"]) + engine.release(new) + engine.release(h2) + engine.replace_model( + version="v2", **dict(zip(("ape", "norm", "cos", "sin"), parameters, strict=True)) + ) + with pytest.raises(CompressorStateError, match="version"): + engine.create_request(10, version="v1", writer_rank=0) + h3 = engine.create_request(10, version="v2", writer_rank=0) + engine.append(h3, x, start_position=0) + assert engine.snapshot(h3)["version"] == "v2" + + +@pytest.mark.parametrize("ratio", [4, 128]) +def test_capture_warmup_does_not_commit_or_write_partial_state(backend, ratio): + engine, x, _ = make(backend, ratio, graph=True) + assert bool((engine._cache == 0).all()) and bool((engine._state == 0).all()) + h = engine.create_request(7, version="v1", writer_rank=0) + engine.append(h, x[: ratio - 1], start_position=0) + assert engine.snapshot(h)["rows"].shape == (0, 512) + assert bool((engine._cache == 0).all()) + engine.append(h, x[ratio - 1 : ratio], start_position=ratio - 1) + snapshot = engine.snapshot(h) + assert snapshot["rows"].shape == (1, 512) + # Every unused cache row must retain its zero initializer. + assert bool((engine._cache.flatten(0, 1)[1:] == 0).all()) + + +@pytest.mark.parametrize("ratio", [4, 128]) +def test_upstream_order_unique_ring_writers_and_no_custom_math(backend, ratio, monkeypatch): + engine, x, _ = make(backend, ratio) + h = engine.create_request(7, version="v1", writer_rank=0) + observed = [] + compress, writer = backend.compress, backend.save_partial + + def compress_spy(**kwargs): + observed.append("compress") + assert kwargs["head_dim"] == 512 + assert kwargs.get("k_split_num_waves", 1) == 1 + return compress(**kwargs) + + def writer_spy(**kwargs): + observed.append("state") + slots = kwargs["slot_mapping"] + active = slots[slots >= 0].cpu().tolist() + assert len(active) == len(set(active)) == engine.contract.state_rows + return writer(**kwargs) + + with monkeypatch.context() as patch: + patch.setattr(backend, "compress", compress_spy) + patch.setattr(backend, "save_partial", writer_spy) + engine.append(h, x, start_position=0) + assert observed == ["compress", "state"] + + +def test_staging_owned_aliases_input_validation_and_writer_identity(backend): + engine, x, _ = make(backend, 4, graph=True, world_size=2) + h = engine.create_request(7, version="v1", writer_rank=0) + expected = x[:4].float().clone() + lease = engine.prepare(h, x[:4], start_position=0) + x[:4].data.zero_() + engine.commit(lease) + byte_equal(engine.snapshot(h)["partial_ring"][:4, :1024], expected[:, :1024]) + other = engine.create_request(8, version="v1", writer_rank=1) + with pytest.raises(CompressorStateError, match="writer"): + engine.prepare(other, x[:4], start_position=0) + # Rejection precedes any staging/cache/state write. + before = engine.snapshot(h) + x[4, 0] = torch.inf + with pytest.raises(CompressorStateError, match="finite"): + engine.prepare(h, x[4:8], start_position=4) + byte_equal(engine.snapshot(h)["partial_ring"], before["partial_ring"]) + byte_equal(engine.snapshot(h)["rows"], before["rows"]) + + +@pytest.mark.parametrize("graph", [False, True]) +def test_cross_stream_handoff_host_race_and_no_early_buffer_reuse(backend, graph): + engine, x, _ = make(backend, 4, graph=graph) + h = engine.create_request(7, version="v1", writer_rank=0) + producer, consumer = torch.cuda.Stream(), torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + lease = engine.prepare(h, x[:8], start_position=0) + with torch.cuda.stream(consumer): + engine.commit(lease) + assert engine.snapshot(h)["tokens"] == 8 + lease = engine.prepare(h, x[8:16], start_position=8) + + def commit(): + with torch.cuda.device(engine.device): + try: + engine.commit(lease) + return 1 + except CompressorStateError: + return 0 + + with ThreadPoolExecutor(2) as pool: + assert sum(pool.map(lambda _: commit(), range(2))) == 1 + assert engine.snapshot(h)["commits"] == 2 + + +def test_static_graph_buffers_stay_stable_across_requests_and_updates(backend, tmp_path): + engine, x, parameters = make(backend, 4, graph=True) + buffers = ( + "_state", + "_cache", + "_input", + "_plan", + "_positions", + "_write_slots", + "_read_kv", + "_read_score", + "_ape", + "_norm", + "_cos", + "_sin", + ) + pointers = {name: getattr(engine, name).data_ptr() for name in buffers} + h = engine.create_request(7, version="v1", writer_rank=0) + engine.append(h, x[:4], start_position=0) + engine.release(h) + gc.collect() + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + for _ in range(40): + h = engine.create_request(7, version="v1", writer_rank=0) + engine.append(h, x[:4], start_position=0) + engine.release(h) + gc.collect() + torch.cuda.synchronize() + after = torch.cuda.memory_allocated() + assert before == after + engine.replace_model( + version="v2", **dict(zip(("ape", "norm", "cos", "sin"), parameters, strict=True)) + ) + assert pointers == {name: getattr(engine, name).data_ptr() for name in buffers} + (tmp_path / "allocation.json").write_text( + json.dumps({"before": before, "after": after, "cycles": 40}) + ) + + +def test_model_weights_are_snapshotted_and_c128_profile_mismatch_is_rejected(backend): + engine, x, (a, n, co, si) = make(backend, 128) + reference, _, _ = make(backend, 128) + a.data.zero_() + n.data.zero_() + co.data.zero_() + si.data.zero_() + actual, _ = run(engine, x, 17) + expected, _ = run(reference, x, len(x)) + byte_equal(actual["rows"], expected["rows"]) + with pytest.raises(CompressorStateError, match="arithmetic mismatch"): + CompressorStateCache( + engine.contract, + ape=a, + norm=n, + cos=co, + sin=si, + version="v1", + writer_rank=0, + expected_arithmetic="aiter.hca.eight_wave.fp32", + backend=backend, + ) + + +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("seed", [0, 9, 73]) +def test_seeded_uneven_graph_updates_equal_full_native_execution(backend, ratio, seed): + baseline, x, _ = make(backend, ratio, seed=seed) + expected, _ = run(baseline, x, len(x)) + engine, x, _ = make(backend, ratio, seed=seed, graph=True) + h = engine.create_request(7, version="v1", writer_rank=0) + sizes = [1, 3, 4, 7, 31, 127] + start, step = 0, seed + while start < len(x): + count = min(sizes[step % len(sizes)], len(x) - start) + engine.append(h, x[start : start + count], start_position=start) + step += 1 + start += count + actual = engine.snapshot(h) + byte_equal(actual["rows"], expected["rows"]) + byte_equal(actual["partial_ring"], expected["partial_ring"]) + + +@pytest.mark.parametrize("ratio", [4, 128]) +def test_graph_model_replacement_uses_new_parameters_and_no_stale_request_data(backend, ratio): + engine, x, (a, n, co, si) = make(backend, ratio, graph=True) + old, h = run(engine, x, 17) + engine.release(h) + new_a, new_n = a * 1.5, n * 0.5 + engine.replace_model(version="v2", ape=new_a, norm=new_n, cos=co, sin=si) + fresh = CompressorStateCache( + engine.contract, + ape=new_a, + norm=new_n, + cos=co, + sin=si, + version="v2", + writer_rank=0, + expected_arithmetic=engine.contract.arithmetic_id, + backend=backend, + ) + snapshots = [] + for cache in (engine, fresh): + handle = cache.create_request(7, version="v2", writer_rank=0) + for start in range(0, len(x), 17): + cache.append(handle, x[start : start + 17], start_position=start) + snapshots.append(cache.snapshot(handle)) + byte_equal(snapshots[0]["rows"], snapshots[1]["rows"]) + byte_equal(snapshots[0]["partial_ring"], snapshots[1]["partial_ring"]) + assert not torch.equal(old["rows"], snapshots[0]["rows"]) + + +def test_zero_length_no_native_launch_and_external_graph_lifecycle_rejection(backend, monkeypatch): + engine, x, _ = make(backend, 4) + h = engine.create_request(7, version="v1", writer_rank=0) + + def forbidden(_): + raise AssertionError("zero-length commit must not launch") + + with monkeypatch.context() as patch: + patch.setattr(backend, "execute", forbidden) + receipt = engine.append(h, x[:0], start_position=0) + assert receipt.first_group == receipt.end_group == 0 + assert engine.snapshot(h)["tokens"] == 0 + with pytest.raises(CompressorStateError, match="external Graph"): + with torch.cuda.graph(torch.cuda.CUDAGraph()): + engine.prepare(h, x[:1], start_position=0) diff --git a/tests/models/deepseek_v4/test_kv_compressor_c4.py b/tests/models/deepseek_v4/test_kv_compressor_c4.py new file mode 100644 index 000000000..3522d8192 --- /dev/null +++ b/tests/models/deepseek_v4/test_kv_compressor_c4.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-safe C4 complete-group, overlap, ownership and VJP tests.""" + +from dataclasses import replace + +import pytest +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract, C4Error +from rl_engine.integrations.common.kv_compressor_c4 import bind_c4_compressor +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from tests.models.deepseek_v4.c4_helpers import ReferenceC4Backend, byte_equal, fixture, reference + +pytestmark = pytest.mark.unit + + +def operator(**kwargs): + return C4Compressor(C4Contract(**kwargs), backend=ReferenceC4Backend()) + + +@pytest.mark.parametrize("groups", [0, 1, 2, 3, 4, 7, 16]) +def test_only_complete_current_groups_emit_and_first_group_pads_previous(groups): + op = operator() + values, ids, ape, norm, cos, sin = fixture(groups) + tail, _, _, _, _, _ = fixture(1, seed=77) + tail_ids = torch.stack( + ( + torch.full((3,), 7, dtype=torch.int64), + torch.arange(groups * 4, groups * 4 + 3, dtype=torch.int64), + ), + dim=-1, + ) + values = torch.cat((values, tail[:3]), dim=0) + ids = torch.cat((ids, tail_ids), dim=0) + bound = op.bind(values, ids, version="v1") + output = op(bound, ape, norm, cos, sin, version="v1") + assert output.values.shape == (groups, 512) + assert bound.logical_keys == tuple((7, group) for group in range(groups)) + assert (7, groups) in bound.incomplete_groups + expected = ( + reference(values[: groups * 4].float(), ape, norm, range(groups), cos, sin, 1e-6).bfloat16() + if groups + else torch.empty((0, 512), dtype=torch.bfloat16) + ) + byte_equal(output.values, expected) + + +def test_overlap_uses_previous_group_first_half_and_current_group_last_half(): + op = operator() + values, ids, ape, norm, cos, sin = fixture(3, start=0) + values[0:4, :1024] = 1 + values[0:4, 1024:] = -1 + values[4:8, :1024] = 2 + values[4:8, 1024:] = 3 + values[8:12, :1024] = 4 + values[8:12, 1024:] = 5 + output = op(op.bind(values, ids, version="v1"), ape, norm, cos, sin, version="v1").values + expected = reference(values, ape, norm, range(3), cos, sin, 1e-6).bfloat16() + byte_equal(output, expected) + assert not torch.equal(output[1], output[0]) + assert not torch.equal(output[2], output[1]) + + +@pytest.mark.parametrize("rank,expected", [(0, ((7, 0),)), (1, ((7, 1),)), (2, ((7, 2),))]) +def test_cp_final_token_ownership_and_previous_context(rank, expected): + op = operator() + values, ids, ape, norm, cos, sin = fixture(3) + owners = torch.arange(12).remainder(3).long() + bound = op.bind(values, ids, version="v1", owners=owners, cp_size=3, rank=rank) + assert bound.logical_keys == expected + assert len(bound.unowned_groups) == 2 + assert bound.context_indices.shape == (1, 8) + if rank == 0: + assert tuple(bound.context_indices_host[0][:4]) == (-1, -1, -1, -1) + else: + assert all(index >= 0 for index in bound.context_indices_host[0]) + + +def test_permuted_rows_padding_duplicate_requests_and_snapshot_alias(): + op = operator() + values, ids, ape, norm, cos, sin = fixture(3) + expected = op(op.bind(values, ids, version="v1"), ape, norm, cos, sin, version="v1").values + second, second_ids, *_ = fixture(3, request=99, seed=29) + merged = torch.cat((values, second, torch.full((3, 2048), float("nan")).bfloat16())) + metadata = torch.cat((ids, second_ids, torch.full((3, 2), -1, dtype=torch.int64))) + permutation = torch.randperm(len(metadata), generator=torch.Generator().manual_seed(4)) + bound = op.bind(merged[permutation].contiguous(), metadata[permutation], version="v1") + merged.data.zero_() + output = op(bound, ape, norm, cos, sin, version="v1").values + byte_equal(output[:3], expected) + bound.context_values.zero_() + with pytest.raises(C4Error, match="storage changed"): + op(bound, ape, norm, cos, sin, version="v1") + + +@pytest.mark.parametrize( + "change", + [ + {"head_dim": 128}, + {"ratio": 128}, + {"overlap": 4}, + {"rope_dim": 32}, + {"rms_eps": 0}, + {"rms_eps": float("nan")}, + {"max_groups": 2**24}, + {"max_input_rows": 2**31}, + {"max_backward_groups": 0}, + {"backend_id": "auto"}, + {"output_dtype": "fp32"}, + {"graph_capture": True}, + ], +) +def test_invalid_contract(change): + with pytest.raises(C4Error): + replace(C4Contract(), **change) + + +@pytest.mark.parametrize( + "bad", + [ + "dtype", + "width", + "ids", + "partial", + "negative", + "duplicate", + "overflow", + "version", + "cp", + "cp_bool", + "owner_dtype", + "owner_range", + "owner_shape", + "nan", + "inf", + "stride", + ], +) +def test_invalid_input_never_reaches_backend(bad): + op = operator() + values, ids, *_ = fixture(2) + args = {"version": "v1"} + if bad == "dtype": + values = values.float() + elif bad == "width": + values = values[:, :1024].contiguous() + elif bad == "ids": + ids = ids.int() + elif bad == "partial": + ids[0, 0] = -1 + elif bad == "negative": + ids[0, 1] = -1 + elif bad == "duplicate": + ids[1] = ids[0] + elif bad == "overflow": + ids[0, 1] = 2**31 + elif bad == "version": + args["version"] = "" + elif bad == "cp": + args["cp_size"] = 2 + elif bad == "cp_bool": + args["cp_size"] = True + elif bad == "owner_dtype": + args["owners"] = torch.zeros(8) + elif bad == "owner_range": + args["owners"] = torch.full((8,), 3, dtype=torch.int64) + elif bad == "owner_shape": + args["owners"] = torch.zeros((8, 1), dtype=torch.int64) + elif bad in ("nan", "inf"): + values[0, 0] = float("nan") if bad == "nan" else float("inf") + else: + wide = torch.empty((values.shape[0], values.shape[1] * 2), dtype=values.dtype) + values = wide[:, ::2] + values.copy_(fixture(2)[0]) + with pytest.raises(C4Error): + op.bind(values, ids, **args) + assert op.backend.calls == 0 + + +@pytest.mark.parametrize( + "bad", + [ + "ape_dtype", + "ape_shape", + "ape_nan", + "norm_shape", + "norm_inf", + "rope_shape", + "rope_dtype", + "rope_short", + "rope_bad", + "version", + ], +) +def test_parameter_rope_and_version_contract(bad): + op = operator() + values, ids, ape, norm, cos, sin = fixture(2) + bound = op.bind(values, ids, version="v1") + version = "v1" + if bad == "ape_dtype": + ape = ape.bfloat16() + elif bad == "ape_shape": + ape = ape[:3] + elif bad == "ape_nan": + ape[0, 0] = float("nan") + elif bad == "norm_shape": + norm = norm[:511] + elif bad == "norm_inf": + norm[0] = float("inf") + elif bad == "rope_shape": + cos = cos[:, :16] + elif bad == "rope_dtype": + cos = cos.float() + elif bad == "rope_short": + cos, sin = cos[:1], sin[:1] + elif bad == "rope_bad": + sin[0, 0] = 2 + else: + version = "v2" + with pytest.raises(C4Error): + op(bound, ape, norm, cos, sin, version=version) + assert op.backend.calls == 0 + + +def train_partition(op, chunks, *, reverse=False): + values, ids, ape, norm, cos, sin = fixture(8, requires_grad=True) + upstream = torch.arange(8 * 512).reshape(8, 512).remainder(17).bfloat16() / 17 + with canonical_backward_session() as session: + calls = [] + for start, end in chunks: + input_start = max(0, start - 1) + bound = op.bind( + values[input_start * 4 : end * 4].contiguous(), + ids[input_start * 4 : end * 4], + version="v1", + emit_group_ids=tuple((7, group) for group in range(start, end)), + ) + output = op(bound, ape, norm, cos, sin, version="v1", parameter_id="layer.c4") + calls.append((output, upstream[start:end])) + order = list(range(len(calls))) + if reverse: + order.reverse() + for index in order: + calls[index][0].values.backward(calls[index][1]) + session.validate_complete() + return values.grad, ape.grad, norm.grad + + +def test_canonical_backward_fixed_order_and_overlap_accumulation(): + op = operator() + baseline = train_partition(op, [(0, 4), (4, 8)]) + for chunks, reverse in [ + ([(0, 2), (2, 4), (4, 8)], False), + ([(4, 8), (0, 4)], True), + ([(0, 1), (1, 3), (3, 8)], True), + ]: + actual = train_partition(operator(), chunks, reverse=reverse) + for got, expected in zip(actual, baseline, strict=True): + byte_equal(got, expected) + + +def test_input_only_backward_and_empty_groups(): + op = operator() + values, ids, ape, norm, cos, sin = fixture(1, requires_grad=True) + with canonical_backward_session() as session: + output = op( + op.bind(values, ids, version="v1"), ape.detach(), norm.detach(), cos, sin, version="v1" + ) + output.values.sum().backward() + session.validate_complete() + assert torch.count_nonzero(values.grad) > 0 + values, ids, ape, norm, cos, sin = fixture(0, requires_grad=True) + output = op(op.bind(values, ids, version="v1"), ape, norm, cos, sin, version="v1") + assert output.values.shape == (0, 512) + + +def test_session_identity_version_and_adapters(): + op = operator() + values, ids, ape, norm, cos, sin = fixture(2, requires_grad=True) + bound = op.bind(values, ids, version="v1") + with pytest.raises(C4Error, match="canonical_backward_session"): + op(bound, ape, norm, cos, sin, version="v1", parameter_id="layer") + with canonical_backward_session(): + op(bound, ape, norm, cos, sin, version="v1", parameter_id="layer") + with pytest.raises(C4Error, match="duplicate"): + op(bound, ape, norm, cos, sin, version="v1", parameter_id="layer") + from rl_engine.integrations.engines.rollout.vllm.kv_compressor_c4 import ( + bind_c4_compressor as rollout, + ) + from rl_engine.integrations.engines.train.megatron.kv_compressor_c4 import ( + bind_c4_compressor as train, + ) + + for factory in (rollout, train, bind_c4_compressor): + assert ( + factory( + contract=op.contract, + backend=ReferenceC4Backend(), + expected_implementation=op.implementation_fingerprint, + ).implementation_fingerprint + == op.implementation_fingerprint + ) diff --git a/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py b/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py new file mode 100644 index 000000000..9248dd777 --- /dev/null +++ b/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MI300X acceptance for the real aiter overlap C4 path and mLite VJP.""" + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract, C4Error +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from tests.models.deepseek_v4.c4_helpers import byte_equal, fixture, reference + +pytestmark = [ + pytest.mark.ws1, + pytest.mark.skipif( + torch.version.hip is None or not torch.cuda.is_available(), reason="ROCm GPU" + ), +] + + +@pytest.fixture(scope="module") +def backend(): + return AiterC4Backend() + + +def op(backend, **kwargs): + return C4Compressor(C4Contract(**kwargs), backend=backend) + + +def numerical_close(actual, expected): + exponent = torch.floor(torch.log2(expected.abs().clamp_min(2.0**-126))) + ulp = torch.exp2(exponent - 7) + assert bool(((actual.double() - expected).abs() <= (2 * ulp + 2e-5)).all()) + + +@pytest.mark.parametrize("groups", [1, 2, 3, 4, 7, 16, 33]) +def test_native_forward_independent_fp64_overlap_oracle(backend, groups): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(groups, device="cuda") + bound = operator.bind(values, ids, version="v1") + output = operator(bound, ape, norm, cos, sin, version="v1").values + expected = reference( + bound.values.double(), ape.double(), norm.double(), bound.group_ids[:, 1], cos, sin, 1e-6 + ) + numerical_close(output, expected) + assert backend.provenance["arch"] == "gfx942" + assert len(bound.incomplete_groups) == 0 + + +def test_native_forward_fixed_overlap_plan_and_cache_isolation(backend, monkeypatch): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(8, device="cuda") + captured = [] + original = backend.aiter.flydsl_fused_compress_attn + + def capture(**kwargs): + assert kwargs["ratio"] == 4 and kwargs["overlap"] is True + assert kwargs.get("k_split_num_waves", 1) == 1 + assert kwargs["k_per_block"] == 64 + captured.append(kwargs["plan_gpu"].clone()) + return original(**kwargs) + + with monkeypatch.context() as patch: + patch.setattr(backend.aiter, "flydsl_fused_compress_attn", capture) + output = operator( + operator.bind(values, ids, version="v1"), ape, norm, cos, sin, version="v1" + ).values + assert captured and captured[0].shape == (8, 4) + assert captured[0][:, 3].cpu().tolist() == [4, 0, 0, 0, 0, 0, 0, 0] + assert torch.equal(captured[0][:, 2].cpu(), torch.tensor([3, 7, 11, 15, 19, 23, 27, 31])) + expected = reference( + values.double(), ape.double(), norm.double(), torch.arange(8), cos, sin, 1e-6 + ) + numerical_close(output, expected) + + +@pytest.mark.parametrize("seed", range(8)) +def test_ragged_permuted_requests_padding_and_cp_ownership(backend, seed): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(5, seed=seed, device="cuda") + other, other_ids, *_ = fixture(5, request=99, seed=seed + 100, device="cuda") + merged = torch.cat( + (values, other, torch.full((3, 2048), float("nan"), dtype=values.dtype, device="cuda")) + ) + metadata = torch.cat((ids, other_ids, torch.full((3, 2), -1, dtype=torch.int64, device="cuda"))) + owners = torch.arange(len(metadata), device="cuda").remainder(3).long() + permutation = torch.randperm( + len(metadata), generator=torch.Generator().manual_seed(seed) + ).cuda() + merged, metadata, owners = merged[permutation], metadata[permutation], owners[permutation] + observed = {} + for rank in range(3): + bound = operator.bind(merged, metadata, owners=owners, cp_size=3, rank=rank, version="v1") + output = operator(bound, ape, norm, cos, sin, version="v1").values + for key, value in zip(bound.logical_keys, output, strict=True): + observed[key] = value + baseline = reference( + values.double(), ape.double(), norm.double(), torch.arange(5), cos, sin, 1e-6 + ) + assert len(observed) == 10 + for group in range(5): + numerical_close(observed[(7, group)], baseline[group]) + + +@pytest.mark.parametrize("content_scale,gate_scale", [(0, 1), (2.0**-25, 1), (2.0**20, 1), (1, 80)]) +def test_extreme_content_gate_scales_are_finite(backend, content_scale, gate_scale): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(4, device="cuda") + values[:, :1024] *= content_scale + values[:, 1024:] *= gate_scale + bound = operator.bind(values, ids, version="v1") + actual = operator(bound, ape, norm, cos, sin, version="v1").values + expected = reference( + bound.values.double(), ape.double(), norm.double(), torch.arange(4), cos, sin, 1e-6 + ) + numerical_close(actual, expected) + assert torch.isfinite(actual).all() + + +def train_chunks(backend, chunks, *, order=None, seed=19): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(8, seed=seed, device="cuda", requires_grad=True) + dy = torch.randn((8, 512), generator=torch.Generator().manual_seed(44)).bfloat16().cuda() + with canonical_backward_session() as session: + calls = [] + for start, end in chunks: + input_start = max(0, start - 1) + bound = operator.bind( + values[input_start * 4 : end * 4].contiguous(), + ids[input_start * 4 : end * 4], + version="v1", + emit_group_ids=tuple((7, group) for group in range(start, end)), + ) + output = operator( + bound, + ape, + norm, + cos, + sin, + version="v1", + parameter_id="layer.c4", + ) + calls.append((output, dy[start:end])) + sequence = list(range(len(calls))) if order is None else order + for index in sequence: + calls[index][0].values.backward(calls[index][1]) + session.validate_complete() + return values.grad, ape.grad, norm.grad + + +@pytest.mark.parametrize("seed", [0, 19, 73]) +def test_native_backward_order_invariance_and_independent_autograd(backend, seed): + baseline = train_chunks(backend, [(0, 4), (4, 8)], seed=seed) + for chunks, order in [ + ([(0, 2), (2, 4), (4, 8)], [0, 1, 2]), + ([(4, 8), (0, 4)], [1, 0]), + ([(0, 1), (1, 3), (3, 8)], [2, 0, 1]), + ]: + actual = train_chunks(backend, chunks, order=order, seed=seed) + for got, expected in zip(actual, baseline, strict=True): + byte_equal(got, expected) + values, ids, ape, norm, cos, sin = fixture(8, seed=seed, device="cuda", requires_grad=True) + dy = torch.randn((8, 512), generator=torch.Generator().manual_seed(44)).bfloat16().cuda() + vd, ad, nd = ( + values.detach().double().requires_grad_(), + ape.detach().double().requires_grad_(), + norm.detach().double().requires_grad_(), + ) + expected = torch.autograd.grad( + (reference(vd, ad, nd, torch.arange(8, device="cuda"), cos, sin, 1e-6) * dy.double()).sum(), + (vd, ad, nd), + ) + for got, want in zip(baseline, expected, strict=True): + torch.testing.assert_close(got.double(), want, rtol=4e-3, atol=3e-5) + + +def test_empty_input_no_native_launch_and_stream_reuse(backend, monkeypatch): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(0, device="cuda") + with monkeypatch.context() as patch: + patch.setattr( + backend.aiter, + "flydsl_fused_compress_attn", + lambda **_: (_ for _ in ()).throw(AssertionError("launch")), + ) + output = operator( + operator.bind(values, ids, version="v1"), ape, norm, cos, sin, version="v1" + ).values + assert output.shape == (0, 512) + values, ids, ape, norm, cos, sin = fixture(8, device="cuda") + bound = operator.bind(values, ids, version="v1") + expected = operator(bound, ape, norm, cos, sin, version="v1").values + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + for stream in streams: + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + actual = operator(bound, ape, norm, cos, sin, version="v1").values + torch.cuda.current_stream().wait_stream(stream) + byte_equal(actual, expected) + + +def test_mutation_and_graph_rejection(backend): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(2, device="cuda", requires_grad=True) + with canonical_backward_session(): + output = operator( + operator.bind(values, ids, version="v1"), + ape, + norm, + cos, + sin, + version="v1", + parameter_id="mutation", + ) + with torch.no_grad(): + values.add_(0.1) + with pytest.raises((RuntimeError, C4Error), match="modified|changed"): + output.values.sum().backward() + bound = operator.bind(values.detach(), ids, version="v1") + graph = torch.cuda.CUDAGraph() + with pytest.raises(C4Error, match="Graph"): + with torch.cuda.graph(graph): + operator(bound, ape.detach(), norm.detach(), cos, sin, version="v1") + + +def test_fresh_process_provider_fingerprint(backend, tmp_path): + del backend + script = Path(__file__).with_name("c4_process_probe.py") + records = [] + for index in range(2): + output = tmp_path / f"probe{index}.json" + env = { + **os.environ, + "PYTHONPATH": str(Path(__file__).resolve().parents[3]), + "TRITON_CACHE_DIR": str(tmp_path / f"cache{index}"), + } + subprocess.run( + [sys.executable, str(script), str(output)], + env=env, + check=True, + capture_output=True, + text=True, + timeout=180, + ) + records.append(json.loads(output.read_text())) + assert records[0] == records[1] + + +@pytest.mark.parametrize("bias", [-(2.0**30), -(2.0**25), 0, 2.0**25]) +def test_first_window_native_mask_excludes_padding_even_for_extreme_scores(backend, bias): + operator = op(backend) + values, ids, ape, norm, cos, sin = fixture(1, device="cuda") + values[:, 1024:] = bias + ape.zero_() + bound = operator.bind(values, ids, version="v1") + actual = operator(bound, ape, norm, cos, sin, version="v1").values + expected = reference( + values.double(), ape.double(), norm.double(), torch.tensor([0]), cos, sin, 1e-6 + ) + numerical_close(actual, expected) + + +def test_nonfinite_predecessor_context_is_not_hidden_by_emit_filter(backend): + operator = op(backend) + values, ids, *_ = fixture(2, device="cuda") + values[0, 0] = float("nan") + with pytest.raises(C4Error, match="finite"): + operator.bind(values, ids, version="v1", emit_group_ids=[(7, 1)]) + + +def test_fp64_oracle_really_preserves_fp64(): + values, ids, ape, norm, cos, sin = fixture(2, device="cuda") + x = values.double().requires_grad_() + a = ape.double().requires_grad_() + n = norm.double().requires_grad_() + y = reference(x, a, n, torch.arange(2, device="cuda"), cos, sin, 1e-6) + assert y.dtype == torch.float64 + eps = 1e-8 + with torch.no_grad(): + perturbed = x.clone() + perturbed[0, 0] += eps + assert not torch.equal( + y, reference(perturbed, a, n, torch.arange(2, device="cuda"), cos, sin, 1e-6) + ) + assert all(t.dtype == torch.float64 for t in torch.autograd.grad(y.sum(), (x, a, n)))