[DSV4][ROCm] Reuse C4 and native quantization for separate Indexer compression (task 20)

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))

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()

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.

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.