From 94bb6d02067ac5c8d74d17eaae7069586d1cee63 Mon Sep 17 00:00:00 2001 From: Takumi <108318514+inaniloquentee@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:40:51 +0000 Subject: [PATCH 1/2] feat(dsv4): add contract-bound C4 KV compressor --- benchmarks/benchmark_kv_compressor_c4.py | 89 +++ docs/models/deepseek_v4/kv-compressor-c4.md | 81 +++ reports/kv-compressor-c4-mi300x-20261028.md | 47 ++ .../rocm/attention/kv_compressor_c4.py | 164 ++++++ .../contracts/operators/kv_compressor_c4.py | 99 ++++ .../integrations/common/kv_compressor_c4.py | 8 + .../engines/rollout/vllm/kv_compressor_c4.py | 6 + .../train/megatron/kv_compressor_c4.py | 7 + .../models/deepseek_v4/kv_compressor_c4.py | 505 ++++++++++++++++++ tests/models/deepseek_v4/c4_helpers.py | 136 +++++ tests/models/deepseek_v4/c4_process_probe.py | 49 ++ .../deepseek_v4/test_kv_compressor_c4.py | 303 +++++++++++ .../deepseek_v4/test_kv_compressor_c4_rocm.py | 257 +++++++++ 13 files changed, 1751 insertions(+) create mode 100644 benchmarks/benchmark_kv_compressor_c4.py create mode 100644 docs/models/deepseek_v4/kv-compressor-c4.md create mode 100644 reports/kv-compressor-c4-mi300x-20261028.md create mode 100644 rl_engine/backends/rocm/attention/kv_compressor_c4.py create mode 100644 rl_engine/contracts/operators/kv_compressor_c4.py create mode 100644 rl_engine/integrations/common/kv_compressor_c4.py create mode 100644 rl_engine/integrations/engines/rollout/vllm/kv_compressor_c4.py create mode 100644 rl_engine/integrations/engines/train/megatron/kv_compressor_c4.py create mode 100644 rl_engine/models/deepseek_v4/kv_compressor_c4.py create mode 100644 tests/models/deepseek_v4/c4_helpers.py create mode 100644 tests/models/deepseek_v4/c4_process_probe.py create mode 100644 tests/models/deepseek_v4/test_kv_compressor_c4.py create mode 100644 tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py 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/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/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/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..33a3afd21 --- /dev/null +++ b/rl_engine/backends/rocm/attention/kv_compressor_c4.py @@ -0,0 +1,164 @@ +# 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, + 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 + require( + source_hash(getattr(self.aiter, "__file__", None)) == AITer_C4_SHA256, + "unvalidated aiter C4 source revision", + ) + 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": AITer_C4_SHA256, + "mlite_backward": MLITE_BACKWARD_SHA256, + "adapter": source_hash(__file__), + } + + def forward(self, bound, ape, norm, cos, sin, contract): + groups = len(bound.logical_keys) + if groups == 0: + return torch.empty((0, 512), dtype=torch.bfloat16, device=bound.values.device) + blocks = max((block for block, _slot in bound.cache_slots_host), default=-1) + 1 + cache = torch.empty((blocks, 64, 512), dtype=torch.bfloat16, device=bound.values.device) + provider_values = bound.context_values.reshape(-1, bound.context_values.shape[-1]).clone() + state_kv = torch.zeros_like(bound.kv_state) + state_score = torch.zeros_like(bound.score_state) + state_slots = bound.state_slots + plan = bound.plan.contiguous() + block_table = bound.block_table + for index, (_request, group) in enumerate(bound.logical_keys): + if group == 0: + # The overlap kernel's input phase intentionally has no + # negative-infinity sentinel path. A finite BF16 floor is + # below the softmax range of the admitted score domain and + # avoids -inf - (-inf) while remaining numerically masked. + provider_values[index * 8 : index * 8 + 4, 1024:] = -(2.0**30) + self.aiter.flydsl_fused_compress_attn( + kv_in=provider_values[:, :1024], + score_in=provider_values[:, 1024:], + kv_state=state_kv, + score_state=state_score, + plan_gpu=plan, + state_slot_mapping=state_slots, + ape=ape, + rms_weight=norm, + rms_eps=contract.rms_eps, + cos_cache=cos, + sin_cache=sin, + kv_cache=cache, + block_tables=block_table, + k_per_block=64, + overlap=True, + ratio=4, + head_dim=512, + rope_head_dim=64, + quant=False, + use_ue8m0=True, + preshuffle=True, + quant_mode="none", + k_split_num_waves=1, + ) + return torch.stack( + [cache[block, slot] for block, slot in bound.cache_slots_host], + dim=0, + ).contiguous() + + def vjp(self, bound, ape, norm, cos, sin, grad, contract): + contributions, ape_rows, norm_rows = [], [], [] + cache = torch.cat((cos, sin), dim=-1) + with torch.inference_mode(False), torch.enable_grad(): + for index, (_request, group) in enumerate(bound.logical_keys): + context = bound.context_values[index].float().detach().requires_grad_() + a = ape.detach().requires_grad_() + n = norm.detach().requires_grad_() + group_ids = torch.tensor( + [group - 1, group] if group > 0 else [0], + dtype=torch.int64, + device=context.device, + ) + source = context if group > 0 else context[4:] + y = self.functional( + source, + a, + n, + group_ids, + cache, + ratio=4, + head_dim=512, + rope_dim=64, + eps=contract.rms_eps, + ) + if group > 0: + grad_out = torch.cat((torch.zeros_like(y[:1]), grad[index : index + 1].float())) + dx, da, dn = torch.autograd.grad(y, (context, a, n), grad_outputs=grad_out) + else: + dx, da, dn = torch.autograd.grad( + y, (source, a, n), grad_outputs=grad[index : index + 1].float() + ) + dx = torch.cat((torch.zeros_like(context[:4]), dx)) + contributions.append((bound.context_indices[index], dx)) + ape_rows.append(da) + norm_rows.append(dn) + record_backward( + "kv_compressor_c4", + kernel_id="megatron.compressed_compact_graph.ratio4+rlk.fixed_order_rows", + impl="upstream_functional_vjp", + family="c4", + ) + return contributions, torch.stack(ape_rows), torch.stack(norm_rows) diff --git a/rl_engine/contracts/operators/kv_compressor_c4.py b/rl_engine/contracts/operators/kv_compressor_c4.py new file mode 100644 index 000000000..1608f19fa --- /dev/null +++ b/rl_engine/contracts/operators/kv_compressor_c4.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Item 19: overlap-window C4 main-KV compressor contract.""" + +from dataclasses import asdict, dataclass +import hashlib +import json +import math +from pathlib import Path + +from rl_engine.contracts.operators.attention import AttentionContractError + +C4_BACKEND_ID = "aiter.csa.c4.gfx942.bf16.overlap.fixed1.v1" +AITer_C4_SHA256 = "53cdd2801b6b2f5866b02d18bff1b2d56a19289a8048cdc14fe8bb687e572a14" +MLITE_BACKWARD_SHA256 = "6a1e78ca89b83f56391b5e5ac409e765338f1e3a625251d0d4c877c1e80b4a68" + + +class C4Error(AttentionContractError): + pass + + +def require(ok, message): + if not ok: + raise C4Error(message) + + +def source_hash(path): + require(isinstance(path, (str, Path)) and bool(path), "provider source is unavailable") + try: + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + except (OSError, ValueError) as exc: + raise C4Error("provider source is unreadable") from exc + + +@dataclass(frozen=True) +class C4Contract: + head_dim: int = 512 + ratio: int = 4 + overlap: int = 8 + rope_dim: int = 64 + rms_eps: float = 1e-6 + max_input_rows: int = 524288 + max_groups: int = 8192 + max_block_table_bytes: int = 16 * 1024 * 1024 + max_backward_groups: int = 128 + backend_id: str = C4_BACKEND_ID + output_dtype: str = "bf16" + graph_capture: bool = False + + def __post_init__(self): + for name in ( + "head_dim", + "ratio", + "overlap", + "rope_dim", + "max_input_rows", + "max_groups", + "max_block_table_bytes", + "max_backward_groups", + ): + require( + type(getattr(self, name)) is int and getattr(self, name) > 0, + f"{name} must be a positive integer", + ) + require( + (self.head_dim, self.ratio, self.overlap, self.rope_dim) == (512, 4, 8, 64), + "validated C4 main-KV profile is D=512, ratio=4, overlap=8, RoPE=64", + ) + require( + type(self.rms_eps) in (int, float) + and math.isfinite(self.rms_eps) + and 1e-12 <= self.rms_eps <= 1, + "invalid RMS epsilon", + ) + require( + self.max_input_rows < 2**31 and self.max_groups < 2**24, + "int32 provider indexing capacity exceeded", + ) + require(self.backend_id == C4_BACKEND_ID, "unknown C4 backend; no fallback") + require(self.output_dtype == "bf16", "C4 emits BF16 independent rows") + require(self.graph_capture is False, "Graph capture belongs to runtime integration") + + @property + def input_width(self): + return 4 * self.head_dim + + @property + def fingerprint(self): + return hashlib.sha256(json.dumps(asdict(self), sort_keys=True).encode()).hexdigest() + + @property + def arithmetic(self): + return { + "pool": ( + "aiter FP32 online softmax;K=8;overlap previous first4/current last4;fixed NW=1" + ), + "norm_rope": "FP32 RMSNorm;BF16 cos/sin;GPT-J tail;BF16 output", + "position": "global group ID * 4; first group previous window is -inf padded", + "backward": "mLite ratio4 compact graph;fixed logical row accumulation", + } diff --git a/rl_engine/integrations/common/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/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/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/kv_compressor_c4.py b/rl_engine/models/deepseek_v4/kv_compressor_c4.py new file mode 100644 index 000000000..a3810e3b0 --- /dev/null +++ b/rl_engine/models/deepseek_v4/kv_compressor_c4.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Complete-group C4 admission, fixed overlap state and canonical VJP wiring.""" + +from dataclasses import dataclass, field +import hashlib +import json + +import torch + +from rl_engine.contracts.operators.kv_compressor_c4 import C4Contract, require, source_hash +from rl_engine.ops.autograd.canonical_backward import active_session + + +def eager(device): + if device.type == "cuda": + with torch.cuda.device(device): + require(not torch.cuda.is_current_stream_capturing(), "Graph capture is not validated") + + +def stamp(tensor): + return ( + tensor.data_ptr(), + None if tensor.is_inference() else tensor._version, + tensor.shape, + tensor.stride(), + tensor.dtype, + tensor.device, + ) + + +@dataclass(frozen=True) +class C4Groups: + source: torch.Tensor = field(repr=False) + values: torch.Tensor = field(repr=False) + context_values: torch.Tensor = field(repr=False) + context_indices: torch.Tensor = field(repr=False) + selector: torch.Tensor = field(repr=False) + group_ids: torch.Tensor = field(repr=False) + plan: torch.Tensor = field(repr=False) + block_table: torch.Tensor = field(repr=False) + kv_state: torch.Tensor = field(repr=False) + score_state: torch.Tensor = field(repr=False) + state_slots: torch.Tensor = field(repr=False) + logical_keys: tuple[tuple[int, int], ...] + request_spans: tuple[tuple[int, int], ...] + incomplete_groups: tuple[tuple[int, int], ...] + unowned_groups: tuple[tuple[int, int], ...] + context_indices_host: tuple[tuple[int, ...], ...] + cache_slots_host: tuple[tuple[int, int], ...] + version: str + owner_rank: int + cp_size: int + contract: C4Contract + _stamps: tuple = field(repr=False) + _source_stamp: tuple = field(repr=False) + + @property + def tensors(self): + return ( + self.values, + self.context_values, + self.context_indices, + self.selector, + self.group_ids, + self.plan, + self.block_table, + self.kv_state, + self.score_state, + self.state_slots, + ) + + def validate(self): + require(tuple(stamp(t) for t in self.tensors) == self._stamps, "bound C4 storage changed") + if self.source.requires_grad: + require(stamp(self.source) == self._source_stamp, "training C4 source changed") + + +@dataclass(frozen=True) +class C4Output: + values: torch.Tensor + group_ids: torch.Tensor + owner_rank: int + version: str + contract_fingerprint: str + implementation_fingerprint: str + + +def bind_complete_groups( + values, + token_ids, + *, + version, + contract, + owners=None, + cp_size=1, + rank=0, + emit_group_ids=None, +): + require(isinstance(contract, C4Contract), "typed C4 contract required") + require( + isinstance(version, str) and bool(version), + "explicit model/position/request version required", + ) + require( + isinstance(values, torch.Tensor) + and values.dtype == torch.bfloat16 + and values.ndim == 2 + and values.shape[1] == contract.input_width + and values.is_contiguous(), + "projected [tokens, 4*head_dim] contiguous BF16 input required", + ) + eager(values.device) + require(values.shape[0] <= contract.max_input_rows, "input capacity exceeded") + require( + isinstance(token_ids, torch.Tensor) + and token_ids.dtype == torch.int64 + and tuple(token_ids.shape) == (values.shape[0], 2) + and token_ids.device == values.device, + "token IDs must be int64 [tokens,2] on the input device", + ) + require( + type(cp_size) is int and cp_size > 0 and type(rank) is int and 0 <= rank < cp_size, + "invalid CP size/rank", + ) + host_ids = tuple(tuple(row) for row in token_ids.cpu().tolist()) + if owners is None: + require(cp_size == 1, "CP token ownership metadata required") + host_owners = (0,) * len(host_ids) + else: + require( + isinstance(owners, torch.Tensor) + and owners.dtype == torch.int64 + and tuple(owners.shape) == (values.shape[0],) + and owners.device == values.device, + "owners must be int64 [tokens] on the input device", + ) + host_owners = tuple(owners.cpu().tolist()) + emit = None if emit_group_ids is None else set(tuple(key) for key in emit_group_ids) + if emit is not None: + require( + all(len(key) == 2 and key[0] >= 0 and key[1] >= 0 for key in emit), + "emit_group_ids must contain nonnegative [request,group] keys", + ) + groups, seen = {}, set() + for physical, (key, owner) in enumerate(zip(host_ids, host_owners, strict=True)): + if key == (-1, -1): + continue + request, position = key + require(request >= 0 and 0 <= position < 2**31, "invalid logical token ID") + require(key not in seen, "duplicate logical token ID") + require(0 <= owner < cp_size, "token owner outside CP group") + seen.add(key) + group, offset = divmod(position, 4) + groups.setdefault((request, group), {})[offset] = (physical, owner) + + complete = {key: entries for key, entries in groups.items() if len(entries) == 4} + incomplete = [key for key, entries in groups.items() if len(entries) != 4] + selected, unowned = [], [] + for key, entries in sorted(complete.items()): + request, group = key + if emit is not None and key not in emit: + continue + if group > 0 and (request, group - 1) not in complete: + incomplete.append(key) + continue + if entries[3][1] != rank: + unowned.append(key) + continue + selected.append(key) + selected = sorted(selected) + require(len(selected) <= contract.max_groups, "complete-group capacity exceeded") + request_order = {request: i for i, request in enumerate(sorted({key[0] for key in selected}))} + max_block_column = max((key[1] // 64 + 1 for key in selected), default=0) + require( + len(request_order) * max_block_column * 4 <= contract.max_block_table_bytes, + "C4 block-table byte budget exceeded", + ) + block_order = { + key: index + for index, key in enumerate(sorted({(request, group // 64) for request, group in selected})) + } + + current_indices = [] + context_indices_host = [] + for request, group in selected: + current = [complete[(request, group)][offset][0] for offset in range(4)] + previous = ( + [complete[(request, group - 1)][offset][0] for offset in range(4)] + if group > 0 + else [-1] * 4 + ) + current_indices.extend(current) + context_indices_host.append(tuple(previous + current)) + cache_slots_host = tuple( + (block_order[(request, group // 64)], group % 64) for request, group in selected + ) + with torch.inference_mode(False), torch.no_grad(): + selector = torch.tensor(current_indices, dtype=torch.int64, device=values.device) + packed = values.index_select(0, selector).detach() + context_indices = torch.tensor( + context_indices_host, dtype=torch.int64, device=values.device + ).reshape(-1, 8) + safe_context = context_indices.clamp_min(0) + context = ( + values.index_select(0, safe_context.reshape(-1)) + .reshape(-1, 8, contract.input_width) + .detach() + ) + if context.numel(): + context[:, :4] = torch.where( + (context_indices[:, :4] >= 0).unsqueeze(-1), + context[:, :4], + torch.zeros_like(context[:, :4]), + ) + require( + bool((torch.isfinite(packed.float()) & (packed.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, 0] + for i, (request, group) in enumerate(selected) + ], + dtype=torch.int32, + device=values.device, + ).reshape(-1, 4) + table = torch.full( + (len(request_order), max_block_column), + -1, + dtype=torch.int32, + device=values.device, + ) + for (request, group), (block, _slot) in zip(selected, cache_slots_host, strict=True): + table[request_order[request], group // 64] = block + slots = torch.arange(len(request_order), dtype=torch.int32, device=values.device) + state_kv = torch.zeros( + (len(request_order), 8, 2 * contract.head_dim), + dtype=torch.float32, + device=values.device, + ) + state_score = torch.zeros_like(state_kv) + spans = [] + cursor = 0 + for request in sorted({key[0] for key in selected}): + length = sum(key[0] == request for key in selected) + spans.append((cursor, cursor + length)) + cursor += length + tensors = ( + packed, + context, + context_indices, + selector, + group_ids, + plan, + table, + state_kv, + state_score, + slots, + ) + return C4Groups( + values, + *tensors, + tuple(selected), + tuple(spans), + tuple(sorted(set(incomplete))), + tuple(unowned), + tuple(context_indices_host), + cache_slots_host, + version, + rank, + cp_size, + contract, + tuple(stamp(tensor) for tensor in tensors), + stamp(values), + ) + + +def canonical_sum(rows): + result = torch.zeros(rows.shape[1:], dtype=rows.dtype, device=rows.device) + for row in rows.unbind(0): + result = torch.add(result, row) + return result + + +def register_parameters(session, prefix, parameters, bound, implementation): + require(session is not None, "trainable C4 parameters require canonical_backward_session") + require(not session.backward_started, "register all C4 calls before backward") + require(isinstance(prefix, str) and bool(prefix), "stable layer parameter_id required") + bindings = getattr(session, "_c4_bindings", {}) + signature = (implementation, bound.version) + for suffix, parameter in parameters.items(): + if not parameter.requires_grad or not bound.logical_keys: + continue + entry = bindings.get(prefix + "." + suffix) + if entry is not None: + require( + entry[0] is parameter and entry[1] == signature, + "canonical parameter identity/version mismatch", + ) + require( + not entry[2].intersection(bound.logical_keys), + "duplicate C4 group across microbatches", + ) + slots = {} + for suffix, parameter in parameters.items(): + if parameter.requires_grad and bound.logical_keys: + key = prefix + "." + suffix + entry = bindings.setdefault(key, (parameter, signature, set())) + entry[2].update(bound.logical_keys) + slots[suffix] = (key, session.register(key, bound.group_ids)) + session._c4_bindings = bindings + return slots + + +class _C4Autograd(torch.autograd.Function): + @staticmethod + def forward(ctx, source, ape, norm, cos, sin, operator, bound, parameter_id): + ctx.operator, ctx.bound = operator, bound + ctx.session = active_session() + ctx.slots = {} + if (ape.requires_grad or norm.requires_grad) and bound.logical_keys: + ctx.slots = register_parameters( + ctx.session, + parameter_id, + {"ape": ape, "norm": norm}, + bound, + operator.implementation_fingerprint, + ) + ctx.save_for_backward(source, ape, norm, cos, sin) + return operator.backend.forward( + bound, ape.detach(), norm.detach(), cos, sin, operator.contract + ) + + @staticmethod + def backward(ctx, grad): + require(not torch.is_grad_enabled(), "higher-order C4 gradients are not validated") + source, ape, norm, cos, sin = ctx.saved_tensors + op, bound = ctx.operator, ctx.bound + eager(bound.values.device) + bound.validate() + if not bound.logical_keys: + return ( + torch.zeros_like(source), + torch.zeros_like(ape), + torch.zeros_like(norm), + None, + None, + None, + None, + None, + ) + contributions, da, dn = op.backend.vjp(bound, ape, norm, cos, sin, grad, op.contract) + dsource = torch.zeros_like(source) + for indices, rows in contributions: + for physical, row in zip(indices, rows, strict=True): + if physical >= 0: + dsource[physical] = dsource[physical] + row.to(source.dtype) + parameter_grads = [] + for suffix, rows, parameter in (("ape", da, ape), ("norm", dn, norm)): + if suffix not in ctx.slots: + parameter_grads.append(None) + continue + key, slot = ctx.slots[suffix] + reduced = ctx.session.submit_rows(key, slot, rows, canonical_sum) + parameter_grads.append(None if reduced is None else reduced.reshape(parameter.shape)) + return dsource, *parameter_grads, None, None, None, None, None + + +class C4Compressor: + def __init__(self, contract, *, backend=None, expected_implementation=None): + require(isinstance(contract, C4Contract), "typed C4 contract required") + if backend is None: + from rl_engine.backends.rocm.attention.kv_compressor_c4 import AiterC4Backend + + backend = AiterC4Backend() + require(backend.backend_id == contract.backend_id, "C4 backend identity mismatch") + self.contract, self.backend = contract, backend + identity = { + "contract": contract.fingerprint, + "arithmetic": contract.arithmetic, + "backend": getattr(backend, "provenance", {"test_provider": type(backend).__name__}), + "bridge": source_hash(__file__), + "canonical": source_hash(active_session.__code__.co_filename), + } + self.implementation_fingerprint = hashlib.sha256( + json.dumps(identity, sort_keys=True).encode() + ).hexdigest() + if expected_implementation is not None: + require( + expected_implementation == self.implementation_fingerprint, + "C4 implementation fingerprint mismatch", + ) + + def bind( + self, + values, + token_ids, + *, + version, + owners=None, + cp_size=1, + rank=0, + emit_group_ids=None, + ): + require( + isinstance(values, torch.Tensor) and values.device.type == self.backend.device_type, + "C4 input/backend device mismatch", + ) + if hasattr(self.backend, "device"): + require(values.device == self.backend.device, "C4 device index mismatch") + return bind_complete_groups( + values, + token_ids, + version=version, + contract=self.contract, + owners=owners, + cp_size=cp_size, + rank=rank, + emit_group_ids=emit_group_ids, + ) + + def __call__(self, bound, ape, norm, cos, sin, *, version, parameter_id=""): + require( + isinstance(bound, C4Groups) and bound.contract == self.contract, + "C4 bound contract mismatch", + ) + require(bound.version == version, "C4 model/position/request version mismatch") + require(bound.values.device.type == self.backend.device_type, "C4 bound backend mismatch") + if hasattr(self.backend, "device"): + require(bound.values.device == self.backend.device, "C4 bound device index mismatch") + eager(bound.values.device) + bound.validate() + require( + isinstance(ape, torch.Tensor) + and tuple(ape.shape) == (4, 1024) + and ape.dtype == torch.float32 + and ape.is_contiguous() + and ape.device == bound.values.device, + "APE must be contiguous FP32 [4,1024]", + ) + require( + isinstance(norm, torch.Tensor) + and tuple(norm.shape) == (512,) + and norm.dtype == torch.float32 + and norm.is_contiguous() + and norm.device == bound.values.device, + "RMS weight must be contiguous FP32 [512]", + ) + for value, name in ((ape, "APE"), (norm, "RMS weight")): + require( + bool((torch.isfinite(value) & (value.abs() <= 2.0**20)).all()), + f"{name} outside finite C4 parameter domain", + ) + require( + isinstance(cos, torch.Tensor) + and isinstance(sin, torch.Tensor) + and cos.shape == sin.shape + and cos.ndim == 2 + and cos.shape[1] == 32 + and cos.dtype == torch.bfloat16 + and sin.dtype == torch.bfloat16 + and cos.is_contiguous() + and sin.is_contiguous() + and cos.device == bound.values.device + and sin.device == bound.values.device + and not cos.requires_grad + and not sin.requires_grad, + "RoPE tables must be matching contiguous BF16 [positions,32]", + ) + if bound.logical_keys: + max_position = max(group * 4 for _request, group in bound.logical_keys) + require(max_position < cos.shape[0], "RoPE table misses a global C4 position") + for table in (cos, sin): + # Do not create a transient GPU gather immediately before the + # FlyDSL launch; the provider applies this global offset in-kernel. + require( + bool((torch.isfinite(table) & (table.abs() <= 1)).all()), "invalid RoPE values" + ) + training = torch.is_grad_enabled() and any( + tensor.requires_grad for tensor in (bound.source, ape, norm) + ) + if training: + require( + len(bound.logical_keys) <= self.contract.max_backward_groups, + "C4 backward group budget exceeded", + ) + output = _C4Autograd.apply(bound.source, ape, norm, cos, sin, self, bound, parameter_id) + else: + with torch.no_grad(): + output = self.backend.forward(bound, ape, norm, cos, sin, self.contract) + require( + output.shape == (len(bound.logical_keys), 512) + and output.dtype == torch.bfloat16 + and output.device == bound.values.device, + "C4 provider returned invalid output", + ) + return C4Output( + output, + bound.group_ids, + bound.owner_rank, + version, + self.contract.fingerprint, + self.implementation_fingerprint, + ) diff --git a/tests/models/deepseek_v4/c4_helpers.py b/tests/models/deepseek_v4/c4_helpers.py new file mode 100644 index 000000000..199bd3934 --- /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): + # 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.float(), dim=0).to(content.dtype) + pooled = (content * weights).sum(0).float() + pooled = pooled * torch.rsqrt(pooled.square().mean(-1, keepdim=True) + eps) + pooled = pooled * norm.float() + selected_cos = cos[group * 4].float() + selected_sin = sin[group * 4].float() + 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.float(), 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/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..a20bf4af0 --- /dev/null +++ b/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py @@ -0,0 +1,257 @@ +# 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["k_split_num_waves"] == 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 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]), + "RLK_MLITE_DSV4_BACKWARD_SOURCE": str( + Path(__file__).resolve().parents[3] + / ".codex-work" + / "upstream" + / "megatron_backward.py" + ), + "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] From 879305251c37e1b8bb0cb44409f1762752d03561 Mon Sep 17 00:00:00 2001 From: Takumi <108318514+inaniloquentee@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:14:23 +0000 Subject: [PATCH 2/2] fix(dsv4): use native C4 padding mask and preserve FP64 test arithmetic --- reports/kv-compressor-c4-new-host-20261029.md | 55 +++++++++++++++++++ .../rocm/attention/kv_compressor_c4.py | 32 ++++++----- .../contracts/operators/kv_compressor_c4.py | 1 + .../models/deepseek_v4/kv_compressor_c4.py | 4 +- tests/models/deepseek_v4/c4_helpers.py | 16 +++--- .../deepseek_v4/test_kv_compressor_c4_rocm.py | 48 +++++++++++++--- 6 files changed, 124 insertions(+), 32 deletions(-) create mode 100644 reports/kv-compressor-c4-new-host-20261029.md 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/kv_compressor_c4.py b/rl_engine/backends/rocm/attention/kv_compressor_c4.py index 33a3afd21..fa00be5b0 100644 --- a/rl_engine/backends/rocm/attention/kv_compressor_c4.py +++ b/rl_engine/backends/rocm/attention/kv_compressor_c4.py @@ -10,6 +10,7 @@ from rl_engine.contracts.operators.kv_compressor_c4 import ( AITer_C4_SHA256, + AITer_C4_LEGACY_SHA256, C4_BACKEND_ID, C4Error, MLITE_BACKWARD_SHA256, @@ -37,10 +38,17 @@ def __init__(self, *, megatron_backward_path=None): ) except (ImportError, AttributeError) as exc: raise C4Error("pinned aiter C4 ABI unavailable") from exc + provider_hash = source_hash(getattr(self.aiter, "__file__", None)) require( - source_hash(getattr(self.aiter, "__file__", None)) == AITer_C4_SHA256, + 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: @@ -65,7 +73,9 @@ def __init__(self, *, megatron_backward_path=None): "hip": torch.version.hip, "triton": importlib.import_module("triton").__version__, "flydsl": importlib.metadata.version("flydsl"), - "aiter_c4": AITer_C4_SHA256, + "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__), } @@ -76,19 +86,12 @@ def forward(self, bound, ape, norm, cos, sin, contract): return torch.empty((0, 512), dtype=torch.bfloat16, device=bound.values.device) blocks = max((block for block, _slot in bound.cache_slots_host), default=-1) + 1 cache = torch.empty((blocks, 64, 512), dtype=torch.bfloat16, device=bound.values.device) - provider_values = bound.context_values.reshape(-1, bound.context_values.shape[-1]).clone() - state_kv = torch.zeros_like(bound.kv_state) - state_score = torch.zeros_like(bound.score_state) + 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 - for index, (_request, group) in enumerate(bound.logical_keys): - if group == 0: - # The overlap kernel's input phase intentionally has no - # negative-infinity sentinel path. A finite BF16 floor is - # below the softmax range of the admitted score domain and - # avoids -inf - (-inf) while remaining numerically masked. - provider_values[index * 8 : index * 8 + 4, 1024:] = -(2.0**30) self.aiter.flydsl_fused_compress_attn( kv_in=provider_values[:, :1024], score_in=provider_values[:, 1024:], @@ -111,8 +114,7 @@ def forward(self, bound, ape, norm, cos, sin, contract): quant=False, use_ue8m0=True, preshuffle=True, - quant_mode="none", - k_split_num_waves=1, + **self.launch_options, ) return torch.stack( [cache[block, slot] for block, slot in bound.cache_slots_host], @@ -152,7 +154,7 @@ def vjp(self, bound, ape, norm, cos, sin, grad, contract): 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[index], dx)) + contributions.append((bound.context_indices_host[index], dx)) ape_rows.append(da) norm_rows.append(dn) record_backward( diff --git a/rl_engine/contracts/operators/kv_compressor_c4.py b/rl_engine/contracts/operators/kv_compressor_c4.py index 1608f19fa..5db416de2 100644 --- a/rl_engine/contracts/operators/kv_compressor_c4.py +++ b/rl_engine/contracts/operators/kv_compressor_c4.py @@ -11,6 +11,7 @@ C4_BACKEND_ID = "aiter.csa.c4.gfx942.bf16.overlap.fixed1.v1" AITer_C4_SHA256 = "53cdd2801b6b2f5866b02d18bff1b2d56a19289a8048cdc14fe8bb687e572a14" +AITer_C4_LEGACY_SHA256 = "0c2019ca3a4a0d6dc7d5fdd20ebcce658ff5f504fdc880dffd61c9c398029451" MLITE_BACKWARD_SHA256 = "6a1e78ca89b83f56391b5e5ac409e765338f1e3a625251d0d4c877c1e80b4a68" diff --git a/rl_engine/models/deepseek_v4/kv_compressor_c4.py b/rl_engine/models/deepseek_v4/kv_compressor_c4.py index a3810e3b0..b1d3bc396 100644 --- a/rl_engine/models/deepseek_v4/kv_compressor_c4.py +++ b/rl_engine/models/deepseek_v4/kv_compressor_c4.py @@ -213,13 +213,13 @@ def bind_complete_groups( torch.zeros_like(context[:, :4]), ) require( - bool((torch.isfinite(packed.float()) & (packed.float().abs() <= 2.0**30)).all()), + 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, 0] + [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, diff --git a/tests/models/deepseek_v4/c4_helpers.py b/tests/models/deepseek_v4/c4_helpers.py index 199bd3934..1e7922153 100644 --- a/tests/models/deepseek_v4/c4_helpers.py +++ b/tests/models/deepseek_v4/c4_helpers.py @@ -7,6 +7,8 @@ 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] @@ -27,12 +29,12 @@ def context_reference(context, ape, norm, group, cos, sin, eps): 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.float(), dim=0).to(content.dtype) - pooled = (content * weights).sum(0).float() + 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.float() - selected_cos = cos[group * 4].float() - selected_sin = sin[group * 4].float() + 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( ( @@ -52,9 +54,7 @@ def reference(values, ape, norm, groups, cos, sin, eps): 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.float(), ape, norm, int(groups[group]), cos, sin, eps) - ) + outputs.append(context_reference(context, ape, norm, int(groups[group]), cos, sin, eps)) return ( torch.stack(outputs).to(values.dtype) if outputs diff --git a/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py b/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py index a20bf4af0..9248dd777 100644 --- a/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py +++ b/tests/models/deepseek_v4/test_kv_compressor_c4_rocm.py @@ -61,7 +61,7 @@ def test_native_forward_fixed_overlap_plan_and_cache_isolation(backend, monkeypa def capture(**kwargs): assert kwargs["ratio"] == 4 and kwargs["overlap"] is True - assert kwargs["k_split_num_waves"] == 1 + assert kwargs.get("k_split_num_waves", 1) == 1 assert kwargs["k_per_block"] == 64 captured.append(kwargs["plan_gpu"].clone()) return original(**kwargs) @@ -72,6 +72,7 @@ def capture(**kwargs): 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 @@ -237,12 +238,6 @@ def test_fresh_process_provider_fingerprint(backend, tmp_path): env = { **os.environ, "PYTHONPATH": str(Path(__file__).resolve().parents[3]), - "RLK_MLITE_DSV4_BACKWARD_SOURCE": str( - Path(__file__).resolve().parents[3] - / ".codex-work" - / "upstream" - / "megatron_backward.py" - ), "TRITON_CACHE_DIR": str(tmp_path / f"cache{index}"), } subprocess.run( @@ -255,3 +250,42 @@ def test_fresh_process_provider_fingerprint(backend, tmp_path): ) 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)))