diff --git a/benchmarks/benchmark_indexer_compressor_c4.py b/benchmarks/benchmark_indexer_compressor_c4.py new file mode 100644 index 000000000..eb2992bcf --- /dev/null +++ b/benchmarks/benchmark_indexer_compressor_c4.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Separate binding, native pipeline and validation timing; no speedup claims.""" + +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.indexer_compressor_c4 import AiterIndexerC4Backend +from rl_engine.contracts.operators.indexer_compressor_c4 import IndexerC4Contract +from rl_engine.models.deepseek_v4.indexer_compressor_c4 import IndexerC4Compressor +from tests.models.deepseek_v4.indexer_c4_helpers import fixture + + +def time_us(fn, repeats): + for _ in range(3): + fn() + samples = [] + for _ in range(5): + torch.cuda.synchronize() + begin = time.perf_counter() + for _ in range(repeats): + fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - begin) * 1e6 / repeats) + return statistics.median(samples) + + +def run_case(backend, precision, groups, repeats): + op = IndexerC4Compressor(IndexerC4Contract(precision=precision), backend=backend) + x, ids, a, n, co, si = fixture(groups, device="cuda") + bound = op.bind(x, ids, version="v1") + rotated = backend.forward(bound, a, n, co, si, op.contract) + return { + "precision": precision, + "groups": groups, + "tokens": groups * 4, + "bind_us": time_us(lambda: op.bind(x, ids, version="v1"), repeats), + "native_compress_rotate_us": time_us( + lambda: backend.forward(bound, a, n, co, si, op.contract), repeats + ), + "quant_only_us": time_us(lambda: backend.quantize(rotated, op.contract), repeats), + "full_operator_us": time_us(lambda: op(bound, a, n, co, si, version="v1"), repeats), + "implementation": op.implementation_fingerprint, + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=30) + args = parser.parse_args() + backend = AiterIndexerC4Backend() + result = { + "backend": backend.provenance, + "metric": "median of five synchronized wall-time loops; Python dispatch included", + "scope": "no cache commit, full framework, full model, quantized VJP or speedup claim", + "cases": [ + run_case(backend, p, g, args.repeats) for p in ("mxfp4", "fp8") for g in (1, 8, 64, 128) + ], + } + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps(result, 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/indexer-compressor-c4.md b/docs/models/deepseek_v4/indexer-compressor-c4.md new file mode 100644 index 000000000..cbc5089a2 --- /dev/null +++ b/docs/models/deepseek_v4/indexer-compressor-c4.md @@ -0,0 +1,100 @@ +# Item 20 — separate 128-dimensional Indexer C4 compressor + +## Scope and dependency + +This task is stacked on task 19 / PR #618 (`8793052`). Its PR still targets +`dsv4-rocm`. Merge #618 first; the task-20 commit adds the Indexer-specific +contract and adapters while reusing the C4 implementation rather than copying +its grouping, ownership, cache-index mapping and autograd machinery. + +The main-KV contract remains D=512. `IndexerC4Contract` independently admits +D=128 with its own backend identity, family fingerprint, parameters, bound +data, version and canonical-gradient namespace. Passing the main-KV contract, +bound data or parameter geometry to the Indexer operator fails before launch. + +## Boundary + +```text +projected BF16 input [tokens,512] = [content256,score256] +APE FP32 [4,256], RMS weight FP32 [128], BF16 cos/sin [positions,32] +logical token IDs int64 [tokens,2] = [request,global_position] + -> complete four-token groups / eight-token overlap context + -> existing aiter compression, RMSNorm, interleaved tail RoPE + -> existing aiter normalized Walsh-Hadamard(128) + -> independent BF16 values [rows,128] + -> explicit MXFP4 or FNUZ FP8 payload and scales +``` + +Default `precision="mxfp4"` returns uint8 packed E2M1 `[rows,64]` and unshuffled +E8M0 bytes `[rows,4]`, group size 32, low nibble first. It reuses aiter +`per_1x32_f4_quant_triton`. The alternative explicit `precision="fp8"` reuses +vLLM `per_token_group_quant_fp8` with FNUZ encoding and power-of-two FP32 +per-row multipliers `[rows,1]`. These are distinct scale ABIs, not interchangeable +bytes. Item 16's independent numerical-domain admission still applies when +consuming these layouts. + +Both train and rollout adapters expose the same BF16 values, payload and scales. +Quantization is detached: this task does **not** define a straight-through +gradient through MXFP4/FP8. A loss on `output.values` receives the explicitly +defined BF16-output compressor/Hadamard VJP. The complete quantized Indexer +training chain belongs to item 26. + +## Reused numerical providers + +- aiter `flydsl_fused_compress_attn`: fixed one-wave C4 overlap forward. +- aiter `rotate_activation`: actual native Hadamard; **no generic matmul fallback**. +- aiter MXFP4 and vLLM FP8 quantizers above. +- Pinned mLite `compressed_compact_graph`: FP32 functional compressor VJP. +- RLK `CanonicalBackwardSession`: logical-group ordered parameter reductions. + +The upstream native D=128 Hadamard in the tested image uses eight physical rows +per wave and does not safely bound every partial last-wave allocation. The +adapter therefore pads **both** input and output buffers to eight rows and +returns only the logical rows. No numerical kernel is modified. Tests check +native buffer dimensions, numerical values and retention of live earlier +outputs/canaries across tail launches. + +Backward uses the same native transform for H-transpose (H is symmetric), with +an explicit BF16 dY/transform boundary, followed by mLite's FP32 compact VJP. +The numerical oracle separately checks this rounding boundary and the FP64 +compressor mathematics. Disjoint main/Indexer parameter namespaces are enforced +even if both adapters use the same caller layer ID. + +## State and lifecycle + +Complete groups, predecessor admission, CP final-token ownership, global RoPE +positions, output page mapping and mutation/version diagnostics are inherited +from task 19. `emit_group_ids` permits context-only predecessor groups at chunk +boundaries. C4 source tensors and parameters must not be mutated through any +alias while in use; version stamps are diagnostics, not storage write protection. +State and gradient storage are not shared between main and Indexer branches. + +This task writes a private temporary paged buffer and returns independent rows. +Request cache lifecycle, persistent partial state and exactly-once MXFP4 cache +commit belong to item 22. It does not claim Graph capture, full framework model +registration, RCCL/multi-GPU CP transport or full-model raw-logprob acceptance. + +## Reproduction on the tested member image + +Tested: MI300X `gfx942`, Torch `2.11.0+gitd0c8b1f`, HIP `7.2.53211`, +Triton 3.6.0, aiter 0.1.16.post3, FlyDSL 0.2.0, vLLM 0.26.0+rocm723. +The existing image was reused; no Torch/ROCm/aiter replacement was necessary. + +Set `RLK_MLITE_DSV4_BACKWARD_SOURCE` to the unmodified `backward.py` from mLite +commit `327a238239448de6771beea185dc6902a8de3eaa`, path +`experimental/lite/megatron/lite/model/deepseek_v4/vllm/primitive/attention/backward.py`. +The required SHA256 is +`6a1e78ca89b83f56391b5e5ac409e765338f1e3a625251d0d4c877c1e80b4a68`. +If mLite is installed, the shared backend can also resolve that module directly. + +```bash +python3 -m pytest tests/contracts/operators tests/models/deepseek_v4 -q \ + --basetemp=/outputs/item20-tests --junitxml=/outputs/item20-tests.xml +python3 benchmarks/benchmark_indexer_compressor_c4.py \ + --repeats 30 --output /outputs/item20-benchmark.json +``` + +Runtime fingerprints include the numerical providers, native rotation binary, +software versions, source hashes, profile and rotation padding policy. Configure +`expected_implementation` to require an exact train/rollout match. Fingerprints +are not a substitute for testing a different dependency stack. 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/indexer-compressor-c4-mi300x-20261029.md b/reports/indexer-compressor-c4-mi300x-20261029.md new file mode 100644 index 000000000..e9b84d2c5 --- /dev/null +++ b/reports/indexer-compressor-c4-mi300x-20261029.md @@ -0,0 +1,78 @@ +# Item 20 — independent Indexer C4 acceptance, 2026-10-29 + +Parent task-19 commit: `879305251c37e1b8bb0cb44409f1762752d03561`. +PR target: `dsv4-rocm` (`5ba769f681b7fa765ea6b29f9939454a858390ba`). +The member environment bootstrap built RLK main +`8588158650aff58c9c1f527674de94bad06c3d10` with ROCm Clang and verified +`rl_engine._C.fused_logp`; task tests ran against the task checkout. + +## Environment + +MI300X `gfx942`; Torch `2.11.0+gitd0c8b1f`; HIP `7.2.53211`; Triton 3.6.0; +aiter 0.1.16.post3; FlyDSL 0.2.0; vLLM 0.26.0+rocm723. + +Existing aiter C4/Hadamard/MXFP4 and vLLM FP8 operators are called directly. +Task 19's grouping/state descriptors and canonical backward session are shared. +There is no newly written numerical GPU kernel, duplicate C4 state machine or +generic matrix-multiply Hadamard fallback. + +## Acceptance + +- **296 tests passed twice**, zero failures/errors/skips. + First run with coverage: 57.56 s; independent repeat: 44.91 s. +- 25 item-20 CPU tests, 44 item-20 real GPU tests, 85 main-C4 regression tests + and 142 existing operator-contract tests. +- Each full run reported 16 warnings; warnings are not counted as passed tests. +- Exact byte comparisons cover BF16 output, packed payload and scales across + chunks, row permutations, padding, CP ownership materializations, training/ + rollout, page boundaries and independent processes. +- Real mLite VJP, native H-transpose, exact gradient comparisons across + microbatch/reverse orders, genuine FP64 compressor reference and parameter + finite differences. Quantized payloads deliberately have no gradient. +- Shared-session main/Indexer gradient isolation and cross-branch misuse + rejection, both CPU and native GPU. +- Native Hadamard tail allocations, previous-output/canary preservation, + empty/inference inputs, streams, mutation/version/Graph rejection. +- 50 repeated calls: live GPU allocation stayed **134,541,824 bytes** before + and after; this is live allocation, not a total process-memory bound. +- Ruff lint and format checks passed for the 14 changed/new Python files. + +Python statement coverage: **463/505 (91.68%)**; branch coverage **82/92 +(89.13%)**, across shared C4 and Indexer wrappers. Native compiled instruction +coverage is not measured. Coverage/autograd tracing gaps are not classified as +all constructor-failure paths. + +## Performance (microseconds) + +Median of five synchronized 30-repetition wall-time loops after warmup. +Python dispatch and validation are included; binding and quantization are +reported separately. These are not kernel-only latencies or speedup claims. + +| Precision | Groups | Bind | Native compress+rotate | Quantize | Full operator | +|---|---:|---:|---:|---:|---:| +| MXFP4 | 1 | 352.45 | 125.13 | 35.04 | 562.92 | +| MXFP4 | 8 | 511.48 | 110.64 | 42.62 | 589.46 | +| MXFP4 | 64 | 1674.79 | 237.77 | 44.43 | 714.58 | +| MXFP4 | 128 | 2997.66 | 381.71 | 43.75 | 890.14 | +| FP8 | 1 | 352.03 | 113.08 | 15.54 | 508.78 | +| FP8 | 8 | 508.17 | 110.00 | 15.53 | 505.23 | +| FP8 | 64 | 1696.73 | 236.97 | 15.96 | 647.97 | +| FP8 | 128 | 2951.27 | 364.55 | 15.74 | 792.66 | + +The small-row operator is Python/validation/launch dominated; no end-to-end +training or generation throughput improvement is claimed. + +Implementation fingerprints: + +- MXFP4: `f4d231272cbbf4d351ddc83244f26fb8a4af20799e3c84b21526b11a7f53df04` +- FP8: `c57b8cf88f38f0911eeddc949678cc1f6f0d3def5e11f366261765c61c666e5d` + +The report/log/JUnit/coverage/provenance and benchmark evidence are stored +under the member's mounted `/outputs/item20-20261029/`. + +## Limits + +No live request/cache commit (item 22), quantized STE/full Indexer training +chain (item 26), Graph/runtime integration (item 27), actual RCCL/GPU CP, +multi-node or full-model acceptance is claimed. The opt-in adapters share the +operator boundary but are not full vLLM/Megatron model registration. 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/indexer_compressor_c4.py b/rl_engine/backends/rocm/attention/indexer_compressor_c4.py new file mode 100644 index 000000000..620b9758c --- /dev/null +++ b/rl_engine/backends/rocm/attention/indexer_compressor_c4.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reuse task 19, aiter native Hadamard/MXFP4 and vLLM FP8; no new GPU kernel.""" + +import hashlib +import importlib +import json + +import torch + +from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend +from rl_engine.contracts.operators.indexer_compressor_c4 import INDEXER_C4_BACKEND_ID +from rl_engine.contracts.operators.kv_compressor_c4 import C4Error, source_hash + + +class AiterIndexerC4Backend(AiterC4Backend): + backend_id = INDEXER_C4_BACKEND_ID + + def __init__(self, *, megatron_backward_path=None): + super().__init__(megatron_backward_path=megatron_backward_path) + try: + quant = importlib.import_module("aiter.ops.quant") + fp8 = importlib.import_module("vllm.model_executor.layers.quantization.utils.fp8_utils") + self.rotate_activation = quant.rotate_activation + self.quant_mxfp4 = quant.per_1x32_f4_quant_triton + self.quant_fp8 = fp8.per_token_group_quant_fp8 + except (ImportError, AttributeError) as exc: + raise C4Error("native aiter rotation/quantization or vLLM FP8 ABI unavailable") from exc + modules = { + "aiter_quant": quant, + "aiter_fp4_utils": importlib.import_module("aiter.utility.fp4_utils"), + "vllm_fp8": fp8, + "aiter_tensor_shim": importlib.import_module("aiter.ops.flydsl.kernels.tensor_shim"), + } + quant_impl = modules["aiter_fp4_utils"].dynamic_mxfp4_quant + modules["mxfp4_implementation"] = importlib.import_module(quant_impl.__module__) + # Track the native rotation binary as well as Python glue. Priming the + # actual native entrypoint happens before deriving the fingerprint. + dummy = torch.zeros((1, 128), dtype=torch.bfloat16, device=self.device) + self._rotate(dummy) + native = importlib.import_module("aiter.jit.module_dsv4_rotate_quant") + modules["native_rotation"] = native + self.provenance["indexer_sources"] = { + label: source_hash(getattr(module, "__file__", None)) + for label, module in modules.items() + } + self.provenance["indexer_bridge"] = source_hash(__file__) + self.provenance["rotation"] = "aiter.rotate_activation, no torch.matmul fallback" + self.provenance["rotation_rows"] = "pad both buffers to 8 rows on gfx942, trim logical rows" + self.provenance["provider_sha256"] = hashlib.sha256( + json.dumps(self.provenance, sort_keys=True).encode() + ).hexdigest() + + def _rotate(self, values): + if not values.numel(): + return torch.empty_like(values) + # This native D=128 kernel executes 8 rows per wave. The installed + # provider's tail buffer extent is a full wave, so NEVER pass shorter + # physical input/output allocations. Padding only changes independent + # row launch geometry, not any within-row reduction. + rows = values.shape[0] + physical_rows = (rows + 7) // 8 * 8 + if physical_rows == rows: + source = values.contiguous() + else: + source = torch.zeros((physical_rows, 128), dtype=values.dtype, device=values.device) + source[:rows].copy_(values) + destination = torch.empty_like(source) + self.rotate_activation(destination, source) + return destination[:rows] + + def forward(self, bound, ape, norm, cos, sin, contract): + # Parent executes exactly the same validated C4 overlap and position + # handling, now with the separately admitted head_dim=128 profile. + compressed = super().forward(bound, ape, norm, cos, sin, contract) + return self._rotate(compressed) + + def quantize(self, values, contract): + if not values.shape[0]: + return ( + torch.empty( + (0, 64 if contract.precision == "mxfp4" else 128), + dtype=torch.uint8 if contract.precision == "mxfp4" else torch.float8_e4m3fnuz, + device=values.device, + ), + torch.empty( + (0, 4 if contract.precision == "mxfp4" else 1), + dtype=torch.uint8 if contract.precision == "mxfp4" else torch.float32, + device=values.device, + ), + ) + if contract.precision == "mxfp4": + data, scale = self.quant_mxfp4(values.contiguous(), shuffle=False) + return data.view(torch.uint8), scale.view(torch.uint8) + return self.quant_fp8( + values.contiguous(), + 128, + dtype=torch.float8_e4m3fnuz, + use_ue8m0=True, + column_major_scales=False, + ) + + def vjp(self, bound, ape, norm, cos, sin, grad, contract): + # H is symmetric; reuse the same native transform for its transpose. + # The BF16 dY boundary is explicit and is tested separately from the + # continuous FP64 mathematical reference. + unrotated_grad = self._rotate(grad.to(torch.bfloat16)) + return super().vjp(bound, ape, norm, cos, sin, unrotated_grad, contract) 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..91e996800 --- /dev/null +++ b/rl_engine/backends/rocm/attention/kv_compressor_c4.py @@ -0,0 +1,167 @@ +# 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) + dim = contract.head_dim + if groups == 0: + return torch.empty((0, dim), 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, dim), 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[:, : 2 * dim], + score_in=provider_values[:, 2 * dim :], + 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=dim, + 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=contract.head_dim, + 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( + "indexer_compressor_c4" if contract.family == "indexer_c4" else "kv_compressor_c4", + kernel_id="megatron.compressed_compact_graph.ratio4+rlk.fixed_order_rows", + impl="upstream_functional_vjp", + family=contract.family, + ) + return contributions, torch.stack(ape_rows), torch.stack(norm_rows) diff --git a/rl_engine/contracts/operators/indexer_compressor_c4.py b/rl_engine/contracts/operators/indexer_compressor_c4.py new file mode 100644 index 000000000..e427f9cb4 --- /dev/null +++ b/rl_engine/contracts/operators/indexer_compressor_c4.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Task 20's separate, typed Indexer profile reuses the C4 grouping contract.""" + +from dataclasses import dataclass +from typing import ClassVar + +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract, C4Error, require + +INDEXER_C4_BACKEND_ID = "aiter.indexer_c4.gfx942.native_hadamard.v1" +IndexerC4Error = C4Error + + +@dataclass(frozen=True) +class IndexerC4Contract(C4Contract): + expected_head_dim: ClassVar[int] = 128 + expected_backend_id: ClassVar[str] = INDEXER_C4_BACKEND_ID + family: ClassVar[str] = "indexer_c4" + head_dim: int = 128 + backend_id: str = INDEXER_C4_BACKEND_ID + precision: str = "mxfp4" + + def __post_init__(self): + super().__post_init__() + require(self.precision in ("mxfp4", "fp8"), "explicit MXFP4 or FNUZ FP8 precision required") + + @property + def layout(self): + return ( + "e2m1.low_nibble_first.row_major.e8m0_group32" + if self.precision == "mxfp4" + else "e4m3.row_major.fp32_per_row_scale" + ) + + @property + def arithmetic(self): + return { + **super().arithmetic, + "rotate": "aiter.rotate_activation;normalized_Walsh_Hadamard128;BF16", + "quantize": ( + "aiter.per_1x32_f4_quant_triton;unshuffled_group32_E8M0" + if self.precision == "mxfp4" + else "vllm.per_token_group_quant_fp8;FNUZ;UE8M0_FP32_multiplier" + ), + "backward_rotation": "native_Hadamard(BF16_dY);mLite_FP32_compact_VJP", + "quantized_backward": "not_defined;train_on_explicit_BF16_values_only", + } 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..5c6246cfb --- /dev/null +++ b/rl_engine/contracts/operators/kv_compressor_c4.py @@ -0,0 +1,108 @@ +# 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 typing import ClassVar + +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: + expected_head_dim: ClassVar[int] = 512 + expected_backend_id: ClassVar[str] = C4_BACKEND_ID + family: ClassVar[str] = "main_kv_c4" + 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) + == (self.expected_head_dim, 4, 8, 64), + f"validated {self.family} profile is D={self.expected_head_dim}, " + "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 == self.expected_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({"family": self.family, **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/indexer_compressor_c4.py b/rl_engine/integrations/common/indexer_compressor_c4.py new file mode 100644 index 000000000..72cc7d622 --- /dev/null +++ b/rl_engine/integrations/common/indexer_compressor_c4.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit opt-in Indexer boundary, not an entire framework registration.""" + +from rl_engine.models.deepseek_v4.indexer_compressor_c4 import IndexerC4Compressor + + +def bind_indexer_c4_compressor(*, contract, backend=None, expected_implementation=None): + return IndexerC4Compressor( + contract, backend=backend, expected_implementation=expected_implementation + ) 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/indexer_compressor_c4.py b/rl_engine/integrations/engines/rollout/vllm/indexer_compressor_c4.py new file mode 100644 index 000000000..3ed8cdb9d --- /dev/null +++ b/rl_engine/integrations/engines/rollout/vllm/indexer_compressor_c4.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Task 20 raw Indexer rows. Task 22 owns MXFP4 cache lifecycle/commit.""" + +from rl_engine.integrations.common.indexer_compressor_c4 import bind_indexer_c4_compressor + +__all__ = ["bind_indexer_c4_compressor"] 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/indexer_compressor_c4.py b/rl_engine/integrations/engines/train/megatron/indexer_compressor_c4.py new file mode 100644 index 000000000..8277546a1 --- /dev/null +++ b/rl_engine/integrations/engines/train/megatron/indexer_compressor_c4.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reuse common Indexer arithmetic and RLK canonical parameter gradients.""" + +from rl_engine.integrations.common.indexer_compressor_c4 import bind_indexer_c4_compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session + +__all__ = ["bind_indexer_c4_compressor", "canonical_backward_session"] 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/indexer_compressor_c4.py b/rl_engine/models/deepseek_v4/indexer_compressor_c4.py new file mode 100644 index 000000000..a0f1d66ed --- /dev/null +++ b/rl_engine/models/deepseek_v4/indexer_compressor_c4.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Small Indexer-specific layer over the shared C4 admission/autograd machinery.""" + +from dataclasses import dataclass +import hashlib +import json + +import torch + +from rl_engine.contracts.operators.indexer_compressor_c4 import IndexerC4Contract +from rl_engine.contracts.operators.kv_compressor_c4 import require, source_hash +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor + + +@dataclass(frozen=True) +class IndexerC4Output: + values: torch.Tensor + packed: torch.Tensor + scales: torch.Tensor + group_ids: torch.Tensor + owner_rank: int + version: str + contract_fingerprint: str + implementation_fingerprint: str + layout: str + precision: str + + +class IndexerC4Compressor(C4Compressor): + contract_type = IndexerC4Contract + + @staticmethod + def _make_backend(): + from rl_engine.backends.rocm.attention.indexer_compressor_c4 import AiterIndexerC4Backend + + return AiterIndexerC4Backend() + + def __init__(self, contract, *, backend=None, expected_implementation=None): + super().__init__(contract, backend=backend) + self.implementation_fingerprint = hashlib.sha256( + json.dumps( + { + "shared": self.implementation_fingerprint, + "indexer_model": source_hash(__file__), + }, + sort_keys=True, + ).encode() + ).hexdigest() + if expected_implementation is not None: + require( + expected_implementation == self.implementation_fingerprint, + "Indexer C4 implementation fingerprint mismatch", + ) + + def __call__(self, bound, ape, norm, cos, sin, *, version, parameter_id=""): + out = super().__call__( + bound, ape, norm, cos, sin, version=version, parameter_id=parameter_id + ) + # Both training and rollout expose identical quantized bytes/scales. + # Quantization is explicitly detached, NOT a hidden straight-through VJP. + with torch.no_grad(): + packed, scales = self.backend.quantize(out.values.detach(), self.contract) + fp4 = self.contract.precision == "mxfp4" + rows = len(bound.logical_keys) + require( + packed.shape == (rows, 64 if fp4 else 128) + and scales.shape == (rows, 4 if fp4 else 1) + and packed.dtype == (torch.uint8 if fp4 else torch.float8_e4m3fnuz) + and scales.dtype == (torch.uint8 if fp4 else torch.float32) + and packed.device == out.values.device + and scales.device == out.values.device + and packed.is_contiguous() + and scales.is_contiguous(), + "Indexer C4 provider returned an invalid quantized layout", + ) + return IndexerC4Output( + out.values, + packed, + scales, + out.group_ids, + out.owner_rank, + out.version, + out.contract_fingerprint, + self.implementation_fingerprint, + self.contract.layout, + self.contract.precision, + ) 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..f1e53ee1d --- /dev/null +++ b/rl_engine/models/deepseek_v4/kv_compressor_c4.py @@ -0,0 +1,524 @@ +# 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", + ) + require( + len(entry[2]) + len(bound.logical_keys) <= bound.contract.max_backward_groups, + "canonical backward group budget exceeded", + ) + 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: + contract_type = C4Contract + + @staticmethod + def _make_backend(): + from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend + + return AiterC4Backend() + + def __init__(self, contract, *, backend=None, expected_implementation=None): + require(type(contract) is self.contract_type, "typed C4 contract required") + if backend is None: + backend = self._make_backend() + 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, 2 * self.contract.head_dim) + and ape.dtype == torch.float32 + and ape.is_contiguous() + and ape.device == bound.values.device, + f"APE must be contiguous FP32 [4,{2 * self.contract.head_dim}]", + ) + require( + isinstance(norm, torch.Tensor) + and tuple(norm.shape) == (self.contract.head_dim,) + and norm.dtype == torch.float32 + and norm.is_contiguous() + and norm.device == bound.values.device, + f"RMS weight must be contiguous FP32 [{self.contract.head_dim}]", + ) + 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, + f"{self.contract.family}:{parameter_id}" if parameter_id else "", + ) + else: + with torch.no_grad(): + output = self.backend.forward(bound, ape, norm, cos, sin, self.contract) + require( + output.shape == (len(bound.logical_keys), self.contract.head_dim) + 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/indexer_c4_helpers.py b/tests/models/deepseek_v4/indexer_c4_helpers.py new file mode 100644 index 000000000..8bc4c4f59 --- /dev/null +++ b/tests/models/deepseek_v4/indexer_c4_helpers.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Test-only independent math; production always reuses upstream providers.""" + +import torch + +from rl_engine.contracts.operators.indexer_compressor_c4 import INDEXER_C4_BACKEND_ID + + +def fixture(groups=4, *, request=7, start=0, seed=20, device="cpu", requires_grad=False): + generator = torch.Generator().manual_seed(seed) + x = torch.randn((groups * 4, 512), generator=generator).bfloat16().to(device) + ape = (torch.randn((4, 256), generator=generator) * 0.1).to(device) + norm = (torch.randn((128,), generator=generator) * 0.1 + 1).to(device) + for t in (x, ape, norm): + t.requires_grad_(requires_grad) + positions = torch.arange(start, start + x.shape[0], dtype=torch.int64, device=device) + ids = torch.stack((torch.full_like(positions, request), positions), -1) + angle = torch.arange(max(4, start + x.shape[0])).double()[:, None] * ( + 10000.0 ** (-torch.arange(0, 64, 2).double() / 64) + ) + return x, ids, ape, norm, angle.cos().bfloat16().to(device), angle.sin().bfloat16().to(device) + + +def hadamard64(device): + h = torch.ones((1, 1), dtype=torch.float64, device=device) + for _ in range(7): + h = torch.cat((torch.cat((h, h), 1), torch.cat((h, -h), 1)), 0) + return h / 128**0.5 + + +def pre_rotation_reference(context, ape, norm, group, cos, sin, eps): + # Float64 callers remain float64 all the way through, including gradients. + x = context if context.dtype == torch.float64 else context.float() + dim = norm.numel() + content = torch.cat((x[:4, :dim], x[4:, dim : 2 * dim]), 0) + gate = torch.cat((x[:4, 2 * dim : 3 * dim] + ape[:, :dim], x[4:, 3 * dim :] + ape[:, dim:]), 0) + if group == 0: + gate = torch.cat((torch.full_like(gate[:4], -torch.inf), gate[4:]), 0) + weights = torch.exp(gate - gate.amax(0, keepdim=True)) + weights = weights / weights.sum(0, keepdim=True) + pooled = (content * weights).sum(0) + normed = pooled * torch.rsqrt(pooled.square().mean() + eps) * norm + c, s = cos[group * 4].to(x.dtype), sin[group * 4].to(x.dtype) + even, odd = normed[-64::2], normed[-63::2] + return torch.cat( + (normed[:-64], torch.stack((even * c - odd * s, odd * c + even * s), -1).flatten()) + ) + + +def rotate_oracle(x): + return (x.double() @ hadamard64(x.device).T).to(x.dtype) + + +class ReferenceIndexerC4Backend: + backend_id = INDEXER_C4_BACKEND_ID + device_type = "cpu" + + def forward(self, bound, ape, norm, cos, sin, contract): + if not bound.logical_keys: + return torch.empty((0, 128), dtype=torch.bfloat16) + out = torch.stack( + [ + pre_rotation_reference(ctx.float(), ape, norm, group, cos, sin, contract.rms_eps) + for ctx, (_req, group) in zip(bound.context_values, bound.logical_keys, strict=True) + ] + ).bfloat16() + return rotate_oracle(out) + + def quantize(self, values, contract): + # CPU contract tests only. Numerical quantization is tested with the + # actual upstream GPU providers, never this dummy layout implementation. + fp4 = contract.precision == "mxfp4" + return ( + torch.zeros( + (values.shape[0], 64 if fp4 else 128), + dtype=torch.uint8 if fp4 else torch.float8_e4m3fnuz, + ), + torch.ones( + (values.shape[0], 4 if fp4 else 1), dtype=torch.uint8 if fp4 else torch.float32 + ), + ) + + def vjp(self, bound, ape, norm, cos, sin, grad, contract): + contributions, aa, nn = [], [], [] + grad = rotate_oracle(grad.bfloat16()).float() + with torch.enable_grad(): + for i, (_req, group) in enumerate(bound.logical_keys): + x = bound.context_values[i].float().detach().requires_grad_() + a, n = ape.detach().requires_grad_(), norm.detach().requires_grad_() + y = pre_rotation_reference(x, a, n, group, cos, sin, contract.rms_eps) + dx, da, dn = torch.autograd.grad(y, (x, a, n), grad[i]) + contributions.append((bound.context_indices_host[i], dx)) + aa.append(da) + nn.append(dn) + return contributions, torch.stack(aa), torch.stack(nn) diff --git a/tests/models/deepseek_v4/indexer_c4_process_probe.py b/tests/models/deepseek_v4/indexer_c4_process_probe.py new file mode 100644 index 000000000..de048ed0a --- /dev/null +++ b/tests/models/deepseek_v4/indexer_c4_process_probe.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +import hashlib +import json +from pathlib import Path +import sys + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from rl_engine.contracts.operators.indexer_compressor_c4 import IndexerC4Contract +from rl_engine.models.deepseek_v4.indexer_compressor_c4 import IndexerC4Compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from tests.models.deepseek_v4.indexer_c4_helpers import fixture + + +def digest(t): + return hashlib.sha256( + t.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + ).hexdigest() + + +records = {} +for precision in ("mxfp4", "fp8"): + op = IndexerC4Compressor(IndexerC4Contract(precision=precision)) + x, ids, a, n, co, si = fixture(8, device="cuda", requires_grad=True) + with canonical_backward_session() as session: + out = op(op.bind(x, ids, version="v1"), a, n, co, si, version="v1", parameter_id="probe") + out.values.sum().backward() + session.validate_complete() + records[precision] = { + "implementation": op.implementation_fingerprint, + "values": digest(out.values), + "packed": digest(out.packed), + "scales": digest(out.scales), + "dx": digest(x.grad), + "da": digest(a.grad), + "dn": digest(n.grad), + } +Path(sys.argv[1]).write_text(json.dumps(records, indent=2)) diff --git a/tests/models/deepseek_v4/test_indexer_compressor_c4.py b/tests/models/deepseek_v4/test_indexer_compressor_c4.py new file mode 100644 index 000000000..f759301fb --- /dev/null +++ b/tests/models/deepseek_v4/test_indexer_compressor_c4.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import replace + +import pytest +import torch + +from rl_engine.contracts.operators.indexer_compressor_c4 import IndexerC4Contract +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract, C4Error +from rl_engine.models.deepseek_v4.kv_compressor_c4 import C4Compressor, C4Groups +from rl_engine.models.deepseek_v4.indexer_compressor_c4 import IndexerC4Compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from tests.models.deepseek_v4.c4_helpers import ReferenceC4Backend, byte_equal +from tests.models.deepseek_v4.indexer_c4_helpers import ( + fixture, + ReferenceIndexerC4Backend, + rotate_oracle, + pre_rotation_reference, +) + +pytestmark = pytest.mark.unit + + +def operator(**changes): + return IndexerC4Compressor(IndexerC4Contract(**changes), backend=ReferenceIndexerC4Backend()) + + +@pytest.mark.parametrize("precision", ["mxfp4", "fp8"]) +@pytest.mark.parametrize("groups", [0, 1, 2, 8, 17]) +def test_shared_binding_and_explicit_quantized_boundary(precision, groups): + op = operator(precision=precision) + x, ids, a, n, co, si = fixture(groups) + b = op.bind(x, ids, version="model1-indexer-epoch1") + assert isinstance(b, C4Groups) + assert b.contract.family == "indexer_c4" + out = op(b, a, n, co, si, version=b.version) + assert out.values.shape == (groups, 128) + assert out.packed.shape == (groups, 64 if precision == "mxfp4" else 128) + assert out.scales.shape == (groups, 4 if precision == "mxfp4" else 1) + assert out.layout == op.contract.layout and not out.packed.requires_grad + if groups: + expected = torch.stack( + [ + pre_rotation_reference(ctx.float(), a, n, g, co, si, 1e-6) + for ctx, (_, g) in zip(b.context_values, b.logical_keys, strict=True) + ] + ).bfloat16() + byte_equal(out.values, rotate_oracle(expected)) + + +@pytest.mark.parametrize( + "change", + [ + {"head_dim": 512}, + {"ratio": 128}, + {"overlap": 4}, + {"rope_dim": 32}, + {"precision": "bf16"}, + {"precision": "auto"}, + {"backend_id": C4Contract().backend_id}, + {"output_dtype": "fp32"}, + {"graph_capture": True}, + {"max_groups": False}, + ], +) +def test_reject_mixed_or_unknown_profiles(change): + with pytest.raises(C4Error): + replace(IndexerC4Contract(), **change) + + +def test_main_and_indexer_contract_state_and_parameter_isolation(): + from tests.models.deepseek_v4.c4_helpers import fixture as main_fixture + + main = C4Compressor(C4Contract(), backend=ReferenceC4Backend()) + index = operator() + mx, mi, ma, mn, co, si = main_fixture(2) + ix, ii, ia, inn, co, si = fixture(2) + mb, ib = main.bind(mx, mi, version="v1"), index.bind(ix, ii, version="v1") + with pytest.raises(C4Error, match="typed"): + C4Compressor(IndexerC4Contract(), backend=ReferenceIndexerC4Backend()) + with pytest.raises(C4Error, match="typed"): + IndexerC4Compressor(C4Contract(), backend=ReferenceC4Backend()) + for op, b, a, n in [ + (main, ib, ia, inn), + (index, mb, ma, mn), + (main, mb, ia, inn), + (index, ib, ma, mn), + ]: + with pytest.raises(C4Error): + op(b, a, n, co, si, version="v1") + before = index(ib, ia, inn, co, si, version="v1").values.clone() + mb.context_values.zero_() + byte_equal(index(ib, ia, inn, co, si, version="v1").values, before) + + +def train_partition(op, chunks, *, order=None): + x, ids, a, n, co, si = fixture(8, requires_grad=True) + dy = torch.randn((8, 128), generator=torch.Generator().manual_seed(37)).bfloat16() + with canonical_backward_session() as session: + calls = [] + for first, last in chunks: + lo = max(0, first - 1) + b = op.bind( + x[lo * 4 : last * 4], + ids[lo * 4 : last * 4], + version="v1", + emit_group_ids=[(7, i) for i in range(first, last)], + ) + out = op(b, a, n, co, si, version="v1", parameter_id="same-layer-prefix") + assert out.packed is not None and not out.packed.requires_grad + calls.append((out.values, dy[first:last])) + for i in range(len(calls)) if order is None else order: + calls[i][0].backward(calls[i][1]) + session.validate_complete() + return x.grad, a.grad, n.grad + + +def test_backward_reuses_canonical_rows_across_chunks_and_orders(): + baseline = train_partition(operator(), [(0, 8)]) + for chunks, order in [([(0, 4), (4, 8)], [1, 0]), ([(4, 8), (0, 2), (2, 4)], [0, 2, 1])]: + for a, b in zip(baseline, train_partition(operator(), chunks, order=order), strict=True): + byte_equal(a, b) + + +def test_shared_session_with_identical_layer_prefix_keeps_branches_separate(): + from tests.models.deepseek_v4.c4_helpers import fixture as main_fixture + + main = C4Compressor(C4Contract(), backend=ReferenceC4Backend()) + index = operator() + m = main_fixture(2, requires_grad=True) + i = fixture(2, requires_grad=True) + with canonical_backward_session() as session: + outputs = [] + for op, (x, ids, a, n, co, si) in [(main, m), (index, i)]: + outputs.append( + op( + op.bind(x, ids, version="v1"), a, n, co, si, version="v1", parameter_id="layer0" + ).values + ) + outputs[1].sum().backward() + assert m[0].grad is None and m[2].grad is None + outputs[0].sum().backward() + session.validate_complete() + assert len(session.uses) == 4 + assert m[2].grad.shape == (4, 1024) and i[2].grad.shape == (4, 256) + + +def test_framework_adapters_are_the_same_operator_not_duplicate_arithmetic(): + from rl_engine.integrations.engines.rollout.vllm.indexer_compressor_c4 import ( + bind_indexer_c4_compressor as rollout, + ) + from rl_engine.integrations.engines.train.megatron.indexer_compressor_c4 import ( + bind_indexer_c4_compressor as train, + ) + + assert train is rollout + op = operator() + assert train( + contract=op.contract, + backend=ReferenceIndexerC4Backend(), + expected_implementation=op.implementation_fingerprint, + ) + with pytest.raises(C4Error, match="fingerprint"): + train( + contract=op.contract, backend=ReferenceIndexerC4Backend(), expected_implementation="old" + ) + + +def test_backward_budget_covers_all_registered_chunks(): + op = operator(max_backward_groups=2) + x, ids, a, n, co, si = fixture(3, requires_grad=True) + with canonical_backward_session(): + b = op.bind(x, ids, version="v1", emit_group_ids=[(7, 0), (7, 1)]) + op(b, a, n, co, si, version="v1", parameter_id="budget") + b2 = op.bind(x, ids, version="v1", emit_group_ids=[(7, 2)]) + with pytest.raises(C4Error, match="budget"): + op(b2, a, n, co, si, version="v1", parameter_id="budget") diff --git a/tests/models/deepseek_v4/test_indexer_compressor_c4_rocm.py b/tests/models/deepseek_v4/test_indexer_compressor_c4_rocm.py new file mode 100644 index 000000000..0713c7083 --- /dev/null +++ b/tests/models/deepseek_v4/test_indexer_compressor_c4_rocm.py @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Actual upstream MI300X kernels; exact BI tests separated from numerical error.""" + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +from rl_engine.backends.rocm.attention.indexer_compressor_c4 import AiterIndexerC4Backend +from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend +from rl_engine.contracts.operators.indexer_compressor_c4 import IndexerC4Contract +from rl_engine.contracts.operators.kv_compressor_c4 import C4Error +from rl_engine.models.deepseek_v4.indexer_compressor_c4 import IndexerC4Compressor +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from tests.models.deepseek_v4.c4_helpers import byte_equal +from tests.models.deepseek_v4.indexer_c4_helpers import fixture, pre_rotation_reference, hadamard64 + +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 AiterIndexerC4Backend() + + +def operator(backend, precision="mxfp4"): + return IndexerC4Compressor(IndexerC4Contract(precision=precision), backend=backend) + + +def assert_bf16_close(actual, exact): + exponent = torch.floor(torch.log2(exact.abs().clamp_min(2.0**-126))) + tolerance = 2 * torch.exp2(exponent - 7) + 1e-5 + error = (actual.double() - exact).abs() + assert bool((error <= tolerance).all()), (error.max().item(), tolerance.max().item()) + + +def assert_quantization(out): + x = out.values.double() + if out.precision == "mxfp4": + assert out.packed.dtype == out.scales.dtype == torch.uint8 + assert bool((out.scales != 255).all()) + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float64, + device=x.device, + ) + code = torch.stack((out.packed & 15, out.packed >> 4), -1).flatten(-2) + scale = torch.exp2(out.scales.double() - 127).repeat_interleave(32, -1) + decoded = table[code.long()] * scale + # Independent nearest-representable check for the emitted E8M0 scale, + # including saturation. Does not call the upstream quantizer again. + distance = (x[..., None] - scale[..., None] * table).abs().amin(-1) + torch.testing.assert_close((decoded - x).abs(), distance, rtol=0, atol=1e-10) + else: + assert out.packed.dtype == torch.float8_e4m3fnuz + assert out.scales.dtype == torch.float32 + assert bool((torch.frexp(out.scales)[0] == 0.5).all()) + assert bool((out.scales > 0).all()) + expected = (out.values.float() / out.scales).clamp(-240, 240).to(torch.float8_e4m3fnuz) + byte_equal(out.packed, expected) + + +@pytest.mark.parametrize("precision", ["mxfp4", "fp8"]) +@pytest.mark.parametrize("groups", [0, 1, 2, 31, 65]) +def test_real_compression_rotation_and_quantized_layout(backend, precision, groups): + op = operator(backend, precision) + x, ids, a, n, co, si = fixture(groups, device="cuda") + b = op.bind(x, ids, version="v1") + out = op(b, a, n, co, si, version="v1") + assert out.values.shape == (groups, 128) + assert not out.packed.requires_grad and not out.scales.requires_grad + if not groups: + assert out.scales.shape == (0, 4 if precision == "mxfp4" else 1) + return + before_h = AiterC4Backend.forward(backend, b, a, n, co, si, op.contract) + exact = torch.stack( + [ + pre_rotation_reference(ctx.double(), a.double(), n.double(), group, co, si, 1e-6) + for ctx, (_r, group) in zip(b.context_values, b.logical_keys, strict=True) + ] + ) + assert_bf16_close(before_h, exact) + assert_bf16_close(out.values, before_h.double() @ hadamard64(x.device).T) + assert_quantization(out) + + +@pytest.mark.parametrize("precision", ["mxfp4", "fp8"]) +@pytest.mark.parametrize("seed", range(5)) +def test_exact_bytes_across_chunks_permutations_padding_and_cp(backend, precision, seed): + op = operator(backend, precision) + x, ids, a, n, co, si = fixture(17, device="cuda", seed=seed) + baseline = op(op.bind(x, ids, version="v1"), a, n, co, si, version="v1") + perm = torch.randperm(x.shape[0], generator=torch.Generator().manual_seed(seed)).cuda() + extra = torch.full((3, 512), torch.nan, dtype=x.dtype, device=x.device) + padded_x = torch.cat((x[perm], extra)) + padded_ids = torch.cat((ids[perm], torch.full((3, 2), -1, dtype=ids.dtype, device=x.device))) + owners = torch.cat( + (ids[perm, 1].remainder(3), torch.zeros(3, dtype=torch.int64, device=x.device)) + ) + seen = set() + for rank in range(3): + b = op.bind(padded_x, padded_ids, owners=owners, cp_size=3, rank=rank, version="v1") + out = op(b, a, n, co, si, version="v1") + for i, key in enumerate(b.logical_keys): + assert key not in seen + seen.add(key) + g = key[1] + for attr in ("values", "packed", "scales"): + byte_equal(getattr(out, attr)[i], getattr(baseline, attr)[g]) + assert len(seen) == 17 + for first, end in [(0, 1), (1, 4), (4, 17)]: + lo = max(0, first - 1) + b = op.bind( + x[lo * 4 : end * 4], + ids[lo * 4 : end * 4], + version="v1", + emit_group_ids=[(7, g) for g in range(first, end)], + ) + out = op(b, a, n, co, si, version="v1") + for attr in ("values", "packed", "scales"): + byte_equal(getattr(out, attr), getattr(baseline, attr)[first:end]) + + +@pytest.mark.parametrize("precision", ["mxfp4", "fp8"]) +def test_train_and_rollout_expose_identical_bytes_and_scales(backend, precision): + op = operator(backend, precision) + x, ids, a, n, co, si = fixture(4, device="cuda", requires_grad=True) + b = op.bind(x, ids, version="v1") + with torch.no_grad(): + rollout = op(b, a, n, co, si, version="v1") + with canonical_backward_session() as session: + train = op(b, a, n, co, si, version="v1", parameter_id="indexer") + for attr in ("values", "packed", "scales"): + byte_equal(getattr(train, attr), getattr(rollout, attr)) + train.values.sum().backward() + session.validate_complete() + assert x.grad is not None and a.grad is not None and n.grad is not None + + +def training_run(backend, chunks, order, seed): + op = operator(backend) + x, ids, a, n, co, si = fixture(8, device="cuda", requires_grad=True, seed=seed) + grad = torch.randn((8, 128), generator=torch.Generator().manual_seed(808)).bfloat16().cuda() + with canonical_backward_session() as session: + calls = [] + for first, end in chunks: + lo = max(0, first - 1) + b = op.bind( + x[lo * 4 : end * 4], + ids[lo * 4 : end * 4], + version="v1", + emit_group_ids=[(7, g) for g in range(first, end)], + ) + out = op(b, a, n, co, si, version="v1", parameter_id="layer1") + calls.append((out.values, grad[first:end])) + for i in order: + calls[i][0].backward(calls[i][1]) + session.validate_complete() + return (x.grad, a.grad, n.grad), (x, ids, a, n, co, si, grad) + + +@pytest.mark.parametrize("seed", [0, 20, 73]) +def test_real_vjp_exact_across_microbatches_and_independent_fp64(backend, seed): + baseline, inputs = training_run(backend, [(0, 8)], [0], seed) + for chunks, order in [([(0, 4), (4, 8)], [1, 0]), ([(4, 8), (0, 2), (2, 4)], [0, 2, 1])]: + actual, _ = training_run(backend, chunks, order, seed) + for a, b in zip(actual, baseline, strict=True): + byte_equal(a, b) + x, ids, a, n, co, si, grad = inputs + # Check the explicit BF16 H-transpose dY boundary independently; keep the + # compressor reference itself FP64 for its forward and gradient arithmetic. + unrotated = backend._rotate(grad) + assert_bf16_close(unrotated, grad.double() @ hadamard64(x.device).T) + xd, ad, nd = [t.detach().double().requires_grad_() for t in (x, a, n)] + outputs = [] + for g in range(8): + ctx = torch.cat( + ( + torch.zeros_like(xd[:4]) if g == 0 else xd[(g - 1) * 4 : g * 4], + xd[g * 4 : (g + 1) * 4], + ), + 0, + ) + outputs.append(pre_rotation_reference(ctx, ad, nd, g, co, si, 1e-6)) + expected = torch.autograd.grad(torch.stack(outputs), (xd, ad, nd), unrotated.double()) + torch.testing.assert_close(baseline[0].double(), expected[0], rtol=0.0045, atol=3e-5) + for actual, exact in zip(baseline[1:], expected[1:], strict=True): + torch.testing.assert_close(actual.double(), exact, rtol=3e-4, atol=5e-5) + + +def test_actual_native_calls_no_matmul_fallback(backend, monkeypatch): + op = operator(backend) + x, ids, a, n, co, si = fixture(4, device="cuda") + b = op.bind(x, ids, version="v1") + observed = {"rotate": 0, "quant": 0} + rotate, quant = backend.rotate_activation, backend.quant_mxfp4 + + def counted_rotate(*args, **kwargs): + observed["rotate"] += 1 + return rotate(*args, **kwargs) + + def counted_quant(*args, **kwargs): + observed["quant"] += 1 + return quant(*args, **kwargs) + + def forbidden(*args, **kwargs): + raise AssertionError("generic matmul is not the native Hadamard provider") + + with monkeypatch.context() as patch: + patch.setattr(backend, "rotate_activation", counted_rotate) + patch.setattr(backend, "quant_mxfp4", counted_quant) + patch.setattr(torch, "matmul", forbidden) + patch.setattr(torch.Tensor, "__matmul__", forbidden) + op(b, a, n, co, si, version="v1") + assert observed == {"rotate": 1, "quant": 1} + + +@pytest.mark.parametrize("rows", [1, 2, 3, 7, 8, 9, 15, 17, 65]) +def test_native_hadamard_tail_uses_full_wave_allocations_and_preserves_live_outputs( + backend, monkeypatch, rows +): + rotate = backend.rotate_activation + shapes = [] + + def checked(out, inp): + assert out.shape == inp.shape and inp.shape[0] % 8 == 0 + shapes.append(tuple(inp.shape)) + return rotate(out, inp) + + generator = torch.Generator().manual_seed(rows) + x = torch.randn((rows, 128), generator=generator).bfloat16().cuda() + with monkeypatch.context() as patch: + patch.setattr(backend, "rotate_activation", checked) + previous = backend._rotate(x) + saved = previous.clone() + canary = torch.full((9, 128), 3.25, dtype=torch.bfloat16, device="cuda") + for count in (1, 7, 9, 17): + backend._rotate(torch.ones((count, 128), dtype=x.dtype, device=x.device)) + torch.cuda.synchronize() + byte_equal(previous, saved) + assert bool((canary == 3.25).all()) + assert shapes[0] == ((rows + 7) // 8 * 8, 128) + assert_bf16_close(previous, x.double() @ hadamard64(x.device).T) + + +@pytest.mark.parametrize("precision", ["mxfp4", "fp8"]) +def test_zero_rows_first_group_extreme_gate_and_stream_reuse(backend, precision): + op = operator(backend, precision) + x, ids, a, n, co, si = fixture(4, device="cuda") + x.zero_() + x[:, 256:] = -(2.0**30) + a.zero_() + b = op.bind(x, ids, version="v1") + baseline = op(b, a, n, co, si, version="v1") + assert bool((baseline.values == 0).all()) + assert_quantization(baseline) + producer = torch.cuda.current_stream() + for stream in [torch.cuda.Stream(), torch.cuda.Stream()]: + stream.wait_stream(producer) + with torch.cuda.stream(stream): + out = op(b, a, n, co, si, version="v1") + producer.wait_stream(stream) + for attr in ("values", "packed", "scales"): + byte_equal(getattr(out, attr), getattr(baseline, attr)) + + +def test_graph_rejected_and_predecessor_mutations_detected(backend): + op = operator(backend) + x, ids, a, n, co, si = fixture(2, device="cuda") + b = op.bind(x, ids, version="v1") + with pytest.raises(C4Error, match="Graph"): + with torch.cuda.graph(torch.cuda.CUDAGraph()): + op(b, a, n, co, si, version="v1") + b.context_values.zero_() + with pytest.raises(C4Error, match="storage changed"): + op(b, a, n, co, si, version="v1") + + +def test_fresh_process_fingerprints(backend, tmp_path): + del backend + script = Path(__file__).with_name("indexer_c4_process_probe.py") + records = [] + for index in range(2): + output = tmp_path / f"record-{index}.json" + env = { + **os.environ, + "PYTHONPATH": str(Path(__file__).resolve().parents[3]), + "TRITON_CACHE_DIR": str(tmp_path / f"triton-{index}"), + } + proc = subprocess.run( + [sys.executable, str(script), str(output)], + env=env, + text=True, + capture_output=True, + timeout=180, + ) + assert proc.returncode == 0, proc.stderr[-8000:] + records.append(json.loads(output.read_text())) + assert records[0] == records[1] + (tmp_path / "process-evidence.json").write_text(json.dumps(records, indent=2)) + + +@pytest.mark.parametrize("precision", ["mxfp4", "fp8"]) +def test_nonzero_global_positions_across_cache_page_boundary(backend, precision): + op = operator(backend, precision) + x, ids, a, n, co, si = fixture(67, device="cuda") + baseline = op(op.bind(x, ids, version="v1"), a, n, co, si, version="v1") + b = op.bind( + x[62 * 4 : 67 * 4], + ids[62 * 4 : 67 * 4], + version="v1", + emit_group_ids=[(7, g) for g in range(63, 67)], + ) + out = op(b, a, n, co, si, version="v1") + assert b.cache_slots_host == ((0, 63), (1, 0), (1, 1), (1, 2)) + for attr in ("values", "packed", "scales"): + byte_equal(getattr(out, attr), getattr(baseline, attr)[63:67]) + + +def test_native_main_and_indexer_gradients_are_isolated_in_one_session(backend): + 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 as main_fixture + + main = C4Compressor(C4Contract(), backend=AiterC4Backend()) + index = operator(backend) + m = main_fixture(2, device="cuda", requires_grad=True) + i = fixture(2, device="cuda", requires_grad=True) + with canonical_backward_session() as session: + results = [] + for op, (x, ids, a, n, co, si) in [(main, m), (index, i)]: + results.append( + op( + op.bind(x, ids, version="v1"), + a, + n, + co, + si, + version="v1", + parameter_id="same-layer", + ) + ) + results[1].values.sum().backward() + assert all(t.grad is None for t in (m[0], m[2], m[3])) + saved = [t.grad.clone() for t in (i[0], i[2], i[3])] + results[0].values.sum().backward() + session.validate_complete() + assert len(session.uses) == 4 + for t, expected in zip((i[0], i[2], i[3]), saved, strict=True): + byte_equal(t.grad, expected) + + +def test_native_vjp_finite_difference_ape_and_norm(backend): + op = operator(backend) + x, ids, a, n, co, si = fixture(1, device="cuda", requires_grad=True) + grad = torch.randn((1, 128), generator=torch.Generator().manual_seed(13)).bfloat16().cuda() + with canonical_backward_session() as session: + out = op(op.bind(x, ids, version="v1"), a, n, co, si, version="v1", parameter_id="fd") + out.values.backward(grad) + session.validate_complete() + context = torch.cat((torch.zeros_like(x), x), 0).detach().double() + ad, nd = a.detach().double(), n.detach().double() + dg = backend._rotate(grad)[0].double() + + def loss(): + return (pre_rotation_reference(context, ad, nd, 0, co, si, 1e-6) * dg).sum() + + for target, idx, analytic in [ + (ad, (0, 128), a.grad[0, 128]), + (ad, (3, 255), a.grad[3, 255]), + (nd, (0,), n.grad[0]), + (nd, (127,), n.grad[127]), + ]: + epsilon = 1e-5 + value = target[idx].item() + target[idx] = value + epsilon + plus = loss() + target[idx] = value - epsilon + minus = loss() + target[idx] = value + torch.testing.assert_close( + analytic.double(), (plus - minus) / (2 * epsilon), rtol=5e-4, atol=3e-5 + ) + + +def test_inference_empty_autograd_and_workspace_reuse(backend, tmp_path): + import gc + + op = operator(backend) + with torch.inference_mode(): + x, ids, a, n, co, si = fixture(0, device="cuda") + out = op(op.bind(x, ids, version="v1"), a, n, co, si, version="v1") + assert out.values.shape == (0, 128) + x, ids, a, n, co, si = fixture(0, device="cuda", requires_grad=True) + with canonical_backward_session() as session: + out = op(op.bind(x, ids, version="v1"), a, n, co, si, version="v1", parameter_id="empty") + out.values.sum().backward() + session.validate_complete() + assert n.grad is not None and bool((n.grad == 0).all()) + x, ids, a, n, co, si = fixture(17, device="cuda") + b = op.bind(x, ids, version="v1") + for _ in range(3): + op(b, a, n, co, si, version="v1") + gc.collect() + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + for _ in range(50): + op(b, a, n, co, si, version="v1") + gc.collect() + torch.cuda.synchronize() + after = torch.cuda.memory_allocated() + assert after <= before + 512 * 1024 + (tmp_path / "allocation-evidence.json").write_text( + json.dumps( + { + "calls": 50, + "live_before": before, + "live_after": after, + }, + indent=2, + ) + ) 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)))