[DSV4][ROCm] Add versioned compressor state and single cache commit (task 21)

benchmarks/benchmark_compressor_state.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stateful append timing includes validation, staging and synchronous commit."""
+
+import argparse
+import json
+from pathlib import Path
+import statistics
+import sys
+import time
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from rl_engine.backends.rocm.attention.compressor_state import AiterCompressorStateBackend
+from rl_engine.contracts.operators.compressor_state import CompressorStateContract
+from rl_engine.models.deepseek_v4.compressor_state import CompressorStateCache
+from tests.models.deepseek_v4.compressor_state_helpers import tensors
+
+
+def case(backend, ratio, graph, rows, repeats):
+    c = CompressorStateContract(
+        ratio=ratio,
+        graph=graph,
+        max_requests=1,
+        max_tokens_per_step=rows,
+        max_sequence_tokens=max(ratio, rows * (5 * repeats + 3)),
+    )
+    x, a, n, co, si = tensors(c, device="cuda")
+    arena = CompressorStateCache(
+        c,
+        ape=a,
+        norm=n,
+        cos=co,
+        sin=si,
+        version="benchmark",
+        writer_rank=0,
+        expected_arithmetic=c.arithmetic_id,
+        backend=backend,
+    )
+    h = arena.create_request(0, version="benchmark", writer_rank=0)
+    start = 0
+    for _ in range(3):
+        arena.append(h, x[start : start + rows], start_position=start)
+        start += rows
+    samples = []
+    for _ in range(5):
+        torch.cuda.synchronize()
+        begin = time.perf_counter()
+        for _ in range(repeats):
+            arena.append(h, x[start : start + rows], start_position=start)
+            start += rows
+        torch.cuda.synchronize()
+        samples.append((time.perf_counter() - begin) * 1e6 / repeats)
+    snapshot = arena.snapshot(h)
+    arena.release(h)
+    return {
+        "ratio": ratio,
+        "graph": graph,
+        "rows_per_append": rows,
+        "append_us": statistics.median(samples),
+        "processed_tokens": start,
+        "complete_cache_rows": len(snapshot["rows"]),
+        "resident_buffer_bytes": c.resident_bytes,
+        "arithmetic_id": c.arithmetic_id,
+        "implementation": arena.implementation_fingerprint,
+    }
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("--repeats", type=int, default=30)
+    p.add_argument("--output", type=Path, required=True)
+    args = p.parse_args()
+    backend = AiterCompressorStateBackend()
+    report = {
+        "provider": backend.provenance,
+        "metric": "median of five synchronized loops; validation and device completion included",
+        "scope": "single writer; not kernel-only latency or full model throughput",
+        "cases": [
+            case(backend, ratio, graph, rows, args.repeats)
+            for ratio in (4, 128)
+            for graph in (False, True)
+            for rows in (1, 256)
+        ],
+    }
+    args.output.write_text(json.dumps(report, indent=2))
+    print(json.dumps(report, indent=2))

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/compressor-state.md

@@ -0,0 +1,99 @@
+# Task 21 — request-owned compressor state and single cache commit
+
+This change is stacked on task 19 / PR #618 at `8793052`; its target remains
+`dsv4-rocm`. It adds lifecycle/control code, not numerical GPU kernels.
+
+## Reused providers and explicit arithmetic
+
+The backend directly calls aiter `flydsl_fused_compress_attn` to read previous
+state and compress/scatter complete rows, then vLLM `save_partial_states` to
+update FP32 `[kv, score+ape]` state. The order is fixed: **compress before
+updating history**. The installed aiter ABI needs contiguous split readers;
+preallocated copies adapt vLLM's packed state without changing its values.
+
+The admitted main-KV profiles are:
+
+- `aiter.main_c4.online_fp32.wave64.vllm_partial_f32.v1`
+- `aiter.main_c128.online_fp32.wave64.vllm_partial_f32.v1`
+
+Both use a fixed single-wave online-FP32 recipe. **C128 is not task 18's
+eight-wave HCA profile.** The required `expected_arithmetic` constructor
+argument rejects a different recipe rather than silently substituting it.
+Callers must select this same profile for the paths they compare; shape
+compatibility alone is not a numerical compatibility claim. Training/rollout
+integration with the eight-wave HCA path is not covered by this profile.
+
+## API and ownership
+
+Construct `CompressorStateCache` with a `CompressorStateContract`, FP32 APE/RMS
+weights, BF16 position tables, a nonempty model/position version and a writer
+rank. The cache owns snapshots of the parameters and a fixed-capacity workspace.
+No model downloads or provider replacement occurs during construction.
+
+1. `create_request(request_id, version=..., writer_rank=...)` returns a handle
+   bound to this arena, a generation, version and private page slot.
+2. `prepare(handle, values, start_position=...)` validates detached BF16 input
+   and owns its staging copy. Only contiguous new logical tokens are accepted.
+3. `commit(lease)` consumes that **exact lease once**, performs ordered native
+   writes and waits for device completion before publishing the token cursor.
+4. `append(...)` combines prepare and commit under the same host lock.
+5. `cancel(lease)` aborts a prepared, not-yet-executed transaction.
+6. `snapshot(handle)` returns caller-owned complete rows and the FP32 ring,
+   plus token/group/version metadata. It does not expose mutable arena storage.
+7. `release(handle)` clears its state/pages before slot reuse. The old handle
+   cannot access a later generation with the same logical request ID.
+8. `replace_model(...)` requires all requests and leases to be released, a
+   different version, and fully validated replacement parameters.
+
+There is one staging lease per arena. Thread races, stale/foreign handles,
+wrong writers, duplicate commits, noncontiguous tokens and capacity violations
+are rejected. Errors after mutable GPU work poison the arena; callers must
+discard it instead of retrying possibly partial writes.
+
+Only the final ring-capacity suffix of an oversized input chunk is eligible
+to write partial state. Earlier writers receive vLLM's invalid slot sentinel,
+so one native launch never has competing writers to the same ring slot.
+
+The request's complete-row interval is determined by the logical cursor:
+`[old_tokens // ratio, new_tokens // ratio)`. Incomplete groups are never
+committed. The ring also retains the predecessor needed for C4 overlap.
+
+## Graph and memory
+
+Graph mode captures the same upstream operations with fixed input, metadata,
+state-reader and cache addresses. Warmup/capture uses only invalid plans and
+write slots and is checked not to mutate state/cache. Per-request lifecycle,
+version validation, metadata preparation and exactly-once checks stay **outside**
+capture. Direct external capture of lifecycle methods is rejected.
+
+Commits synchronously finish their device event. This intentionally conservative
+policy prevents early staging/page reuse between threads or streams; asynchronous
+communication/compute overlap is a later task. Input producers must order their
+writes before `prepare`; private arena storage must not be modified externally.
+
+The byte budget bounds explicitly owned resident tensors, not PyTorch/FlyDSL
+compiler caches, Graph memory pools, admission temporaries or caller snapshots.
+Provider-addressing bounds are checked before allocation.
+
+## Integration and limits
+
+The vLLM and Megatron opt-in factories share this exact lifecycle implementation.
+This is not complete framework model registration. Cache writes are deliberately
+non-differentiable; detached inputs are required. Native CP transport must occur
+before ingestion; a writer-rank check is not a distributed lock or collective.
+
+No multi-GPU/RCCL, multi-node, full model or eight-wave C128 training/rollout
+equivalence is claimed. Indexer MXFP4 state and cache belong to task 22.
+
+Run the targeted suite and benchmark in the member container, keeping results
+in the mounted output directory:
+
+```bash
+python3 -m pytest tests/contracts/operators tests/models/deepseek_v4 -q \
+  --basetemp=/outputs/state-tests --junitxml=/outputs/state-tests.xml
+python3 benchmarks/benchmark_compressor_state.py \
+  --repeats 30 --output /outputs/state-benchmark.json
+```
+
+The existing task-19 regression tests require the pinned mLite source through
+`RLK_MLITE_DSV4_BACKWARD_SOURCE`; the state backend itself has no mLite dependency.

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.