[DSV4][ROCm] Reuse native operators for Indexer state and MXFP4 cache (task 22)

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_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_indexer_state.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reuse the state benchmark shape: one synchronous Indexer cache append."""
+
+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_state import AiterIndexerStateBackend
+from tests.models.deepseek_v4.indexer_state_helpers import make
+
+
+def case(backend, graph, rows, repeats):
+    arena, x, _, _ = make(
+        backend,
+        graph=graph,
+        step=rows,
+        tokens=max(4, rows * (5 * repeats + 3)),
+        max_requests=1,
+        device="cuda",
+    )
+    h = arena.create_request(0, version="v1", 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 {
+        "graph": graph,
+        "rows_per_append": rows,
+        "append_us": statistics.median(samples),
+        "processed_tokens": start,
+        "complete_rows": snapshot["packed"].shape[0],
+        "resident_buffer_bytes": arena.contract.resident_bytes,
+        "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 = AiterIndexerStateBackend()
+    report = {
+        "provider": backend.provenance,
+        "metric": "median of five synchronized loops, validation/staging/commit included",
+        "scope": "single-writer complete cache API; not kernel-only latency or model throughput",
+        "cases": [
+            case(backend, graph, rows, args.repeats) 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()