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 01/10] 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 02/10] 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 03/10] 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 7054d99b1483817d0eeda84d840ff87a817f47fc Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Mon, 3 Aug 2026 16:55:48 +0800 Subject: [PATCH 04/10] feat(attention): add decode KV-cache replay harness Signed-off-by: JLiu4Coding --- docs/design/ws2-attention-decode-replay.md | 61 +++ rl_engine/testing/__init__.py | 10 + rl_engine/testing/attention_comparison.py | 425 +++++++++++++++++++++ tests/test_attention_comparison.py | 230 +++++++++++ 4 files changed, 726 insertions(+) create mode 100644 docs/design/ws2-attention-decode-replay.md diff --git a/docs/design/ws2-attention-decode-replay.md b/docs/design/ws2-attention-decode-replay.md new file mode 100644 index 00000000..0deca953 --- /dev/null +++ b/docs/design/ws2-attention-decode-replay.md @@ -0,0 +1,61 @@ +# WS2 attention decode replay + +PR6 of issue #235 extends the single-GPU attention attribution harness with a +correctness-first decode-stage KV-cache replay. The reference compares each +decode query with the same query evaluated over a fully materialized logical KV +sequence. Both paths export attention output and attention-domain LSE. + +## Cache identity + +`DecodeKVCacheMetadata` separates logical sequence identity from physical cache +layout: + +- `cache_position` and `query_position_ids` identify each decode query; +- `kv_seq_lens` records the active cached sequence length; +- `block_table` maps logical blocks to physical pages; +- `global_token_positions` identifies every populated physical cache slot; +- `key_position_ids` binds cached K to its RoPE positions; +- `q_rope_state` and `k_cache_rope_state` declare whether tensors are before or + after RoPE; +- `prefix_cache_key` identifies a reused prefix when prefix caching is enabled; +- `cp_block_owners` records logical CP ownership without changing merge order. + +Metadata is validated before attention runs. Missing pages, duplicated active +pages, out-of-range positions, mismatched RoPE positions, or inconsistent +prefix-cache identity fail with an explicit error. + +## Deterministic replay + +The replay restores logical KV order from the block table, computes one FP32 +partial `(out, lse)` state per logical block, and merges partial states in +`global_block_index` order. CP ownership and physical page order never determine +the numerical reduction order. Downcast occurs only at final write. + +For each decode query at logical position `t`, only cached positions less than +or equal to `t` participate. This supports `Sq=1` and few-query replay. + +The reference path removes physical page boundaries after reconstructing the +same logical KV sequence. Therefore the primary check is: + +```text +decode_paged_kv(query_t) == full_logical_kv(query_t) +``` + +The report reuses the PR2 drift format and includes maximum, mean, p95, and p99 +absolute drift for output and LSE, plus decode/cache/RoPE provenance. + +## Transformer Engine + +Transformer Engine remains optional. When requested, the replay passes the same +sorted partial states to the capability-probed TE context-parallel correction +helpers already used by the PR2 harness. TE validates merge arithmetic only; it +does not interpret cache metadata, choose logical order, or replace RL-Kernel's +reference semantics. + +## Current boundary + +The harness models CP block ownership on one device so cache construction, +logical ordering, RoPE identity, and deterministic merging can be attributed +without communication. A distributed caller can gather the same partial states +using the PR3 transport layer; numerical merging must still happen in the fixed +logical order described above. diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 792b36d0..559a1026 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -8,11 +8,16 @@ AttentionComparisonReport, AttentionPathDrift, AttentionPathResult, + DecodeAttentionInputs, + DecodeKVCacheMetadata, DriftStats, TransformerEngineUnavailable, + compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, 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 +40,23 @@ "AttentionComparisonReport", "AttentionPathDrift", "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", "DriftStats", "SyntheticRLKernelBatch", "TransformerEngineUnavailable", "active_token_count", "compare_single_gpu_rope_attention", "compare_single_gpu_attention", + "compare_decode_kv_replay", "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..fde2773d 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -24,6 +24,7 @@ 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 +72,49 @@ 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 + q_rope_state: RoPEState = "post_rope" + k_cache_rope_state: RoPEState = "post_rope" + cp_block_owners: torch.Tensor | None = None + + +@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" + 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 + + @dataclass(frozen=True) class AttentionPathResult: """One materialized attention path result.""" @@ -210,6 +254,177 @@ def compare_single_gpu_rope_attention( return AttentionComparisonReport(reference_name=reference.name, drifts=drifts) +def compare_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend = "rl_kernel", +) -> AttentionComparisonReport: + """Compare paged decode replay with a logical full-KV teacher-forcing view.""" + + _validate_decode_inputs(inputs) + reference = run_decode_full_prefill_reference(inputs) + candidate = run_decode_kv_replay(inputs, merge_backend=merge_backend) + 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) + drift = AttentionPathDrift( + candidate_name=candidate.name, + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=dlogp, + provenance=candidate.provenance, + ) + return AttentionComparisonReport(reference_name=reference.name, drifts=(drift,)) + + +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) + 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) + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + merge_orders: list[list[int]] = [] + cp_block_owners: list[list[int]] = [] + 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] = [] + 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] = [] + for block_index, (block_start, block_end) in enumerate( + _chunk_bounds(k.size(2), inputs.metadata.page_size) + ): + 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) + merge_orders.append(order) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + + provenance: dict[str, Any] = { + "attention_mode": "decode", + "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, + "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, + "rope_output_dtype": str(_decode_rope_output_dtype(inputs)).replace("torch.", ""), + "cp_block_owners": cp_block_owners, + "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.""" @@ -464,6 +679,71 @@ 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 + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + physical_slots: list[int] = [] + logical_positions: list[int] = [] + logical_block_count = math.ceil(sequence_length / metadata.page_size) + for logical_block in range(logical_block_count): + physical_page = int(metadata.block_table[batch_index, logical_block].item()) + tokens_in_block = min( + metadata.page_size, + sequence_length - logical_block * metadata.page_size, + ) + for page_offset in range(tokens_in_block): + slot = physical_page * metadata.page_size + page_offset + physical_slots.append(slot) + logical_positions.append(int(metadata.global_token_positions[batch_index, slot].item())) + + slot_index = torch.tensor(physical_slots, device=inputs.k_cache.device, dtype=torch.long) + logical_position_tensor = torch.tensor( + logical_positions, + device=inputs.k_cache.device, + dtype=torch.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() + output_dtype = _decode_rope_output_dtype(inputs) + 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(output_dtype) + 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(output_dtype) + return q, k, v, logical_position_tensor + + +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_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.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 +1023,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 +1151,124 @@ 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) + 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") + 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 and not metadata.prefix_cache_key: + raise ValueError("prefix_cache_key is required when prefix cache is enabled") + if not metadata.prefix_cache_enabled and metadata.prefix_cache_key is not None: + raise ValueError("prefix_cache_key must be None 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.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") + 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") + slots: list[int] = [] + for logical_block, page in enumerate(pages.tolist()): + token_count = min(page_size, sequence_length - logical_block * page_size) + slots.extend(int(page) * page_size + offset for offset in range(token_count)) + slot_index = torch.tensor(slots, device=inputs.q.device, dtype=torch.long) + 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] + expected_positions = torch.arange( + 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 " + "0..kv_seq_len-1 exactly" + ) + 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 bool((cache_positions < 0).any()) or bool((cache_positions >= 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") + + 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]") @@ -894,13 +1314,18 @@ 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", "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..9251513f 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -8,16 +8,23 @@ 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.testing.attention_comparison import ( AttentionComparisonInputs, + DecodeAttentionInputs, + DecodeKVCacheMetadata, + compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + run_decode_kv_replay, run_paged_kv_attention, ) @@ -57,6 +64,54 @@ 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, + ) + return 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, + prefix_cache_enabled=prefix_cache_enabled, + prefix_cache_key="shared-prefix" if prefix_cache_enabled else None, + 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), + ), + 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), + ) + + def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): report = compare_single_gpu_attention( _comparison_inputs(), @@ -155,6 +210,181 @@ 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" + + 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, + ), + ) + 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" + + +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_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, + ) + + 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, + ) + + 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), + ), + output_dtype=torch.bfloat16, + ) + + report = compare_decode_kv_replay(inputs) + assert report.drifts[0].out.max_abs <= 1.0e-6 + assert report.drifts[0].lse.max_abs <= 1.0e-6 + + def test_transformer_engine_merge_oracle_can_be_reused_when_available(monkeypatch): calls = {"lse": 0, "out": 0} From 2eb4362d04ae0184d4112a4dba79029aa5492db7 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Tue, 4 Aug 2026 10:40:52 +0800 Subject: [PATCH 05/10] fix(attention): address decode replay review feedback Signed-off-by: JLiu4Coding --- docs/design/ws2-attention-decode-replay.md | 14 +- rl_engine/testing/__init__.py | 2 + rl_engine/testing/attention_comparison.py | 231 ++++++++++++++++----- tests/test_attention_comparison.py | 140 ++++++++++++- 4 files changed, 323 insertions(+), 64 deletions(-) diff --git a/docs/design/ws2-attention-decode-replay.md b/docs/design/ws2-attention-decode-replay.md index 0deca953..16b96e28 100644 --- a/docs/design/ws2-attention-decode-replay.md +++ b/docs/design/ws2-attention-decode-replay.md @@ -17,12 +17,18 @@ layout: - `key_position_ids` binds cached K to its RoPE positions; - `q_rope_state` and `k_cache_rope_state` declare whether tensors are before or after RoPE; -- `prefix_cache_key` identifies a reused prefix when prefix caching is enabled; +- `prefix_cache_key`, `prefix_length`, and `prefix_cache_fingerprint` identify a + reused logical prefix and bind it to the cached K/V content; - `cp_block_owners` records logical CP ownership without changing merge order. Metadata is validated before attention runs. Missing pages, duplicated active pages, out-of-range positions, mismatched RoPE positions, or inconsistent -prefix-cache identity fail with an explicit error. +prefix-cache identity fail with an explicit error. Prefix identity is verified +by recomputing a physical-layout-invariant SHA-256 fingerprint over logical +prefix positions and cached K/V content. The current reference requires +`rope_cast_at="after_rope"` and `rope_rotary_dim == head_dim`; partial rotary is +not supported yet. Q and cached K retain separate RoPE output dtype contracts so +mixed rollout query/KV storage dtypes are represented faithfully. ## Deterministic replay @@ -50,7 +56,9 @@ Transformer Engine remains optional. When requested, the replay passes the same sorted partial states to the capability-probed TE context-parallel correction helpers already used by the PR2 harness. TE validates merge arithmetic only; it does not interpret cache metadata, choose logical order, or replace RL-Kernel's -reference semantics. +reference semantics. The comparison always reports the RL-Kernel result; when +TE is requested but unavailable or incompatible, it records the reason in the +report instead of failing the core decode harness. ## Current boundary diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 559a1026..3897e009 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -15,6 +15,7 @@ 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, @@ -49,6 +50,7 @@ "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", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index fde2773d..5edec9ae 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 @@ -91,6 +92,8 @@ class DecodeKVCacheMetadata: 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 @@ -109,7 +112,8 @@ class DecodeAttentionInputs: 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 + 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 @@ -257,26 +261,25 @@ def compare_single_gpu_rope_attention( def compare_decode_kv_replay( inputs: DecodeAttentionInputs, *, - merge_backend: MergeBackend = "rl_kernel", + 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) - candidate = run_decode_kv_replay(inputs, merge_backend=merge_backend) - 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) - drift = AttentionPathDrift( - candidate_name=candidate.name, - out=_drift_stats(candidate.out, reference.out), - lse=_drift_stats(candidate.lse, reference.lse), - dlogp=dlogp, - provenance=candidate.provenance, + 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), ) - return AttentionComparisonReport(reference_name=reference.name, drifts=(drift,)) def run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: @@ -287,6 +290,10 @@ def run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> Attentio """ _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)): @@ -334,9 +341,17 @@ def run_decode_kv_replay( """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[int]] = [] + merge_orders: list[list[list[int]]] = [] cp_block_owners: list[list[int]] = [] for batch_index in range(inputs.q.size(0)): q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) @@ -344,6 +359,7 @@ def run_decode_kv_replay( cp_block_owners.append(owners) batch_out: list[torch.Tensor] = [] batch_lse: list[torch.Tensor] = [] + batch_orders: list[list[int]] = [] for query_index in range(q.size(2)): query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) states: list[_PartialAttentionState] = [] @@ -383,7 +399,8 @@ def run_decode_kv_replay( out, lse = _merge_partial_states(states, backend=merge_backend) batch_out.append(out.to(inputs.output_dtype)) batch_lse.append(lse) - merge_orders.append(order) + batch_orders.append(order) + merge_orders.append(batch_orders) outs.append(torch.cat(batch_out, dim=2)) lses.append(torch.cat(batch_lse, dim=2)) @@ -400,12 +417,15 @@ def run_decode_kv_replay( "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, - "rope_output_dtype": str(_decode_rope_output_dtype(inputs)).replace("torch.", ""), + "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, "merge_order": "global_block_index", "logical_merge_orders": merge_orders, @@ -639,6 +659,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, @@ -669,6 +722,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 @@ -686,45 +758,46 @@ def _decode_logical_qkv( """Restore one batch's cache to logical order and materialize RoPE state.""" metadata = inputs.metadata - sequence_length = int(metadata.kv_seq_lens[batch_index].item()) - physical_slots: list[int] = [] - logical_positions: list[int] = [] - logical_block_count = math.ceil(sequence_length / metadata.page_size) - for logical_block in range(logical_block_count): - physical_page = int(metadata.block_table[batch_index, logical_block].item()) - tokens_in_block = min( - metadata.page_size, - sequence_length - logical_block * metadata.page_size, - ) - for page_offset in range(tokens_in_block): - slot = physical_page * metadata.page_size + page_offset - physical_slots.append(slot) - logical_positions.append(int(metadata.global_token_positions[batch_index, slot].item())) - - slot_index = torch.tensor(physical_slots, device=inputs.k_cache.device, dtype=torch.long) - logical_position_tensor = torch.tensor( - logical_positions, - device=inputs.k_cache.device, - dtype=torch.long, - ) + 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() - output_dtype = _decode_rope_output_dtype(inputs) 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(output_dtype) + ).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(output_dtype) + k = rope.forward_fp32(k, key_positions, theta=inputs.rope_theta).to( + _decode_k_rope_output_dtype(inputs) + ) 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 @@ -736,8 +809,16 @@ def _logical_block_owners(inputs: DecodeAttentionInputs, batch_index: int) -> li ] -def _decode_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: - return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype +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: @@ -1153,6 +1234,8 @@ def _validate_rope_inputs(inputs: AttentionComparisonInputs) -> None: 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) @@ -1199,10 +1282,21 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: 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 and not metadata.prefix_cache_key: - raise ValueError("prefix_cache_key is required when prefix cache is enabled") - if not metadata.prefix_cache_enabled and metadata.prefix_cache_key is not None: - raise ValueError("prefix_cache_key must be None when prefix cache is disabled") + 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: @@ -1211,10 +1305,26 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: _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.rope_output_dtype is not None and not isinstance( - inputs.rope_output_dtype, torch.dtype + if inputs.q_rope_output_dtype is not None and not isinstance( + inputs.q_rope_output_dtype, torch.dtype ): - raise ValueError("rope_output_dtype must be a torch.dtype when provided") + 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 ( + 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): @@ -1237,11 +1347,7 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: 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") - slots: list[int] = [] - for logical_block, page in enumerate(pages.tolist()): - token_count = min(page_size, sequence_length - logical_block * page_size) - slots.extend(int(page) * page_size + offset for offset in range(token_count)) - slot_index = torch.tensor(slots, device=inputs.q.device, dtype=torch.long) + 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()): @@ -1268,6 +1374,16 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: 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: @@ -1321,6 +1437,7 @@ def _positive_int(value: int, name: str) -> int: "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", diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 9251513f..0733a321 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -24,6 +24,7 @@ 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, ) @@ -86,7 +87,7 @@ def _decode_inputs( logical_page * page_size, (logical_page + 1) * page_size, ) - return DecodeAttentionInputs( + inputs = DecodeAttentionInputs( q=q, k_cache=physical_k, v_cache=physical_v, @@ -98,8 +99,6 @@ def _decode_inputs( query_position_ids=torch.tensor([[4, 5], [4, 5]], dtype=torch.long), key_position_ids=positions.clone(), page_size=page_size, - prefix_cache_enabled=prefix_cache_enabled, - prefix_cache_key="shared-prefix" if prefix_cache_enabled else None, 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), @@ -110,6 +109,20 @@ def _decode_inputs( 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(): @@ -226,6 +239,10 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): 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:, :], @@ -255,6 +272,10 @@ def test_decode_replay_is_invariant_to_physical_page_and_prefix_layout(): 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(): @@ -305,6 +326,51 @@ def test_decode_replay_pre_rope_cache_matches_equivalent_post_rope_cache(): 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() @@ -381,10 +447,76 @@ def test_decode_replay_covers_qwen3_gqa_head_layout(): ) report = compare_decode_kv_replay(inputs) - assert report.drifts[0].out.max_abs <= 1.0e-6 + 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_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 be6ed39f19c794b0b7cfdd8fa8f4c9d267d0ed64 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 00:51:49 +0800 Subject: [PATCH 06/10] fix(attention): bind decode replay to cache execution identity Signed-off-by: lamentropetion <3051000145@qq.com> --- docs/design/ws2-attention-decode-replay.md | 23 ++++-- rl_engine/testing/attention_comparison.py | 92 +++++++++++++++++++--- tests/test_attention_comparison.py | 91 +++++++++++++++++++++ 3 files changed, 187 insertions(+), 19 deletions(-) diff --git a/docs/design/ws2-attention-decode-replay.md b/docs/design/ws2-attention-decode-replay.md index 16b96e28..c0c1fff1 100644 --- a/docs/design/ws2-attention-decode-replay.md +++ b/docs/design/ws2-attention-decode-replay.md @@ -22,10 +22,12 @@ layout: - `cp_block_owners` records logical CP ownership without changing merge order. Metadata is validated before attention runs. Missing pages, duplicated active -pages, out-of-range positions, mismatched RoPE positions, or inconsistent -prefix-cache identity fail with an explicit error. Prefix identity is verified -by recomputing a physical-layout-invariant SHA-256 fingerprint over logical -prefix positions and cached K/V content. The current reference requires +pages, non-canonical `-1` page/owner tails, out-of-range positions, mismatched +RoPE positions, or inconsistent prefix-cache identity fail with an explicit +error. Prefix identity is verified by recomputing a physical-layout-invariant +SHA-256 fingerprint over logical prefix positions, cached K/V content, storage +dtypes, and the cache-side RoPE materialization configuration (`theta`, rotary +dimension, cast boundary, and output dtype). The current reference requires `rope_cast_at="after_rope"` and `rope_rotary_dim == head_dim`; partial rotary is not supported yet. Q and cached K retain separate RoPE output dtype contracts so mixed rollout query/KV storage dtypes are represented faithfully. @@ -35,7 +37,9 @@ mixed rollout query/KV storage dtypes are represented faithfully. The replay restores logical KV order from the block table, computes one FP32 partial `(out, lse)` state per logical block, and merges partial states in `global_block_index` order. CP ownership and physical page order never determine -the numerical reduction order. Downcast occurs only at final write. +the numerical reduction order. Both the full logical-KV reference and paged +replay accumulate and export LSE in FP32; output downcast occurs only at final +write. For each decode query at logical position `t`, only cached positions less than or equal to `t` participate. This supports `Sq=1` and few-query replay. @@ -64,6 +68,9 @@ report instead of failing the core decode harness. The harness models CP block ownership on one device so cache construction, logical ordering, RoPE identity, and deterministic merging can be attributed -without communication. A distributed caller can gather the same partial states -using the PR3 transport layer; numerical merging must still happen in the fixed -logical order described above. +without communication. It does not implement P2P NCCL transport, production +custom CUDA all-gather/reduce-scatter, FlashInfer runtime execution, or training +backward. A distributed caller can gather the same partial states using the PR3 +transport layer; numerical merging must still happen in the fixed logical order +described above. This PR targets the shared `test` integration branch and has a +logical dependency on PR2/#253. diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 5edec9ae..30fb4838 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -17,6 +17,7 @@ import inspect import math from dataclasses import dataclass +from numbers import Integral from typing import Any, Literal import torch @@ -329,6 +330,13 @@ def _run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> Attenti "materialization": "full_logical_kv", "lse_domain": "attention", "accum_dtype": "fp32", + "lse_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(inputs.output_dtype).replace("torch.", ""), + "scale": _decode_attention_scale(inputs), + "q_dtype": str(inputs.q.dtype).replace("torch.", ""), + "k_cache_dtype": str(inputs.k_cache.dtype).replace("torch.", ""), + "v_cache_dtype": str(inputs.v_cache.dtype).replace("torch.", ""), }, ) @@ -433,7 +441,13 @@ def _run_decode_kv_replay( "lse_domain": "attention", "lse_exported": True, "accum_dtype": "fp32", + "lse_dtype": "fp32", "downcast_at": "final_write", + "output_dtype": str(inputs.output_dtype).replace("torch.", ""), + "scale": _decode_attention_scale(inputs), + "q_dtype": str(inputs.q.dtype).replace("torch.", ""), + "k_cache_dtype": str(inputs.k_cache.dtype).replace("torch.", ""), + "v_cache_dtype": str(inputs.v_cache.dtype).replace("torch.", ""), } if merge_backend == "transformer_engine": provenance.update(_te_context_parallel_provenance()) @@ -667,16 +681,24 @@ def decode_prefix_cache_fingerprint( """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. + are first restored to logical token order. It includes every cache-side + RoPE materialization fact that can change the replayed K values, so a + prefix cannot be reused under a different rotary configuration. """ 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_rope_state={inputs.metadata.k_cache_rope_state};" + f"rope_theta={float(inputs.rope_theta):.17g};" + f"rotary_dim={_decode_rope_rotary_dim(inputs)};" + f"rope_cast_at={inputs.rope_cast_at};" + f"k_rope_output_dtype={_decode_k_rope_output_dtype(inputs)}\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] @@ -825,6 +847,14 @@ 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 _decode_attention_scale(inputs: DecodeAttentionInputs) -> float: + return ( + 1.0 / math.sqrt(inputs.q.size(-1)) + if inputs.scale is None + else float(inputs.scale) + ) + + def _rope_output_dtype(inputs: AttentionComparisonInputs) -> torch.dtype: return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype @@ -1191,9 +1221,12 @@ 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: + if ( + not isinstance(inputs.rope_theta, (float, int)) + or isinstance(inputs.rope_theta, bool) + or not math.isfinite(float(inputs.rope_theta)) + or 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 @@ -1236,6 +1269,25 @@ 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") + if ( + not inputs.q.is_floating_point() + or not inputs.k_cache.is_floating_point() + or not inputs.v_cache.is_floating_point() + ): + raise ValueError("q, k_cache, and v_cache must use floating-point dtypes") + if ( + not isinstance(inputs.output_dtype, torch.dtype) + or not inputs.output_dtype.is_floating_point + ): + raise ValueError("output_dtype must be a floating-point torch.dtype") + if inputs.scale is not None: + if ( + not isinstance(inputs.scale, (float, int)) + or isinstance(inputs.scale, bool) + or not math.isfinite(float(inputs.scale)) + or float(inputs.scale) <= 0 + ): + raise ValueError("scale must be a positive finite number") metadata = inputs.metadata batch, _, sq, head_dim = inputs.q.shape cache_capacity = inputs.k_cache.size(2) @@ -1274,8 +1326,6 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: 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") 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"}: @@ -1303,7 +1353,12 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> 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: + if ( + not isinstance(inputs.rope_theta, (float, int)) + or isinstance(inputs.rope_theta, bool) + or not math.isfinite(float(inputs.rope_theta)) + or 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 @@ -1313,6 +1368,12 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: inputs.k_cache_rope_output_dtype, torch.dtype ): raise ValueError("k_cache_rope_output_dtype must be a torch.dtype when provided") + for name, dtype in ( + ("q_rope_output_dtype", inputs.q_rope_output_dtype), + ("k_cache_rope_output_dtype", inputs.k_cache_rope_output_dtype), + ): + if dtype is not None and not dtype.is_floating_point: + raise ValueError(f"{name} must be a floating-point torch.dtype") if ( metadata.q_rope_state == "post_rope" and inputs.q_rope_output_dtype is not None @@ -1347,6 +1408,15 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: 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") + if bool((metadata.block_table[batch_index, block_count:] != -1).any()): + raise ValueError("unused block_table entries must be -1") + if metadata.cp_block_owners is not None: + active_owners = metadata.cp_block_owners[batch_index, :block_count] + inactive_owners = metadata.cp_block_owners[batch_index, block_count:] + if bool((active_owners < 0).any()): + raise ValueError("active cp_block_owners must be non-negative") + if bool((inactive_owners != -1).any()): + raise ValueError("unused cp_block_owners entries must be -1") 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 @@ -1420,7 +1490,7 @@ def _chunk_bounds(length: int, chunk_size: int) -> list[tuple[int, int]]: def _positive_int(value: int, name: str) -> int: - if isinstance(value, bool) or value <= 0: + if isinstance(value, bool) or not isinstance(value, Integral) or value <= 0: raise ValueError(f"{name} must be a positive integer") return int(value) diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 0733a321..bc21bd11 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -239,6 +239,15 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): 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["accum_dtype"] == "fp32" + assert drift.provenance["lse_dtype"] == "fp32" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.provenance["output_dtype"] == "float32" + assert drift.provenance["scale"] == pytest.approx(1.0 / (8.0**0.5)) + assert drift.provenance["q_dtype"] == "float32" + assert drift.provenance["k_cache_dtype"] == "float32" + assert drift.provenance["v_cache_dtype"] == "float32" + assert drift.provenance["scale"] == pytest.approx(1.0 / (8.0**0.5)) assert drift.provenance["logical_merge_orders"] == [ [[0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2]], @@ -371,6 +380,88 @@ def test_decode_replay_rejects_stale_prefix_cache_content(): run_decode_kv_replay(replace(inputs, k_cache=stale_k)) +def test_decode_replay_rejects_prefix_cache_theta_drift(): + inputs = _decode_inputs(prefix_cache_enabled=True) + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, rope_theta=10_000.0)) + + +def test_decode_replay_rejects_prefix_cache_storage_dtype_drift(): + inputs = _decode_inputs(prefix_cache_enabled=True) + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, k_cache=inputs.k_cache.to(torch.bfloat16))) + + +def test_decode_replay_rejects_unsupported_rope_cast_boundary(): + inputs = _decode_inputs(prefix_cache_enabled=True) + + with pytest.raises(ValueError, match="rope_cast_at"): + run_decode_kv_replay(replace(inputs, rope_cast_at="before_rope")) + + +@pytest.mark.parametrize("theta", [float("nan"), float("inf"), float("-inf")]) +def test_decode_replay_rejects_non_finite_rope_theta(theta): + with pytest.raises(ValueError, match="rope_theta"): + run_decode_kv_replay(replace(_decode_inputs(), rope_theta=theta)) + + +def test_decode_replay_rejects_noncanonical_inactive_page_metadata(): + inputs = _decode_inputs() + bad_metadata = replace( + inputs.metadata, + block_table=torch.cat( + [inputs.metadata.block_table, torch.zeros((2, 1), dtype=torch.long)], dim=1 + ), + cp_block_owners=torch.cat( + [inputs.metadata.cp_block_owners, torch.full((2, 1), -1, dtype=torch.long)], dim=1 + ), + ) + + with pytest.raises(ValueError, match="unused block_table"): + run_decode_kv_replay(replace(inputs, metadata=bad_metadata)) + + +def test_decode_replay_rejects_negative_active_cp_owner(): + inputs = _decode_inputs() + bad_owners = inputs.metadata.cp_block_owners.clone() + bad_owners[0, 1] = -1 + + with pytest.raises(ValueError, match="active cp_block_owners"): + run_decode_kv_replay( + replace(inputs, metadata=replace(inputs.metadata, cp_block_owners=bad_owners)) + ) + + +def test_decode_replay_rejects_noncanonical_inactive_cp_owner(): + inputs = _decode_inputs() + bad_metadata = replace( + inputs.metadata, + block_table=torch.cat( + [inputs.metadata.block_table, torch.full((2, 1), -1, dtype=torch.long)], dim=1 + ), + cp_block_owners=torch.cat( + [inputs.metadata.cp_block_owners, torch.zeros((2, 1), dtype=torch.long)], dim=1 + ), + ) + + with pytest.raises(ValueError, match="unused cp_block_owners"): + run_decode_kv_replay(replace(inputs, metadata=bad_metadata)) + + +@pytest.mark.parametrize("scale", [float("nan"), float("inf"), 0.0, -1.0]) +def test_decode_replay_rejects_invalid_scale(scale): + with pytest.raises(ValueError, match="scale"): + run_decode_kv_replay(replace(_decode_inputs(), scale=scale)) + + +def test_decode_replay_rejects_non_floating_cache_dtypes(): + inputs = _decode_inputs() + with pytest.raises(ValueError, match="floating-point dtypes"): + run_decode_kv_replay(replace(inputs, q=inputs.q.to(torch.int32))) + + def test_decode_replay_fails_loudly_on_position_identity_mismatch(): inputs = _decode_inputs() bad_query_positions = inputs.metadata.query_position_ids.clone() From f432b658ddeec85df36a5abbf1af66b7c6f85b77 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:17:00 +0800 Subject: [PATCH 07/10] fix(attention): harden decode cache execution identity Signed-off-by: lamentropetion <3051000145@qq.com> --- docs/design/ws2-attention-decode-replay.md | 13 +- rl_engine/testing/__init__.py | 2 + rl_engine/testing/attention_comparison.py | 103 +++++++++--- tests/test_attention_comparison.py | 182 ++++++++++++++++++++- 4 files changed, 270 insertions(+), 30 deletions(-) diff --git a/docs/design/ws2-attention-decode-replay.md b/docs/design/ws2-attention-decode-replay.md index c0c1fff1..60583d5b 100644 --- a/docs/design/ws2-attention-decode-replay.md +++ b/docs/design/ws2-attention-decode-replay.md @@ -24,10 +24,15 @@ layout: Metadata is validated before attention runs. Missing pages, duplicated active pages, non-canonical `-1` page/owner tails, out-of-range positions, mismatched RoPE positions, or inconsistent prefix-cache identity fail with an explicit -error. Prefix identity is verified by recomputing a physical-layout-invariant +error. Logical positions must be non-negative and strictly increasing, but may +start at a nonzero global offset (for example after a sliding-window eviction). +Prefix identity is verified by recomputing a physical-layout-invariant SHA-256 fingerprint over logical prefix positions, cached K/V content, storage dtypes, and the cache-side RoPE materialization configuration (`theta`, rotary -dimension, cast boundary, and output dtype). The current reference requires +dimension, cast boundary, and output dtype). The prefix key and enabled state are +part of that fingerprint. Every replay report also includes a physical-layout- +invariant full cache execution fingerprint over logical positions, page size, CP +ownership, prefix/RoPE identity, dtypes, and active K/V content. The current reference requires `rope_cast_at="after_rope"` and `rope_rotary_dim == head_dim`; partial rotary is not supported yet. Q and cached K retain separate RoPE output dtype contracts so mixed rollout query/KV storage dtypes are represented faithfully. @@ -70,7 +75,9 @@ The harness models CP block ownership on one device so cache construction, logical ordering, RoPE identity, and deterministic merging can be attributed without communication. It does not implement P2P NCCL transport, production custom CUDA all-gather/reduce-scatter, FlashInfer runtime execution, or training -backward. A distributed caller can gather the same partial states using the PR3 +backward. Reports explicitly record `single_device_logical_reference`, +`runtime_verified=false`, `supports_backward=false`, and `communication=none`; +they are not valid Megatron/vLLM runtime readbacks. A distributed caller can gather the same partial states using the PR3 transport layer; numerical merging must still happen in the fixed logical order described above. This PR targets the shared `test` integration branch and has a logical dependency on PR2/#253. diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 3897e009..0edbb5ef 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -15,6 +15,7 @@ compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + decode_kv_cache_fingerprint, decode_prefix_cache_fingerprint, run_chunked_query_attention, run_decode_full_prefill_reference, @@ -50,6 +51,7 @@ "compare_single_gpu_rope_attention", "compare_single_gpu_attention", "compare_decode_kv_replay", + "decode_kv_cache_fingerprint", "decode_prefix_cache_fingerprint", "compute_policy_ratio", "compute_reference_kl", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 30fb4838..95879517 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -317,6 +317,7 @@ def _run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> Attenti total_kv_len=int(visible.sum().item()), output_dtype=inputs.output_dtype, ) + _require_finite_decode_result(out, lse, path="full logical-KV reference") batch_out.append(out) batch_lse.append(lse) outs.append(torch.cat(batch_out, dim=2)) @@ -328,6 +329,10 @@ def _run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> Attenti provenance={ "attention_mode": "decode", "materialization": "full_logical_kv", + "execution_scope": "single_device_logical_reference", + "runtime_verified": False, + "supports_backward": False, + "communication": "none", "lse_domain": "attention", "accum_dtype": "fp32", "lse_dtype": "fp32", @@ -405,6 +410,7 @@ def _run_decode_kv_replay( 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) + _require_finite_decode_result(out, lse, path=f"{merge_backend} paged-KV replay") batch_out.append(out.to(inputs.output_dtype)) batch_lse.append(lse) batch_orders.append(order) @@ -415,6 +421,10 @@ def _run_decode_kv_replay( provenance: dict[str, Any] = { "attention_mode": "decode", "materialization": "paged_kv_replay", + "execution_scope": "single_device_logical_reference", + "runtime_verified": False, + "supports_backward": False, + "communication": "none", "sq": inputs.q.size(2), "page_size": inputs.metadata.page_size, "cache_position": inputs.metadata.cache_position.tolist(), @@ -448,6 +458,11 @@ def _run_decode_kv_replay( "q_dtype": str(inputs.q.dtype).replace("torch.", ""), "k_cache_dtype": str(inputs.k_cache.dtype).replace("torch.", ""), "v_cache_dtype": str(inputs.v_cache.dtype).replace("torch.", ""), + "cache_execution_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="full_logical_kv_cache", + ), } if merge_backend == "transformer_engine": provenance.update(_te_context_parallel_provenance()) @@ -689,19 +704,54 @@ def decode_prefix_cache_fingerprint( 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") + return _decode_cache_fingerprint( + inputs, + token_counts=(prefix_length,) * inputs.q.size(0), + domain="shared_prefix", + ) + + +def decode_kv_cache_fingerprint(inputs: DecodeAttentionInputs) -> str: + """Fingerprint complete logical cache execution identity after validation.""" + + _validate_decode_inputs(inputs) + return _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="full_logical_kv_cache", + ) + + +def _decode_cache_fingerprint( + inputs: DecodeAttentionInputs, + *, + token_counts: tuple[int, ...], + domain: str, +) -> str: + if len(token_counts) != inputs.q.size(0): + raise ValueError("cache fingerprint requires one token count per batch item") digest = hashlib.sha256() digest.update( ( + f"domain={domain};page_size={inputs.metadata.page_size};" + f"prefix_cache_enabled={inputs.metadata.prefix_cache_enabled};" + f"prefix_cache_key={inputs.metadata.prefix_cache_key!r};" + f"prefix_length={inputs.metadata.prefix_length};" f"k_rope_state={inputs.metadata.k_cache_rope_state};" f"rope_theta={float(inputs.rope_theta):.17g};" f"rotary_dim={_decode_rope_rotary_dim(inputs)};" f"rope_cast_at={inputs.rope_cast_at};" - f"k_rope_output_dtype={_decode_k_rope_output_dtype(inputs)}\n" + f"k_rope_output_dtype={_decode_k_rope_output_dtype(inputs)};" + f"k_dtype={inputs.k_cache.dtype};v_dtype={inputs.v_cache.dtype}\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 batch_index, token_count in enumerate(token_counts): + if token_count <= 0 or token_count > int(inputs.metadata.kv_seq_lens[batch_index].item()): + raise ValueError("cache fingerprint token count is outside kv_seq_lens") + slots = _decode_logical_slot_index(inputs, batch_index)[:token_count] + owner_count = math.ceil(token_count / inputs.metadata.page_size) + owners = _logical_block_owners(inputs, batch_index)[:owner_count] + digest.update(f"batch={batch_index};tokens={token_count};owners={owners}\n".encode()) for tensor in ( inputs.metadata.global_token_positions[batch_index, slots], inputs.metadata.key_position_ids[batch_index, slots], @@ -848,11 +898,7 @@ def _decode_rope_rotary_dim(inputs: DecodeAttentionInputs) -> int: def _decode_attention_scale(inputs: DecodeAttentionInputs) -> float: - return ( - 1.0 / math.sqrt(inputs.q.size(-1)) - if inputs.scale is None - else float(inputs.scale) - ) + return 1.0 / math.sqrt(inputs.q.size(-1)) if inputs.scale is None else float(inputs.scale) def _rope_output_dtype(inputs: AttentionComparisonInputs) -> torch.dtype: @@ -1418,6 +1464,12 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: if bool((inactive_owners != -1).any()): raise ValueError("unused cp_block_owners entries must be -1") slot_index = _decode_logical_slot_index(inputs, batch_index) + if not torch.isfinite(inputs.q[batch_index]).all(): + raise ValueError("active decode q values must be finite") + if not torch.isfinite(inputs.k_cache[batch_index, :, slot_index, :]).all(): + raise ValueError("active logical k_cache values must be finite") + if not torch.isfinite(inputs.v_cache[batch_index, :, slot_index, :]).all(): + raise ValueError("active logical v_cache values must be finite") 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()): @@ -1425,21 +1477,19 @@ def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: 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] - expected_positions = torch.arange( - sequence_length, - device=inputs.q.device, - dtype=global_positions.dtype, - ) - if not torch.equal(global_positions, expected_positions): + if bool((global_positions < 0).any()) or ( + sequence_length > 1 and bool((global_positions[1:] <= global_positions[:-1]).any()) + ): raise ValueError( "block_table/global_token_positions must reconstruct logical positions " - "0..kv_seq_len-1 exactly" + "in strictly increasing global order" ) 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 bool((cache_positions < 0).any()) or bool((cache_positions >= sequence_length).any()): + query_is_cached = (cache_positions[:, None] == global_positions[None, :]).any(dim=1) + if not bool(query_is_cached.all()): 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") @@ -1468,15 +1518,29 @@ def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: def _validate_partial_states(states: list[_PartialAttentionState]) -> None: first = states[0] + if first.block_start < 0 or first.block_end <= first.block_start: + raise ValueError("partial state ranges must satisfy 0 <= block_start < block_end") 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_end <= state.block_start: + raise ValueError("partial state block ranges must have positive width") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be contiguous and non-overlapping") previous_end = state.block_end +def _require_finite_decode_result( + out: torch.Tensor, + lse: torch.Tensor, + *, + path: str, +) -> None: + if not torch.isfinite(out).all() or not torch.isfinite(lse).all(): + raise ValueError(f"{path} produced non-finite attention output or LSE") + + def _chunk_bounds(length: int, chunk_size: int) -> list[tuple[int, int]]: if length <= 0: raise ValueError("sequence length must be positive") @@ -1507,6 +1571,7 @@ def _positive_int(value: int, name: str) -> int: "compare_single_gpu_rope_attention", "compare_single_gpu_attention", "compare_decode_kv_replay", + "decode_kv_cache_fingerprint", "decode_prefix_cache_fingerprint", "run_chunked_query_attention", "run_fused_like_rope_attention", diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index bc21bd11..10619101 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -24,6 +24,7 @@ compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + decode_kv_cache_fingerprint, decode_prefix_cache_fingerprint, run_decode_kv_replay, run_paged_kv_attention, @@ -71,6 +72,7 @@ def _decode_inputs( 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", + position_offset: int = 0, ) -> DecodeAttentionInputs: q, logical_k, logical_v = _qkv(seed=17) q = q[:, :, 4:6, :] @@ -84,19 +86,25 @@ def _decode_inputs( 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, + position_offset + logical_page * page_size, + position_offset + (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), + cache_position=torch.tensor( + [[position_offset + 4, position_offset + 5]] * 2, + 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), + query_position_ids=torch.tensor( + [[position_offset + 4, position_offset + 5]] * 2, + dtype=torch.long, + ), key_position_ids=positions.clone(), page_size=page_size, q_rope_state=q_rope_state, @@ -112,14 +120,23 @@ def _decode_inputs( if not prefix_cache_enabled: return inputs prefix_length = 4 - fingerprint = decode_prefix_cache_fingerprint(inputs, prefix_length=prefix_length) - return replace( + with_prefix_identity = replace( inputs, metadata=replace( inputs.metadata, prefix_cache_enabled=True, prefix_cache_key="shared-prefix", prefix_length=prefix_length, + ), + ) + fingerprint = decode_prefix_cache_fingerprint( + with_prefix_identity, + prefix_length=prefix_length, + ) + return replace( + with_prefix_identity, + metadata=replace( + with_prefix_identity.metadata, prefix_cache_fingerprint=fingerprint, ), ) @@ -247,11 +264,17 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): assert drift.provenance["q_dtype"] == "float32" assert drift.provenance["k_cache_dtype"] == "float32" assert drift.provenance["v_cache_dtype"] == "float32" + assert drift.provenance["execution_scope"] == "single_device_logical_reference" + assert drift.provenance["runtime_verified"] is False + assert drift.provenance["supports_backward"] is False + assert drift.provenance["communication"] == "none" + assert drift.provenance["cache_execution_fingerprint"] == decode_kv_cache_fingerprint(inputs) assert drift.provenance["scale"] == pytest.approx(1.0 / (8.0**0.5)) assert drift.provenance["logical_merge_orders"] == [ [[0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2]], ] + json.dumps(report.to_dict()) single_query = DecodeAttentionInputs( q=inputs.q[:, :, -1:, :], @@ -275,7 +298,8 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): 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)) + permuted_inputs = _decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True) + permuted = run_decode_kv_replay(permuted_inputs) 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) @@ -283,8 +307,150 @@ def test_decode_replay_is_invariant_to_physical_page_and_prefix_layout(): 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 + permuted_inputs, + prefix_length=permuted_inputs.metadata.prefix_length, + ) + + +def test_decode_replay_supports_nonzero_global_position_origin(): + base = compare_decode_kv_replay(_decode_inputs()) + offset = compare_decode_kv_replay(_decode_inputs(position_offset=128)) + + assert offset.drifts[0].out.max_abs <= 1.0e-6 + assert offset.drifts[0].lse.max_abs <= 1.0e-6 + assert base.drifts[0].out.max_abs <= 1.0e-6 + + +def test_decode_cache_execution_fingerprint_is_layout_invariant_and_content_bound(): + contiguous = _decode_inputs() + permuted = _decode_inputs(page_order=(2, 0, 1)) + + assert decode_kv_cache_fingerprint(contiguous) == decode_kv_cache_fingerprint(permuted) + assert run_decode_kv_replay(contiguous).provenance[ + "cache_execution_fingerprint" + ] == decode_kv_cache_fingerprint(contiguous) + + changed_v = contiguous.v_cache.clone() + first_slot = int(contiguous.metadata.block_table[0, 0].item()) * contiguous.metadata.page_size + changed_v[0, 0, first_slot, 0] += 1.0 + changed = replace(contiguous, v_cache=changed_v) + assert decode_kv_cache_fingerprint(changed) != decode_kv_cache_fingerprint(contiguous) + + contiguous_prefix = _decode_inputs(prefix_cache_enabled=True) + permuted_prefix = _decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True) + assert contiguous_prefix.metadata.prefix_cache_fingerprint == ( + permuted_prefix.metadata.prefix_cache_fingerprint + ) + + +def test_decode_cache_fingerprint_ignores_inactive_physical_page_content(): + inputs = _decode_inputs() + generator = torch.Generator().manual_seed(91) + extra_k = torch.randn( + inputs.k_cache.size(0), + inputs.k_cache.size(1), + inputs.metadata.page_size, + inputs.k_cache.size(3), + generator=generator, + ) + extra_v = torch.randn(extra_k.shape, generator=generator) + extended = replace( + inputs, + k_cache=torch.cat((inputs.k_cache, extra_k), dim=2), + v_cache=torch.cat((inputs.v_cache, extra_v), dim=2), + metadata=replace( + inputs.metadata, + block_table=torch.cat( + (inputs.metadata.block_table, torch.full((2, 1), -1, dtype=torch.long)), + dim=1, + ), + global_token_positions=torch.cat( + ( + inputs.metadata.global_token_positions, + torch.full((2, inputs.metadata.page_size), -1, dtype=torch.long), + ), + dim=1, + ), + key_position_ids=torch.cat( + ( + inputs.metadata.key_position_ids, + torch.full((2, inputs.metadata.page_size), -1, dtype=torch.long), + ), + dim=1, + ), + cp_block_owners=torch.cat( + ( + inputs.metadata.cp_block_owners, + torch.full((2, 1), -1, dtype=torch.long), + ), + dim=1, + ), + ), + ) + changed_inactive = replace( + extended, + k_cache=extended.k_cache.clone(), + v_cache=extended.v_cache.clone(), ) + changed_inactive.k_cache[:, :, -2:, :].fill_(123.0) + changed_inactive.v_cache[:, :, -2:, :].fill_(-456.0) + + assert decode_kv_cache_fingerprint(extended) == decode_kv_cache_fingerprint(changed_inactive) + extended_result = run_decode_kv_replay(extended) + changed_result = run_decode_kv_replay(changed_inactive) + torch.testing.assert_close(extended_result.out, changed_result.out) + torch.testing.assert_close(extended_result.lse, changed_result.lse) + + +def test_decode_replay_supports_per_batch_lengths_and_partial_final_page(): + inputs = _decode_inputs() + metadata = replace( + inputs.metadata, + cache_position=torch.tensor([[3, 4], [4, 5]], dtype=torch.long), + query_position_ids=torch.tensor([[3, 4], [4, 5]], dtype=torch.long), + kv_seq_lens=torch.tensor([5, 6], dtype=torch.long), + global_token_positions=inputs.metadata.global_token_positions.clone(), + key_position_ids=inputs.metadata.key_position_ids.clone(), + ) + first_unused_slot = int(metadata.block_table[0, 2].item()) * metadata.page_size + 1 + metadata.global_token_positions[0, first_unused_slot] = -1 + metadata.key_position_ids[0, first_unused_slot] = -1 + shortened = replace(inputs, metadata=metadata) + + report = compare_decode_kv_replay(shortened) + + assert report.drifts[0].out.max_abs <= 1.0e-6 + assert report.drifts[0].lse.max_abs <= 1.0e-6 + assert report.drifts[0].provenance["kv_seq_lens"] == [5, 6] + + +@pytest.mark.parametrize("field_name", ["q", "k_cache", "v_cache"]) +def test_decode_replay_rejects_non_finite_active_values(field_name): + inputs = _decode_inputs() + tensor = getattr(inputs, field_name).clone() + tensor[0, 0, 0, 0] = float("nan") + + with pytest.raises(ValueError, match="must be finite"): + run_decode_kv_replay(replace(inputs, **{field_name: tensor})) + + +def test_decode_replay_rejects_finite_inputs_that_overflow_attention(): + inputs = _decode_inputs() + huge = torch.full_like(inputs.q, torch.finfo(inputs.q.dtype).max) + + with pytest.raises(ValueError, match="produced non-finite"): + compare_decode_kv_replay(replace(inputs, q=huge)) + + +def test_decode_replay_rejects_prefix_cache_key_drift(): + inputs = _decode_inputs(prefix_cache_enabled=True) + changed_key = replace( + inputs, + metadata=replace(inputs.metadata, prefix_cache_key="different-prefix"), + ) + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(changed_key) def test_decode_replay_is_invariant_to_equivalent_cp_block_ownership(): From a9c7506256e1a56a7d75dd8fac4a4298d2a01f4e Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:18:35 +0800 Subject: [PATCH 08/10] test(attention): cover decode cache identity edge cases Signed-off-by: lamentropetion <3051000145@qq.com> --- tests/test_attention_comparison.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 10619101..628778e8 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -269,7 +269,6 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): assert drift.provenance["supports_backward"] is False assert drift.provenance["communication"] == "none" assert drift.provenance["cache_execution_fingerprint"] == decode_kv_cache_fingerprint(inputs) - assert drift.provenance["scale"] == pytest.approx(1.0 / (8.0**0.5)) assert drift.provenance["logical_merge_orders"] == [ [[0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2]], From 8730fd2b8908eb10f463364c77b9abd29819439a Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:48:04 +0800 Subject: [PATCH 09/10] test(attention): bind decode cache and rope identity --- rl_engine/testing/__init__.py | 2 ++ rl_engine/testing/attention_comparison.py | 44 +++++++++++++++++++++++ tests/test_attention_comparison.py | 5 +++ 3 files changed, 51 insertions(+) diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 0edbb5ef..0386e374 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -16,6 +16,7 @@ compare_single_gpu_attention, compare_single_gpu_rope_attention, decode_kv_cache_fingerprint, + decode_rope_identity_fingerprint, decode_prefix_cache_fingerprint, run_chunked_query_attention, run_decode_full_prefill_reference, @@ -52,6 +53,7 @@ "compare_single_gpu_attention", "compare_decode_kv_replay", "decode_kv_cache_fingerprint", + "decode_rope_identity_fingerprint", "decode_prefix_cache_fingerprint", "compute_policy_ratio", "compute_reference_kl", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 95879517..2a9216c8 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -463,6 +463,26 @@ def _run_decode_kv_replay( token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), domain="full_logical_kv_cache", ), + "kv_cache_identity_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="full_logical_kv_cache", + ), + "rope_identity_fingerprint": decode_rope_identity_fingerprint(inputs), + "preprocess_policy": "reference_only_not_production", + "preprocess_backends": { + "qk_rmsnorm": "not_executed_projected_qk_input", + "rope": ( + "rlkernel.pytorch.rope_reference" + if inputs.metadata.q_rope_state == "pre_rope" + or inputs.metadata.k_cache_rope_state == "pre_rope" + else "external_post_rope_input" + ), + }, + "preprocess_fallback": True, + "preprocess_fallback_reason": ( + "decode logical replay is a single-device correctness reference, not native vLLM" + ), } if merge_backend == "transformer_engine": provenance.update(_te_context_parallel_provenance()) @@ -722,6 +742,29 @@ def decode_kv_cache_fingerprint(inputs: DecodeAttentionInputs) -> str: ) +def decode_rope_identity_fingerprint(inputs: DecodeAttentionInputs) -> str: + """Fingerprint every RoPE fact that changes decode Q/K interpretation.""" + + _validate_decode_inputs(inputs) + digest = hashlib.sha256() + digest.update( + ( + f"q_rope_state={inputs.metadata.q_rope_state};" + f"k_cache_rope_state={inputs.metadata.k_cache_rope_state};" + f"rope_theta={float(inputs.rope_theta):.17g};" + f"rotary_dim={_decode_rope_rotary_dim(inputs)};" + f"rope_cast_at={inputs.rope_cast_at};" + f"q_rope_output_dtype={_decode_q_rope_output_dtype(inputs)};" + f"k_rope_output_dtype={_decode_k_rope_output_dtype(inputs)}\n" + ).encode() + ) + for tensor in (inputs.metadata.query_position_ids, inputs.metadata.key_position_ids): + 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 _decode_cache_fingerprint( inputs: DecodeAttentionInputs, *, @@ -1572,6 +1615,7 @@ def _positive_int(value: int, name: str) -> int: "compare_single_gpu_attention", "compare_decode_kv_replay", "decode_kv_cache_fingerprint", + "decode_rope_identity_fingerprint", "decode_prefix_cache_fingerprint", "run_chunked_query_attention", "run_fused_like_rope_attention", diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 628778e8..cc8a92ac 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -25,6 +25,7 @@ compare_single_gpu_attention, compare_single_gpu_rope_attention, decode_kv_cache_fingerprint, + decode_rope_identity_fingerprint, decode_prefix_cache_fingerprint, run_decode_kv_replay, run_paged_kv_attention, @@ -269,6 +270,10 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): assert drift.provenance["supports_backward"] is False assert drift.provenance["communication"] == "none" assert drift.provenance["cache_execution_fingerprint"] == decode_kv_cache_fingerprint(inputs) + assert drift.provenance["kv_cache_identity_fingerprint"] == decode_kv_cache_fingerprint(inputs) + assert drift.provenance["rope_identity_fingerprint"] == decode_rope_identity_fingerprint(inputs) + assert drift.provenance["preprocess_policy"] == "reference_only_not_production" + assert drift.provenance["preprocess_fallback"] is True assert drift.provenance["logical_merge_orders"] == [ [[0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2]], From 3755c5caa3ebfde162378c1bc7d6c3801d47328d Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 23:25:29 +0800 Subject: [PATCH 10/10] test(attention): verify strict decode layout invariance --- rl_engine/testing/__init__.py | 2 + rl_engine/testing/attention_comparison.py | 121 +++++++++++++++++++++- tests/test_attention_comparison.py | 17 +++ 3 files changed, 138 insertions(+), 2 deletions(-) diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 24c0a75a..602c875b 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -21,6 +21,7 @@ run_chunked_query_attention, run_decode_full_prefill_reference, run_decode_kv_replay, + run_decode_strict_shared_core, run_full_attention, run_fused_like_rope_attention, run_paged_kv_attention, @@ -90,6 +91,7 @@ "run_chunked_query_attention", "run_decode_full_prefill_reference", "run_decode_kv_replay", + "run_decode_strict_shared_core", "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 2a9216c8..def7a449 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -263,12 +263,19 @@ def compare_decode_kv_replay( inputs: DecodeAttentionInputs, *, include_transformer_engine: bool = False, + strict_bitwise: 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")] + if strict_bitwise: + reference = _run_decode_strict_shared_core(inputs, materialization="logical_prefill") + candidates = [ + _run_decode_strict_shared_core(inputs, materialization="paged_kv_layout") + ] + else: + 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: @@ -283,6 +290,86 @@ def compare_decode_kv_replay( ) +def run_decode_strict_shared_core( + inputs: DecodeAttentionInputs, +) -> AttentionPathResult: + """Run decode through the same fixed one-query arithmetic schedule.""" + + _validate_decode_inputs(inputs) + return _run_decode_strict_shared_core(inputs, materialization="paged_kv_layout") + + +def _run_decode_strict_shared_core( + inputs: DecodeAttentionInputs, + *, + materialization: str, +) -> 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 + if not bool(visible.any()): + raise ValueError("each decode query must have at least one visible KV token") + out, lse = _strict_decode_attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, visible, :], + v[:, :, visible, :], + 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=f"strict_shared_core_{materialization}", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance={ + "attention_mode": "decode", + "materialization": materialization, + "execution_scope": "single_device_strict_shared_core", + "runtime_verified": False, + "actual_backend": "rlkernel.pytorch.strict_attention_reference", + "communication_backend": "none", + "production_ready": False, + "supports_backward": False, + "communication": "none", + "strict_mode": True, + "strict_core_id": "rlkernel.attention.deterministic_core.v1", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "native_attention_arithmetic": False, + "fallback": False, + "fallback_reason": None, + "split_kv_policy": "disabled", + "merge_order": "global_block_index", + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + "cache_execution_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="strict_shared_core_decode", + ), + "kv_cache_identity_fingerprint": _decode_cache_fingerprint( + inputs, + token_counts=tuple(int(length) for length in inputs.metadata.kv_seq_lens.tolist()), + domain="strict_shared_core_decode", + ), + "rope_identity_fingerprint": decode_rope_identity_fingerprint(inputs), + "query_position_ids": inputs.metadata.query_position_ids.tolist(), + "key_position_ids": inputs.metadata.key_position_ids.tolist(), + "block_table": inputs.metadata.block_table.tolist(), + "page_size": inputs.metadata.page_size, + }, + ) + + def run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: """Materialize the full logical KV sequence for each decode query. @@ -978,6 +1065,36 @@ def _rope_attention_provenance( } +def _strict_decode_attention_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + output_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Canonical one-query score/softmax/value schedule for decode replay.""" + + qf, kf, vf = q.float(), k.float(), v.float() + hq, hkv = qf.size(1), kf.size(1) + if hq % hkv: + raise ValueError("strict decode Attention requires GQA-compatible head counts") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + vf = vf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) / math.sqrt(qf.size(-1)) + row_max = scores.amax(dim=-1, keepdim=True) + finite = torch.isfinite(row_max) + exp_scores = torch.where(finite, torch.exp(scores - row_max), torch.zeros_like(scores)) + row_sum = exp_scores.sum(dim=-1, keepdim=True) + lse = torch.where( + row_sum > 0, + row_max + torch.log(row_sum), + torch.full_like(row_sum, float("-inf")), + ) + weights = torch.where(row_sum > 0, exp_scores / row_sum, torch.zeros_like(exp_scores)) + return torch.matmul(weights, vf).to(output_dtype), lse.squeeze(-1) + + def _attention_with_lse( q: torch.Tensor, k: torch.Tensor, diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index cc8a92ac..f26f5f66 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -300,6 +300,23 @@ def test_decode_replay_matches_full_prefill_for_single_and_few_query(): assert single_report.drifts[0].lse.max_abs <= 1.0e-6 +def test_strict_decode_replay_is_bitwise_across_logical_and_paged_materialization(): + report = compare_decode_kv_replay(_decode_inputs(page_order=(2, 0, 1)), strict_bitwise=True) + assert report.reference_name == "strict_shared_core_logical_prefill" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.candidate_name == "strict_shared_core_paged_kv_layout" + assert drift.out.max_abs == 0.0 + assert drift.lse.max_abs == 0.0 + assert drift.provenance["strict_core_id"] == ( + "rlkernel.attention.deterministic_core.v1" + ) + assert drift.provenance["strict_schedule"] == ( + "single_batch_single_query_global_kv_blocks" + ) + assert drift.provenance["split_kv_policy"] == "disabled" + + def test_decode_replay_is_invariant_to_physical_page_and_prefix_layout(): contiguous = run_decode_kv_replay(_decode_inputs()) permuted_inputs = _decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True)