From aa15aa84c0485e6d7a68578a8693197591bc0c85 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Thu, 30 Jul 2026 16:43:32 +0800 Subject: [PATCH 1/5] feat(attention): add single-gpu comparison harness Signed-off-by: inaniloquentee <3051000145@qq.com> --- .../ws2-attention-single-gpu-harness.md | 95 +++ rl_engine/testing/__init__.py | 24 + rl_engine/testing/attention_comparison.py | 578 ++++++++++++++++++ tests/test_attention_comparison.py | 197 ++++++ 4 files changed, 894 insertions(+) create mode 100644 docs/design/ws2-attention-single-gpu-harness.md create mode 100644 rl_engine/testing/attention_comparison.py create mode 100644 tests/test_attention_comparison.py diff --git a/docs/design/ws2-attention-single-gpu-harness.md b/docs/design/ws2-attention-single-gpu-harness.md new file mode 100644 index 00000000..32eba5da --- /dev/null +++ b/docs/design/ws2-attention-single-gpu-harness.md @@ -0,0 +1,95 @@ +# WS2 Attention Single-GPU Comparison Harness + +Status: PR2 harness for [#235](https://github.com/RL-Align/RL-Kernel/issues/235) + +## Scope + +This harness compares attention materializations on one device before CP +communication is introduced. It is diagnostic infrastructure: it does not launch +collectives and does not replace the deterministic CP reference planned in PR3. + +Implemented paths: + +- `full_prefill`: training-style full-sequence softmax attention; +- `chunked_prefill`: rollout-style query chunk replay over full KV; +- `rl_kernel_paged_kv`: rollout-style KV page replay with fp32 attention-domain + LSE merge by logical KV block order; +- `transformer_engine_paged_kv`: optional oracle that reuses NVIDIA Transformer + Engine's context-parallel PyTorch correction helpers when TE is installed. + +## Report + +`rl_engine.testing.attention_comparison.compare_single_gpu_attention` emits a +structured report with: + +- `out` max / mean / p95 / p99 absolute drift; +- attention-domain `lse` max / mean / p95 / p99 absolute drift; +- optional active-token-only `dlogp` drift when `lm_head_weight`, `target_ids`, + and an active token mask are provided; +- per-path provenance including chunk/page sizes, KV page bounds, merge backend, + merge order, and LSE domain; +- optional-backend unavailability reasons. + +The selected-logprob convention follows #207: + +```text +dlogp = candidate selected logp - full_prefill selected logp +``` + +## Transformer Engine Reuse + +The harness does not make Transformer Engine a runtime dependency. When +available, it lazily imports: + +```text +transformer_engine.pytorch.attention.dot_product_attention.context_parallel +``` + +and calls: + +```text +flash_attn_fwd_softmax_lse_correction +flash_attn_fwd_out_correction_init +flash_attn_fwd_out_correction +``` + +Those helpers provide an industrial implementation oracle for the same fp32 +`(out, lse)` online-softmax merge policy that later CP/fused paths must match. +When TE is not installed, the TE path is reported as unavailable and the local +RL-Kernel paths still run. + +## CLI Registration + +The existing generic operator harness now registers `attention`, so a local +candidate smoke can run with: + +```bash +python scripts/check_operator.py --op attention --candidate pytorch --dtype fp32 +``` + +The attention-specific WS2 comparison entry point is Python-first for now: + +```python +from rl_engine.testing.attention_comparison import ( + AttentionComparisonInputs, + compare_single_gpu_attention, +) + +report = compare_single_gpu_attention( + AttentionComparisonInputs(q=q, k=k, v=v, target_ids=target_ids, lm_head_weight=w), + query_chunk_size=512, + kv_page_size=512, + include_transformer_engine=True, +) +print(report.to_dict()) +``` + +## Validation + +```bash +python -m pytest tests/test_attention_comparison.py -q +``` + +The tests cover full vs chunked/paged equivalence, active-token `dlogp` drift, +optional TE correction-helper reuse through a fake TE module, JSON-compatible +reports, and `attention` registration in the generic operator comparison specs. diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 42be8c1b..b9bffee0 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,19 @@ """Testing helpers for RL-shaped kernel validation.""" +from .attention_comparison import ( + AttentionComparisonInputs, + AttentionComparisonReport, + AttentionPathDrift, + AttentionPathResult, + DriftStats, + TransformerEngineUnavailable, + compare_single_gpu_attention, + run_chunked_query_attention, + run_full_attention, + run_paged_kv_attention, + transformer_engine_context_parallel_available, +) from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -15,13 +28,24 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ + "AttentionComparisonInputs", + "AttentionComparisonReport", + "AttentionPathDrift", + "AttentionPathResult", + "DriftStats", "SyntheticRLKernelBatch", + "TransformerEngineUnavailable", "active_token_count", + "compare_single_gpu_attention", "compute_policy_ratio", "compute_reference_kl", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", + "run_chunked_query_attention", + "run_full_attention", + "run_paged_kv_attention", "selected_logprobs_reference", "summarize_kernel_drift", + "transformer_engine_context_parallel_available", ] diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py new file mode 100644 index 00000000..350ecdf2 --- /dev/null +++ b/rl_engine/testing/attention_comparison.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Single-GPU WS2 attention cross-implementation comparison harness. + +This module compares logically equivalent attention materializations before CP +communication is introduced. The full path is the training-style reference. +The chunked-query and paged-KV paths emulate rollout-style prefill layouts on a +single device while preserving global causal positions and attention-domain LSE. +""" + +from __future__ import annotations + +import importlib +import math +from dataclasses import dataclass +from typing import Any, Literal + +import torch + +from rl_engine.testing.reference_ops import selected_logprobs_reference + +MergeBackend = Literal["rl_kernel", "transformer_engine"] + +_TE_CONTEXT_PARALLEL_MODULE = ( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" +) + + +class TransformerEngineUnavailable(RuntimeError): + """Raised when the optional Transformer Engine oracle cannot be imported.""" + + +@dataclass(frozen=True) +class AttentionComparisonInputs: + """Inputs shared by every single-GPU attention comparison path.""" + + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + causal: bool = True + scale: float | None = None + key_padding_mask: torch.Tensor | None = None + lm_head_weight: torch.Tensor | None = None + target_ids: torch.Tensor | None = None + active_token_mask: torch.Tensor | None = None + output_dtype: torch.dtype = torch.float32 + + +@dataclass(frozen=True) +class AttentionPathResult: + """One materialized attention path result.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, Any] + + +@dataclass(frozen=True) +class DriftStats: + """Shape-aware absolute drift summary.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, Any]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class AttentionPathDrift: + """Candidate-vs-reference drift for one attention path.""" + + candidate_name: str + out: DriftStats + lse: DriftStats + dlogp: DriftStats | None + provenance: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return { + "candidate_name": self.candidate_name, + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "dlogp": None if self.dlogp is None else self.dlogp.to_dict(), + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionComparisonReport: + """Structured report for PR2 single-GPU attention attribution.""" + + reference_name: str + drifts: tuple[AttentionPathDrift, ...] + unavailable: tuple[str, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + "unavailable": list(self.unavailable), + } + + +@dataclass(frozen=True) +class _PartialAttentionState: + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + +def compare_single_gpu_attention( + inputs: AttentionComparisonInputs, + *, + query_chunk_size: int | None = None, + kv_page_size: int | None = None, + include_transformer_engine: bool = False, +) -> AttentionComparisonReport: + """Compare full attention with chunked/paged single-GPU materializations. + + If ``lm_head_weight`` and ``target_ids`` are provided, the report also + includes active-token selected-logprob drift using the #207 convention: + candidate logp minus reference logp. + """ + + _validate_comparison_inputs(inputs) + reference = run_full_attention(inputs) + candidates = [ + run_chunked_query_attention(inputs, query_chunk_size=query_chunk_size), + run_paged_kv_attention(inputs, kv_page_size=kv_page_size, merge_backend="rl_kernel"), + ] + unavailable: list[str] = [] + if include_transformer_engine: + try: + candidates.append( + run_paged_kv_attention( + inputs, + kv_page_size=kv_page_size, + merge_backend="transformer_engine", + ) + ) + except TransformerEngineUnavailable as exc: + unavailable.append(f"transformer_engine_paged_kv: {exc}") + + drifts = tuple(_compare_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport( + reference_name=reference.name, + drifts=drifts, + unavailable=tuple(unavailable), + ) + + +def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Training-style full-sequence attention with exported attention-domain LSE.""" + + out, lse = _attention_with_lse( + inputs.q, + inputs.k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=inputs.q.size(2), + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="full_prefill", + out=out, + lse=lse, + provenance={ + "attention_mode": "prefill", + "materialization": "full_sequence", + "lse_domain": "attention", + }, + ) + + +def run_chunked_query_attention( + inputs: AttentionComparisonInputs, + *, + query_chunk_size: int | None, +) -> AttentionPathResult: + """Rollout-style chunked prefill replay over full KV on one device.""" + + sq = inputs.q.size(2) + chunk_size = ( + sq if query_chunk_size is None else _positive_int(query_chunk_size, "query_chunk_size") + ) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + chunk_bounds = _chunk_bounds(sq, chunk_size) + for q_start, q_end in chunk_bounds: + out, lse = _attention_with_lse( + inputs.q[:, :, q_start:q_end, :], + inputs.k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=q_start, + k_start=0, + total_query_len=sq, + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + out_chunks.append(out) + lse_chunks.append(lse) + + return AttentionPathResult( + name="chunked_prefill", + out=torch.cat(out_chunks, dim=2), + lse=torch.cat(lse_chunks, dim=2), + provenance={ + "attention_mode": "chunked_prefill", + "materialization": "query_chunks", + "query_chunk_size": chunk_size, + "chunk_bounds": [list(bound) for bound in chunk_bounds], + "lse_domain": "attention", + }, + ) + + +def run_paged_kv_attention( + inputs: AttentionComparisonInputs, + *, + kv_page_size: int | None, + merge_backend: MergeBackend = "rl_kernel", +) -> AttentionPathResult: + """Rollout-style paged-KV prefill replay with explicit LSE merge.""" + + skv = inputs.k.size(2) + page_size = skv if kv_page_size is None else _positive_int(kv_page_size, "kv_page_size") + states: list[_PartialAttentionState] = [] + page_bounds = _chunk_bounds(skv, page_size) + for k_start, k_end in page_bounds: + key_mask = ( + None if inputs.key_padding_mask is None else inputs.key_padding_mask[:, k_start:k_end] + ) + out, lse = _attention_with_lse( + inputs.q, + inputs.k[:, :, k_start:k_end, :], + inputs.v[:, :, k_start:k_end, :], + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=key_mask, + q_start=0, + k_start=k_start, + total_query_len=inputs.q.size(2), + total_kv_len=skv, + output_dtype=torch.float32, + ) + states.append( + _PartialAttentionState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_end, + ) + ) + + out, lse = _merge_partial_states(states, backend=merge_backend) + return AttentionPathResult( + name=f"{merge_backend}_paged_kv", + out=out.to(inputs.output_dtype), + lse=lse, + provenance={ + "attention_mode": "prefill", + "materialization": "paged_kv", + "kv_page_size": page_size, + "kv_page_bounds": [list(bound) for bound in page_bounds], + "merge_backend": merge_backend, + "merge_order": "global_block_index", + "lse_domain": "attention", + }, + ) + + +def transformer_engine_context_parallel_available() -> bool: + """Return whether the optional TE context-parallel helper module imports.""" + + try: + _load_te_context_parallel() + except TransformerEngineUnavailable: + return False + return True + + +def _compare_path( + candidate: AttentionPathResult, + reference: AttentionPathResult, + inputs: AttentionComparisonInputs, +) -> AttentionPathDrift: + dlogp = None + if inputs.lm_head_weight is not None and inputs.target_ids is not None: + candidate_logp = _selected_logps_from_attention(candidate.out, inputs) + reference_logp = _selected_logps_from_attention(reference.out, inputs) + dlogp = _drift_stats(candidate_logp, reference_logp, mask=inputs.active_token_mask) + + return AttentionPathDrift( + candidate_name=candidate.name, + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=dlogp, + provenance=candidate.provenance, + ) + + +def _attention_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: float | None, + key_padding_mask: torch.Tensor | None, + q_start: int, + k_start: int, + total_query_len: int, + total_kv_len: int, + output_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + + qf, kf, vf = q.float(), k.float(), v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + scale_value = scale if scale is not None else 1.0 / math.sqrt(dim) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_offset = total_kv_len - total_query_len + q_pos = torch.arange(sq, device=q.device) + q_start + query_offset + k_pos = torch.arange(skv, device=q.device) + k_start + scores = scores.masked_fill(k_pos[None, :] > q_pos[:, None], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return out.to(output_dtype), lse + + +def _merge_partial_states( + states: list[_PartialAttentionState], + *, + backend: MergeBackend, +) -> tuple[torch.Tensor, torch.Tensor]: + if not states: + raise ValueError("at least one partial state is required") + ordered = sorted(states, key=lambda state: (state.block_start, state.block_end)) + _validate_partial_states(ordered) + if backend == "rl_kernel": + return _merge_partial_states_rl_kernel(ordered) + if backend == "transformer_engine": + return _merge_partial_states_transformer_engine(ordered) + raise ValueError(f"unsupported merge backend: {backend}") + + +def _merge_partial_states_rl_kernel( + states: list[_PartialAttentionState], +) -> tuple[torch.Tensor, torch.Tensor]: + merged_out = states[0].out.float() + merged_lse = states[0].lse.float() + for state in states[1:]: + next_lse = torch.logaddexp(merged_lse, state.lse.float()) + finite = torch.isfinite(next_lse) + weight_prev = torch.where( + finite, + torch.exp(merged_lse - next_lse), + torch.zeros_like(next_lse), + ) + weight_next = torch.where( + finite, + torch.exp(state.lse.float() - next_lse), + torch.zeros_like(next_lse), + ) + merged_out = ( + weight_prev.unsqueeze(-1) * merged_out + weight_next.unsqueeze(-1) * state.out.float() + ) + merged_lse = next_lse + return merged_out, merged_lse + + +def _merge_partial_states_transformer_engine( + states: list[_PartialAttentionState], +) -> tuple[torch.Tensor, torch.Tensor]: + te_cp = _load_te_context_parallel() + merged_out = states[0].out.float() + merged_lse = states[0].lse.float() + for state in states[1:]: + previous_lse = merged_lse + merged_lse = previous_lse.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(merged_lse, state.lse.float()) + merged_out = te_cp.flash_attn_fwd_out_correction_init( + merged_out, + merged_lse, + previous_lse, + seq_dim=2, + ) + te_cp.flash_attn_fwd_out_correction( + merged_out, + state.out.float(), + merged_lse, + state.lse.float(), + seq_dim=2, + ) + return merged_out, merged_lse + + +def _load_te_context_parallel() -> Any: + try: + return importlib.import_module(_TE_CONTEXT_PARALLEL_MODULE) + except (ImportError, OSError, RuntimeError) as exc: + raise TransformerEngineUnavailable(str(exc)) from exc + + +def _selected_logps_from_attention( + out: torch.Tensor, + inputs: AttentionComparisonInputs, +) -> torch.Tensor: + if inputs.lm_head_weight is None or inputs.target_ids is None: + raise ValueError("lm_head_weight and target_ids are required for dlogp drift") + batch, heads, seq, dim = out.shape + hidden = out.transpose(1, 2).reshape(batch, seq, heads * dim) + if inputs.lm_head_weight.shape[1] != hidden.size(-1): + raise ValueError( + "lm_head_weight hidden dimension must equal Hq * D; " + f"got {inputs.lm_head_weight.shape[1]} and {hidden.size(-1)}" + ) + logits = torch.matmul(hidden.float(), inputs.lm_head_weight.float().transpose(0, 1)) + return selected_logprobs_reference( + logits, + inputs.target_ids, + mask=inputs.active_token_mask, + output_dtype=torch.float32, + ) + + +def _drift_stats( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + mask: torch.Tensor | None = None, +) -> DriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs() + values = _active_values(diff, mask) + active_count = int(values.numel()) + if active_count == 0: + return DriftStats(0.0, 0.0, 0.0, 0.0, 0) + return DriftStats( + max_abs=float(values.max().item()), + mean_abs=float(values.mean().item()), + p95_abs=float(torch.quantile(values, 0.95).item()), + p99_abs=float(torch.quantile(values, 0.99).item()), + active_count=active_count, + ) + + +def _active_values(diff: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: + if mask is None: + return diff.reshape(-1) + if mask.shape == diff.shape: + return diff[mask.to(device=diff.device, dtype=torch.bool)] + if mask.ndim == 2 and diff.ndim == 4 and mask.shape == (diff.size(0), diff.size(2)): + expanded = mask[:, None, :, None].expand_as(diff) + return diff[expanded.to(device=diff.device, dtype=torch.bool)] + if mask.ndim == 2 and diff.ndim == 3 and mask.shape == (diff.size(0), diff.size(2)): + expanded = mask[:, None, :].expand_as(diff) + return diff[expanded.to(device=diff.device, dtype=torch.bool)] + raise ValueError(f"mask shape {tuple(mask.shape)} cannot select diff shape {tuple(diff.shape)}") + + +def _validate_comparison_inputs(inputs: AttentionComparisonInputs) -> None: + _validate_qkv(inputs.q, inputs.k, inputs.v) + if inputs.key_padding_mask is not None: + if inputs.key_padding_mask.shape != (inputs.q.size(0), inputs.k.size(2)): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if inputs.key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + if (inputs.lm_head_weight is None) != (inputs.target_ids is None): + raise ValueError("lm_head_weight and target_ids must be provided together") + if inputs.target_ids is not None and inputs.target_ids.shape != ( + inputs.q.size(0), + inputs.q.size(2), + ): + raise ValueError("target_ids must have shape [B, Sq]") + if inputs.active_token_mask is not None: + if inputs.active_token_mask.shape != (inputs.q.size(0), inputs.q.size(2)): + raise ValueError("active_token_mask must have shape [B, Sq]") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have matching shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} must be divisible by Hkv={k.size(1)}") + + +def _validate_partial_states(states: list[_PartialAttentionState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching shapes") + if state.block_start < previous_end: + raise ValueError("partial state block ranges must not overlap") + previous_end = state.block_end + + +def _chunk_bounds(length: int, chunk_size: int) -> list[tuple[int, int]]: + if length <= 0: + raise ValueError("sequence length must be positive") + bounds: list[tuple[int, int]] = [] + cursor = 0 + while cursor < length: + end = min(cursor + chunk_size, length) + bounds.append((cursor, end)) + cursor = end + return bounds + + +def _positive_int(value: int, name: str) -> int: + if isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + return int(value) + + +__all__ = [ + "AttentionComparisonInputs", + "AttentionComparisonReport", + "AttentionPathDrift", + "AttentionPathResult", + "DriftStats", + "TransformerEngineUnavailable", + "compare_single_gpu_attention", + "run_chunked_query_attention", + "run_full_attention", + "run_paged_kv_attention", + "transformer_engine_context_parallel_available", +] diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py new file mode 100644 index 00000000..3b7a5edc --- /dev/null +++ b/tests/test_attention_comparison.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import argparse +import importlib +import json +import sys +import types + +import torch + +from rl_engine.kernels.gtest import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.testing.attention_comparison import ( + AttentionComparisonInputs, + compare_single_gpu_attention, +) + +_TE_CONTEXT_PARALLEL_MODULE = ( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" +) + + +def _qkv(*, seed: int = 1): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(2, 4, 6, 8, generator=gen) + k = torch.randn(2, 2, 6, 8, generator=gen) + v = torch.randn(2, 2, 6, 8, generator=gen) + return q, k, v + + +def _comparison_inputs() -> AttentionComparisonInputs: + q, k, v = _qkv() + gen = torch.Generator().manual_seed(2) + lm_head_weight = torch.randn(13, q.size(1) * q.size(3), generator=gen) + target_ids = torch.randint(0, 13, (q.size(0), q.size(2)), generator=gen) + active_mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, True, True, True, True, False], + ], + dtype=torch.bool, + ) + return AttentionComparisonInputs( + q=q, + k=k, + v=v, + causal=True, + lm_head_weight=lm_head_weight, + target_ids=target_ids, + active_token_mask=active_mask, + ) + + +def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=2, + kv_page_size=3, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert set(by_name) == {"chunked_prefill", "rl_kernel_paged_kv"} + for drift in by_name.values(): + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.active_count == 7 + assert drift.dlogp.p95_abs <= 1.0e-6 + + payload = report.to_dict() + assert payload["reference_name"] == "full_prefill" + assert payload["drifts"][0]["out"]["p99_abs"] >= 0.0 + json.dumps(payload) + + +def test_single_gpu_attention_harness_preserves_key_padding_mask(): + q, k, v = _qkv(seed=3) + key_padding_mask = torch.tensor( + [ + [True, True, True, False, False, False], + [True, False, True, True, False, False], + ], + dtype=torch.bool, + ) + + report = compare_single_gpu_attention( + AttentionComparisonInputs( + q=q, + k=k, + v=v, + causal=True, + key_padding_mask=key_padding_mask, + ), + query_chunk_size=4, + kv_page_size=2, + ) + + assert report.unavailable == () + for drift in report.drifts: + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + + +def test_transformer_engine_merge_oracle_can_be_reused_when_available(monkeypatch): + calls = {"lse": 0, "out": 0} + + def lse_correction(softmax_lse, softmax_lse_per_step): + calls["lse"] += 1 + softmax_lse.copy_(torch.logaddexp(softmax_lse, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + calls["out"] += 1 + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert "transformer_engine_paged_kv" in by_name + assert by_name["transformer_engine_paged_kv"].out.max_abs <= 1.0e-6 + assert by_name["transformer_engine_paged_kv"].lse.max_abs <= 1.0e-6 + assert calls["lse"] > 0 + assert calls["out"] > 0 + assert report.unavailable == () + + +def test_transformer_engine_path_reports_unavailable_without_failing(monkeypatch): + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == _TE_CONTEXT_PARALLEL_MODULE: + raise ImportError("test TE unavailable") + return real_import_module(name, package) + + monkeypatch.setattr(importlib, "import_module", fake_import_module) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + assert {drift.candidate_name for drift in report.drifts} == { + "chunked_prefill", + "rl_kernel_paged_kv", + } + assert report.unavailable == ("transformer_engine_paged_kv: test TE unavailable",) + + +def test_operator_comparison_specs_register_attention(): + args = argparse.Namespace( + op="attention", + candidate="pytorch", + arch_key=None, + batch=1, + seq=3, + vocab=17, + seed=7, + input_mode="random", + constant_value=0.5, + token_value=3, + normalized_dim=128, + k_dim=16, + n_dim=32, + theta=1.0e6, + eps=1.0e-6, + ) + + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(args) + report = run_operator_suite("attention", candidates=[candidate], cases=[case]) + + assert report.passed + assert report.candidates[0].cases[0].op_class == "attention" From 6d57478aff25fea0e683152c0b8f9add059e3177 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 2 Aug 2026 23:39:40 +0800 Subject: [PATCH 2/5] feat(attention): add rope comparison harness Signed-off-by: inaniloquentee <3051000145@qq.com> --- .../ws2-attention-single-gpu-harness.md | 36 +++- rl_engine/testing/__init__.py | 6 + rl_engine/testing/attention_comparison.py | 188 ++++++++++++++++++ tests/test_attention_comparison.py | 50 +++++ 4 files changed, 279 insertions(+), 1 deletion(-) diff --git a/docs/design/ws2-attention-single-gpu-harness.md b/docs/design/ws2-attention-single-gpu-harness.md index 32eba5da..acca4101 100644 --- a/docs/design/ws2-attention-single-gpu-harness.md +++ b/docs/design/ws2-attention-single-gpu-harness.md @@ -17,6 +17,17 @@ Implemented paths: - `transformer_engine_paged_kv`: optional oracle that reuses NVIDIA Transformer Engine's context-parallel PyTorch correction helpers when TE is installed. +RoPE scope: + +- `unfused_rope_attention`: canonical `RoPE -> Attention` path; +- `fused_like_rope_attention`: semantic `RoPE+Attention` path that applies the + same canonical RoPE rules before attention, then records the fused boundary in + provenance. + +The RoPE path is still single-GPU attribution. It proves that both sides agree +on post-RoPE Q/K, `out`, attention-domain `lse`, and optional active-token +`dlogp` before CP communication or production fused kernels are introduced. + ## Report `rl_engine.testing.attention_comparison.compare_single_gpu_attention` emits a @@ -30,6 +41,15 @@ structured report with: merge order, and LSE domain; - optional-backend unavailability reasons. +`compare_single_gpu_rope_attention` emits the same drift schema and additionally +reports post-RoPE Q/K drift. Its provenance records: + +- Q/K state as `post_rope`; +- `position_ids` shape and range; +- `rope_theta`, `rotary_dim`, `rope_cast_at`, and `rope_output_dtype`; +- `fusion_boundary` as either `unfused_rope_attention` or + `fused_rope_attention`. + The selected-logprob convention follows #207: ```text @@ -72,6 +92,7 @@ The attention-specific WS2 comparison entry point is Python-first for now: ```python from rl_engine.testing.attention_comparison import ( AttentionComparisonInputs, + compare_single_gpu_rope_attention, compare_single_gpu_attention, ) @@ -82,6 +103,18 @@ report = compare_single_gpu_attention( include_transformer_engine=True, ) print(report.to_dict()) + +rope_report = compare_single_gpu_rope_attention( + AttentionComparisonInputs( + q=q, + k=k, + v=v, + rope_positions=torch.arange(q.size(2), device=q.device), + target_ids=target_ids, + lm_head_weight=w, + ) +) +print(rope_report.to_dict()) ``` ## Validation @@ -92,4 +125,5 @@ python -m pytest tests/test_attention_comparison.py -q The tests cover full vs chunked/paged equivalence, active-token `dlogp` drift, optional TE correction-helper reuse through a fake TE module, JSON-compatible -reports, and `attention` registration in the generic operator comparison specs. +reports, RoPE+Attention post-RoPE Q/K attribution, and `attention` registration +in the generic operator comparison specs. diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index b9bffee0..792b36d0 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -11,9 +11,12 @@ DriftStats, TransformerEngineUnavailable, compare_single_gpu_attention, + compare_single_gpu_rope_attention, run_chunked_query_attention, run_full_attention, + run_fused_like_rope_attention, run_paged_kv_attention, + run_unfused_rope_attention, transformer_engine_context_parallel_available, ) from .reference_ops import ( @@ -36,6 +39,7 @@ "SyntheticRLKernelBatch", "TransformerEngineUnavailable", "active_token_count", + "compare_single_gpu_rope_attention", "compare_single_gpu_attention", "compute_policy_ratio", "compute_reference_kl", @@ -43,8 +47,10 @@ "masked_mean", "masked_sum", "run_chunked_query_attention", + "run_fused_like_rope_attention", "run_full_attention", "run_paged_kv_attention", + "run_unfused_rope_attention", "selected_logprobs_reference", "summarize_kernel_drift", "transformer_engine_context_parallel_available", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 350ecdf2..ebf772f7 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -18,6 +18,7 @@ import torch +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.testing.reference_ops import selected_logprobs_reference MergeBackend = Literal["rl_kernel", "transformer_engine"] @@ -45,6 +46,11 @@ class AttentionComparisonInputs: target_ids: torch.Tensor | None = None active_token_mask: torch.Tensor | None = None output_dtype: torch.dtype = torch.float32 + rope_positions: torch.Tensor | None = None + rope_theta: float = 1_000_000.0 + rope_rotary_dim: int | None = None + rope_cast_at: str = "after_rope" + rope_output_dtype: torch.dtype | None = None @dataclass(frozen=True) @@ -55,6 +61,8 @@ class AttentionPathResult: out: torch.Tensor lse: torch.Tensor provenance: dict[str, Any] + post_rope_q: torch.Tensor | None = None + post_rope_k: torch.Tensor | None = None @dataclass(frozen=True) @@ -86,6 +94,8 @@ class AttentionPathDrift: lse: DriftStats dlogp: DriftStats | None provenance: dict[str, Any] + post_rope_q: DriftStats | None = None + post_rope_k: DriftStats | None = None def to_dict(self) -> dict[str, Any]: return { @@ -93,6 +103,8 @@ def to_dict(self) -> dict[str, Any]: "out": self.out.to_dict(), "lse": self.lse.to_dict(), "dlogp": None if self.dlogp is None else self.dlogp.to_dict(), + "post_rope_q": (None if self.post_rope_q is None else self.post_rope_q.to_dict()), + "post_rope_k": (None if self.post_rope_k is None else self.post_rope_k.to_dict()), "provenance": self.provenance, } @@ -162,6 +174,24 @@ def compare_single_gpu_attention( ) +def compare_single_gpu_rope_attention( + inputs: AttentionComparisonInputs, +) -> AttentionComparisonReport: + """Compare canonical unfused RoPE+Attention with fused-like materialization. + + This attribution path keeps the computation on one device and checks the + boundary that matters before CP communication: post-RoPE Q/K identity and + the resulting attention ``out`` / attention-domain ``lse``. + """ + + _validate_comparison_inputs(inputs) + _validate_rope_inputs(inputs) + reference = run_unfused_rope_attention(inputs) + candidates = [run_fused_like_rope_attention(inputs)] + drifts = tuple(_compare_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport(reference_name=reference.name, drifts=drifts) + + def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: """Training-style full-sequence attention with exported attention-domain LSE.""" @@ -190,6 +220,68 @@ def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult ) +def run_unfused_rope_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Canonical ``RoPE -> Attention`` reference materialization.""" + + post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) + out, lse = _attention_with_lse( + post_rope_q, + post_rope_k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=post_rope_q.size(2), + total_kv_len=post_rope_k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="unfused_rope_attention", + out=out, + lse=lse, + provenance=_rope_attention_provenance( + inputs, + materialization="rope_then_attention", + fusion_boundary="unfused_rope_attention", + ), + post_rope_q=post_rope_q, + post_rope_k=post_rope_k, + ) + + +def run_fused_like_rope_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Semantic fused ``RoPE+Attention`` path using the same canonical RoPE rules.""" + + post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) + out, lse = _attention_with_lse( + post_rope_q, + post_rope_k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=post_rope_q.size(2), + total_kv_len=post_rope_k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="fused_like_rope_attention", + out=out, + lse=lse, + provenance=_rope_attention_provenance( + inputs, + materialization="fused_like_rope_attention", + fusion_boundary="fused_rope_attention", + ), + post_rope_q=post_rope_q, + post_rope_k=post_rope_k, + ) + + def run_chunked_query_attention( inputs: AttentionComparisonInputs, *, @@ -317,9 +409,63 @@ def _compare_path( lse=_drift_stats(candidate.lse, reference.lse), dlogp=dlogp, provenance=candidate.provenance, + post_rope_q=( + None + if candidate.post_rope_q is None or reference.post_rope_q is None + else _drift_stats(candidate.post_rope_q, reference.post_rope_q) + ), + post_rope_k=( + None + if candidate.post_rope_k is None or reference.post_rope_k is None + else _drift_stats(candidate.post_rope_k, reference.post_rope_k) + ), ) +def _apply_rope_to_qk(inputs: AttentionComparisonInputs) -> tuple[torch.Tensor, torch.Tensor]: + _validate_rope_inputs(inputs) + assert inputs.rope_positions is not None + rope = NativeRoPEOp() + output_dtype = _rope_output_dtype(inputs) + q = rope.forward_fp32(inputs.q, inputs.rope_positions, theta=inputs.rope_theta).to(output_dtype) + k = rope.forward_fp32(inputs.k, inputs.rope_positions, theta=inputs.rope_theta).to(output_dtype) + return q, k + + +def _rope_output_dtype(inputs: AttentionComparisonInputs) -> torch.dtype: + return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype + + +def _rope_rotary_dim(inputs: AttentionComparisonInputs) -> int: + return inputs.q.size(-1) if inputs.rope_rotary_dim is None else inputs.rope_rotary_dim + + +def _rope_attention_provenance( + inputs: AttentionComparisonInputs, + *, + materialization: str, + fusion_boundary: str, +) -> dict[str, Any]: + assert inputs.rope_positions is not None + return { + "attention_mode": "prefill", + "materialization": materialization, + "rope_state": "post_rope", + "q_rope_state": "post_rope", + "k_rope_state": "post_rope", + "position_kind": "position_ids", + "position_ids_shape": list(inputs.rope_positions.shape), + "position_ids_min": int(inputs.rope_positions.min().item()), + "position_ids_max": int(inputs.rope_positions.max().item()), + "rope_theta": float(inputs.rope_theta), + "rotary_dim": _rope_rotary_dim(inputs), + "rope_cast_at": inputs.rope_cast_at, + "rope_output_dtype": str(_rope_output_dtype(inputs)).replace("torch.", ""), + "fusion_boundary": fusion_boundary, + "lse_domain": "attention", + } + + def _attention_with_lse( q: torch.Tensor, k: torch.Tensor, @@ -521,6 +667,45 @@ def _validate_comparison_inputs(inputs: AttentionComparisonInputs) -> None: raise ValueError("active_token_mask must have shape [B, Sq]") if inputs.active_token_mask.dtype != torch.bool: raise ValueError("active_token_mask must be bool") + if not isinstance(inputs.rope_theta, (float, int)) or isinstance(inputs.rope_theta, bool): + raise ValueError("rope_theta must be a positive number") + if float(inputs.rope_theta) <= 0: + raise ValueError("rope_theta must be a positive number") + if inputs.rope_output_dtype is not None and not isinstance( + inputs.rope_output_dtype, torch.dtype + ): + raise ValueError("rope_output_dtype must be a torch.dtype when provided") + + +def _validate_rope_inputs(inputs: AttentionComparisonInputs) -> None: + if inputs.rope_positions is None: + raise ValueError("rope_positions are required for RoPE+Attention comparison") + if inputs.q.size(2) != inputs.k.size(2): + raise ValueError("RoPE+Attention comparison currently requires Sq == Skv") + if inputs.rope_rotary_dim is not None: + if isinstance(inputs.rope_rotary_dim, bool) or inputs.rope_rotary_dim <= 0: + raise ValueError("rope_rotary_dim must be a positive integer when provided") + if inputs.rope_rotary_dim != inputs.q.size(-1): + raise ValueError( + "rope_rotary_dim must equal head_dim until partial-rotary RoPE is supported" + ) + if inputs.rope_cast_at != "after_rope": + raise ValueError("rope_cast_at must be 'after_rope' for the current fp32 RoPE reference") + if ( + inputs.rope_positions.device != inputs.q.device + or inputs.rope_positions.device != inputs.k.device + ): + raise ValueError("rope_positions must be on the same device as q/k") + if inputs.rope_positions.dtype not in {torch.int32, torch.int64, torch.long}: + raise ValueError("rope_positions must contain integer token positions") + if inputs.rope_positions.ndim == 1: + if inputs.rope_positions.numel() != inputs.q.size(2): + raise ValueError("1D rope_positions must have length Sq") + elif inputs.rope_positions.ndim == 2: + if inputs.rope_positions.shape != (inputs.q.size(0), inputs.q.size(2)): + raise ValueError("2D rope_positions must have shape [B, Sq]") + else: + raise ValueError("rope_positions must have shape [Sq] or [B, Sq]") def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: @@ -570,9 +755,12 @@ def _positive_int(value: int, name: str) -> int: "AttentionPathResult", "DriftStats", "TransformerEngineUnavailable", + "compare_single_gpu_rope_attention", "compare_single_gpu_attention", "run_chunked_query_attention", + "run_fused_like_rope_attention", "run_full_attention", "run_paged_kv_attention", + "run_unfused_rope_attention", "transformer_engine_context_parallel_available", ] diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 3b7a5edc..100a52fb 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -9,6 +9,7 @@ import sys import types +import pytest import torch from rl_engine.kernels.gtest import run_operator_suite @@ -16,6 +17,7 @@ from rl_engine.testing.attention_comparison import ( AttentionComparisonInputs, compare_single_gpu_attention, + compare_single_gpu_rope_attention, ) _TE_CONTEXT_PARALLEL_MODULE = ( @@ -104,6 +106,54 @@ def test_single_gpu_attention_harness_preserves_key_padding_mask(): assert drift.lse.max_abs <= 1.0e-6 +def test_single_gpu_rope_attention_harness_reports_rope_and_attention_drift(): + base = _comparison_inputs() + report = compare_single_gpu_rope_attention( + AttentionComparisonInputs( + q=base.q, + k=base.k, + v=base.v, + causal=True, + lm_head_weight=base.lm_head_weight, + target_ids=base.target_ids, + active_token_mask=base.active_token_mask, + rope_positions=torch.arange(base.q.size(2), dtype=torch.long), + rope_theta=1_000_000.0, + rope_rotary_dim=base.q.size(-1), + rope_output_dtype=torch.float32, + ) + ) + + assert report.reference_name == "unfused_rope_attention" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.candidate_name == "fused_like_rope_attention" + assert drift.post_rope_q is not None + assert drift.post_rope_k is not None + assert drift.post_rope_q.max_abs <= 1.0e-6 + assert drift.post_rope_k.max_abs <= 1.0e-6 + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.max_abs <= 1.0e-6 + assert drift.provenance["position_kind"] == "position_ids" + assert drift.provenance["position_ids_shape"] == [base.q.size(2)] + assert drift.provenance["rotary_dim"] == base.q.size(-1) + assert drift.provenance["rope_cast_at"] == "after_rope" + assert drift.provenance["fusion_boundary"] == "fused_rope_attention" + + payload = report.to_dict() + assert payload["drifts"][0]["post_rope_q"]["active_count"] == base.q.numel() + json.dumps(payload) + + +def test_single_gpu_rope_attention_requires_position_metadata(): + base = _comparison_inputs() + + with pytest.raises(ValueError, match="rope_positions are required"): + compare_single_gpu_rope_attention(AttentionComparisonInputs(q=base.q, k=base.k, v=base.v)) + + def test_transformer_engine_merge_oracle_can_be_reused_when_available(monkeypatch): calls = {"lse": 0, "out": 0} From cbb5e2ceb54947ebfed578db55ca0be1a5c22b44 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 3 Aug 2026 16:18:01 +0800 Subject: [PATCH 3/5] fix(attention): harden transformer engine merge oracle Signed-off-by: inaniloquentee <3051000145@qq.com> --- ...attention-transformer-engine-reuse-plan.md | 106 +++++++++++ rl_engine/testing/attention_comparison.py | 173 ++++++++++++++++-- tests/test_attention_comparison.py | 117 ++++++++++++ 3 files changed, 380 insertions(+), 16 deletions(-) create mode 100644 docs/design/ws2-attention-transformer-engine-reuse-plan.md diff --git a/docs/design/ws2-attention-transformer-engine-reuse-plan.md b/docs/design/ws2-attention-transformer-engine-reuse-plan.md new file mode 100644 index 00000000..76c5bb40 --- /dev/null +++ b/docs/design/ws2-attention-transformer-engine-reuse-plan.md @@ -0,0 +1,106 @@ +# WS2 Attention Transformer Engine 复用方案 + +Status: #235 设计补充 + +## 设计结论 + +Transformer Engine（TE）在 #235 中只能是显式 opt-in 的 validation oracle +或 backend candidate，不是 RL-Kernel attention 语义的可信源。可信源仍然是 +RL-Kernel 自己的 `AttentionContract`、RoPE/cache metadata、 +attention-domain `lse`、固定 `global_block_index` merge 顺序、 +deterministic reference 和 drift report。 + +TE 复用分为三层： + +| 层级 | TE 角色 | 允许范围 | +| --- | --- | --- | +| Merge oracle | 复用 TE context-parallel correction helpers 校验 `(out, lse)` online-softmax merge | PR2、PR3、PR5、PR6 | +| Fused forward candidate | 评估 `DotProductAttention` 作为 opt-in 生产后端候选 | 仅 PR7 | +| Backward oracle | 仅在 TE 暴露兼容 saved forward state 时，通过 autograd/backward 对比 `dq/dk/dv` | 仅 PR8 | + +## Merge Oracle Contract + +对任意 Q row，RL-Kernel 先按逻辑 KV block 生成 partial states： + +```text +state_i = (out_i, lse_i, global_block_index_i) +``` + +其中 `out_i` 是本地 KV block 内已经归一化的 attention output，`lse_i` +是 attention-domain LSE，shape 为 `[B, Hq, Sq]`。所有 state 必须按 +`global_block_index` 排序后再合并： + +```text +lse_new = logaddexp(lse_prev, lse_i) +out_new = exp(lse_prev - lse_new) * out_prev + + exp(lse_i - lse_new) * out_i +``` + +TE helper 可以负责 correction arithmetic，但语义输入必须由 RL-Kernel 提供： + +```text +TE_merge(sorted(RL-Kernel partial states)) == RL-Kernel_merge(sorted(partial states)) +``` + +调用 TE 前，RL-Kernel 必须保证： + +- merge accumulation 使用 FP32，只有 `final_write` 才 downcast； +- merge 顺序来自逻辑 `global_block_index`，不是通信 arrival order； +- all-masked / empty-KV row 保持 `lse = -inf`、`out = 0`，不能产生 NaN； +- TE adapter 启用前必须完成 capability probe：module/symbol 存在、helper + signature 兼容、tiny numeric merge smoke 通过； +- RoPE state、causal/padding mask、packed/varlen boundary、cache position 已经对齐。 + +## PR-level TE Plan + +| PR | RL-Kernel 核心功能 | TE 复用方式 | 精确 TE API | RL-Kernel 必须准备 | Gate / fallback | +| --- | --- | --- | --- | --- | --- | +| PR1 / #236 | 定义 attention contract、sharding/reduction metadata、RoPE/cache 字段 | 不调用 TE；只预留 `transformer_engine` 作为未来显式 backend 名称 | 无 | backend、reduction、`lse_domain`、`merge_order`、RoPE/cache identity 字段 | 不依赖 TE；metadata 缺失仍由 RL-Kernel contract fail | +| PR2 / #253 | 单 GPU full/chunked/paged-KV attention comparison harness | optional paged-KV merge oracle | `transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py`；`transformer_engine.pytorch.attention.dot_product_attention.context_parallel`；`flash_attn_fwd_softmax_lse_correction`；`flash_attn_fwd_out_correction_init`；`flash_attn_fwd_out_correction` | 相同 Q/K/V、相同 causal/padding metadata、相同 KV page order、RL-Kernel partial states `(out_i, lse_i)` | 对比 `TE_merge(partials)` 和 `RL-Kernel_merge(partials)` 的 `out/lse`；TE 不可用时 report `unavailable` | +| PR3 / #238 | post-RoPE Q/K 上的 deterministic CP attention reference | optional CP merge oracle test | 同 PR2 的 `context_parallel.py` module/functions | post-RoPE Q/K boundary、CP partial states、不重叠 global KV block ranges、固定 merge order | TE 不可用时 skip；TE 不定义 reference path | +| PR4 | Qwen3-8B TP=2 CP=2 BF16 cross-config 集成和 backend provenance | policy/provenance only | 不新增 TE 调用 | runtime descriptor 可 request `transformer_engine`，但默认执行仍是 deterministic reference | 记录 requested backend、actual backend、fallback reason、TE availability；禁止 silent fallback | +| PR5 | 分布式 prefill/chunked-prefill drift benchmark 和 report artifacts | benchmark merge oracle | 通过 `TEContextParallelMergeAdapter` 调用同 PR2 的 `context_parallel.py` module/functions | 与 RL-Kernel merge 完全相同的 gathered CP partial states、per-rank block metadata hash、FP32 merge dtype | 报告 `merge_drift = drift(TE_merge(partials), RL-Kernel_merge(partials))`；benchmark 可 provenance fallback | +| PR6 | decode-stage KV-cache CP attention replay | decode / paged-KV merge oracle only | 通过 decode TE merge adapter 调用同 PR2 的 `context_parallel.py` module/functions | `cache_position`、`kv_seq_lens`、page table、prefix-cache identity、global token positions、RoPE cache state、sorted logical page/block order | TE 只验证 `(out, lse)` merge；cache/page identity 不一致时，在调用 TE 前 fail | +| PR7 | deterministic reference 稳定后的 fused prefill/decode backend alignment | full fused forward backend candidate | `transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py`；`transformer_engine.pytorch.DotProductAttention`；actual backend 可观测时记录为 `FlashAttention` / `FusedAttention` / `UnfusedDotProductAttention` | 精确 layout / `qkv_format`、mask mode、RoPE fusion boundary、dtype、scale placement、dropout=0 correctness mode、deterministic controls、LSE export capability、actual-backend 观测方式 | 只有 TE output、attention-domain LSE、actual backend provenance 都能对齐/记录时才可作为 production candidate；如果不能导出 LSE 或不能观测 actual backend，只能算 exploratory，并记录原因 | +| PR8 | training backward CP attention reference 和 gradient drift validation | optional backward oracle | `DotProductAttention` autograd/backward path，仅当 compatible saved forward state 暴露时使用 | 与 RL-Kernel reference 相同的 forward inputs/metadata：`out`、attention-domain `lse`、masks、RoPE state、sequence/cache metadata、CP block ownership | 对比 `dq/dk/dv`；没有兼容 TE backward state 时明确写 `not used`，不能宣称复用 TE backward | + +## Capability / Provenance Checklist + +任何 PR 只要提到 TE，都必须写清： + +```text +te_available, te_version, te_module, te_symbols +te_capability_probe, te_signature_checked, te_numeric_selftest +requested_backend, actual_backend, actual_backend_source +fallback, fallback_reason +attention_mode, dtype, layout/qkv_format, mask_alignment +lse_domain, lse_exported, merge_order, accum_dtype, downcast_at +split_kv_policy, paged_kv_policy, cp_block_metadata_hash +scale_placement, deterministic_controls, dropout_policy, te_env_controls +``` + +fallback 策略： + +| 场景 | TE 不可用 / capability 不匹配时 | +| --- | --- | +| optional oracle test | skip / report unavailable | +| benchmark exploration | provenance fallback 到 deterministic reference | +| correctness gate | fail closed | +| production backend | fail closed 或显式 provenance fallback；禁止 silent fallback | + +## 不宣称的事 + +- 不把 TE 设为 #235 的硬依赖。 +- 不用 TE API 反向定义 RL-Kernel contract。 +- 不在 metadata 不完整时 silent fallback 到 TE。 +- 不用 NCCL / TE arrival order 决定 attention merge 数值顺序。 +- 不在 PR7 前把 TE fused path 宣称为默认生产路径。 +- PR7 如果拿不到 attention-domain LSE，不宣称完整 correctness closure。 +- PR8 如果拿不到兼容 backward state，不宣称复用 TE backward。 + +## 最终判断标准 + +TE 可以帮助验证和加速，但 #235 的正确性仍由 RL-Kernel 自己的 contract、 +metadata、deterministic reference 和 drift report 保证。当前最值得复用的是 +TE context-parallel correction helper；完整 `DotProductAttention` 路径只有在 +显式声明 capability 并满足 RL-Kernel 语义契约后，才允许作为生产候选后端。 diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index ebf772f7..dbee69af 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -12,6 +12,8 @@ from __future__ import annotations import importlib +import importlib.metadata as importlib_metadata +import inspect import math from dataclasses import dataclass from typing import Any, Literal @@ -26,6 +28,22 @@ _TE_CONTEXT_PARALLEL_MODULE = ( "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" ) +_TE_CONTEXT_PARALLEL_HELPERS = { + "flash_attn_fwd_softmax_lse_correction": ("softmax_lse", "softmax_lse_per_step"), + "flash_attn_fwd_out_correction_init": ( + "out_init_step", + "softmax_lse", + "softmax_lse_init_step", + "seq_dim", + ), + "flash_attn_fwd_out_correction": ( + "out", + "out_per_step", + "softmax_lse", + "softmax_lse_per_step", + "seq_dim", + ), +} class TransformerEngineUnavailable(RuntimeError): @@ -366,19 +384,33 @@ def run_paged_kv_attention( ) out, lse = _merge_partial_states(states, backend=merge_backend) + provenance = { + "attention_mode": "prefill", + "materialization": "paged_kv", + "kv_page_size": page_size, + "kv_page_bounds": [list(bound) for bound in page_bounds], + "merge_backend": merge_backend, + "requested_backend": merge_backend, + "actual_backend": ( + "te_context_parallel_merge_helpers" + if merge_backend == "transformer_engine" + else "rl_kernel" + ), + "fallback": False, + "fallback_reason": None, + "merge_order": "global_block_index", + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + } + if merge_backend == "transformer_engine": + provenance.update(_te_context_parallel_provenance()) return AttentionPathResult( name=f"{merge_backend}_paged_kv", out=out.to(inputs.output_dtype), lse=lse, - provenance={ - "attention_mode": "prefill", - "materialization": "paged_kv", - "kv_page_size": page_size, - "kv_page_bounds": [list(bound) for bound in page_bounds], - "merge_backend": merge_backend, - "merge_order": "global_block_index", - "lse_domain": "attention", - }, + provenance=provenance, ) @@ -562,29 +594,131 @@ def _merge_partial_states_transformer_engine( merged_lse = states[0].lse.float() for state in states[1:]: previous_lse = merged_lse - merged_lse = previous_lse.clone() - te_cp.flash_attn_fwd_softmax_lse_correction(merged_lse, state.lse.float()) + state_out = state.out.float() + state_lse = state.lse.float() + both_masked = torch.isneginf(previous_lse) & torch.isneginf(state_lse) + te_previous_lse = torch.where(both_masked, torch.zeros_like(previous_lse), previous_lse) + te_state_lse = torch.where(both_masked, torch.zeros_like(state_lse), state_lse) + merged_lse = te_previous_lse.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(merged_lse, te_state_lse) merged_out = te_cp.flash_attn_fwd_out_correction_init( merged_out, merged_lse, - previous_lse, + te_previous_lse, seq_dim=2, ) te_cp.flash_attn_fwd_out_correction( merged_out, - state.out.float(), + state_out, merged_lse, - state.lse.float(), + te_state_lse, seq_dim=2, ) + if both_masked.any(): + merged_lse = torch.where(both_masked, previous_lse, merged_lse) + merged_out = torch.where( + both_masked.unsqueeze(-1), + torch.zeros_like(merged_out), + merged_out, + ) return merged_out, merged_lse def _load_te_context_parallel() -> Any: try: - return importlib.import_module(_TE_CONTEXT_PARALLEL_MODULE) + module = importlib.import_module(_TE_CONTEXT_PARALLEL_MODULE) except (ImportError, OSError, RuntimeError) as exc: raise TransformerEngineUnavailable(str(exc)) from exc + _probe_te_context_parallel(module) + return module + + +def _probe_te_context_parallel(module: Any) -> None: + missing = [ + name for name in _TE_CONTEXT_PARALLEL_HELPERS if not callable(getattr(module, name, None)) + ] + if missing: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} missing required helpers: {', '.join(missing)}" + ) + + for name, expected in _TE_CONTEXT_PARALLEL_HELPERS.items(): + helper = getattr(module, name) + try: + parameters = tuple(inspect.signature(helper).parameters) + except (TypeError, ValueError) as exc: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE}.{name} signature is not inspectable" + ) from exc + if parameters[: len(expected)] != expected: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE}.{name} has incompatible signature " + f"{parameters}; expected prefix {expected}" + ) + + try: + lse_a = torch.tensor([[[0.0, -1.0]]], dtype=torch.float32) + lse_b = torch.tensor([[[1.0, -3.0]]], dtype=torch.float32) + out_a = torch.tensor([[[[1.0, -2.0], [0.5, 2.0]]]], dtype=torch.float32) + out_b = torch.tensor([[[[-1.0, 4.0], [3.0, -0.5]]]], dtype=torch.float32) + expected_lse = torch.logaddexp(lse_a, lse_b) + expected_out = ( + torch.exp(lse_a - expected_lse).unsqueeze(-1) * out_a + + torch.exp(lse_b - expected_lse).unsqueeze(-1) * out_b + ) + + probed_lse = lse_a.clone() + module.flash_attn_fwd_softmax_lse_correction(probed_lse, lse_b) + probed_out = module.flash_attn_fwd_out_correction_init( + out_a.clone(), + probed_lse, + lse_a, + seq_dim=2, + ) + module.flash_attn_fwd_out_correction( + probed_out, + out_b, + probed_lse, + lse_b, + seq_dim=2, + ) + except Exception as exc: + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} helper numeric self-test failed: {exc}" + ) from exc + + if not torch.allclose(probed_lse, expected_lse, atol=1.0e-6, rtol=0.0): + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} LSE helper numeric self-test failed" + ) + if not torch.allclose(probed_out, expected_out, atol=1.0e-6, rtol=0.0): + raise TransformerEngineUnavailable( + f"{_TE_CONTEXT_PARALLEL_MODULE} out helper numeric self-test failed" + ) + + +def _te_context_parallel_provenance() -> dict[str, Any]: + return { + "te_available": True, + "te_version": _te_version(), + "te_module": _TE_CONTEXT_PARALLEL_MODULE, + "te_symbols": list(_TE_CONTEXT_PARALLEL_HELPERS), + "te_capability_probe": "passed", + "te_signature_checked": True, + "te_numeric_selftest": "passed", + "actual_backend_source": "rl_kernel_te_context_parallel_adapter", + "deterministic_controls": "not_applicable_merge_only", + "dropout_policy": "not_applicable_merge_only", + } + + +def _te_version() -> str | None: + for package_name in ("transformer-engine", "transformer_engine"): + try: + return importlib_metadata.version(package_name) + except importlib_metadata.PackageNotFoundError: + continue + return None def _selected_logps_from_attention( @@ -620,7 +754,14 @@ def _drift_stats( f"candidate shape {tuple(candidate.shape)} must match " f"reference shape {tuple(reference.shape)}" ) - diff = (candidate.float() - reference.float()).abs() + candidate_fp32 = candidate.float() + reference_fp32 = reference.float() + raw_diff = (candidate_fp32 - reference_fp32).abs() + diff = torch.where( + candidate_fp32 == reference_fp32, + torch.zeros_like(raw_diff), + raw_diff, + ) values = _active_values(diff, mask) active_count = int(values.numel()) if active_count == 0: diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 100a52fb..061e8edd 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -18,6 +18,7 @@ AttentionComparisonInputs, compare_single_gpu_attention, compare_single_gpu_rope_attention, + run_paged_kv_attention, ) _TE_CONTEXT_PARALLEL_MODULE = ( @@ -191,11 +192,74 @@ def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim assert "transformer_engine_paged_kv" in by_name assert by_name["transformer_engine_paged_kv"].out.max_abs <= 1.0e-6 assert by_name["transformer_engine_paged_kv"].lse.max_abs <= 1.0e-6 + provenance = by_name["transformer_engine_paged_kv"].provenance + assert provenance["te_available"] is True + assert provenance["te_module"] == _TE_CONTEXT_PARALLEL_MODULE + assert provenance["te_capability_probe"] == "passed" + assert provenance["te_signature_checked"] is True + assert provenance["te_numeric_selftest"] == "passed" + assert provenance["actual_backend"] == "te_context_parallel_merge_helpers" + assert provenance["actual_backend_source"] == "rl_kernel_te_context_parallel_adapter" + assert provenance["accum_dtype"] == "fp32" + assert provenance["downcast_at"] == "final_write" assert calls["lse"] > 0 assert calls["out"] > 0 assert report.unavailable == () +def test_transformer_engine_merge_oracle_keeps_all_masked_rows_stable(monkeypatch): + def lse_correction(softmax_lse, softmax_lse_per_step): + max_scale = torch.max(softmax_lse, softmax_lse_per_step) + min_scale = torch.min(softmax_lse, softmax_lse_per_step) + softmax_lse.copy_(max_scale + torch.log1p(torch.exp(min_scale - max_scale))) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + q, k, v = _qkv(seed=11) + inputs = AttentionComparisonInputs( + q=q, + k=k, + v=v, + causal=False, + key_padding_mask=torch.zeros(q.size(0), k.size(2), dtype=torch.bool), + ) + + te_result = run_paged_kv_attention( + inputs, + kv_page_size=2, + merge_backend="transformer_engine", + ) + assert torch.equal(te_result.out, torch.zeros_like(te_result.out)) + assert torch.isneginf(te_result.lse).all() + + report = compare_single_gpu_attention( + inputs, + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert by_name["transformer_engine_paged_kv"].out.max_abs == 0.0 + assert by_name["transformer_engine_paged_kv"].lse.max_abs == 0.0 + assert report.unavailable == () + + def test_transformer_engine_path_reports_unavailable_without_failing(monkeypatch): real_import_module = importlib.import_module @@ -220,6 +284,59 @@ def fake_import_module(name, package=None): assert report.unavailable == ("transformer_engine_paged_kv: test TE unavailable",) +def test_transformer_engine_path_reports_missing_helpers(monkeypatch): + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lambda softmax_lse, per_step: None, + ), + ) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + assert len(report.unavailable) == 1 + assert "missing required helpers" in report.unavailable[0] + + +def test_transformer_engine_path_reports_incompatible_helper_signature(monkeypatch): + def lse_correction(wrong_name, softmax_lse_per_step): + wrong_name.copy_(torch.logaddexp(wrong_name, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=3, + kv_page_size=2, + include_transformer_engine=True, + ) + + assert len(report.unavailable) == 1 + assert "incompatible signature" in report.unavailable[0] + + def test_operator_comparison_specs_register_attention(): args = argparse.Namespace( op="attention", From f5dd5cabfba9c997e5ae26f5c609c4a8fafeb566 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:26:07 +0800 Subject: [PATCH 4/5] feat(attention): harden single-gpu split-kv harness --- rl_engine/kernels/attention_contract.py | 1472 +++++++++++++++++++++ rl_engine/testing/__init__.py | 12 + rl_engine/testing/attention_comparison.py | 652 ++++++++- tests/test_attention_comparison.py | 444 +++++++ 4 files changed, 2578 insertions(+), 2 deletions(-) create mode 100644 rl_engine/kernels/attention_contract.py diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..eb4994b3 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1472 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class SplitKVMode(str, Enum): + DISABLED = "disabled" + FIXED = "fixed" + AUTO = "auto" + + +class RoPEState(str, Enum): + PRE_ROPE = "pre_rope" + POST_ROPE = "post_rope" + + +class RoPECastPoint(str, Enum): + NONE = "none" + BEFORE_ROPE = "before_rope" + AFTER_ROPE = "after_rope" + + +class RoPEFusionBoundary(str, Enum): + UNFUSED_ROPE_ATTENTION = "unfused_rope_attention" + FUSED_ROPE_ATTENTION = "fused_rope_attention" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise AttentionContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _integer_tuple(values: Iterable[int], field: str) -> tuple[int, ...]: + try: + result = tuple(values) + except TypeError as exc: + raise AttentionContractError(f"{field} must be an iterable of integers") from exc + for index, value in enumerate(result): + if isinstance(value, bool) or not isinstance(value, int): + raise AttentionContractError(f"{field}[{index}] must be an integer; got {value!r}") + return result + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical TP/CP ownership for one attention invocation. + + TP head shards are currently required to be equal and contiguous. CP + sequence ownership may be uneven, but every local block must carry a stable + logical global index so a later implementation can merge by logical order + instead of collective arrival order. + """ + + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + global_q_heads: int + global_kv_heads: int + local_q_head_start: int + local_q_heads: int + local_kv_head_start: int + local_kv_heads: int + global_sequence_length: int + local_sequence_length: int + global_block_indices: tuple[int, ...] + global_block_token_starts: tuple[int, ...] + local_block_offsets: tuple[int, ...] + packed_sequence_offsets: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if tp_rank >= tp_world_size: + raise AttentionContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + if cp_rank >= cp_world_size: + raise AttentionContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") + global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") + if global_q_heads % global_kv_heads != 0: + raise AttentionContractError( + f"global_q_heads={global_q_heads} must be divisible by " + f"global_kv_heads={global_kv_heads} for GQA" + ) + if global_q_heads % tp_world_size != 0 or global_kv_heads % tp_world_size != 0: + raise AttentionContractError( + "global Q/KV heads must be evenly divisible by tp_world_size; " + f"got Hq={global_q_heads}, Hkv={global_kv_heads}, TP={tp_world_size}" + ) + + local_q_heads = _positive_int(self.local_q_heads, "local_q_heads") + local_kv_heads = _positive_int(self.local_kv_heads, "local_kv_heads") + expected_q_heads = global_q_heads // tp_world_size + expected_kv_heads = global_kv_heads // tp_world_size + if local_q_heads != expected_q_heads or local_kv_heads != expected_kv_heads: + raise AttentionContractError( + "local TP head counts do not preserve the global Q/KV mapping; " + f"expected ({expected_q_heads}, {expected_kv_heads}), got " + f"({local_q_heads}, {local_kv_heads})" + ) + + local_q_head_start = _non_negative_int(self.local_q_head_start, "local_q_head_start") + local_kv_head_start = _non_negative_int(self.local_kv_head_start, "local_kv_head_start") + expected_q_start = tp_rank * expected_q_heads + expected_kv_start = tp_rank * expected_kv_heads + if local_q_head_start != expected_q_start or local_kv_head_start != expected_kv_start: + raise AttentionContractError( + "local TP head starts do not match contiguous rank ownership; " + f"expected ({expected_q_start}, {expected_kv_start}), got " + f"({local_q_head_start}, {local_kv_head_start})" + ) + + global_sequence_length = _positive_int( + self.global_sequence_length, "global_sequence_length" + ) + local_sequence_length = _positive_int(self.local_sequence_length, "local_sequence_length") + + block_indices = _integer_tuple(self.global_block_indices, "global_block_indices") + if not block_indices: + raise AttentionContractError("global_block_indices must not be empty") + if any(index < 0 for index in block_indices): + raise AttentionContractError("global_block_indices must be non-negative") + if any( + left >= right for left, right in zip(block_indices, block_indices[1:], strict=False) + ): + raise AttentionContractError( + "global_block_indices must be unique and strictly increasing" + ) + object.__setattr__(self, "global_block_indices", block_indices) + + block_token_starts = _integer_tuple( + self.global_block_token_starts, "global_block_token_starts" + ) + local_block_offsets = _integer_tuple(self.local_block_offsets, "local_block_offsets") + if len(block_token_starts) != len(block_indices): + raise AttentionContractError( + "global_block_token_starts must contain one entry per global_block_indices entry" + ) + if any(start < 0 for start in block_token_starts): + raise AttentionContractError("global_block_token_starts must be non-negative") + if len(local_block_offsets) != len(block_indices) + 1: + raise AttentionContractError( + "local_block_offsets must contain one boundary more than global_block_indices" + ) + if local_block_offsets[0] != 0 or local_block_offsets[-1] != local_sequence_length: + raise AttentionContractError( + "local_block_offsets must start at 0 and end at local_sequence_length" + ) + if any( + left >= right + for left, right in zip(local_block_offsets, local_block_offsets[1:], strict=False) + ): + raise AttentionContractError("local_block_offsets must be strictly increasing") + + previous_global_end = 0 + for index, (global_start, local_start, local_end) in enumerate( + zip( + block_token_starts, + local_block_offsets[:-1], + local_block_offsets[1:], + strict=True, + ) + ): + global_end = global_start + (local_end - local_start) + if global_end > global_sequence_length: + raise AttentionContractError( + f"global block {block_indices[index]} exceeds global_sequence_length" + ) + if index > 0 and global_start < previous_global_end: + raise AttentionContractError( + "global block token ranges must be non-overlapping and ordered" + ) + previous_global_end = global_end + object.__setattr__(self, "global_block_token_starts", block_token_starts) + object.__setattr__(self, "local_block_offsets", local_block_offsets) + + if self.packed_sequence_offsets is not None: + offsets = _integer_tuple(self.packed_sequence_offsets, "packed_sequence_offsets") + if len(offsets) < 2 or offsets[0] != 0: + raise AttentionContractError( + "packed_sequence_offsets must start at 0 and contain an end offset" + ) + if any(left >= right for left, right in zip(offsets, offsets[1:], strict=False)): + raise AttentionContractError("packed_sequence_offsets must be strictly increasing") + if offsets[-1] != local_sequence_length: + raise AttentionContractError( + "the final packed_sequence_offsets value must equal local_sequence_length; " + f"got {offsets[-1]} and {local_sequence_length}" + ) + object.__setattr__(self, "packed_sequence_offsets", offsets) + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic CP ``(out, lse)`` merge semantics.""" + + merge: AttentionMerge = AttentionMerge.ONLINE_SOFTMAX_LSE + acc_dtype: AttentionDType = AttentionDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(AttentionMerge, self.merge, "merge")) + object.__setattr__( + self, "acc_dtype", _enum_value(AttentionDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError( + f"CP attention accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class SplitKVExecutionPlan: + """Actual backend-local Split-KV schedule emitted by a runtime. + + CP ownership is intentionally separate from this plan. ``boundaries`` + describe the canonical logical KV ranges reduced by one backend invocation; + CP may transport those partial states between ranks but must not change the + FP32 merge contract recorded here. + """ + + requested_mode: SplitKVMode + requested_split_size: int | None + actual_mode: SplitKVMode | None + actual_split_size: int | None + boundaries: tuple[tuple[int, int], ...] + merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + acc_dtype: AttentionDType = AttentionDType.FP32 + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + backend: str = "reference" + source: str = "contract_exact" + fallback: bool = False + fallback_reason: str | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "requested_mode", + _enum_value(SplitKVMode, self.requested_mode, "requested_mode"), + ) + if self.actual_mode is not None: + object.__setattr__( + self, + "actual_mode", + _enum_value(SplitKVMode, self.actual_mode, "actual_mode"), + ) + for field_name in ("requested_split_size", "actual_split_size"): + value = getattr(self, field_name) + if value is not None: + _positive_int(value, field_name) + if self.requested_mode is SplitKVMode.FIXED: + if self.requested_split_size is None: + raise AttentionContractError( + "requested fixed Split-KV mode requires requested_split_size" + ) + elif self.requested_split_size is not None: + raise AttentionContractError( + "requested_split_size is only valid for requested fixed Split-KV mode" + ) + try: + boundaries = tuple(tuple(boundary) for boundary in self.boundaries) + except TypeError as exc: + raise AttentionContractError( + "Split-KV boundaries must be an iterable of (start, end) pairs" + ) from exc + previous_end = 0 + for index, boundary in enumerate(boundaries): + if len(boundary) != 2: + raise AttentionContractError( + f"Split-KV boundary {index} must contain exactly start and end" + ) + start, end = boundary + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError( + "Split-KV boundaries must satisfy 0 <= start < end" + ) + if index > 0 and start != previous_end: + raise AttentionContractError( + "Split-KV boundaries must be contiguous and in logical KV order" + ) + previous_end = end + if self.actual_mode is None and boundaries: + raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries") + if self.actual_mode is not None and not boundaries: + raise AttentionContractError("actual Split-KV plan must declare logical boundaries") + if self.actual_mode is SplitKVMode.FIXED: + if self.actual_split_size is None: + raise AttentionContractError( + "actual fixed Split-KV mode requires actual_split_size" + ) + widths = tuple(end - start for start, end in boundaries) + if any(width != self.actual_split_size for width in widths[:-1]) or ( + widths and widths[-1] > self.actual_split_size + ): + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_size except " + "for a shorter final split" + ) + elif self.actual_split_size is not None: + raise AttentionContractError( + "actual_split_size is only valid for actual fixed Split-KV mode" + ) + if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1: + raise AttentionContractError( + "disabled Split-KV execution must contain exactly one boundary" + ) + object.__setattr__(self, "boundaries", boundaries) + object.__setattr__( + self, + "merge_order", + _enum_value(ReductionOrder, self.merge_order, "merge_order"), + ) + object.__setattr__( + self, + "acc_dtype", + _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"), + ) + object.__setattr__( + self, + "downcast_at", + _enum_value(DowncastPoint, self.downcast_at, "downcast_at"), + ) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Split-KV partial states must be merged in fp32") + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Split-KV backend must be a non-empty string") + if not isinstance(self.source, str) or not self.source.strip(): + raise AttentionContractError("Split-KV plan source must be a non-empty string") + if not isinstance(self.fallback, bool): + raise AttentionContractError("Split-KV fallback must be a bool") + if self.fallback and not self.fallback_reason: + raise AttentionContractError("Split-KV fallback_reason is required for a fallback") + if not self.fallback and self.fallback_reason is not None: + raise AttentionContractError( + "Split-KV fallback_reason must be None when fallback=False" + ) + if not self.fallback and self.actual_mode is not None: + if self.actual_mode is not self.requested_mode: + raise AttentionContractError( + "actual Split-KV mode may differ from requested mode only for a fallback" + ) + if self.actual_split_size != self.requested_split_size: + raise AttentionContractError( + "actual Split-KV size may differ from requested size only for a fallback" + ) + + @property + def actual_split_count(self) -> int | None: + return None if self.actual_mode is None else len(self.boundaries) + + def to_dict(self) -> dict[str, Any]: + return { + "requested_split_kv_policy": self.requested_mode.value, + "requested_split_kv_size": self.requested_split_size, + "actual_split_kv_policy": ( + None if self.actual_mode is None else self.actual_mode.value + ), + "actual_split_kv_size": self.actual_split_size, + "actual_split_kv_count": self.actual_split_count, + "actual_split_boundaries": [list(boundary) for boundary in self.boundaries], + "split_kv_merge_order": self.merge_order.value, + "split_kv_accum_dtype": self.acc_dtype.value, + "split_kv_downcast_at": self.downcast_at.value, + "split_kv_backend": self.backend, + "split_kv_plan_source": self.source, + "split_kv_fallback": self.fallback, + "split_kv_fallback_reason": self.fallback_reason, + } + + +@dataclass(frozen=True) +class SplitKVSpec: + """Requested Split-KV policy shared by training and rollout paths.""" + + mode: SplitKVMode = SplitKVMode.DISABLED + fixed_split_size: int | None = None + strict_consistency: bool = True + + @classmethod + def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency) + + @classmethod + def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls( + mode=SplitKVMode.FIXED, + fixed_split_size=split_size, + strict_consistency=strict_consistency, + ) + + @classmethod + def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec": + return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency) + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode")) + if not isinstance(self.strict_consistency, bool): + raise AttentionContractError("split_kv.strict_consistency must be a bool") + if self.mode is SplitKVMode.FIXED: + if self.fixed_split_size is None: + raise AttentionContractError( + "fixed Split-KV policy requires fixed_split_size" + ) + _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") + elif self.fixed_split_size is not None: + raise AttentionContractError( + "fixed_split_size is only valid for fixed Split-KV policy" + ) + if self.strict_consistency and self.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" + ) + + def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan: + """Resolve policies whose logical schedule is fully known by contract.""" + + total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens") + if self.mode is SplitKVMode.AUTO: + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=None, + actual_mode=None, + actual_split_size=None, + boundaries=(), + backend=backend, + source="runtime_required", + ) + split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size + assert split_size is not None + boundaries = tuple( + (start, min(start + split_size, total_kv_tokens)) + for start in range(0, total_kv_tokens, split_size) + ) + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=self.fixed_split_size, + actual_mode=self.mode, + actual_split_size=self.fixed_split_size, + boundaries=boundaries, + backend=backend, + source="contract_exact", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "fixed_split_size": self.fixed_split_size, + "strict_consistency": self.strict_consistency, + } + + +def validate_split_kv_alignment( + training: SplitKVExecutionPlan, + rollout: SplitKVExecutionPlan, +) -> None: + """Fail closed unless train and rollout executed the same logical plan.""" + + if training.actual_mode is None or rollout.actual_mode is None: + raise AttentionContractError( + "strict Split-KV alignment requires actual runtime plans from both sides" + ) + fields = ( + "requested_mode", + "requested_split_size", + "actual_mode", + "actual_split_size", + "boundaries", + "merge_order", + "acc_dtype", + "downcast_at", + "fallback", + ) + mismatches = [ + field_name + for field_name in fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if mismatches: + raise AttentionContractError( + "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches) + ) + + +@dataclass(frozen=True, order=True) +class SplitKVRuntimeCoordinate: + """Identity of one batch/rank/owner Split-KV runtime plan.""" + + batch_index: int + tp_rank: int + cp_rank: int + owner_cp_rank: int + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + ) -> None: + batch_index = _non_negative_int(self.batch_index, "batch_index") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank") + if batch_index >= batch_size: + raise AttentionContractError("Split-KV batch_index is outside batch_size") + if tp_rank >= tp_world_size: + raise AttentionContractError("Split-KV tp_rank is outside tp_world_size") + if cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV cp_rank is outside cp_world_size") + if owner_cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size") + + def to_dict(self) -> dict[str, int]: + return { + "batch_index": self.batch_index, + "tp_rank": self.tp_rank, + "cp_rank": self.cp_rank, + "owner_cp_rank": self.owner_cp_rank, + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanEntry: + """Actual plan for one batch/TP/CP consumer and logical KV owner.""" + + coordinate: SplitKVRuntimeCoordinate + expected_kv_range: tuple[int, int] + execution: SplitKVExecutionPlan + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + total_kv_tokens: int, + ) -> None: + self.coordinate.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + try: + start, end = self.expected_kv_range + except (TypeError, ValueError) as exc: + raise AttentionContractError( + "expected_kv_range must contain exactly (start, end)" + ) from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + or end > total_kv_tokens + ): + raise AttentionContractError( + "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" + ) + if self.execution.actual_mode is None: + raise AttentionContractError( + "complete Split-KV plan sets require actual runtime plans" + ) + if ( + self.execution.boundaries[0][0] != start + or self.execution.boundaries[-1][1] != end + ): + raise AttentionContractError( + "Split-KV execution boundaries must exactly cover expected_kv_range" + ) + if any( + boundary_start < start or boundary_end > end + for boundary_start, boundary_end in self.execution.boundaries + ): + raise AttentionContractError( + "Split-KV execution boundary escapes expected_kv_range" + ) + + def to_dict(self) -> dict[str, Any]: + return { + **self.coordinate.to_dict(), + "expected_kv_range": list(self.expected_kv_range), + **self.execution.to_dict(), + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanSet: + """Complete actual Split-KV plans across batch, TP, CP, and KV owners. + + Every CP consumer must report the plan used for every CP-owned KV range. + This duplicates owner plans across consumers intentionally: it detects one + rank silently choosing a different Split-K schedule or merge policy. + """ + + batch_size: int + tp_world_size: int + cp_world_size: int + total_kv_tokens: tuple[int, ...] + entries: tuple[SplitKVRuntimePlanEntry, ...] + + def __post_init__(self) -> None: + batch_size = _positive_int(self.batch_size, "batch_size") + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens") + if len(totals) != batch_size or any(total <= 0 for total in totals): + raise AttentionContractError( + "total_kv_tokens must contain one positive length per batch item" + ) + object.__setattr__(self, "total_kv_tokens", totals) + entries = tuple(self.entries) + object.__setattr__(self, "entries", entries) + expected_coordinates = { + SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(batch_size) + for tp_rank in range(tp_world_size) + for cp_rank in range(cp_world_size) + for owner_cp_rank in range(cp_world_size) + } + actual_coordinates = [entry.coordinate for entry in entries] + if len(set(actual_coordinates)) != len(actual_coordinates): + raise AttentionContractError( + "Split-KV runtime plan set contains duplicate coordinates" + ) + missing = expected_coordinates.difference(actual_coordinates) + extra = set(actual_coordinates).difference(expected_coordinates) + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan set coordinate coverage is incomplete; " + f"missing={_format_split_kv_coordinates(missing)}, " + f"extra={_format_split_kv_coordinates(extra)}" + ) + for entry in entries: + entry.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals[entry.coordinate.batch_index], + ) + self._validate_owner_coverage() + self._validate_rank_invariance() + + def _validate_owner_coverage(self) -> None: + for batch_index, total in enumerate(self.total_kv_tokens): + for tp_rank in range(self.tp_world_size): + ranges = [] + for owner_cp_rank in range(self.cp_world_size): + matches = [ + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.tp_rank == tp_rank + and entry.coordinate.cp_rank == 0 + and entry.coordinate.owner_cp_rank == owner_cp_rank + ] + ranges.append(matches[0].expected_kv_range) + previous_end = 0 + for start, end in ranges: + if start != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = end + if previous_end != total: + raise AttentionContractError( + "Split-KV owner ranges do not cover total_kv_tokens" + ) + + def _validate_rank_invariance(self) -> None: + for batch_index in range(self.batch_size): + for owner_cp_rank in range(self.cp_world_size): + entries = sorted( + ( + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.owner_cp_rank == owner_cp_rank + ), + key=lambda entry: ( + entry.coordinate.tp_rank, + entry.coordinate.cp_rank, + ), + ) + reference = entries[0] + for entry in entries[1:]: + if entry.expected_kv_range != reference.expected_kv_range: + raise AttentionContractError( + "Split-KV owner range differs across TP/CP consumers" + ) + try: + validate_split_kv_alignment( + reference.execution, + entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return { + "batch_size": self.batch_size, + "tp_world_size": self.tp_world_size, + "cp_world_size": self.cp_world_size, + "total_kv_tokens": list(self.total_kv_tokens), + "entries": [ + entry.to_dict() + for entry in sorted(self.entries, key=lambda entry: entry.coordinate) + ], + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + } + + +def validate_split_kv_plan_set_alignment( + training: SplitKVRuntimePlanSet, + rollout: SplitKVRuntimePlanSet, +) -> None: + """Fail closed unless complete train/rollout runtime plan sets align.""" + + topology_fields = ( + "batch_size", + "tp_world_size", + "cp_world_size", + "total_kv_tokens", + ) + topology_mismatches = [ + field_name + for field_name in topology_fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if topology_mismatches: + raise AttentionContractError( + "training/rollout Split-KV plan-set topology differs: " + + ", ".join(topology_mismatches) + ) + training_by_coordinate = { + entry.coordinate: entry for entry in training.entries + } + rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} + if training_by_coordinate.keys() != rollout_by_coordinate.keys(): + raise AttentionContractError( + "training/rollout Split-KV plan-set coordinates differ" + ) + for coordinate in sorted(training_by_coordinate): + train_entry = training_by_coordinate[coordinate] + rollout_entry = rollout_by_coordinate[coordinate] + if train_entry.expected_kv_range != rollout_entry.expected_kv_range: + raise AttentionContractError( + f"training/rollout expected KV range differs at {coordinate}" + ) + try: + validate_split_kv_alignment( + train_entry.execution, + rollout_entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + f"training/rollout Split-KV plan differs at {coordinate}: {exc}" + ) from exc + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, field), bool): + raise AttentionContractError(f"{field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append( + f"Split-KV policy={contract.split_kv.mode.value} is unsupported" + ) + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 792b36d0..3897e009 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -8,11 +8,17 @@ AttentionComparisonReport, AttentionPathDrift, AttentionPathResult, + DecodeAttentionInputs, + DecodeKVCacheMetadata, DriftStats, TransformerEngineUnavailable, + compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + decode_prefix_cache_fingerprint, run_chunked_query_attention, + run_decode_full_prefill_reference, + run_decode_kv_replay, run_full_attention, run_fused_like_rope_attention, run_paged_kv_attention, @@ -35,18 +41,24 @@ "AttentionComparisonReport", "AttentionPathDrift", "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", "DriftStats", "SyntheticRLKernelBatch", "TransformerEngineUnavailable", "active_token_count", "compare_single_gpu_rope_attention", "compare_single_gpu_attention", + "compare_decode_kv_replay", + "decode_prefix_cache_fingerprint", "compute_policy_ratio", "compute_reference_kl", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", "run_chunked_query_attention", + "run_decode_full_prefill_reference", + "run_decode_kv_replay", "run_fused_like_rope_attention", "run_full_attention", "run_paged_kv_attention", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index dbee69af..41d5abfc 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -11,6 +11,7 @@ from __future__ import annotations +import hashlib import importlib import importlib.metadata as importlib_metadata import inspect @@ -20,10 +21,16 @@ import torch +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVSpec, +) from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.testing.reference_ops import selected_logprobs_reference MergeBackend = Literal["rl_kernel", "transformer_engine"] +RoPEState = Literal["pre_rope", "post_rope"] _TE_CONTEXT_PARALLEL_MODULE = ( "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" @@ -71,6 +78,56 @@ class AttentionComparisonInputs: rope_output_dtype: torch.dtype | None = None +@dataclass(frozen=True) +class DecodeKVCacheMetadata: + """Logical identity and physical layout for decode-stage cached KV. + + ``block_table`` maps logical KV blocks to physical cache pages. Positions + are stored per physical cache slot; unused slots must contain ``-1``. + Keeping both mappings explicit lets the harness distinguish layout changes + from changes to the logical token sequence. + """ + + cache_position: torch.Tensor + kv_seq_lens: torch.Tensor + block_table: torch.Tensor + global_token_positions: torch.Tensor + query_position_ids: torch.Tensor + key_position_ids: torch.Tensor + page_size: int + prefix_cache_key: str | None = None + prefix_cache_enabled: bool = False + prefix_length: int = 0 + prefix_cache_fingerprint: str | None = None + q_rope_state: RoPEState = "post_rope" + k_cache_rope_state: RoPEState = "post_rope" + cp_block_owners: torch.Tensor | None = None + cp_world_size: int = 1 + + +@dataclass(frozen=True) +class DecodeAttentionInputs: + """Decode queries and physically paged KV cache used by the PR6 harness.""" + + q: torch.Tensor + k_cache: torch.Tensor + v_cache: torch.Tensor + metadata: DecodeKVCacheMetadata + scale: float | None = None + output_dtype: torch.dtype = torch.float32 + rope_theta: float = 1_000_000.0 + rope_rotary_dim: int | None = None + rope_cast_at: str = "after_rope" + q_rope_output_dtype: torch.dtype | None = None + k_cache_rope_output_dtype: torch.dtype | None = None + lm_head_weight: torch.Tensor | None = None + target_ids: torch.Tensor | None = None + active_token_mask: torch.Tensor | None = None + k_new: torch.Tensor | None = None + v_new: torch.Tensor | None = None + split_kv: SplitKVSpec | None = None + + @dataclass(frozen=True) class AttentionPathResult: """One materialized attention path result.""" @@ -210,6 +267,218 @@ def compare_single_gpu_rope_attention( return AttentionComparisonReport(reference_name=reference.name, drifts=drifts) +def compare_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + include_transformer_engine: bool = False, +) -> AttentionComparisonReport: + """Compare paged decode replay with a logical full-KV teacher-forcing view.""" + + _validate_decode_inputs(inputs) + reference = _run_decode_full_prefill_reference(inputs) + candidates = [_run_decode_kv_replay(inputs, merge_backend="rl_kernel")] + unavailable: list[str] = [] + if include_transformer_engine: + try: + candidates.append(_run_decode_kv_replay(inputs, merge_backend="transformer_engine")) + except TransformerEngineUnavailable as exc: + unavailable.append(f"transformer_engine_decode_kv_replay: {exc}") + drifts = tuple(_compare_decode_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport( + reference_name=reference.name, + drifts=drifts, + unavailable=tuple(unavailable), + ) + + +def run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: + """Materialize the full logical KV sequence for each decode query. + + This is the teacher-forcing side of the PR6 comparison. It deliberately + ignores physical page boundaries after restoring logical token order. + """ + + _validate_decode_inputs(inputs) + return _run_decode_full_prefill_reference(inputs) + + +def _run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + visible = logical_positions <= query_position + out, lse = _attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, visible, :], + v[:, :, visible, :], + causal=False, + scale=inputs.scale, + key_padding_mask=None, + q_start=0, + k_start=0, + total_query_len=1, + total_kv_len=int(visible.sum().item()), + output_dtype=inputs.output_dtype, + ) + batch_out.append(out) + batch_lse.append(lse) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + return AttentionPathResult( + name="full_prefill_decode_reference", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance={ + "attention_mode": "decode", + "materialization": "full_logical_kv", + "lse_domain": "attention", + "accum_dtype": "fp32", + }, + ) + + +def run_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend = "rl_kernel", +) -> AttentionPathResult: + """Replay decode over physical KV pages and merge by logical block index.""" + + _validate_decode_inputs(inputs) + return _run_decode_kv_replay(inputs, merge_backend=merge_backend) + + +def _run_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend, +) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + merge_orders: list[list[list[int]]] = [] + actual_split_plans: list[list[dict[str, Any]]] = [] + cp_block_owners: list[list[int]] = [] + split_kv = _resolved_decode_split_kv(inputs) + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + owners = _logical_block_owners(inputs, batch_index) + cp_block_owners.append(owners) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + batch_orders: list[list[int]] = [] + batch_split_plans: list[dict[str, Any]] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + states: list[_PartialAttentionState] = [] + order: list[int] = [] + visible_count = int((logical_positions <= query_position).sum().item()) + split_bounds = _decode_split_bounds(visible_count, split_kv) + for block_index, (block_start, block_end) in enumerate( + split_bounds + ): + block_positions = logical_positions[block_start:block_end] + visible = block_positions <= query_position + if not bool(visible.any()): + continue + visible_end = block_start + int(visible.sum().item()) + out, lse = _attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, block_start:visible_end, :], + v[:, :, block_start:visible_end, :], + causal=False, + scale=inputs.scale, + key_padding_mask=None, + q_start=0, + k_start=block_start, + total_query_len=1, + total_kv_len=visible_end, + output_dtype=torch.float32, + ) + states.append( + _PartialAttentionState( + out=out, + lse=lse, + block_start=block_index, + block_end=block_index + 1, + ) + ) + order.append(block_index) + if not states: + raise ValueError("each decode query must have at least one visible cached KV token") + out, lse = _merge_partial_states(states, backend=merge_backend) + batch_out.append(out.to(inputs.output_dtype)) + batch_lse.append(lse) + batch_orders.append(order) + plan = SplitKVExecutionPlan( + requested_mode=split_kv.mode, + requested_split_size=split_kv.fixed_split_size, + actual_mode=split_kv.mode, + actual_split_size=split_kv.fixed_split_size, + boundaries=tuple(split_bounds), + backend=f"{merge_backend}_decode_kv_replay", + source="reference_execution", + ) + batch_split_plans.append(plan.to_dict()) + merge_orders.append(batch_orders) + actual_split_plans.append(batch_split_plans) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + + provenance: dict[str, Any] = { + "attention_mode": "decode", + "decode_semantics": ( + "past_kv_plus_new_kv_append" if inputs.k_new is not None else "cache_replay" + ), + "past_kv_lengths": inputs.metadata.kv_seq_lens.tolist(), + "new_kv_length": (0 if inputs.k_new is None else inputs.k_new.size(2)), + "materialization": "paged_kv_replay", + "sq": inputs.q.size(2), + "page_size": inputs.metadata.page_size, + "cache_position": inputs.metadata.cache_position.tolist(), + "kv_seq_lens": inputs.metadata.kv_seq_lens.tolist(), + "block_table": inputs.metadata.block_table.tolist(), + "global_token_positions": inputs.metadata.global_token_positions.tolist(), + "query_position_ids": inputs.metadata.query_position_ids.tolist(), + "key_position_ids": inputs.metadata.key_position_ids.tolist(), + "prefix_cache_enabled": inputs.metadata.prefix_cache_enabled, + "prefix_cache_key": inputs.metadata.prefix_cache_key, + "prefix_length": inputs.metadata.prefix_length, + "prefix_cache_fingerprint": inputs.metadata.prefix_cache_fingerprint, + "q_rope_state": inputs.metadata.q_rope_state, + "k_cache_rope_state": inputs.metadata.k_cache_rope_state, + "rope_theta": float(inputs.rope_theta), + "rotary_dim": _decode_rope_rotary_dim(inputs), + "rope_cast_at": inputs.rope_cast_at, + "q_rope_output_dtype": str(_decode_q_rope_output_dtype(inputs)).replace("torch.", ""), + "k_cache_rope_output_dtype": str(_decode_k_rope_output_dtype(inputs)).replace("torch.", ""), + "cp_block_owners": cp_block_owners, + "cp_world_size": inputs.metadata.cp_world_size, + "requested_split_kv_policy": split_kv.mode.value, + "requested_split_kv_size": split_kv.fixed_split_size, + "actual_split_kv_plans": actual_split_plans, + "merge_order": "global_block_index", + "logical_merge_orders": merge_orders, + "merge_backend": merge_backend, + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + } + if merge_backend == "transformer_engine": + provenance.update(_te_context_parallel_provenance()) + return AttentionPathResult( + name=f"{merge_backend}_decode_kv_replay", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance=provenance, + ) + + def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: """Training-style full-sequence attention with exported attention-domain LSE.""" @@ -424,6 +693,39 @@ def transformer_engine_context_parallel_available() -> bool: return True +def decode_prefix_cache_fingerprint( + inputs: DecodeAttentionInputs, + *, + prefix_length: int, +) -> str: + """Fingerprint logical prefix positions and cached K/V content. + + The fingerprint is invariant to physical page placement because cache slots + are first restored to logical token order. It intentionally includes the + cached-K RoPE state and tensor dtypes so it identifies the actual replay + boundary rather than only the token positions. + """ + + prefix_length = _positive_int(prefix_length, "prefix_length") + if bool((inputs.metadata.kv_seq_lens < prefix_length).any()): + raise ValueError("prefix_length must not exceed any kv_seq_lens entry") + digest = hashlib.sha256() + digest.update(f"k_rope_state={inputs.metadata.k_cache_rope_state}\n".encode()) + digest.update(f"k_dtype={inputs.k_cache.dtype};v_dtype={inputs.v_cache.dtype}\n".encode()) + for batch_index in range(inputs.q.size(0)): + slots = _decode_logical_slot_index(inputs, batch_index)[:prefix_length] + for tensor in ( + inputs.metadata.global_token_positions[batch_index, slots], + inputs.metadata.key_position_ids[batch_index, slots], + inputs.k_cache[batch_index, :, slots, :], + inputs.v_cache[batch_index, :, slots, :], + ): + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + def _compare_path( candidate: AttentionPathResult, reference: AttentionPathResult, @@ -454,6 +756,25 @@ def _compare_path( ) +def _compare_decode_path( + candidate: AttentionPathResult, + reference: AttentionPathResult, + inputs: DecodeAttentionInputs, +) -> AttentionPathDrift: + dlogp = None + if inputs.lm_head_weight is not None and inputs.target_ids is not None: + candidate_logp = _selected_logps_from_decode_attention(candidate.out, inputs) + reference_logp = _selected_logps_from_decode_attention(reference.out, inputs) + dlogp = _drift_stats(candidate_logp, reference_logp, mask=inputs.active_token_mask) + return AttentionPathDrift( + candidate_name=candidate.name, + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=dlogp, + provenance=candidate.provenance, + ) + + def _apply_rope_to_qk(inputs: AttentionComparisonInputs) -> tuple[torch.Tensor, torch.Tensor]: _validate_rope_inputs(inputs) assert inputs.rope_positions is not None @@ -464,6 +785,97 @@ def _apply_rope_to_qk(inputs: AttentionComparisonInputs) -> tuple[torch.Tensor, return q, k +def _decode_logical_qkv( + inputs: DecodeAttentionInputs, + batch_index: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Restore one batch's cache to logical order and materialize RoPE state.""" + + metadata = inputs.metadata + slot_index = _decode_logical_slot_index(inputs, batch_index) + logical_position_tensor = metadata.global_token_positions[batch_index, slot_index].long() + k = inputs.k_cache[batch_index : batch_index + 1, :, slot_index, :] + v = inputs.v_cache[batch_index : batch_index + 1, :, slot_index, :] + q = inputs.q[batch_index : batch_index + 1] + + rope = NativeRoPEOp() + if metadata.q_rope_state == "pre_rope": + q = rope.forward_fp32( + q, + metadata.query_position_ids[batch_index : batch_index + 1], + theta=inputs.rope_theta, + ).to(_decode_q_rope_output_dtype(inputs)) + if metadata.k_cache_rope_state == "pre_rope": + key_positions = metadata.key_position_ids[batch_index, slot_index].unsqueeze(0) + k = rope.forward_fp32(k, key_positions, theta=inputs.rope_theta).to( + _decode_k_rope_output_dtype(inputs) + ) + if inputs.k_new is not None: + assert inputs.v_new is not None + k_new = inputs.k_new[batch_index : batch_index + 1] + if metadata.k_cache_rope_state == "pre_rope": + k_new = rope.forward_fp32( + k_new, + metadata.query_position_ids[batch_index : batch_index + 1], + theta=inputs.rope_theta, + ).to(_decode_k_rope_output_dtype(inputs)) + k = torch.cat((k, k_new), dim=2) + v = torch.cat((v, inputs.v_new[batch_index : batch_index + 1]), dim=2) + logical_position_tensor = torch.cat( + ( + logical_position_tensor, + metadata.query_position_ids[batch_index].long(), + ) + ) + return q, k, v, logical_position_tensor + + +def _decode_logical_slot_index( + inputs: DecodeAttentionInputs, + batch_index: int, +) -> torch.Tensor: + metadata = inputs.metadata + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + logical_block_count = math.ceil(sequence_length / metadata.page_size) + logical_index = torch.arange( + sequence_length, + device=inputs.k_cache.device, + dtype=torch.long, + ) + pages = metadata.block_table[batch_index, :logical_block_count].long() + return ( + pages[logical_index // metadata.page_size] * metadata.page_size + + logical_index % metadata.page_size + ) + + +def _logical_block_owners(inputs: DecodeAttentionInputs, batch_index: int) -> list[int]: + block_count = math.ceil( + int(inputs.metadata.kv_seq_lens[batch_index].item()) / inputs.metadata.page_size + ) + if inputs.metadata.cp_block_owners is None: + return [0] * block_count + return [ + int(owner) for owner in inputs.metadata.cp_block_owners[batch_index, :block_count].tolist() + ] + + +def _decode_q_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return inputs.q.dtype if inputs.q_rope_output_dtype is None else inputs.q_rope_output_dtype + + +def _decode_k_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return ( + inputs.k_cache.dtype + if inputs.k_cache_rope_output_dtype is None + else inputs.k_cache_rope_output_dtype + ) + + +def _decode_rope_rotary_dim(inputs: DecodeAttentionInputs) -> int: + return inputs.q.size(-1) if inputs.rope_rotary_dim is None else inputs.rope_rotary_dim + + def _rope_output_dtype(inputs: AttentionComparisonInputs) -> torch.dtype: return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype @@ -743,6 +1155,28 @@ def _selected_logps_from_attention( ) +def _selected_logps_from_decode_attention( + out: torch.Tensor, + inputs: DecodeAttentionInputs, +) -> torch.Tensor: + if inputs.lm_head_weight is None or inputs.target_ids is None: + raise ValueError("lm_head_weight and target_ids are required for decode dlogp drift") + batch, heads, seq, dim = out.shape + hidden = out.transpose(1, 2).reshape(batch, seq, heads * dim) + if inputs.lm_head_weight.shape[1] != hidden.size(-1): + raise ValueError( + "lm_head_weight hidden dimension must equal Hq * D; " + f"got {inputs.lm_head_weight.shape[1]} and {hidden.size(-1)}" + ) + logits = torch.matmul(hidden.float(), inputs.lm_head_weight.float().transpose(0, 1)) + return selected_logprobs_reference( + logits, + inputs.target_ids, + mask=inputs.active_token_mask, + output_dtype=torch.float32, + ) + + def _drift_stats( candidate: torch.Tensor, reference: torch.Tensor, @@ -849,6 +1283,193 @@ def _validate_rope_inputs(inputs: AttentionComparisonInputs) -> None: raise ValueError("rope_positions must have shape [Sq] or [B, Sq]") +def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: + _validate_qkv(inputs.q, inputs.k_cache, inputs.v_cache) + if inputs.q.device != inputs.k_cache.device or inputs.q.device != inputs.v_cache.device: + raise ValueError("q, k_cache, and v_cache must be on the same device") + metadata = inputs.metadata + batch, _, sq, head_dim = inputs.q.shape + cache_capacity = inputs.k_cache.size(2) + page_size = _positive_int(metadata.page_size, "page_size") + if cache_capacity % page_size != 0: + raise ValueError("physical KV cache capacity must be divisible by page_size") + physical_page_count = cache_capacity // page_size + if metadata.cache_position.shape != (batch, sq): + raise ValueError("cache_position must have shape [B, Sq]") + if metadata.query_position_ids.shape != (batch, sq): + raise ValueError("query_position_ids must have shape [B, Sq]") + if metadata.kv_seq_lens.shape != (batch,): + raise ValueError("kv_seq_lens must have shape [B]") + if metadata.block_table.ndim != 2 or metadata.block_table.size(0) != batch: + raise ValueError("block_table must have shape [B, max_blocks]") + expected_cache_shape = (batch, cache_capacity) + if metadata.global_token_positions.shape != expected_cache_shape: + raise ValueError("global_token_positions must have shape [B, cache_capacity]") + if metadata.key_position_ids.shape != expected_cache_shape: + raise ValueError("key_position_ids must have shape [B, cache_capacity]") + integer_tensors = { + "cache_position": metadata.cache_position, + "query_position_ids": metadata.query_position_ids, + "kv_seq_lens": metadata.kv_seq_lens, + "block_table": metadata.block_table, + "global_token_positions": metadata.global_token_positions, + "key_position_ids": metadata.key_position_ids, + } + if metadata.cp_block_owners is not None: + integer_tensors["cp_block_owners"] = metadata.cp_block_owners + for name, tensor in integer_tensors.items(): + if tensor.device != inputs.q.device: + raise ValueError(f"{name} must be on the same device as q/k/v") + if tensor.dtype not in {torch.int32, torch.int64, torch.long}: + raise ValueError(f"{name} must contain integers") + if metadata.cp_block_owners is not None: + if metadata.cp_block_owners.shape != metadata.block_table.shape: + raise ValueError("cp_block_owners must have the same shape as block_table") + if bool((metadata.cp_block_owners < 0).any()): + raise ValueError("cp_block_owners must be non-negative") + cp_world_size = _positive_int(metadata.cp_world_size, "cp_world_size") + if bool((metadata.cp_block_owners >= cp_world_size).any()): + raise ValueError("cp_block_owners must be smaller than cp_world_size") + if not torch.equal(metadata.cache_position, metadata.query_position_ids): + raise ValueError("cache_position and query_position_ids must identify the same positions") + if metadata.q_rope_state not in {"pre_rope", "post_rope"}: + raise ValueError("q_rope_state must be 'pre_rope' or 'post_rope'") + if metadata.k_cache_rope_state not in {"pre_rope", "post_rope"}: + raise ValueError("k_cache_rope_state must be 'pre_rope' or 'post_rope'") + if metadata.prefix_cache_enabled: + if not metadata.prefix_cache_key: + raise ValueError("prefix_cache_key is required when prefix cache is enabled") + _positive_int(metadata.prefix_length, "prefix_length") + if not metadata.prefix_cache_fingerprint: + raise ValueError("prefix_cache_fingerprint is required when prefix cache is enabled") + elif ( + metadata.prefix_cache_key is not None + or metadata.prefix_length != 0 + or metadata.prefix_cache_fingerprint is not None + ): + raise ValueError( + "prefix cache key/fingerprint must be None and prefix_length must be 0 " + "when prefix cache is disabled" + ) + if inputs.rope_cast_at != "after_rope": + raise ValueError("rope_cast_at must be 'after_rope' for the current fp32 RoPE reference") + if inputs.rope_rotary_dim is not None: + if inputs.rope_rotary_dim != head_dim: + raise ValueError("rope_rotary_dim must equal head_dim") + _positive_int(inputs.rope_rotary_dim, "rope_rotary_dim") + if float(inputs.rope_theta) <= 0: + raise ValueError("rope_theta must be a positive number") + if inputs.q_rope_output_dtype is not None and not isinstance( + inputs.q_rope_output_dtype, torch.dtype + ): + raise ValueError("q_rope_output_dtype must be a torch.dtype when provided") + if inputs.k_cache_rope_output_dtype is not None and not isinstance( + inputs.k_cache_rope_output_dtype, torch.dtype + ): + raise ValueError("k_cache_rope_output_dtype must be a torch.dtype when provided") + if (inputs.k_new is None) != (inputs.v_new is None): + raise ValueError("k_new and v_new must be provided together") + append_mode = inputs.k_new is not None + if append_mode: + assert inputs.k_new is not None and inputs.v_new is not None + if inputs.k_new.shape != inputs.v_new.shape: + raise ValueError("k_new and v_new must have matching shapes") + expected_new_shape = (batch, inputs.k_cache.size(1), sq, head_dim) + if inputs.k_new.shape != expected_new_shape: + raise ValueError("k_new and v_new must have shape [B, Hkv, Sq, D]") + if inputs.k_new.device != inputs.q.device or inputs.v_new.device != inputs.q.device: + raise ValueError("k_new and v_new must be on the same device as q") + if inputs.split_kv is not None and not isinstance(inputs.split_kv, SplitKVSpec): + raise ValueError("split_kv must be a SplitKVSpec when provided") + if inputs.split_kv is not None and inputs.split_kv.mode is SplitKVMode.AUTO: + raise ValueError("decode replay requires disabled or fixed Split-KV, not auto") + if ( + metadata.q_rope_state == "post_rope" + and inputs.q_rope_output_dtype is not None + and inputs.q.dtype != inputs.q_rope_output_dtype + ): + raise ValueError("post-RoPE q dtype must match q_rope_output_dtype") + if ( + metadata.k_cache_rope_state == "post_rope" + and inputs.k_cache_rope_output_dtype is not None + and inputs.k_cache.dtype != inputs.k_cache_rope_output_dtype + ): + raise ValueError("post-RoPE k_cache dtype must match k_cache_rope_output_dtype") + if (inputs.lm_head_weight is None) != (inputs.target_ids is None): + raise ValueError("lm_head_weight and target_ids must be provided together") + if inputs.target_ids is not None and inputs.target_ids.shape != (batch, sq): + raise ValueError("target_ids must have shape [B, Sq]") + if inputs.active_token_mask is not None: + if inputs.active_token_mask.shape != (batch, sq): + raise ValueError("active_token_mask must have shape [B, Sq]") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + + for batch_index in range(batch): + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + if sequence_length <= 0 or sequence_length > cache_capacity: + raise ValueError("each kv_seq_lens entry must be in [1, cache_capacity]") + block_count = math.ceil(sequence_length / page_size) + if block_count > metadata.block_table.size(1): + raise ValueError("block_table does not contain enough logical KV blocks") + pages = metadata.block_table[batch_index, :block_count] + if bool(((pages < 0) | (pages >= physical_page_count)).any()): + raise ValueError("block_table contains an out-of-range physical page") + if torch.unique(pages).numel() != block_count: + raise ValueError("active block_table entries must not contain duplicate pages") + slot_index = _decode_logical_slot_index(inputs, batch_index) + active_slot_mask = torch.zeros(cache_capacity, device=inputs.q.device, dtype=torch.bool) + active_slot_mask[slot_index] = True + if bool((metadata.global_token_positions[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused global_token_positions entries must be -1") + if bool((metadata.key_position_ids[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused key_position_ids entries must be -1") + global_positions = metadata.global_token_positions[batch_index, slot_index] + position_offset = int(global_positions[0].item()) + expected_positions = torch.arange( + position_offset, + position_offset + sequence_length, + device=inputs.q.device, + dtype=global_positions.dtype, + ) + if not torch.equal(global_positions, expected_positions): + raise ValueError( + "block_table/global_token_positions must reconstruct logical positions " + "as one contiguous global range" + ) + key_positions = metadata.key_position_ids[batch_index, slot_index] + if not torch.equal(key_positions, global_positions): + raise ValueError("key_position_ids must match cached global token positions") + cache_positions = metadata.cache_position[batch_index] + if append_mode: + expected_new_positions = torch.arange( + position_offset + sequence_length, + position_offset + sequence_length + sq, + device=inputs.q.device, + dtype=cache_positions.dtype, + ) + if not torch.equal(cache_positions, expected_new_positions): + raise ValueError( + "append cache_position must identify the contiguous new-token suffix" + ) + elif bool((cache_positions < position_offset).any()) or bool( + (cache_positions >= position_offset + sequence_length).any() + ): + raise ValueError("cache_position must refer to a token present in the KV cache") + if sq > 1 and bool((cache_positions[1:] <= cache_positions[:-1]).any()): + raise ValueError("few-query cache_position values must be strictly increasing") + + if metadata.prefix_cache_enabled: + actual_fingerprint = decode_prefix_cache_fingerprint( + inputs, + prefix_length=metadata.prefix_length, + ) + if actual_fingerprint != metadata.prefix_cache_fingerprint: + raise ValueError( + "prefix_cache_fingerprint does not match the logical prefix positions/content" + ) + + def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: raise ValueError("q, k, and v must have shape [B, H, S, D]") @@ -861,13 +1482,17 @@ def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: def _validate_partial_states(states: list[_PartialAttentionState]) -> None: + if not states: + raise ValueError("at least one partial attention state is required") first = states[0] + if first.block_start != 0: + raise ValueError("partial state coverage must start at logical KV token 0") previous_end = first.block_end for state in states[1:]: if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: raise ValueError("all partial states must have matching shapes") - if state.block_start < previous_end: - raise ValueError("partial state block ranges must not overlap") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") previous_end = state.block_end @@ -883,6 +1508,23 @@ def _chunk_bounds(length: int, chunk_size: int) -> list[tuple[int, int]]: return bounds +def _resolved_decode_split_kv(inputs: DecodeAttentionInputs) -> SplitKVSpec: + # Existing replay behavior used one partial state per logical cache page. + # Keep that as the explicit default while allowing disabled/fixed sweeps on + # the same physical page layout. + return inputs.split_kv or SplitKVSpec.fixed(inputs.metadata.page_size) + + +def _decode_split_bounds(length: int, split_kv: SplitKVSpec) -> list[tuple[int, int]]: + if split_kv.mode is SplitKVMode.AUTO: + raise ValueError("decode replay cannot materialize an unknown auto Split-KV plan") + chunk_size = length + if split_kv.mode is SplitKVMode.FIXED: + assert split_kv.fixed_split_size is not None + chunk_size = split_kv.fixed_split_size + return _chunk_bounds(length, chunk_size) + + def _positive_int(value: int, name: str) -> int: if isinstance(value, bool) or value <= 0: raise ValueError(f"{name} must be a positive integer") @@ -894,13 +1536,19 @@ def _positive_int(value: int, name: str) -> int: "AttentionComparisonReport", "AttentionPathDrift", "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", "DriftStats", "TransformerEngineUnavailable", "compare_single_gpu_rope_attention", "compare_single_gpu_attention", + "compare_decode_kv_replay", + "decode_prefix_cache_fingerprint", "run_chunked_query_attention", "run_fused_like_rope_attention", "run_full_attention", + "run_decode_full_prefill_reference", + "run_decode_kv_replay", "run_paged_kv_attention", "run_unfused_rope_attention", "transformer_engine_context_parallel_available", diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 061e8edd..5ae5d5e3 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -8,16 +8,25 @@ import json import sys import types +from dataclasses import replace +from typing import Literal import pytest import torch from rl_engine.kernels.gtest import run_operator_suite from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.testing.attention_comparison import ( AttentionComparisonInputs, + DecodeAttentionInputs, + DecodeKVCacheMetadata, + compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + decode_prefix_cache_fingerprint, + run_decode_kv_replay, run_paged_kv_attention, ) @@ -57,6 +66,67 @@ def _comparison_inputs() -> AttentionComparisonInputs: ) +def _decode_inputs( + *, + page_order: tuple[int, ...] = (0, 1, 2), + prefix_cache_enabled: bool = False, + q_rope_state: Literal["pre_rope", "post_rope"] = "post_rope", + k_cache_rope_state: Literal["pre_rope", "post_rope"] = "post_rope", +) -> DecodeAttentionInputs: + q, logical_k, logical_v = _qkv(seed=17) + q = q[:, :, 4:6, :] + page_size = 2 + physical_k = torch.empty_like(logical_k) + physical_v = torch.empty_like(logical_v) + positions = torch.full((2, 6), -1, dtype=torch.long) + for logical_page, physical_page in enumerate(page_order): + logical_slice = slice(logical_page * page_size, (logical_page + 1) * page_size) + physical_slice = slice(physical_page * page_size, (physical_page + 1) * page_size) + physical_k[:, :, physical_slice, :] = logical_k[:, :, logical_slice, :] + physical_v[:, :, physical_slice, :] = logical_v[:, :, logical_slice, :] + positions[:, physical_slice] = torch.arange( + logical_page * page_size, + (logical_page + 1) * page_size, + ) + inputs = DecodeAttentionInputs( + q=q, + k_cache=physical_k, + v_cache=physical_v, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[4, 5], [4, 5]], dtype=torch.long), + kv_seq_lens=torch.tensor([6, 6], dtype=torch.long), + block_table=torch.tensor([page_order, page_order], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor([[4, 5], [4, 5]], dtype=torch.long), + key_position_ids=positions.clone(), + page_size=page_size, + q_rope_state=q_rope_state, + k_cache_rope_state=k_cache_rope_state, + cp_block_owners=torch.tensor([[0, 1, 0], [0, 1, 0]], dtype=torch.long), + cp_world_size=2, + ), + lm_head_weight=torch.randn( + 11, q.size(1) * q.size(3), generator=torch.Generator().manual_seed(18) + ), + target_ids=torch.tensor([[1, 2], [3, 4]], dtype=torch.long), + active_token_mask=torch.tensor([[True, True], [False, True]], dtype=torch.bool), + ) + if not prefix_cache_enabled: + return inputs + prefix_length = 4 + fingerprint = decode_prefix_cache_fingerprint(inputs, prefix_length=prefix_length) + return replace( + inputs, + metadata=replace( + inputs.metadata, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + prefix_length=prefix_length, + prefix_cache_fingerprint=fingerprint, + ), + ) + + def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): report = compare_single_gpu_attention( _comparison_inputs(), @@ -155,6 +225,380 @@ def test_single_gpu_rope_attention_requires_position_metadata(): compare_single_gpu_rope_attention(AttentionComparisonInputs(q=base.q, k=base.k, v=base.v)) +def test_decode_replay_matches_full_prefill_for_single_and_few_query(): + inputs = _decode_inputs() + report = compare_decode_kv_replay(inputs) + + assert report.reference_name == "full_prefill_decode_reference" + drift = report.drifts[0] + assert drift.candidate_name == "rl_kernel_decode_kv_replay" + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.max_abs <= 3.0e-6 + assert drift.dlogp.active_count == 3 + assert drift.provenance["attention_mode"] == "decode" + assert drift.provenance["cache_position"] == [[4, 5], [4, 5]] + assert drift.provenance["cp_block_owners"] == [[0, 1, 0], [0, 1, 0]] + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["logical_merge_orders"] == [ + [[0, 1, 2], [0, 1, 2]], + [[0, 1, 2], [0, 1, 2]], + ] + + single_query = DecodeAttentionInputs( + q=inputs.q[:, :, -1:, :], + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position[:, -1:], + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=inputs.metadata.global_token_positions, + query_position_ids=inputs.metadata.query_position_ids[:, -1:], + key_position_ids=inputs.metadata.key_position_ids, + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + cp_world_size=inputs.metadata.cp_world_size, + ), + ) + single_report = compare_decode_kv_replay(single_query) + assert single_report.drifts[0].out.max_abs <= 1.0e-6 + assert single_report.drifts[0].lse.max_abs <= 1.0e-6 + + +def test_decode_replay_is_invariant_to_physical_page_and_prefix_layout(): + contiguous = run_decode_kv_replay(_decode_inputs()) + permuted = run_decode_kv_replay(_decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True)) + + torch.testing.assert_close(permuted.out, contiguous.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(permuted.lse, contiguous.lse, atol=1.0e-6, rtol=0.0) + assert permuted.provenance["prefix_cache_enabled"] is True + assert permuted.provenance["prefix_cache_key"] == "shared-prefix" + assert permuted.provenance["prefix_length"] == 4 + assert permuted.provenance["prefix_cache_fingerprint"] == decode_prefix_cache_fingerprint( + _decode_inputs(), prefix_length=4 + ) + + +def test_decode_replay_is_invariant_to_equivalent_cp_block_ownership(): + cp2_inputs = _decode_inputs() + cp1_inputs = replace( + cp2_inputs, + metadata=replace( + cp2_inputs.metadata, + cp_block_owners=torch.zeros_like(cp2_inputs.metadata.cp_block_owners), + ), + ) + + cp1 = run_decode_kv_replay(cp1_inputs) + cp2 = run_decode_kv_replay(cp2_inputs) + torch.testing.assert_close(cp2.out, cp1.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(cp2.lse, cp1.lse, atol=1.0e-6, rtol=0.0) + assert cp1.provenance["cp_block_owners"] == [[0, 0, 0], [0, 0, 0]] + assert cp2.provenance["cp_block_owners"] == [[0, 1, 0], [0, 1, 0]] + + +def test_decode_replay_pre_rope_cache_matches_equivalent_post_rope_cache(): + pre_rope = _decode_inputs(q_rope_state="pre_rope", k_cache_rope_state="pre_rope") + rope = NativeRoPEOp() + post_q = rope.forward_fp32( + pre_rope.q, + pre_rope.metadata.query_position_ids, + theta=pre_rope.rope_theta, + ) + post_k = rope.forward_fp32( + pre_rope.k_cache, + pre_rope.metadata.key_position_ids, + theta=pre_rope.rope_theta, + ) + post_rope = replace( + pre_rope, + q=post_q, + k_cache=post_k, + metadata=replace( + pre_rope.metadata, + q_rope_state="post_rope", + k_cache_rope_state="post_rope", + ), + ) + + pre_result = run_decode_kv_replay(pre_rope) + post_result = run_decode_kv_replay(post_rope) + torch.testing.assert_close(post_result.out, pre_result.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(post_result.lse, pre_result.lse, atol=1.0e-6, rtol=0.0) + + +def test_decode_replay_preserves_separate_q_and_k_rope_output_dtypes(): + base = _decode_inputs(q_rope_state="pre_rope", k_cache_rope_state="pre_rope") + mixed = replace( + base, + k_cache=base.k_cache.to(torch.bfloat16), + v_cache=base.v_cache.to(torch.bfloat16), + ) + rope = NativeRoPEOp() + post = replace( + mixed, + q=rope.forward_fp32( + mixed.q, + mixed.metadata.query_position_ids, + theta=mixed.rope_theta, + ).to(torch.float32), + k_cache=rope.forward_fp32( + mixed.k_cache, + mixed.metadata.key_position_ids, + theta=mixed.rope_theta, + ).to(torch.bfloat16), + metadata=replace( + mixed.metadata, + q_rope_state="post_rope", + k_cache_rope_state="post_rope", + ), + ) + + pre_result = run_decode_kv_replay(mixed) + post_result = run_decode_kv_replay(post) + torch.testing.assert_close(post_result.out, pre_result.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(post_result.lse, pre_result.lse, atol=1.0e-6, rtol=0.0) + assert pre_result.provenance["q_rope_output_dtype"] == "float32" + assert pre_result.provenance["k_cache_rope_output_dtype"] == "bfloat16" + + +def test_decode_replay_rejects_stale_prefix_cache_content(): + inputs = _decode_inputs(prefix_cache_enabled=True) + stale_k = inputs.k_cache.clone() + first_prefix_slot = int(inputs.metadata.block_table[0, 0].item()) * inputs.metadata.page_size + stale_k[0, 0, first_prefix_slot, 0] += 1.0 + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, k_cache=stale_k)) + + +def test_decode_replay_fails_loudly_on_position_identity_mismatch(): + inputs = _decode_inputs() + bad_query_positions = inputs.metadata.query_position_ids.clone() + bad_query_positions[0, -1] = 4 + bad_metadata = DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position, + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=inputs.metadata.global_token_positions, + query_position_ids=bad_query_positions, + key_position_ids=inputs.metadata.key_position_ids, + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + cp_world_size=inputs.metadata.cp_world_size, + ) + + with pytest.raises(ValueError, match="cache_position and query_position_ids"): + run_decode_kv_replay( + DecodeAttentionInputs( + q=inputs.q, + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=bad_metadata, + ) + ) + + +def test_decode_replay_fails_loudly_on_invalid_page_identity(): + inputs = _decode_inputs() + bad_positions = inputs.metadata.global_token_positions.clone() + bad_positions[:, 0] = 1 + bad_metadata = DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position, + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=bad_positions, + query_position_ids=inputs.metadata.query_position_ids, + key_position_ids=bad_positions.clone(), + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + cp_world_size=inputs.metadata.cp_world_size, + ) + + with pytest.raises(ValueError, match="reconstruct logical positions"): + compare_decode_kv_replay( + DecodeAttentionInputs( + q=inputs.q, + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=bad_metadata, + ) + ) + + +def test_decode_replay_covers_qwen3_gqa_head_layout(): + generator = torch.Generator().manual_seed(23) + q = torch.randn(1, 32, 1, 128, generator=generator, dtype=torch.bfloat16) + k = torch.randn(1, 8, 4, 128, generator=generator, dtype=torch.bfloat16) + v = torch.randn(1, 8, 4, 128, generator=generator, dtype=torch.bfloat16) + positions = torch.arange(4, dtype=torch.long).unsqueeze(0) + inputs = DecodeAttentionInputs( + q=q, + k_cache=k, + v_cache=v, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[3]], dtype=torch.long), + kv_seq_lens=torch.tensor([4], dtype=torch.long), + block_table=torch.tensor([[0, 1]], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor([[3]], dtype=torch.long), + key_position_ids=positions.clone(), + page_size=2, + cp_block_owners=torch.tensor([[0, 1]], dtype=torch.long), + cp_world_size=2, + ), + output_dtype=torch.bfloat16, + ) + + report = compare_decode_kv_replay(inputs) + assert report.drifts[0].out.max_abs <= 2 * torch.finfo(torch.bfloat16).eps + assert report.drifts[0].lse.max_abs <= 1.0e-6 + + +def test_decode_append_matches_full_prefill_suffix(): + generator = torch.Generator().manual_seed(71) + q = torch.randn(1, 4, 2, 8, generator=generator) + k_past = torch.randn(1, 2, 4, 8, generator=generator) + v_past = torch.randn(1, 2, 4, 8, generator=generator) + k_new = torch.randn(1, 2, 2, 8, generator=generator) + v_new = torch.randn(1, 2, 2, 8, generator=generator) + inputs = DecodeAttentionInputs( + q=q, + k_cache=k_past, + v_cache=v_past, + k_new=k_new, + v_new=v_new, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[104, 105]], dtype=torch.long), + kv_seq_lens=torch.tensor([4], dtype=torch.long), + block_table=torch.tensor([[0, 1]], dtype=torch.long), + global_token_positions=torch.tensor([[100, 101, 102, 103]], dtype=torch.long), + query_position_ids=torch.tensor([[104, 105]], dtype=torch.long), + key_position_ids=torch.tensor([[100, 101, 102, 103]], dtype=torch.long), + page_size=2, + cp_block_owners=torch.tensor([[0, 1]], dtype=torch.long), + cp_world_size=2, + ), + split_kv=SplitKVSpec.fixed(2), + ) + + report = compare_decode_kv_replay(inputs) + drift = report.drifts[0] + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.provenance["decode_semantics"] == "past_kv_plus_new_kv_append" + assert drift.provenance["past_kv_lengths"] == [4] + assert drift.provenance["new_kv_length"] == 2 + assert drift.provenance["actual_split_kv_plans"][0][1][ + "actual_split_boundaries" + ] == [[0, 2], [2, 4], [4, 6]] + + +def test_decode_replay_supports_nonzero_global_position_offset(): + base = _decode_inputs() + offset = 4096 + active = base.metadata.global_token_positions >= 0 + positions = torch.where( + active, + base.metadata.global_token_positions + offset, + base.metadata.global_token_positions, + ) + inputs = replace( + base, + metadata=replace( + base.metadata, + cache_position=base.metadata.cache_position + offset, + query_position_ids=base.metadata.query_position_ids + offset, + global_token_positions=positions, + key_position_ids=positions.clone(), + ), + ) + + report = compare_decode_kv_replay(inputs) + assert report.drifts[0].out.max_abs <= 1.0e-6 + assert report.drifts[0].provenance["global_token_positions"][0][0] >= offset + + +def test_decode_split_k_disabled_and_fixed_share_cache_layout(): + base = _decode_inputs() + disabled = run_decode_kv_replay(replace(base, split_kv=SplitKVSpec.disabled())) + fixed = run_decode_kv_replay(replace(base, split_kv=SplitKVSpec.fixed(2))) + + torch.testing.assert_close(fixed.out, disabled.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(fixed.lse, disabled.lse, atol=1.0e-6, rtol=0.0) + assert disabled.provenance["requested_split_kv_policy"] == "disabled" + assert fixed.provenance["requested_split_kv_policy"] == "fixed" + assert disabled.provenance["block_table"] == fixed.provenance["block_table"] + + +def test_decode_transformer_engine_oracle_reuses_sorted_partial_states(monkeypatch): + calls = {"lse": 0, "out": 0} + + def lse_correction(softmax_lse, softmax_lse_per_step): + calls["lse"] += 1 + softmax_lse.copy_(torch.logaddexp(softmax_lse, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + calls["out"] += 1 + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_decode_kv_replay( + _decode_inputs(page_order=(2, 0, 1)), + include_transformer_engine=True, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert set(by_name) == { + "rl_kernel_decode_kv_replay", + "transformer_engine_decode_kv_replay", + } + assert by_name["transformer_engine_decode_kv_replay"].out.max_abs <= 1.0e-6 + assert by_name["transformer_engine_decode_kv_replay"].lse.max_abs <= 1.0e-6 + assert by_name["transformer_engine_decode_kv_replay"].provenance["logical_merge_orders"] == [ + [[0, 1, 2], [0, 1, 2]], + [[0, 1, 2], [0, 1, 2]], + ] + assert calls["lse"] > 0 + assert calls["out"] > 0 + assert report.unavailable == () + + +def test_decode_transformer_engine_unavailable_is_reported(monkeypatch): + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == _TE_CONTEXT_PARALLEL_MODULE: + raise ImportError("decode TE unavailable") + return real_import_module(name, package) + + monkeypatch.setattr(importlib, "import_module", fake_import_module) + + report = compare_decode_kv_replay( + _decode_inputs(), + include_transformer_engine=True, + ) + + assert {drift.candidate_name for drift in report.drifts} == {"rl_kernel_decode_kv_replay"} + assert report.unavailable == ("transformer_engine_decode_kv_replay: decode TE unavailable",) + + def test_transformer_engine_merge_oracle_can_be_reused_when_available(monkeypatch): calls = {"lse": 0, "out": 0} From d73fe374e0480f42c69558ea9ceaf3471a2d8036 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 18:19:56 +0800 Subject: [PATCH 5/5] style(attention): satisfy comparison harness lint Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/attention_contract.py | 68 ++++++++--------------- rl_engine/testing/attention_comparison.py | 10 +--- tests/test_attention_comparison.py | 10 ++-- 3 files changed, 32 insertions(+), 56 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index eb4994b3..1750476d 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -359,9 +359,7 @@ def __post_init__(self) -> None: or start < 0 or end <= start ): - raise AttentionContractError( - "Split-KV boundaries must satisfy 0 <= start < end" - ) + raise AttentionContractError("Split-KV boundaries must satisfy 0 <= start < end") if index > 0 and start != previous_end: raise AttentionContractError( "Split-KV boundaries must be contiguous and in logical KV order" @@ -486,14 +484,10 @@ def __post_init__(self) -> None: raise AttentionContractError("split_kv.strict_consistency must be a bool") if self.mode is SplitKVMode.FIXED: if self.fixed_split_size is None: - raise AttentionContractError( - "fixed Split-KV policy requires fixed_split_size" - ) + raise AttentionContractError("fixed Split-KV policy requires fixed_split_size") _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") elif self.fixed_split_size is not None: - raise AttentionContractError( - "fixed_split_size is only valid for fixed Split-KV policy" - ) + raise AttentionContractError("fixed_split_size is only valid for fixed Split-KV policy") if self.strict_consistency and self.mode is SplitKVMode.AUTO: raise AttentionContractError( "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" @@ -647,13 +641,8 @@ def validate( "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" ) if self.execution.actual_mode is None: - raise AttentionContractError( - "complete Split-KV plan sets require actual runtime plans" - ) - if ( - self.execution.boundaries[0][0] != start - or self.execution.boundaries[-1][1] != end - ): + raise AttentionContractError("complete Split-KV plan sets require actual runtime plans") + if self.execution.boundaries[0][0] != start or self.execution.boundaries[-1][1] != end: raise AttentionContractError( "Split-KV execution boundaries must exactly cover expected_kv_range" ) @@ -661,9 +650,7 @@ def validate( boundary_start < start or boundary_end > end for boundary_start, boundary_end in self.execution.boundaries ): - raise AttentionContractError( - "Split-KV execution boundary escapes expected_kv_range" - ) + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") def to_dict(self) -> dict[str, Any]: return { @@ -709,9 +696,7 @@ def __post_init__(self) -> None: } actual_coordinates = [entry.coordinate for entry in entries] if len(set(actual_coordinates)) != len(actual_coordinates): - raise AttentionContractError( - "Split-KV runtime plan set contains duplicate coordinates" - ) + raise AttentionContractError("Split-KV runtime plan set contains duplicate coordinates") missing = expected_coordinates.difference(actual_coordinates) extra = set(actual_coordinates).difference(expected_coordinates) if missing or extra: @@ -821,17 +806,12 @@ def validate_split_kv_plan_set_alignment( ] if topology_mismatches: raise AttentionContractError( - "training/rollout Split-KV plan-set topology differs: " - + ", ".join(topology_mismatches) + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) ) - training_by_coordinate = { - entry.coordinate: entry for entry in training.entries - } + training_by_coordinate = {entry.coordinate: entry for entry in training.entries} rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} if training_by_coordinate.keys() != rollout_by_coordinate.keys(): - raise AttentionContractError( - "training/rollout Split-KV plan-set coordinates differ" - ) + raise AttentionContractError("training/rollout Split-KV plan-set coordinates differ") for coordinate in sorted(training_by_coordinate): train_entry = training_by_coordinate[coordinate] rollout_entry = rollout_by_coordinate[coordinate] @@ -1075,14 +1055,16 @@ def __post_init__(self) -> None: not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() ): raise AttentionContractError("rope_scaling must be a non-empty string when provided") - for field in ("position_ids", "query_position_offsets", "key_position_offsets"): - values = getattr(self, field) + for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, position_field) if values is None: continue - normalized = _integer_tuple(values, field) + normalized = _integer_tuple(values, position_field) if not normalized or any(value < 0 for value in normalized): - raise AttentionContractError(f"{field} must contain non-negative positions") - object.__setattr__(self, field, normalized) + raise AttentionContractError( + f"{position_field} must contain non-negative positions" + ) + object.__setattr__(self, position_field, normalized) object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) object.__setattr__( self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") @@ -1186,11 +1168,11 @@ def __post_init__(self) -> None: "position_ids must describe the local query sequence or full local " "sequence length" ) - for field in ("query_position_offsets", "key_position_offsets"): - offsets = getattr(self.rope, field) + for position_field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, position_field) if offsets is not None and len(offsets) != batch_size: raise AttentionContractError( - f"{field} must contain one entry per logical batch entry" + f"{position_field} must contain one entry per logical batch entry" ) if self.mode is AttentionMode.DECODE and self.kv_cache is not None: if len(self.kv_cache.kv_seq_lens) != batch_size: @@ -1336,7 +1318,7 @@ def __post_init__(self) -> None: raise AttentionContractError("tp_world_sizes must contain positive values") if len(set(tp_world_sizes)) != len(tp_world_sizes): raise AttentionContractError("tp_world_sizes must not contain duplicates") - for field in ( + for capability_field in ( "exports_attention_lse", "deterministic_cp_merge", "supports_packed_varlen", @@ -1348,8 +1330,8 @@ def __post_init__(self) -> None: "supports_split_kv_auto", "reports_actual_split_kv_plan", ): - if not isinstance(getattr(self, field), bool): - raise AttentionContractError(f"{field} must be a bool") + if not isinstance(getattr(self, capability_field), bool): + raise AttentionContractError(f"{capability_field} must be a bool") if self.implementation_kind not in {"production", "reference", "deterministic"}: raise AttentionContractError( "implementation_kind must be production, reference, or deterministic" @@ -1401,9 +1383,7 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: SplitKVMode.AUTO: self.supports_split_kv_auto, } if not split_support[contract.split_kv.mode]: - reasons.append( - f"Split-KV policy={contract.split_kv.mode.value} is unsupported" - ) + reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported") if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: reasons.append("actual Split-KV execution-plan provenance is unsupported") return tuple(reasons) diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 41d5abfc..06b39f78 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -21,11 +21,7 @@ import torch -from rl_engine.kernels.attention_contract import ( - SplitKVExecutionPlan, - SplitKVMode, - SplitKVSpec, -) +from rl_engine.kernels.attention_contract import SplitKVExecutionPlan, SplitKVMode, SplitKVSpec from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.testing.reference_ops import selected_logprobs_reference @@ -378,9 +374,7 @@ def _run_decode_kv_replay( order: list[int] = [] visible_count = int((logical_positions <= query_position).sum().item()) split_bounds = _decode_split_bounds(visible_count, split_kv) - for block_index, (block_start, block_end) in enumerate( - split_bounds - ): + for block_index, (block_start, block_end) in enumerate(split_bounds): block_positions = logical_positions[block_start:block_end] visible = block_positions <= query_position if not bool(visible.any()): diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 5ae5d5e3..86936760 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -14,10 +14,10 @@ import pytest import torch +from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.kernels.gtest import run_operator_suite from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp -from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.testing.attention_comparison import ( AttentionComparisonInputs, DecodeAttentionInputs, @@ -491,9 +491,11 @@ def test_decode_append_matches_full_prefill_suffix(): assert drift.provenance["decode_semantics"] == "past_kv_plus_new_kv_append" assert drift.provenance["past_kv_lengths"] == [4] assert drift.provenance["new_kv_length"] == 2 - assert drift.provenance["actual_split_kv_plans"][0][1][ - "actual_split_boundaries" - ] == [[0, 2], [2, 4], [4, 6]] + assert drift.provenance["actual_split_kv_plans"][0][1]["actual_split_boundaries"] == [ + [0, 2], + [2, 4], + [4, 6], + ] def test_decode_replay_supports_nonzero_global_position_offset():