From ed05a09f66f12fa85e2cb7f5dbf68410178ecb41 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 20 Jul 2026 22:59:48 +0800 Subject: [PATCH 01/16] feat(attention): add deterministic CP reference Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 22 + rl_engine/kernels/gtest/operator_inputs.py | 22 + .../ops/pytorch/attention/cp_attention.py | 548 ++++++++++++++++++ rl_engine/kernels/registry.py | 9 + tests/test_cp_attention.py | 385 ++++++++++++ tests/test_operator_inputs.py | 1 + 6 files changed, 987 insertions(+) create mode 100644 rl_engine/kernels/ops/pytorch/attention/cp_attention.py create mode 100644 tests/test_cp_attention.py diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3cb5f9b..97dfed29 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -83,6 +83,15 @@ path. When fused attention kernels land, they are prepended to the priority list op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. +`kernel_registry.get_op("cp_attention")` resolves to +`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel +reference. It emulates CP prefill and chunked-prefill by splitting logical query +and KV sequence blocks, computing per-block `(out, lse)` partial states, and +merging them in fp32 by global KV block index. This path is not a production +fused backend; it defines the CP/LSE merge behavior that downstream fused paths +must match. Optional per-batch `query_position_offsets` / `key_position_offsets` +cover varlen causal-mask metadata while keeping the dense tensor layout. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -135,6 +144,7 @@ memory. ```bash python -m pytest tests/test_attention.py -v +python -m pytest tests/test_cp_attention.py -v ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -144,15 +154,27 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below), gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. +`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard +attention, CP=2 prefill vs CP=1, chunked-prefill replay, global-position causal +masking across CP boundaries, order-independent LSE merge by global block index, +padding/all-masked stability, BF16 final-write behavior, input purity, argument +validation, and registry dispatch. +`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill +synthetic case for local harnesses. + ## Implementation Files - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` +- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` - `rl_engine/kernels/registry.py` - `tests/test_attention.py` +- `tests/test_cp_attention.py` ## Known Limitations - PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). +- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, + not a distributed runtime or fused kernel. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - The naive path materializes the full `[B, Hq, Sq, Skv]` scores tensor — no query-chunking, so the LARGE load point is memory-heavy and GPU-only. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f124cafb..8d26ac04 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -28,6 +28,7 @@ def make_operator_inputs( "rms_norm": _make_rms_norm_inputs, "matmul": _make_matmul_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "rope": _make_rope_inputs, @@ -50,6 +51,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", @@ -107,6 +109,26 @@ def _make_attention_inputs( } +def _make_cp_attention_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + return { + "q": _floating_tensor( + (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0 + ), + "k": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1 + ), + "v": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2 + ), + "causal": True, + "cp_world_size": 2, + "kv_chunk_size": max(1, seq // 2), + } + + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..c658b134 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,548 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(q.dtype if output_dtype is None else output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + """ + + _validate_qkv(q, k, v) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = 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) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + 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_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, 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 AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> 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 out/lse shapes") + if state.block_start < previous_end: + raise ValueError("partial state block ranges must not overlap") + previous_end = state.block_end + + +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 the same 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)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionPartialState", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "merge_attention_partial_states", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..d463bb5b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -72,6 +72,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + # WS2 correctness-first context-parallel attention reference. It emulates + # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. + PYTORCH_CP_ATTENTION = ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ) # WS1 pure-PyTorch ground-truth linear ops PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops @@ -165,6 +171,7 @@ def __init__(self): ], "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], @@ -188,6 +195,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], @@ -206,6 +214,7 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..dd0f53f7 --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + merge_attention_partial_states, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 + + +@contextlib.contextmanager +def _single_thread(): + prev = torch.get_num_threads() + torch.set_num_threads(1) + try: + yield + finally: + torch.set_num_threads(prev) + + +def _qkv( + batch, + sq, + skv, + *, + seed, + dtype=torch.float32, + heads=_N_HEADS, + kv_heads=_N_KV, + dim=_HEAD_DIM, +): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(batch, heads, sq, dim, generator=gen, dtype=dtype) + k = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + v = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + return q, k, v + + +def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): + qf, kf = q.float(), k.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hq % hkv != 0: + raise ValueError("invalid GQA shape") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * ( + scale if scale is not None else 1.0 / math.sqrt(dim) + ) + if causal: + query_pos = torch.arange(skv - sq, skv) + key_pos = torch.arange(skv) + scores = scores.masked_fill( + (key_pos[None, :] > query_pos[:, None])[None, None, :, :], + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + return torch.logsumexp(scores, dim=-1) + + +def test_cp1_matches_native_attention_and_exports_lse(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 8, 8, seed=1) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + want = native.forward_fp32(q, k, v, causal=True) + want_lse = _full_lse(q, k, causal=True) + + torch.testing.assert_close(out, want, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse, want_lse, atol=_ATOL, rtol=0.0) + assert lse.dtype == torch.float32 + assert lse.shape == q.shape[:3] + + +def test_cp2_prefill_matches_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 9, 9, seed=2) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + out2, lse2 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=2) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_chunked_prefill_replay_matches_unchunked_cp2(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 10, 10, seed=3) + + with _single_thread(): + unchunked_out, unchunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + ) + chunked_out, chunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=3, + ) + + torch.testing.assert_close(chunked_out, unchunked_out, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(chunked_lse, unchunked_lse, atol=_ATOL, rtol=0.0) + + +def test_causal_mask_uses_global_positions_across_cp_boundary(): + op = DeterministicCPAttentionReferenceOp() + batch, heads, kv_heads, seq, dim = 1, 2, 1, 5, 3 + q = torch.zeros(batch, heads, seq, dim) + k = torch.zeros(batch, kv_heads, seq, dim) + v = torch.arange(seq * dim, dtype=torch.float32).reshape(1, 1, seq, dim) + out = op.forward_fp32(q, k, v, causal=True, cp_world_size=2) + + expected = torch.stack([v[0, 0, : index + 1].mean(dim=0) for index in range(seq)]) + expected = expected.reshape(1, 1, seq, dim).repeat(1, heads, 1, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + + +def test_position_offsets_apply_varlen_causal_metadata_per_batch_row(): + op = DeterministicCPAttentionReferenceOp() + q = torch.zeros(2, 2, 2, 1) + k = torch.zeros(2, 1, 4, 1) + v = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4, 1) + + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=torch.tensor([0, 11]), + key_position_offsets=torch.tensor([0, 10]), + cp_world_size=2, + kv_chunk_size=1, + ) + + expected = torch.tensor([0.0, 0.5, 4.5, 5.0]).reshape(2, 1, 2, 1).repeat(1, 2, 1, 1) + expected_lse = torch.log(torch.tensor([1.0, 2.0, 2.0, 3.0])).reshape(2, 1, 2) + expected_lse = expected_lse.repeat(1, 2, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(lse, expected_lse, atol=1.0e-6, rtol=0.0) + + +def test_merge_order_uses_global_block_index_not_arrival_order(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 6, 6, seed=4) + first = op.local_partial_state( + q, + k[:, :, :3], + v[:, :, :3], + q_start=0, + k_start=0, + total_kv_len=6, + causal=True, + ) + second = op.local_partial_state( + q, + k[:, :, 3:], + v[:, :, 3:], + q_start=0, + k_start=3, + total_kv_len=6, + causal=True, + ) + + forward = merge_attention_partial_states([first, second]) + reversed_arrival = merge_attention_partial_states([second, first]) + assert torch.equal(forward.out, reversed_arrival.out) + assert torch.equal(forward.lse, reversed_arrival.lse) + + +def test_key_padding_mask_and_all_masked_rows_are_stable(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 6, 6, seed=5) + mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, False, False, False, False, False], + ], + dtype=torch.bool, + ) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=False, + key_padding_mask=mask, + cp_world_size=2, + kv_chunk_size=2, + ) + want = native.forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + torch.testing.assert_close(out[:1], want[:1], atol=_ATOL, rtol=0.0) + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + assert torch.isfinite(out).all() + + +def test_empty_query_and_empty_kv_edges_are_stable(): + op = DeterministicCPAttentionReferenceOp() + q_empty = torch.randn(1, 2, 0, 4, requires_grad=True) + k_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + v_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + out, lse = op.forward_fp32_with_lse(q_empty, k_empty, v_empty, cp_world_size=2) + assert out.shape == (1, 2, 0, 4) + assert lse.shape == (1, 2, 0) + assert out.requires_grad + out.sum().backward() + assert torch.equal(q_empty.grad, torch.zeros_like(q_empty)) + assert torch.equal(k_empty.grad, torch.zeros_like(k_empty)) + assert torch.equal(v_empty.grad, torch.zeros_like(v_empty)) + + q = torch.randn(1, 2, 3, 4) + out, lse = op.forward_fp32_with_lse(q, k_empty, v_empty, causal=False, cp_world_size=4) + assert torch.equal(out, torch.zeros_like(out)) + assert torch.isneginf(lse).all() + + +def test_empty_kv_backward_returns_zero_grads(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 2, 3, 4, requires_grad=True) + k = torch.randn(1, 1, 0, 4, requires_grad=True) + v = torch.randn(1, 1, 0, 4, requires_grad=True) + + out = op.forward_fp32(q, k, v, causal=False, cp_world_size=4) + assert out.requires_grad + out.sum().backward() + + assert torch.equal(q.grad, torch.zeros_like(q)) + assert torch.equal(k.grad, torch.zeros_like(k)) + assert torch.equal(v.grad, torch.zeros_like(v)) + + +def test_bf16_forward_uses_fp32_merge_then_final_write(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 8, 8, seed=6, dtype=torch.bfloat16) + + out, lse = op.forward_with_lse(q, k, v, causal=True, cp_world_size=2, kv_chunk_size=2) + fp32_out, fp32_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert torch.equal(out, fp32_out.to(torch.bfloat16)) + assert torch.equal(lse, fp32_lse) + + +def test_cp2_chunked_gradients_match_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 5, 5, seed=12, heads=4, kv_heads=2, dim=8) + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + q_cp = q.detach().clone().requires_grad_(True) + k_cp = k.detach().clone().requires_grad_(True) + v_cp = v.detach().clone().requires_grad_(True) + gen = torch.Generator().manual_seed(13) + dy = torch.randn(1, 4, 5, 8, generator=gen) + + with _single_thread(): + out_ref = op.forward_fp32(q_ref, k_ref, v_ref, causal=True, cp_world_size=1) + out_cp = op.forward_fp32( + q_cp, + k_cp, + v_cp, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + out_ref.backward(dy) + out_cp.backward(dy) + + torch.testing.assert_close(out_cp, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(q_cp.grad, q_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(k_cp.grad, k_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) + + +def test_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index bb1a2220..9e3eac34 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -36,6 +36,7 @@ def _args(**overrides): "rms_norm", "matmul", "attention", + "cp_attention", "logp", "linear_logp", "rope", From 0480ce81fe2cf8a543f6aa6613a6ec2cf8550ede Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 2 Aug 2026 23:40:00 +0800 Subject: [PATCH 02/16] docs(attention): clarify rope boundary for cp reference Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 13 ++++++- .../ops/pytorch/attention/cp_attention.py | 8 ++++ tests/test_cp_attention.py | 38 +++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 9aac608e..988fad8c 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -94,6 +94,12 @@ fused backend; it defines the CP/LSE merge behavior that downstream fused paths must match. Optional per-batch `query_position_offsets` / `key_position_offsets` cover varlen causal-mask metadata while keeping the dense tensor layout. +For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not +call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The +position offsets passed to CP attention must describe the same absolute token +positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary +while PR7 can later validate production fused `RoPE+Attention` kernels. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -158,8 +164,9 @@ gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. `tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard -attention, CP=2 prefill vs CP=1, chunked-prefill replay, global-position causal -masking across CP boundaries, order-independent LSE merge by global block index, +attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared +global position metadata, chunked-prefill replay, global-position causal masking +across CP boundaries, order-independent LSE merge by global block index, padding/all-masked stability, BF16 final-write behavior, input purity, argument validation, and registry dispatch. `make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill @@ -246,6 +253,8 @@ for measured peak memory at representative shapes. - Full materialization of scores/P limits practical sequence length. - `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, not a distributed runtime or fused kernel. +- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused + `RoPE+Attention` backend alignment are outside PR3. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 734d97b9..302ad73f 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -84,6 +84,12 @@ def merge_attention_partial_states( class DeterministicCPAttentionReferenceOp: """Correctness-first CP attention reference for prefill and chunked prefill. + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + The op emulates CP by splitting query and KV sequence dimensions into logical CP shards. Each query shard computes one partial attention state per KV block, then merges those states in fixed global-block order using fp32 @@ -267,6 +273,8 @@ def local_partial_state( ``query_position_offsets`` and ``key_position_offsets`` are optional per-batch-row base positions. They let the reference express varlen or packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. """ _validate_qkv(q, k, v) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index dd0f53f7..03c861dd 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -19,6 +19,7 @@ merge_attention_partial_states, ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.kernels.registry import kernel_registry _N_HEADS = 32 @@ -106,6 +107,43 @@ def test_cp2_prefill_matches_cp1_reference(): torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) +def test_cp2_consumes_post_rope_qk_with_shared_global_position_metadata(): + op = DeterministicCPAttentionReferenceOp() + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 7, 7, seed=14, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([17, 103], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + + assert not torch.equal(q, pre_rope_q.float()) + assert not torch.equal(k, pre_rope_k.float()) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=1, + ) + out2, lse2 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + def test_chunked_prefill_replay_matches_unchunked_cp2(): op = DeterministicCPAttentionReferenceOp() q, k, v = _qkv(2, 10, 10, seed=3) From 9f46efff100e8b2067af979ec964c898bfdeeb12 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:27:36 +0800 Subject: [PATCH 03/16] fix(attention): harden CP reference Split-KV validation --- docs/operators/attention.md | 64 +- rl_engine/kernels/attention_contract.py | 1472 +++++++++++++++++ .../ops/pytorch/attention/cp_attention.py | 486 +++++- tests/test_cp_attention.py | 213 +++ 4 files changed, 2198 insertions(+), 37 deletions(-) create mode 100644 rl_engine/kernels/attention_contract.py diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 988fad8c..bc7f4a38 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -85,20 +85,31 @@ Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_ the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. -`kernel_registry.get_op("cp_attention")` resolves to -`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel -reference. It emulates CP prefill and chunked-prefill by splitting logical query -and KV sequence blocks, computing per-block `(out, lse)` partial states, and -merging them in fp32 by global KV block index. This path is not a production -fused backend; it defines the CP/LSE merge behavior that downstream fused paths -must match. Optional per-batch `query_position_offsets` / `key_position_offsets` -cover varlen causal-mask metadata while keeping the dense tensor layout. - -For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not -call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The -position offsets passed to CP attention must describe the same absolute token -positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary -while PR7 can later validate production fused `RoPE+Attention` kernels. +### WS2 CP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed +`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity +before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. + +Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic +CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as +a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). + +Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow +`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block +boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. +Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and +validate the same actual plan. + +The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ + --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ + --output artifacts/ws2-cp-attention-drift.json +``` ## Accuracy @@ -153,7 +164,6 @@ memory. ```bash python -m pytest tests/test_attention.py -v python -m pytest tests/test_cp_attention.py -v -python -m pytest tests/test_cp_attention_transformer_engine.py -v # optional TE oracle ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -163,30 +173,14 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below), gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. -`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard -attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared -global position metadata, chunked-prefill replay, global-position causal masking -across CP boundaries, order-independent LSE merge by global block index, -padding/all-masked stability, BF16 final-write behavior, input purity, argument -validation, and registry dispatch. -`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill -synthetic case for local harnesses. -`tests/test_cp_attention_transformer_engine.py` optionally imports NVIDIA -Transformer Engine's context-parallel PyTorch correction helpers and checks that -RL-Kernel's fp32 `(out, lse)` merge matches those helpers; the test skips when -Transformer Engine is not installed. - ## Implementation Files - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference -- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` — CP prefill/chunked reference - `rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` — CUDA deterministic op - `csrc/cuda/attention/deterministic_attention.cu` — CUDA kernels - `rl_engine/kernels/registry.py` - `tests/test_attention.py` - `tests/test_deterministic_attention_cuda.py` -- `tests/test_cp_attention.py` -- `tests/test_cp_attention_transformer_engine.py` ## Fixed Reduction Order (CUDA Deterministic Backend) @@ -227,6 +221,10 @@ Hooks: - `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable. - `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging, and future KV-cache / training integration. +- `backward_reference(q, k, v, dout, ...)` — runs the deterministic training backward + validation path and returns `dq`, `dk`, `dv`, `out`, `lse`, and provenance. +- `compare_cp_attention_backward(q, k, v, dout, ...)` — compares CP=1 backward against + CP/chunked-prefill backward and emits whole-tensor plus per-logical-rank drift stats. ## Tolerance @@ -251,10 +249,6 @@ for measured peak memory at representative shapes. - First version: `D=128` only (Qwen3-8B alignment). - Supported dtypes: BF16, FP16. - Full materialization of scores/P limits practical sequence length. -- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, - not a distributed runtime or fused kernel. -- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused - `RoPE+Attention` backend alignment are outside PR3. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..eb4994b3 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1472 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class SplitKVMode(str, Enum): + DISABLED = "disabled" + FIXED = "fixed" + AUTO = "auto" + + +class RoPEState(str, Enum): + PRE_ROPE = "pre_rope" + POST_ROPE = "post_rope" + + +class RoPECastPoint(str, Enum): + NONE = "none" + BEFORE_ROPE = "before_rope" + AFTER_ROPE = "after_rope" + + +class RoPEFusionBoundary(str, Enum): + UNFUSED_ROPE_ATTENTION = "unfused_rope_attention" + FUSED_ROPE_ATTENTION = "fused_rope_attention" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise AttentionContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _integer_tuple(values: Iterable[int], field: str) -> tuple[int, ...]: + try: + result = tuple(values) + except TypeError as exc: + raise AttentionContractError(f"{field} must be an iterable of integers") from exc + for index, value in enumerate(result): + if isinstance(value, bool) or not isinstance(value, int): + raise AttentionContractError(f"{field}[{index}] must be an integer; got {value!r}") + return result + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical TP/CP ownership for one attention invocation. + + TP head shards are currently required to be equal and contiguous. CP + sequence ownership may be uneven, but every local block must carry a stable + logical global index so a later implementation can merge by logical order + instead of collective arrival order. + """ + + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + global_q_heads: int + global_kv_heads: int + local_q_head_start: int + local_q_heads: int + local_kv_head_start: int + local_kv_heads: int + global_sequence_length: int + local_sequence_length: int + global_block_indices: tuple[int, ...] + global_block_token_starts: tuple[int, ...] + local_block_offsets: tuple[int, ...] + packed_sequence_offsets: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if tp_rank >= tp_world_size: + raise AttentionContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + if cp_rank >= cp_world_size: + raise AttentionContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") + global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") + if global_q_heads % global_kv_heads != 0: + raise AttentionContractError( + f"global_q_heads={global_q_heads} must be divisible by " + f"global_kv_heads={global_kv_heads} for GQA" + ) + if global_q_heads % tp_world_size != 0 or global_kv_heads % tp_world_size != 0: + raise AttentionContractError( + "global Q/KV heads must be evenly divisible by tp_world_size; " + f"got Hq={global_q_heads}, Hkv={global_kv_heads}, TP={tp_world_size}" + ) + + local_q_heads = _positive_int(self.local_q_heads, "local_q_heads") + local_kv_heads = _positive_int(self.local_kv_heads, "local_kv_heads") + expected_q_heads = global_q_heads // tp_world_size + expected_kv_heads = global_kv_heads // tp_world_size + if local_q_heads != expected_q_heads or local_kv_heads != expected_kv_heads: + raise AttentionContractError( + "local TP head counts do not preserve the global Q/KV mapping; " + f"expected ({expected_q_heads}, {expected_kv_heads}), got " + f"({local_q_heads}, {local_kv_heads})" + ) + + local_q_head_start = _non_negative_int(self.local_q_head_start, "local_q_head_start") + local_kv_head_start = _non_negative_int(self.local_kv_head_start, "local_kv_head_start") + expected_q_start = tp_rank * expected_q_heads + expected_kv_start = tp_rank * expected_kv_heads + if local_q_head_start != expected_q_start or local_kv_head_start != expected_kv_start: + raise AttentionContractError( + "local TP head starts do not match contiguous rank ownership; " + f"expected ({expected_q_start}, {expected_kv_start}), got " + f"({local_q_head_start}, {local_kv_head_start})" + ) + + global_sequence_length = _positive_int( + self.global_sequence_length, "global_sequence_length" + ) + local_sequence_length = _positive_int(self.local_sequence_length, "local_sequence_length") + + block_indices = _integer_tuple(self.global_block_indices, "global_block_indices") + if not block_indices: + raise AttentionContractError("global_block_indices must not be empty") + if any(index < 0 for index in block_indices): + raise AttentionContractError("global_block_indices must be non-negative") + if any( + left >= right for left, right in zip(block_indices, block_indices[1:], strict=False) + ): + raise AttentionContractError( + "global_block_indices must be unique and strictly increasing" + ) + object.__setattr__(self, "global_block_indices", block_indices) + + block_token_starts = _integer_tuple( + self.global_block_token_starts, "global_block_token_starts" + ) + local_block_offsets = _integer_tuple(self.local_block_offsets, "local_block_offsets") + if len(block_token_starts) != len(block_indices): + raise AttentionContractError( + "global_block_token_starts must contain one entry per global_block_indices entry" + ) + if any(start < 0 for start in block_token_starts): + raise AttentionContractError("global_block_token_starts must be non-negative") + if len(local_block_offsets) != len(block_indices) + 1: + raise AttentionContractError( + "local_block_offsets must contain one boundary more than global_block_indices" + ) + if local_block_offsets[0] != 0 or local_block_offsets[-1] != local_sequence_length: + raise AttentionContractError( + "local_block_offsets must start at 0 and end at local_sequence_length" + ) + if any( + left >= right + for left, right in zip(local_block_offsets, local_block_offsets[1:], strict=False) + ): + raise AttentionContractError("local_block_offsets must be strictly increasing") + + previous_global_end = 0 + for index, (global_start, local_start, local_end) in enumerate( + zip( + block_token_starts, + local_block_offsets[:-1], + local_block_offsets[1:], + strict=True, + ) + ): + global_end = global_start + (local_end - local_start) + if global_end > global_sequence_length: + raise AttentionContractError( + f"global block {block_indices[index]} exceeds global_sequence_length" + ) + if index > 0 and global_start < previous_global_end: + raise AttentionContractError( + "global block token ranges must be non-overlapping and ordered" + ) + previous_global_end = global_end + object.__setattr__(self, "global_block_token_starts", block_token_starts) + object.__setattr__(self, "local_block_offsets", local_block_offsets) + + if self.packed_sequence_offsets is not None: + offsets = _integer_tuple(self.packed_sequence_offsets, "packed_sequence_offsets") + if len(offsets) < 2 or offsets[0] != 0: + raise AttentionContractError( + "packed_sequence_offsets must start at 0 and contain an end offset" + ) + if any(left >= right for left, right in zip(offsets, offsets[1:], strict=False)): + raise AttentionContractError("packed_sequence_offsets must be strictly increasing") + if offsets[-1] != local_sequence_length: + raise AttentionContractError( + "the final packed_sequence_offsets value must equal local_sequence_length; " + f"got {offsets[-1]} and {local_sequence_length}" + ) + object.__setattr__(self, "packed_sequence_offsets", offsets) + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic CP ``(out, lse)`` merge semantics.""" + + merge: AttentionMerge = AttentionMerge.ONLINE_SOFTMAX_LSE + acc_dtype: AttentionDType = AttentionDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(AttentionMerge, self.merge, "merge")) + object.__setattr__( + self, "acc_dtype", _enum_value(AttentionDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError( + f"CP attention accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class SplitKVExecutionPlan: + """Actual backend-local Split-KV schedule emitted by a runtime. + + CP ownership is intentionally separate from this plan. ``boundaries`` + describe the canonical logical KV ranges reduced by one backend invocation; + CP may transport those partial states between ranks but must not change the + FP32 merge contract recorded here. + """ + + requested_mode: SplitKVMode + requested_split_size: int | None + actual_mode: SplitKVMode | None + actual_split_size: int | None + boundaries: tuple[tuple[int, int], ...] + merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + acc_dtype: AttentionDType = AttentionDType.FP32 + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + backend: str = "reference" + source: str = "contract_exact" + fallback: bool = False + fallback_reason: str | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "requested_mode", + _enum_value(SplitKVMode, self.requested_mode, "requested_mode"), + ) + if self.actual_mode is not None: + object.__setattr__( + self, + "actual_mode", + _enum_value(SplitKVMode, self.actual_mode, "actual_mode"), + ) + for field_name in ("requested_split_size", "actual_split_size"): + value = getattr(self, field_name) + if value is not None: + _positive_int(value, field_name) + if self.requested_mode is SplitKVMode.FIXED: + if self.requested_split_size is None: + raise AttentionContractError( + "requested fixed Split-KV mode requires requested_split_size" + ) + elif self.requested_split_size is not None: + raise AttentionContractError( + "requested_split_size is only valid for requested fixed Split-KV mode" + ) + try: + boundaries = tuple(tuple(boundary) for boundary in self.boundaries) + except TypeError as exc: + raise AttentionContractError( + "Split-KV boundaries must be an iterable of (start, end) pairs" + ) from exc + previous_end = 0 + for index, boundary in enumerate(boundaries): + if len(boundary) != 2: + raise AttentionContractError( + f"Split-KV boundary {index} must contain exactly start and end" + ) + start, end = boundary + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError( + "Split-KV boundaries must satisfy 0 <= start < end" + ) + if index > 0 and start != previous_end: + raise AttentionContractError( + "Split-KV boundaries must be contiguous and in logical KV order" + ) + previous_end = end + if self.actual_mode is None and boundaries: + raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries") + if self.actual_mode is not None and not boundaries: + raise AttentionContractError("actual Split-KV plan must declare logical boundaries") + if self.actual_mode is SplitKVMode.FIXED: + if self.actual_split_size is None: + raise AttentionContractError( + "actual fixed Split-KV mode requires actual_split_size" + ) + widths = tuple(end - start for start, end in boundaries) + if any(width != self.actual_split_size for width in widths[:-1]) or ( + widths and widths[-1] > self.actual_split_size + ): + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_size except " + "for a shorter final split" + ) + elif self.actual_split_size is not None: + raise AttentionContractError( + "actual_split_size is only valid for actual fixed Split-KV mode" + ) + if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1: + raise AttentionContractError( + "disabled Split-KV execution must contain exactly one boundary" + ) + object.__setattr__(self, "boundaries", boundaries) + object.__setattr__( + self, + "merge_order", + _enum_value(ReductionOrder, self.merge_order, "merge_order"), + ) + object.__setattr__( + self, + "acc_dtype", + _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"), + ) + object.__setattr__( + self, + "downcast_at", + _enum_value(DowncastPoint, self.downcast_at, "downcast_at"), + ) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Split-KV partial states must be merged in fp32") + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Split-KV backend must be a non-empty string") + if not isinstance(self.source, str) or not self.source.strip(): + raise AttentionContractError("Split-KV plan source must be a non-empty string") + if not isinstance(self.fallback, bool): + raise AttentionContractError("Split-KV fallback must be a bool") + if self.fallback and not self.fallback_reason: + raise AttentionContractError("Split-KV fallback_reason is required for a fallback") + if not self.fallback and self.fallback_reason is not None: + raise AttentionContractError( + "Split-KV fallback_reason must be None when fallback=False" + ) + if not self.fallback and self.actual_mode is not None: + if self.actual_mode is not self.requested_mode: + raise AttentionContractError( + "actual Split-KV mode may differ from requested mode only for a fallback" + ) + if self.actual_split_size != self.requested_split_size: + raise AttentionContractError( + "actual Split-KV size may differ from requested size only for a fallback" + ) + + @property + def actual_split_count(self) -> int | None: + return None if self.actual_mode is None else len(self.boundaries) + + def to_dict(self) -> dict[str, Any]: + return { + "requested_split_kv_policy": self.requested_mode.value, + "requested_split_kv_size": self.requested_split_size, + "actual_split_kv_policy": ( + None if self.actual_mode is None else self.actual_mode.value + ), + "actual_split_kv_size": self.actual_split_size, + "actual_split_kv_count": self.actual_split_count, + "actual_split_boundaries": [list(boundary) for boundary in self.boundaries], + "split_kv_merge_order": self.merge_order.value, + "split_kv_accum_dtype": self.acc_dtype.value, + "split_kv_downcast_at": self.downcast_at.value, + "split_kv_backend": self.backend, + "split_kv_plan_source": self.source, + "split_kv_fallback": self.fallback, + "split_kv_fallback_reason": self.fallback_reason, + } + + +@dataclass(frozen=True) +class SplitKVSpec: + """Requested Split-KV policy shared by training and rollout paths.""" + + mode: SplitKVMode = SplitKVMode.DISABLED + fixed_split_size: int | None = None + strict_consistency: bool = True + + @classmethod + def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency) + + @classmethod + def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls( + mode=SplitKVMode.FIXED, + fixed_split_size=split_size, + strict_consistency=strict_consistency, + ) + + @classmethod + def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec": + return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency) + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode")) + if not isinstance(self.strict_consistency, bool): + raise AttentionContractError("split_kv.strict_consistency must be a bool") + if self.mode is SplitKVMode.FIXED: + if self.fixed_split_size is None: + raise AttentionContractError( + "fixed Split-KV policy requires fixed_split_size" + ) + _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") + elif self.fixed_split_size is not None: + raise AttentionContractError( + "fixed_split_size is only valid for fixed Split-KV policy" + ) + if self.strict_consistency and self.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" + ) + + def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan: + """Resolve policies whose logical schedule is fully known by contract.""" + + total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens") + if self.mode is SplitKVMode.AUTO: + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=None, + actual_mode=None, + actual_split_size=None, + boundaries=(), + backend=backend, + source="runtime_required", + ) + split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size + assert split_size is not None + boundaries = tuple( + (start, min(start + split_size, total_kv_tokens)) + for start in range(0, total_kv_tokens, split_size) + ) + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=self.fixed_split_size, + actual_mode=self.mode, + actual_split_size=self.fixed_split_size, + boundaries=boundaries, + backend=backend, + source="contract_exact", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "fixed_split_size": self.fixed_split_size, + "strict_consistency": self.strict_consistency, + } + + +def validate_split_kv_alignment( + training: SplitKVExecutionPlan, + rollout: SplitKVExecutionPlan, +) -> None: + """Fail closed unless train and rollout executed the same logical plan.""" + + if training.actual_mode is None or rollout.actual_mode is None: + raise AttentionContractError( + "strict Split-KV alignment requires actual runtime plans from both sides" + ) + fields = ( + "requested_mode", + "requested_split_size", + "actual_mode", + "actual_split_size", + "boundaries", + "merge_order", + "acc_dtype", + "downcast_at", + "fallback", + ) + mismatches = [ + field_name + for field_name in fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if mismatches: + raise AttentionContractError( + "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches) + ) + + +@dataclass(frozen=True, order=True) +class SplitKVRuntimeCoordinate: + """Identity of one batch/rank/owner Split-KV runtime plan.""" + + batch_index: int + tp_rank: int + cp_rank: int + owner_cp_rank: int + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + ) -> None: + batch_index = _non_negative_int(self.batch_index, "batch_index") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank") + if batch_index >= batch_size: + raise AttentionContractError("Split-KV batch_index is outside batch_size") + if tp_rank >= tp_world_size: + raise AttentionContractError("Split-KV tp_rank is outside tp_world_size") + if cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV cp_rank is outside cp_world_size") + if owner_cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size") + + def to_dict(self) -> dict[str, int]: + return { + "batch_index": self.batch_index, + "tp_rank": self.tp_rank, + "cp_rank": self.cp_rank, + "owner_cp_rank": self.owner_cp_rank, + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanEntry: + """Actual plan for one batch/TP/CP consumer and logical KV owner.""" + + coordinate: SplitKVRuntimeCoordinate + expected_kv_range: tuple[int, int] + execution: SplitKVExecutionPlan + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + total_kv_tokens: int, + ) -> None: + self.coordinate.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + try: + start, end = self.expected_kv_range + except (TypeError, ValueError) as exc: + raise AttentionContractError( + "expected_kv_range must contain exactly (start, end)" + ) from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + or end > total_kv_tokens + ): + raise AttentionContractError( + "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" + ) + if self.execution.actual_mode is None: + raise AttentionContractError( + "complete Split-KV plan sets require actual runtime plans" + ) + if ( + self.execution.boundaries[0][0] != start + or self.execution.boundaries[-1][1] != end + ): + raise AttentionContractError( + "Split-KV execution boundaries must exactly cover expected_kv_range" + ) + if any( + boundary_start < start or boundary_end > end + for boundary_start, boundary_end in self.execution.boundaries + ): + raise AttentionContractError( + "Split-KV execution boundary escapes expected_kv_range" + ) + + def to_dict(self) -> dict[str, Any]: + return { + **self.coordinate.to_dict(), + "expected_kv_range": list(self.expected_kv_range), + **self.execution.to_dict(), + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanSet: + """Complete actual Split-KV plans across batch, TP, CP, and KV owners. + + Every CP consumer must report the plan used for every CP-owned KV range. + This duplicates owner plans across consumers intentionally: it detects one + rank silently choosing a different Split-K schedule or merge policy. + """ + + batch_size: int + tp_world_size: int + cp_world_size: int + total_kv_tokens: tuple[int, ...] + entries: tuple[SplitKVRuntimePlanEntry, ...] + + def __post_init__(self) -> None: + batch_size = _positive_int(self.batch_size, "batch_size") + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens") + if len(totals) != batch_size or any(total <= 0 for total in totals): + raise AttentionContractError( + "total_kv_tokens must contain one positive length per batch item" + ) + object.__setattr__(self, "total_kv_tokens", totals) + entries = tuple(self.entries) + object.__setattr__(self, "entries", entries) + expected_coordinates = { + SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(batch_size) + for tp_rank in range(tp_world_size) + for cp_rank in range(cp_world_size) + for owner_cp_rank in range(cp_world_size) + } + actual_coordinates = [entry.coordinate for entry in entries] + if len(set(actual_coordinates)) != len(actual_coordinates): + raise AttentionContractError( + "Split-KV runtime plan set contains duplicate coordinates" + ) + missing = expected_coordinates.difference(actual_coordinates) + extra = set(actual_coordinates).difference(expected_coordinates) + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan set coordinate coverage is incomplete; " + f"missing={_format_split_kv_coordinates(missing)}, " + f"extra={_format_split_kv_coordinates(extra)}" + ) + for entry in entries: + entry.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals[entry.coordinate.batch_index], + ) + self._validate_owner_coverage() + self._validate_rank_invariance() + + def _validate_owner_coverage(self) -> None: + for batch_index, total in enumerate(self.total_kv_tokens): + for tp_rank in range(self.tp_world_size): + ranges = [] + for owner_cp_rank in range(self.cp_world_size): + matches = [ + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.tp_rank == tp_rank + and entry.coordinate.cp_rank == 0 + and entry.coordinate.owner_cp_rank == owner_cp_rank + ] + ranges.append(matches[0].expected_kv_range) + previous_end = 0 + for start, end in ranges: + if start != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = end + if previous_end != total: + raise AttentionContractError( + "Split-KV owner ranges do not cover total_kv_tokens" + ) + + def _validate_rank_invariance(self) -> None: + for batch_index in range(self.batch_size): + for owner_cp_rank in range(self.cp_world_size): + entries = sorted( + ( + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.owner_cp_rank == owner_cp_rank + ), + key=lambda entry: ( + entry.coordinate.tp_rank, + entry.coordinate.cp_rank, + ), + ) + reference = entries[0] + for entry in entries[1:]: + if entry.expected_kv_range != reference.expected_kv_range: + raise AttentionContractError( + "Split-KV owner range differs across TP/CP consumers" + ) + try: + validate_split_kv_alignment( + reference.execution, + entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return { + "batch_size": self.batch_size, + "tp_world_size": self.tp_world_size, + "cp_world_size": self.cp_world_size, + "total_kv_tokens": list(self.total_kv_tokens), + "entries": [ + entry.to_dict() + for entry in sorted(self.entries, key=lambda entry: entry.coordinate) + ], + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + } + + +def validate_split_kv_plan_set_alignment( + training: SplitKVRuntimePlanSet, + rollout: SplitKVRuntimePlanSet, +) -> None: + """Fail closed unless complete train/rollout runtime plan sets align.""" + + topology_fields = ( + "batch_size", + "tp_world_size", + "cp_world_size", + "total_kv_tokens", + ) + topology_mismatches = [ + field_name + for field_name in topology_fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if topology_mismatches: + raise AttentionContractError( + "training/rollout Split-KV plan-set topology differs: " + + ", ".join(topology_mismatches) + ) + training_by_coordinate = { + entry.coordinate: entry for entry in training.entries + } + rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} + if training_by_coordinate.keys() != rollout_by_coordinate.keys(): + raise AttentionContractError( + "training/rollout Split-KV plan-set coordinates differ" + ) + for coordinate in sorted(training_by_coordinate): + train_entry = training_by_coordinate[coordinate] + rollout_entry = rollout_by_coordinate[coordinate] + if train_entry.expected_kv_range != rollout_entry.expected_kv_range: + raise AttentionContractError( + f"training/rollout expected KV range differs at {coordinate}" + ) + try: + validate_split_kv_alignment( + train_entry.execution, + rollout_entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + f"training/rollout Split-KV plan differs at {coordinate}: {exc}" + ) from exc + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, field), bool): + raise AttentionContractError(f"{field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append( + f"Split-KV policy={contract.split_kv.mode.value} is unsupported" + ) + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 302ad73f..c334f49a 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -16,6 +16,13 @@ import torch +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, +) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp @@ -45,6 +52,104 @@ def __post_init__(self) -> None: raise ValueError("block_end must be >= block_start") +@dataclass(frozen=True) +class AttentionBackwardGradients: + """Training-side gradients emitted by the CP attention backward reference.""" + + dq: torch.Tensor + dk: torch.Tensor + dv: torch.Tensor + + +@dataclass(frozen=True) +class AttentionBackwardPathResult: + """One materialized CP attention backward path.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + gradients: AttentionBackwardGradients + provenance: dict[str, object] + + +@dataclass(frozen=True) +class GradientDriftStats: + """Shape-aware absolute drift summary for backward validation reports.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, object]: + 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 AttentionBackwardRankDrift: + """Backward drift for one logical CP rank's sequence ownership.""" + + rank: int + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + + def to_dict(self) -> dict[str, object]: + return { + "rank": self.rank, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + } + + +@dataclass(frozen=True) +class AttentionBackwardPathDrift: + """Candidate-vs-reference backward drift for one CP path.""" + + candidate_name: str + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + out: GradientDriftStats + lse: GradientDriftStats + per_rank: tuple[AttentionBackwardRankDrift, ...] + provenance: dict[str, object] + + def to_dict(self) -> dict[str, object]: + return { + "candidate_name": self.candidate_name, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "per_rank": [item.to_dict() for item in self.per_rank], + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionBackwardComparisonReport: + """Structured PR8 report for CP attention gradient drift validation.""" + + reference_name: str + drifts: tuple[AttentionBackwardPathDrift, ...] + + def to_dict(self) -> dict[str, object]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + } + + def merge_attention_partial_states( states: Sequence[AttentionPartialState], ) -> AttentionPartialState: @@ -97,6 +202,24 @@ class DeterministicCPAttentionReferenceOp: ``forward_fp32`` keeps the fp32 merged output. """ + op_class = "attention" + + @staticmethod + def split_kv_execution_plans( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> list[dict[str, object]]: + """Export the actual logical Split-KV plan before execution.""" + + return split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + def __call__( self, q: torch.Tensor, @@ -252,6 +375,108 @@ def forward_fp32_with_lse( output_dtype=torch.float32, ) + def backward_reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, + name: Optional[str] = None, + ) -> AttentionBackwardPathResult: + """Run the deterministic training-side backward validation path. + + The semantic backward input is ``dout`` plus the forward attention state + produced from the same Q/K/V, masks, position offsets, CP world, and KV + block order. The reference keeps the softmax/merge math in fp32 and + records the final-write dtype in provenance; decode backward is + intentionally out of scope for PR8. + """ + + _validate_qkv(q, k, v) + if dout.shape != q.shape: + raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") + if not torch.is_floating_point(dout) or torch.is_complex(dout): + raise ValueError("dout must be a real floating-point tensor") + q_leaf = q.detach().clone().requires_grad_(True) + k_leaf = k.detach().clone().requires_grad_(True) + v_leaf = v.detach().clone().requires_grad_(True) + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=resolved_output_dtype, + ) + torch.autograd.backward(out, dout.to(device=out.device, dtype=out.dtype)) + if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: + raise RuntimeError("CP attention backward did not produce dq/dk/dv") + + return AttentionBackwardPathResult( + name=name + or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), + out=out.detach(), + lse=lse.detach(), + gradients=AttentionBackwardGradients( + dq=q_leaf.grad.detach(), + dk=k_leaf.grad.detach(), + dv=v_leaf.grad.detach(), + ), + provenance={ + "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", + "gradient_mode": "training_backward", + "gradient_inputs": ["q", "k", "v"], + "gradient_outputs": ["out"], + "saved_forward_state": [ + "out", + "attention_lse", + "causal_mask", + "key_padding_mask", + "query_position_offsets", + "key_position_offsets", + "global_block_index", + ], + "cp_world_size": cp_world_size, + "kv_chunk_size": kv_chunk_size, + "requested_split_kv_policy": ( + "disabled" if kv_chunk_size is None else "fixed" + ), + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_backward_reference", + ), + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(resolved_output_dtype).replace("torch.", ""), + "q_dtype": str(q.dtype).replace("torch.", ""), + "k_dtype": str(k.dtype).replace("torch.", ""), + "v_dtype": str(v.dtype).replace("torch.", ""), + "dout_dtype": str(dout.dtype).replace("torch.", ""), + "te_backward_oracle": "not_used", + "decode_backward": "not_supported", + }, + ) + def local_partial_state( self, q: torch.Tensor, @@ -456,6 +681,143 @@ def _forward_impl( return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) +def compare_cp_attention_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + candidate_cp_world_size: int = 2, + candidate_kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, +) -> AttentionBackwardComparisonReport: + """Compare CP=1 backward with a CP/chunked-prefill candidate. + + The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank + slices. It is a validation/reporting helper, not a separate production + backward kernel. + """ + + op = DeterministicCPAttentionReferenceOp() + reference = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=1, + kv_chunk_size=None, + output_dtype=output_dtype, + name="cp1_backward_reference", + ) + candidate = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=candidate_cp_world_size, + kv_chunk_size=candidate_kv_chunk_size, + output_dtype=output_dtype, + ) + return AttentionBackwardComparisonReport( + reference_name=reference.name, + drifts=(_compare_backward_path(candidate, reference),), + ) + + +def _compare_backward_path( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, +) -> AttentionBackwardPathDrift: + cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") + return AttentionBackwardPathDrift( + candidate_name=candidate.name, + dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), + dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), + dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), + provenance=candidate.provenance, + ) + + +def _per_rank_backward_drifts( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, + cp_world_size: int, +) -> tuple[AttentionBackwardRankDrift, ...]: + q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) + kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) + per_rank = [] + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + per_rank.append( + AttentionBackwardRankDrift( + rank=rank, + dq=_drift_stats( + candidate.gradients.dq[:, :, q_start:q_end, :], + reference.gradients.dq[:, :, q_start:q_end, :], + ), + dk=_drift_stats( + candidate.gradients.dk[:, :, kv_start:kv_end, :], + reference.gradients.dk[:, :, kv_start:kv_end, :], + ), + dv=_drift_stats( + candidate.gradients.dv[:, :, kv_start:kv_end, :], + reference.gradients.dv[:, :, kv_start:kv_end, :], + ), + ) + ) + return tuple(per_rank) + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: + 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().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return GradientDriftStats( + max_abs=float(diff.max().item()), + mean_abs=float(diff.mean().item()), + p95_abs=float(torch.quantile(diff, 0.95).item()), + p99_abs=float(torch.quantile(diff, 0.99).item()), + active_count=active_count, + ) + + +def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: + prefix = f"cp{cp_world_size}" + if kv_chunk_size is None: + return f"{prefix}_backward" + return f"{prefix}_chunked_backward" + + +def _provenance_int(provenance: dict[str, object], key: str) -> int: + value = provenance[key] + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"provenance field {key!r} must be an int") + return value + + def _merge_two_states( out_a: torch.Tensor, lse_a: torch.Tensor, @@ -476,8 +838,8 @@ def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) - 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 out/lse shapes") - if state.block_start < previous_end: - raise ValueError("partial state block ranges must not overlap") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") previous_end = state.block_end @@ -546,8 +908,128 @@ def _kv_block_bounds( return bounds +def split_kv_execution_plan_provenance( + length: int, + *, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str, +) -> list[dict[str, object]]: + """Return the actual backend-local Split-KV plan for every CP owner.""" + + if length < 1: + raise ValueError("Split-KV sequence length must be >= 1") + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + result: list[dict[str, object]] = [] + for owner_cp_rank, (rank_start, rank_end) in enumerate( + _split_bounds(length, cp_world_size) + ): + if rank_start == rank_end: + continue + if kv_chunk_size is None: + boundaries = ((rank_start, rank_end),) + mode = SplitKVMode.DISABLED + else: + boundaries = tuple( + (start, min(start + kv_chunk_size, rank_end)) + for start in range(rank_start, rank_end, kv_chunk_size) + ) + mode = SplitKVMode.FIXED + plan = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) + return result + + +def build_reference_split_kv_runtime_plan_set( + total_kv_tokens: Sequence[int], + *, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str = "deterministic_cp_reference", +) -> SplitKVRuntimePlanSet: + """Build complete per-batch/TP/CP/owner plans for the reference path.""" + + totals = tuple(total_kv_tokens) + if not totals or any(total < cp_world_size for total in totals): + raise ValueError( + "reference runtime plan sets require at least one KV token per CP owner" + ) + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("TP and CP world sizes must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + owner_ranges = _split_bounds(total, cp_world_size) + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + if kv_chunk_size is None: + mode = SplitKVMode.DISABLED + boundaries = ((owner_start, owner_end),) + else: + mode = SplitKVMode.FIXED + boundaries = tuple( + (start, min(start + kv_chunk_size, owner_end)) + for start in range(owner_start, owner_end, kv_chunk_size) + ) + execution = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + __all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", "merge_attention_partial_states", + "split_kv_execution_plan_provenance", ] diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 03c861dd..196cfcaa 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -8,6 +8,7 @@ """ import contextlib +import json import math import pytest @@ -16,7 +17,9 @@ from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, merge_attention_partial_states, + split_kv_execution_plan_provenance, ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp @@ -26,6 +29,7 @@ _N_KV = 8 _HEAD_DIM = 128 _ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 @contextlib.contextmanager @@ -350,6 +354,203 @@ def test_cp2_chunked_gradients_match_cp1_reference(): torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) +def test_backward_report_cp2_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 5, 5, seed=15, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 5, 8, generator=torch.Generator().manual_seed(16)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + output_dtype=torch.float32, + ) + + assert report.reference_name == "cp1_backward_reference" + drift = report.drifts[0] + assert drift.candidate_name == "cp2_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.out.max_abs <= _ATOL + assert drift.lse.max_abs <= _ATOL + assert len(drift.per_rank) == 2 + assert drift.per_rank[0].dq.active_count > 0 + assert drift.per_rank[1].dk.active_count > 0 + assert drift.provenance["saved_forward_state"][0] == "out" + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["te_backward_oracle"] == "not_used" + assert drift.provenance["decode_backward"] == "not_supported" + json.dumps(report.to_dict()) + + +def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 6, 6, seed=17, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 6, 8, generator=torch.Generator().manual_seed(18)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.candidate_name == "cp2_chunked_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.provenance["attention_mode"] == "chunked_prefill" + assert drift.provenance["kv_chunk_size"] == 2 + assert drift.provenance["requested_split_kv_policy"] == "fixed" + assert drift.provenance["actual_split_kv_plans"] == [ + { + "owner_cp_rank": 0, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[0, 2], [2, 3]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + { + "owner_cp_rank": 1, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[3, 5], [5, 6]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + ] + + +def test_split_kv_plan_never_crosses_cp_owner_boundaries(): + plans = split_kv_execution_plan_provenance( + 10, + cp_world_size=3, + kv_chunk_size=3, + backend="test-reference", + ) + + assert [plan["actual_split_boundaries"] for plan in plans] == [ + [[0, 3], [3, 4]], + [[4, 7]], + [[7, 10]], + ] + assert [plan["owner_cp_rank"] for plan in plans] == [0, 1, 2] + + +def test_backward_report_preserves_post_rope_position_metadata(): + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 5, 5, seed=19, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([23, 101], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + dout = torch.randn(2, 4, 5, 8, generator=torch.Generator().manual_seed(20)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + + +def test_qwen3_8b_local_tp2_cp2_bf16_backward_report_smoke(): + # Qwen3-8B global Hq/Hkv is 32/8. A TP=2 local shard owns 16/4 heads. + q, k, v = _qkv( + 1, + 4, + 4, + seed=21, + dtype=torch.bfloat16, + heads=16, + kv_heads=4, + dim=_HEAD_DIM, + ) + dout = torch.randn( + 1, + 16, + 4, + _HEAD_DIM, + generator=torch.Generator().manual_seed(22), + dtype=torch.bfloat16, + ) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.bfloat16, + ) + + drift = report.drifts[0] + assert drift.provenance["q_dtype"] == "bfloat16" + assert drift.provenance["output_dtype"] == "bfloat16" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.dq.max_abs <= 5.0e-2 + assert drift.dk.max_abs <= 5.0e-2 + assert drift.dv.max_abs <= 5.0e-2 + + +def test_backward_report_validates_dout_shape_and_dtype(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=23, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="dout must have shape"): + op.backward_reference(q, k, v, torch.randn(1, 4, 3, 8), cp_world_size=2) + + with pytest.raises(ValueError, match="dout must be a real floating-point tensor"): + op.backward_reference( + q, + k, + v, + torch.ones(1, 4, 4, 8, dtype=torch.long), + cp_world_size=2, + ) + + def test_inputs_are_not_mutated(): op = DeterministicCPAttentionReferenceOp() q, k, v = _qkv(2, 6, 6, seed=7) @@ -419,5 +620,17 @@ def test_overlapping_partial_ranges_raise(): ) +def test_gapped_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="gap-free"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=2), + AttentionPartialState(out=out, lse=lse, block_start=3, block_end=4), + ] + ) + + def test_registry_dispatches_cp_attention_reference(): assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) From f98279793efce11129f7e9bcf487aad5bb2d0d89 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:22:05 +0800 Subject: [PATCH 04/16] fix(attention): validate CP reference numeric inputs Signed-off-by: lamentropetion <3051000145@qq.com> --- .../ops/pytorch/attention/cp_attention.py | 29 +++++++++++++++++-- tests/test_cp_attention.py | 28 ++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index c334f49a..6f80ea32 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -503,6 +503,7 @@ def local_partial_state( """ _validate_qkv(q, k, v) + _validate_scale(scale) if q_start < 0 or k_start < 0: raise ValueError("q_start and k_start must be non-negative") if total_kv_len < k_start + k.size(2): @@ -598,9 +599,14 @@ def _forward_impl( kv_chunk_size: Optional[int], ) -> tuple[torch.Tensor, torch.Tensor]: _validate_qkv(q, k, v) - if cp_world_size < 1: + _validate_scale(scale) + if isinstance(cp_world_size, bool) or not isinstance(cp_world_size, int) or cp_world_size < 1: raise ValueError("cp_world_size must be >= 1") - if kv_chunk_size is not None and kv_chunk_size < 1: + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): raise ValueError("kv_chunk_size must be >= 1 when provided") batch, hq, sq, dim = q.shape @@ -850,10 +856,29 @@ def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: raise ValueError("k and v must have the same 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) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") if q.size(1) % k.size(1) != 0: raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") +def _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: total = torch.tensor(0.0, device=tensors[0].device) for tensor in tensors: diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 196cfcaa..0db59caf 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -608,6 +608,34 @@ def test_invalid_gqa_and_mask_shapes_raise(): ) +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf"), True, "bad"]) +def test_invalid_scale_fails_before_attention_math(scale): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=24, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="scale must be a positive finite number"): + op.forward_fp32_with_lse(q, k, v, scale=scale) + + +def test_qkv_dtype_and_floating_contract_fails_closed(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=25, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="same dtype"): + op.forward_fp32_with_lse(q, k.to(torch.bfloat16), v) + with pytest.raises(ValueError, match="real floating-point"): + op.forward_fp32_with_lse(q.to(torch.long), k.to(torch.long), v.to(torch.long)) + + +@pytest.mark.parametrize("kwargs", [{"cp_world_size": True}, {"kv_chunk_size": True}]) +def test_boolean_parallelism_arguments_fail_closed(kwargs): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=26, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError): + op.forward_fp32_with_lse(q, k, v, **kwargs) + + def test_overlapping_partial_ranges_raise(): out = torch.zeros(1, 1, 1, 1) lse = torch.zeros(1, 1, 1) From 3c8b734e196207f335b49568ed9847f57dd5efcd Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:28:58 +0800 Subject: [PATCH 05/16] fix(attention): bind backward gradient dtype and device Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 6 +++++- tests/test_cp_attention.py | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 6f80ea32..897b23fd 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -406,6 +406,10 @@ def backward_reference( raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") if not torch.is_floating_point(dout) or torch.is_complex(dout): raise ValueError("dout must be a real floating-point tensor") + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") q_leaf = q.detach().clone().requires_grad_(True) k_leaf = k.detach().clone().requires_grad_(True) v_leaf = v.detach().clone().requires_grad_(True) @@ -424,7 +428,7 @@ def backward_reference( kv_chunk_size=kv_chunk_size, output_dtype=resolved_output_dtype, ) - torch.autograd.backward(out, dout.to(device=out.device, dtype=out.dtype)) + torch.autograd.backward(out, dout.to(dtype=out.dtype)) if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: raise RuntimeError("CP attention backward did not produce dq/dk/dv") diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 0db59caf..8a84ac20 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -550,6 +550,15 @@ def test_backward_report_validates_dout_shape_and_dtype(): cp_world_size=2, ) + with pytest.raises(ValueError, match="dout must have the same dtype"): + op.backward_reference( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + torch.ones_like(q), + cp_world_size=2, + ) + def test_inputs_are_not_mutated(): op = DeterministicCPAttentionReferenceOp() From 88872f3491658acf6ce8bd372db76568a9f92aff Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:31:37 +0800 Subject: [PATCH 06/16] fix(attention): enforce FP32 CP merge state Signed-off-by: lamentropetion <3051000145@qq.com> --- .../ops/pytorch/attention/cp_attention.py | 17 ++++++++++++++++- tests/test_cp_attention.py | 17 +++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 897b23fd..adfcfb53 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -46,6 +46,10 @@ def __post_init__(self) -> None: raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") if self.lse.shape != self.out.shape[:3]: raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") if self.block_start < 0: raise ValueError("block_start must be non-negative") if self.block_end < self.block_start: @@ -330,6 +334,8 @@ def forward_with_lse( ``output_dtype`` defaults to the input dtype. """ + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) out, lse = self._forward_impl( q, k, @@ -342,7 +348,7 @@ def forward_with_lse( cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, ) - out = out.to(q.dtype if output_dtype is None else output_dtype) + out = out.to(resolved_output_dtype) return out, lse def forward_fp32_with_lse( @@ -415,6 +421,7 @@ def backward_reference( v_leaf = v.detach().clone().requires_grad_(True) resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) out, lse = self.forward_with_lse( q_leaf, k_leaf, @@ -883,6 +890,14 @@ def _validate_scale(scale: Optional[float]) -> None: raise ValueError("scale must be a positive finite number") +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: total = torch.tensor(0.0, device=tensors[0].device) for tensor in tensors: diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 8a84ac20..cc93874b 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -645,6 +645,23 @@ def test_boolean_parallelism_arguments_fail_closed(kwargs): op.forward_fp32_with_lse(q, k, v, **kwargs) +@pytest.mark.parametrize("output_dtype", [torch.long, torch.complex64, "fp32"]) +def test_nonfloating_output_dtype_fails_closed(output_dtype): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=27, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="output_dtype must be a real floating-point"): + op.forward_with_lse(q, k, v, output_dtype=output_dtype) + + +def test_partial_states_must_remain_fp32_and_colocated(): + out = torch.zeros(1, 1, 1, 1, dtype=torch.bfloat16) + lse = torch.zeros(1, 1, 1, dtype=torch.float32) + + with pytest.raises(ValueError, match="must remain FP32"): + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=1) + + def test_overlapping_partial_ranges_raise(): out = torch.zeros(1, 1, 1, 1) lse = torch.zeros(1, 1, 1) From 5b567e14b223ddb0263381d9690260575d1e224b Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 21:19:13 +0800 Subject: [PATCH 07/16] style(attention): satisfy full PR lint Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/attention_contract.py | 68 +++++++------------ .../ops/pytorch/attention/cp_attention.py | 18 +++-- 2 files changed, 32 insertions(+), 54 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index eb4994b3..1750476d 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -359,9 +359,7 @@ def __post_init__(self) -> None: or start < 0 or end <= start ): - raise AttentionContractError( - "Split-KV boundaries must satisfy 0 <= start < end" - ) + raise AttentionContractError("Split-KV boundaries must satisfy 0 <= start < end") if index > 0 and start != previous_end: raise AttentionContractError( "Split-KV boundaries must be contiguous and in logical KV order" @@ -486,14 +484,10 @@ def __post_init__(self) -> None: raise AttentionContractError("split_kv.strict_consistency must be a bool") if self.mode is SplitKVMode.FIXED: if self.fixed_split_size is None: - raise AttentionContractError( - "fixed Split-KV policy requires fixed_split_size" - ) + raise AttentionContractError("fixed Split-KV policy requires fixed_split_size") _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") elif self.fixed_split_size is not None: - raise AttentionContractError( - "fixed_split_size is only valid for fixed Split-KV policy" - ) + raise AttentionContractError("fixed_split_size is only valid for fixed Split-KV policy") if self.strict_consistency and self.mode is SplitKVMode.AUTO: raise AttentionContractError( "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" @@ -647,13 +641,8 @@ def validate( "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" ) if self.execution.actual_mode is None: - raise AttentionContractError( - "complete Split-KV plan sets require actual runtime plans" - ) - if ( - self.execution.boundaries[0][0] != start - or self.execution.boundaries[-1][1] != end - ): + raise AttentionContractError("complete Split-KV plan sets require actual runtime plans") + if self.execution.boundaries[0][0] != start or self.execution.boundaries[-1][1] != end: raise AttentionContractError( "Split-KV execution boundaries must exactly cover expected_kv_range" ) @@ -661,9 +650,7 @@ def validate( boundary_start < start or boundary_end > end for boundary_start, boundary_end in self.execution.boundaries ): - raise AttentionContractError( - "Split-KV execution boundary escapes expected_kv_range" - ) + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") def to_dict(self) -> dict[str, Any]: return { @@ -709,9 +696,7 @@ def __post_init__(self) -> None: } actual_coordinates = [entry.coordinate for entry in entries] if len(set(actual_coordinates)) != len(actual_coordinates): - raise AttentionContractError( - "Split-KV runtime plan set contains duplicate coordinates" - ) + raise AttentionContractError("Split-KV runtime plan set contains duplicate coordinates") missing = expected_coordinates.difference(actual_coordinates) extra = set(actual_coordinates).difference(expected_coordinates) if missing or extra: @@ -821,17 +806,12 @@ def validate_split_kv_plan_set_alignment( ] if topology_mismatches: raise AttentionContractError( - "training/rollout Split-KV plan-set topology differs: " - + ", ".join(topology_mismatches) + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) ) - training_by_coordinate = { - entry.coordinate: entry for entry in training.entries - } + training_by_coordinate = {entry.coordinate: entry for entry in training.entries} rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} if training_by_coordinate.keys() != rollout_by_coordinate.keys(): - raise AttentionContractError( - "training/rollout Split-KV plan-set coordinates differ" - ) + raise AttentionContractError("training/rollout Split-KV plan-set coordinates differ") for coordinate in sorted(training_by_coordinate): train_entry = training_by_coordinate[coordinate] rollout_entry = rollout_by_coordinate[coordinate] @@ -1075,14 +1055,16 @@ def __post_init__(self) -> None: not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() ): raise AttentionContractError("rope_scaling must be a non-empty string when provided") - for field in ("position_ids", "query_position_offsets", "key_position_offsets"): - values = getattr(self, field) + for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, position_field) if values is None: continue - normalized = _integer_tuple(values, field) + normalized = _integer_tuple(values, position_field) if not normalized or any(value < 0 for value in normalized): - raise AttentionContractError(f"{field} must contain non-negative positions") - object.__setattr__(self, field, normalized) + raise AttentionContractError( + f"{position_field} must contain non-negative positions" + ) + object.__setattr__(self, position_field, normalized) object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) object.__setattr__( self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") @@ -1186,11 +1168,11 @@ def __post_init__(self) -> None: "position_ids must describe the local query sequence or full local " "sequence length" ) - for field in ("query_position_offsets", "key_position_offsets"): - offsets = getattr(self.rope, field) + for position_field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, position_field) if offsets is not None and len(offsets) != batch_size: raise AttentionContractError( - f"{field} must contain one entry per logical batch entry" + f"{position_field} must contain one entry per logical batch entry" ) if self.mode is AttentionMode.DECODE and self.kv_cache is not None: if len(self.kv_cache.kv_seq_lens) != batch_size: @@ -1336,7 +1318,7 @@ def __post_init__(self) -> None: raise AttentionContractError("tp_world_sizes must contain positive values") if len(set(tp_world_sizes)) != len(tp_world_sizes): raise AttentionContractError("tp_world_sizes must not contain duplicates") - for field in ( + for capability_field in ( "exports_attention_lse", "deterministic_cp_merge", "supports_packed_varlen", @@ -1348,8 +1330,8 @@ def __post_init__(self) -> None: "supports_split_kv_auto", "reports_actual_split_kv_plan", ): - if not isinstance(getattr(self, field), bool): - raise AttentionContractError(f"{field} must be a bool") + if not isinstance(getattr(self, capability_field), bool): + raise AttentionContractError(f"{capability_field} must be a bool") if self.implementation_kind not in {"production", "reference", "deterministic"}: raise AttentionContractError( "implementation_kind must be production, reference, or deterministic" @@ -1401,9 +1383,7 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: SplitKVMode.AUTO: self.supports_split_kv_auto, } if not split_support[contract.split_kv.mode]: - reasons.append( - f"Split-KV policy={contract.split_kv.mode.value} is unsupported" - ) + reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported") if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: reasons.append("actual Split-KV execution-plan provenance is unsupported") return tuple(reasons) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index adfcfb53..e81ab8b7 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -465,9 +465,7 @@ def backward_reference( ], "cp_world_size": cp_world_size, "kv_chunk_size": kv_chunk_size, - "requested_split_kv_policy": ( - "disabled" if kv_chunk_size is None else "fixed" - ), + "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), "requested_split_kv_size": kv_chunk_size, "actual_split_kv_plans": split_kv_execution_plan_provenance( k.size(2), @@ -611,7 +609,11 @@ def _forward_impl( ) -> tuple[torch.Tensor, torch.Tensor]: _validate_qkv(q, k, v) _validate_scale(scale) - if isinstance(cp_world_size, bool) or not isinstance(cp_world_size, int) or cp_world_size < 1: + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): raise ValueError("cp_world_size must be >= 1") if kv_chunk_size is not None and ( isinstance(kv_chunk_size, bool) @@ -968,9 +970,7 @@ def split_kv_execution_plan_provenance( if kv_chunk_size is not None and kv_chunk_size < 1: raise ValueError("kv_chunk_size must be >= 1 when provided") result: list[dict[str, object]] = [] - for owner_cp_rank, (rank_start, rank_end) in enumerate( - _split_bounds(length, cp_world_size) - ): + for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): if rank_start == rank_end: continue if kv_chunk_size is None: @@ -1007,9 +1007,7 @@ def build_reference_split_kv_runtime_plan_set( totals = tuple(total_kv_tokens) if not totals or any(total < cp_world_size for total in totals): - raise ValueError( - "reference runtime plan sets require at least one KV token per CP owner" - ) + raise ValueError("reference runtime plan sets require at least one KV token per CP owner") if tp_world_size < 1 or cp_world_size < 1: raise ValueError("TP and CP world sizes must be >= 1") if kv_chunk_size is not None and kv_chunk_size < 1: From 19bc3436198f46da5db19a2b7f11281708e0db66 Mon Sep 17 00:00:00 2001 From: Takumi <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:44:11 +0800 Subject: [PATCH 08/16] fix(attention): type split-k runtime boundaries --- .../ops/pytorch/attention/cp_attention.py | 2142 +++++++++-------- 1 file changed, 1072 insertions(+), 1070 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index e81ab8b7..2394b98d 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -1,1077 +1,1079 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors -"""Deterministic context-parallel attention reference. - -This module is the correctness-first WS2 reference for CP-aware standard -softmax attention. It intentionally stays in PyTorch and uses fp32 partial -states so fused CUDA/Triton backends can validate their CP/LSE merge semantics -against a small, inspectable implementation. -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import Optional, Sequence - -import torch - -from rl_engine.kernels.attention_contract import ( - SplitKVExecutionPlan, - SplitKVMode, - SplitKVRuntimeCoordinate, - SplitKVRuntimePlanEntry, - SplitKVRuntimePlanSet, -) -from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp - - -@dataclass(frozen=True) -class AttentionPartialState: - """One KV block's attention state before deterministic LSE merge. - - ``out`` is already normalized within the local KV block and has shape - ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with - shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV - positions and define the canonical merge order. - """ - - out: torch.Tensor - lse: torch.Tensor - block_start: int - block_end: int - - def __post_init__(self) -> None: - if self.out.ndim != 4: - raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") - if self.lse.shape != self.out.shape[:3]: - raise ValueError("partial attention lse must have shape [B, Hq, Sq]") - if self.out.device != self.lse.device: - raise ValueError("partial attention out/lse must be on the same device") - if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: - raise ValueError("partial attention out/lse must remain FP32 before merge") - if self.block_start < 0: - raise ValueError("block_start must be non-negative") - if self.block_end < self.block_start: - raise ValueError("block_end must be >= block_start") - - -@dataclass(frozen=True) -class AttentionBackwardGradients: - """Training-side gradients emitted by the CP attention backward reference.""" - - dq: torch.Tensor - dk: torch.Tensor - dv: torch.Tensor - - -@dataclass(frozen=True) -class AttentionBackwardPathResult: - """One materialized CP attention backward path.""" - - name: str - out: torch.Tensor - lse: torch.Tensor - gradients: AttentionBackwardGradients - provenance: dict[str, object] - - -@dataclass(frozen=True) -class GradientDriftStats: - """Shape-aware absolute drift summary for backward validation reports.""" - - max_abs: float - mean_abs: float - p95_abs: float - p99_abs: float - active_count: int - - def to_dict(self) -> dict[str, object]: - 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 AttentionBackwardRankDrift: - """Backward drift for one logical CP rank's sequence ownership.""" - - rank: int - dq: GradientDriftStats - dk: GradientDriftStats - dv: GradientDriftStats - - def to_dict(self) -> dict[str, object]: - return { - "rank": self.rank, - "dq": self.dq.to_dict(), - "dk": self.dk.to_dict(), - "dv": self.dv.to_dict(), - } - - -@dataclass(frozen=True) -class AttentionBackwardPathDrift: - """Candidate-vs-reference backward drift for one CP path.""" - - candidate_name: str - dq: GradientDriftStats - dk: GradientDriftStats - dv: GradientDriftStats - out: GradientDriftStats - lse: GradientDriftStats - per_rank: tuple[AttentionBackwardRankDrift, ...] - provenance: dict[str, object] - - def to_dict(self) -> dict[str, object]: - return { - "candidate_name": self.candidate_name, - "dq": self.dq.to_dict(), - "dk": self.dk.to_dict(), - "dv": self.dv.to_dict(), - "out": self.out.to_dict(), - "lse": self.lse.to_dict(), - "per_rank": [item.to_dict() for item in self.per_rank], - "provenance": self.provenance, - } - - -@dataclass(frozen=True) -class AttentionBackwardComparisonReport: - """Structured PR8 report for CP attention gradient drift validation.""" - - reference_name: str - drifts: tuple[AttentionBackwardPathDrift, ...] - - def to_dict(self) -> dict[str, object]: - return { - "reference_name": self.reference_name, - "drifts": [drift.to_dict() for drift in self.drifts], - } - - -def merge_attention_partial_states( - states: Sequence[AttentionPartialState], -) -> AttentionPartialState: - """Merge CP/chunk partial states in logical block order. - - The merge is the online-softmax/LSE merge used by attention, not a plain - sum. The input order is deliberately ignored: states are sorted by logical - ``block_start`` so the result depends on global block indices rather than - arrival order. - """ - - if not states: - raise ValueError("at least one attention partial state is required") - - ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) - _validate_merge_shapes_and_ranges(ordered) - - merged = ordered[0] - merged_out = merged.out.float() - merged_lse = merged.lse.float() - for state in ordered[1:]: - merged_out, merged_lse = _merge_two_states( - merged_out, - merged_lse, - state.out.float(), - state.lse.float(), - ) - - return AttentionPartialState( - out=merged_out, - lse=merged_lse, - block_start=ordered[0].block_start, - block_end=ordered[-1].block_end, - ) - - -class DeterministicCPAttentionReferenceOp: - """Correctness-first CP attention reference for prefill and chunked prefill. - - The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K - have already passed QK-Norm and RoPE unless an outer contract explicitly - marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge - implementation so fused and unfused ``RoPE+Attention`` paths can compare the - same post-RoPE Q/K boundary before validating CP communication. - - The op emulates CP by splitting query and KV sequence dimensions into - logical CP shards. Each query shard computes one partial attention state per - KV block, then merges those states in fixed global-block order using fp32 - LSE arithmetic. ``forward`` returns the input dtype after the final write; - ``forward_fp32`` keeps the fp32 merged output. - """ - - op_class = "attention" - - @staticmethod - def split_kv_execution_plans( - total_kv_tokens: int, - *, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - ) -> list[dict[str, object]]: - """Export the actual logical Split-KV plan before execution.""" - - return split_kv_execution_plan_provenance( - total_kv_tokens, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - backend="deterministic_cp_reference", - ) - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - ) -> torch.Tensor: - return self.forward( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - ) - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - ) -> torch.Tensor: - """Compute CP attention with fp32 accumulation and final input-dtype write.""" - - out, _ = self.forward_with_lse( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - output_dtype=q.dtype, - ) - return out - - def forward_fp32( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - ) -> torch.Tensor: - """Compute CP attention with fp32 accumulation and fp32 output.""" - - out, _ = self.forward_fp32_with_lse( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - ) - return out - - def forward_with_lse( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - output_dtype: Optional[torch.dtype] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Return ``(out, lse)`` for the CP reference path. - - ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 - until the final write, then downcast to ``output_dtype``. When omitted, - ``output_dtype`` defaults to the input dtype. - """ - - resolved_output_dtype = q.dtype if output_dtype is None else output_dtype - _validate_output_dtype(resolved_output_dtype) - out, lse = self._forward_impl( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - ) - out = out.to(resolved_output_dtype) - return out, lse - - def forward_fp32_with_lse( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Return fp32 ``(out, lse)`` for the CP reference path.""" - - return self.forward_with_lse( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - output_dtype=torch.float32, - ) - - def backward_reference( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dout: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - cp_world_size: int = 1, - kv_chunk_size: Optional[int] = None, - output_dtype: Optional[torch.dtype] = torch.float32, - name: Optional[str] = None, - ) -> AttentionBackwardPathResult: - """Run the deterministic training-side backward validation path. - - The semantic backward input is ``dout`` plus the forward attention state - produced from the same Q/K/V, masks, position offsets, CP world, and KV - block order. The reference keeps the softmax/merge math in fp32 and - records the final-write dtype in provenance; decode backward is - intentionally out of scope for PR8. - """ - - _validate_qkv(q, k, v) - if dout.shape != q.shape: - raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") - if not torch.is_floating_point(dout) or torch.is_complex(dout): - raise ValueError("dout must be a real floating-point tensor") - if dout.device != q.device: - raise ValueError("dout must be on the same device as q, k, and v") - if dout.dtype != q.dtype: - raise ValueError("dout must have the same dtype as q") - q_leaf = q.detach().clone().requires_grad_(True) - k_leaf = k.detach().clone().requires_grad_(True) - v_leaf = v.detach().clone().requires_grad_(True) - - resolved_output_dtype = q.dtype if output_dtype is None else output_dtype - _validate_output_dtype(resolved_output_dtype) - out, lse = self.forward_with_lse( - q_leaf, - k_leaf, - v_leaf, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - output_dtype=resolved_output_dtype, - ) - torch.autograd.backward(out, dout.to(dtype=out.dtype)) - if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: - raise RuntimeError("CP attention backward did not produce dq/dk/dv") - - return AttentionBackwardPathResult( - name=name - or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), - out=out.detach(), - lse=lse.detach(), - gradients=AttentionBackwardGradients( - dq=q_leaf.grad.detach(), - dk=k_leaf.grad.detach(), - dv=v_leaf.grad.detach(), - ), - provenance={ - "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", - "gradient_mode": "training_backward", - "gradient_inputs": ["q", "k", "v"], - "gradient_outputs": ["out"], - "saved_forward_state": [ - "out", - "attention_lse", - "causal_mask", - "key_padding_mask", - "query_position_offsets", - "key_position_offsets", - "global_block_index", - ], - "cp_world_size": cp_world_size, - "kv_chunk_size": kv_chunk_size, - "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), - "requested_split_kv_size": kv_chunk_size, - "actual_split_kv_plans": split_kv_execution_plan_provenance( - k.size(2), - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - backend="deterministic_cp_backward_reference", - ), - "merge_order": "global_block_index", - "accum_dtype": "fp32", - "downcast_at": "final_write", - "output_dtype": str(resolved_output_dtype).replace("torch.", ""), - "q_dtype": str(q.dtype).replace("torch.", ""), - "k_dtype": str(k.dtype).replace("torch.", ""), - "v_dtype": str(v.dtype).replace("torch.", ""), - "dout_dtype": str(dout.dtype).replace("torch.", ""), - "te_backward_oracle": "not_used", - "decode_backward": "not_supported", - }, - ) - - def local_partial_state( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - q_start: int, - k_start: int, - total_kv_len: int, - total_query_len: Optional[int] = None, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - ) -> AttentionPartialState: - """Compute one query shard against one logical KV block. - - ``query_position_offsets`` and ``key_position_offsets`` are optional - per-batch-row base positions. They let the reference express varlen or - packed metadata while retaining the dense [B, H, S, D] tensor layout. - For post-RoPE Q/K, these offsets must describe the same absolute token - positions used when RoPE was applied. - """ - - _validate_qkv(q, k, v) - _validate_scale(scale) - if q_start < 0 or k_start < 0: - raise ValueError("q_start and k_start must be non-negative") - if total_kv_len < k_start + k.size(2): - raise ValueError("total_kv_len must cover the local KV block") - if total_query_len is None: - total_query_len = q.size(2) - if total_query_len < q_start + q.size(2): - raise ValueError("total_query_len must cover the local query block") - if key_padding_mask is not None: - if key_padding_mask.shape != (q.size(0), k.size(2)): - raise ValueError("local key_padding_mask must have shape [B, local_skv]") - if key_padding_mask.dtype != torch.bool: - raise ValueError("local key_padding_mask must be bool") - query_offsets = _normalize_position_offsets( - query_position_offsets, - q.size(0), - q.device, - default=total_kv_len - total_query_len, - name="query_position_offsets", - ) - key_offsets = _normalize_position_offsets( - key_position_offsets, - q.size(0), - q.device, - default=0, - name="key_position_offsets", - ) - - ctx = NativeAttentionOp._strict_fp32_math(q.device.type) - with ctx: - qf = q.float() - kf = k.float() - vf = 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) - - if skv == 0: - zero_dep = _zero_dependency(qf, kf, vf) - return AttentionPartialState( - out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) - + zero_dep, - lse=torch.full( - (q.size(0), hq, sq), - float("-inf"), - device=q.device, - dtype=torch.float32, - ) - + zero_dep, - block_start=k_start, - block_end=k_start, - ) - - 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_base = query_offsets[:, None] + q_start - key_base = key_offsets[:, None] + k_start - q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base - k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base - causal_mask = k_pos[:, None, :] > q_pos[:, :, None] - scores = scores.masked_fill(causal_mask[:, 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 AttentionPartialState( - out=out, - lse=lse, - block_start=k_start, - block_end=k_start + skv, - ) - - def _forward_impl( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - causal: bool, - scale: Optional[float], - key_padding_mask: Optional[torch.Tensor], - query_position_offsets: Optional[torch.Tensor], - key_position_offsets: Optional[torch.Tensor], - cp_world_size: int, - kv_chunk_size: Optional[int], - ) -> tuple[torch.Tensor, torch.Tensor]: - _validate_qkv(q, k, v) - _validate_scale(scale) - if ( - isinstance(cp_world_size, bool) - or not isinstance(cp_world_size, int) - or cp_world_size < 1 - ): - raise ValueError("cp_world_size must be >= 1") - if kv_chunk_size is not None and ( - isinstance(kv_chunk_size, bool) - or not isinstance(kv_chunk_size, int) - or kv_chunk_size < 1 - ): - raise ValueError("kv_chunk_size must be >= 1 when provided") - - batch, hq, sq, dim = q.shape - skv = k.size(2) - if key_padding_mask is not None: - if key_padding_mask.shape != (batch, skv): - raise ValueError("key_padding_mask must have shape [B, Skv]") - if key_padding_mask.dtype != torch.bool: - raise ValueError("key_padding_mask must be bool") - query_offsets = _normalize_position_offsets( - query_position_offsets, - batch, - q.device, - default=skv - sq, - name="query_position_offsets", - ) - key_offsets = _normalize_position_offsets( - key_position_offsets, - batch, - q.device, - default=0, - name="key_position_offsets", - ) - - q_bounds = _split_bounds(sq, cp_world_size) - kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) - out_chunks: list[torch.Tensor] = [] - lse_chunks: list[torch.Tensor] = [] - for q_start, q_end in q_bounds: - if q_start == q_end: - continue - q_block = q[:, :, q_start:q_end, :] - states = [ - self.local_partial_state( - q_block, - k[:, :, k_start:k_end, :], - v[:, :, k_start:k_end, :], - q_start=q_start, - k_start=k_start, - total_kv_len=skv, - total_query_len=sq, - causal=causal, - scale=scale, - key_padding_mask=( - None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] - ), - query_position_offsets=query_offsets, - key_position_offsets=key_offsets, - ) - for k_start, k_end in kv_bounds - if k_start != k_end - ] - if states: - merged = merge_attention_partial_states(states) - out_chunks.append(merged.out) - lse_chunks.append(merged.lse) - else: - zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) - out_chunks.append( - torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep - ) - lse_chunks.append( - torch.full( - (batch, hq, q_end - q_start), - float("-inf"), - device=q.device, - dtype=torch.float32, - ) - + zero_dep - ) - - if not out_chunks: - zero_dep = _zero_dependency(q.float(), k.float(), v.float()) - return ( - torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, - torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, - ) - return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) - - -def compare_cp_attention_backward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - dout: torch.Tensor, - *, - causal: bool = True, - scale: Optional[float] = None, - key_padding_mask: Optional[torch.Tensor] = None, - query_position_offsets: Optional[torch.Tensor] = None, - key_position_offsets: Optional[torch.Tensor] = None, - candidate_cp_world_size: int = 2, - candidate_kv_chunk_size: Optional[int] = None, - output_dtype: Optional[torch.dtype] = torch.float32, -) -> AttentionBackwardComparisonReport: - """Compare CP=1 backward with a CP/chunked-prefill candidate. - - The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank - slices. It is a validation/reporting helper, not a separate production - backward kernel. - """ - - op = DeterministicCPAttentionReferenceOp() - reference = op.backward_reference( - q, - k, - v, - dout, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=1, - kv_chunk_size=None, - output_dtype=output_dtype, - name="cp1_backward_reference", - ) - candidate = op.backward_reference( - q, - k, - v, - dout, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=candidate_cp_world_size, - kv_chunk_size=candidate_kv_chunk_size, - output_dtype=output_dtype, - ) - return AttentionBackwardComparisonReport( - reference_name=reference.name, - drifts=(_compare_backward_path(candidate, reference),), - ) - - -def _compare_backward_path( - candidate: AttentionBackwardPathResult, - reference: AttentionBackwardPathResult, -) -> AttentionBackwardPathDrift: - cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") - return AttentionBackwardPathDrift( - candidate_name=candidate.name, - dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), - dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), - dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), - out=_drift_stats(candidate.out, reference.out), - lse=_drift_stats(candidate.lse, reference.lse), - per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), - provenance=candidate.provenance, - ) - - -def _per_rank_backward_drifts( - candidate: AttentionBackwardPathResult, - reference: AttentionBackwardPathResult, - cp_world_size: int, -) -> tuple[AttentionBackwardRankDrift, ...]: - q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) - kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) - per_rank = [] - for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): - per_rank.append( - AttentionBackwardRankDrift( - rank=rank, - dq=_drift_stats( - candidate.gradients.dq[:, :, q_start:q_end, :], - reference.gradients.dq[:, :, q_start:q_end, :], - ), - dk=_drift_stats( - candidate.gradients.dk[:, :, kv_start:kv_end, :], - reference.gradients.dk[:, :, kv_start:kv_end, :], - ), - dv=_drift_stats( - candidate.gradients.dv[:, :, kv_start:kv_end, :], - reference.gradients.dv[:, :, kv_start:kv_end, :], - ), - ) - ) - return tuple(per_rank) - - -def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: - 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().reshape(-1) - active_count = int(diff.numel()) - if active_count == 0: - return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) - return GradientDriftStats( - max_abs=float(diff.max().item()), - mean_abs=float(diff.mean().item()), - p95_abs=float(torch.quantile(diff, 0.95).item()), - p99_abs=float(torch.quantile(diff, 0.99).item()), - active_count=active_count, - ) - - -def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: - prefix = f"cp{cp_world_size}" - if kv_chunk_size is None: - return f"{prefix}_backward" - return f"{prefix}_chunked_backward" - - -def _provenance_int(provenance: dict[str, object], key: str) -> int: - value = provenance[key] - if isinstance(value, bool) or not isinstance(value, int): - raise ValueError(f"provenance field {key!r} must be an int") - return value - - -def _merge_two_states( - out_a: torch.Tensor, - lse_a: torch.Tensor, - out_b: torch.Tensor, - lse_b: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - merged_lse = torch.logaddexp(lse_a, lse_b) - finite = torch.isfinite(merged_lse) - weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) - weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) - merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b - return merged_out, merged_lse - - -def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> 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 out/lse shapes") - if state.block_start != previous_end: - raise ValueError("partial state block ranges must be gap-free and non-overlapping") - previous_end = state.block_end - - -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 the same 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) < 1 or k.size(1) < 1 or q.size(3) < 1: - raise ValueError("q, k, and v must have positive head counts and head dim") - if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( - torch.is_complex(tensor) for tensor in (q, k, v) - ): - raise ValueError("q, k, and v must be real floating-point tensors") - if q.dtype != k.dtype or q.dtype != v.dtype: - raise ValueError("q, k, and v must have the same dtype") - if q.device != k.device or q.device != v.device: - raise ValueError("q, k, and v must be on the same device") - if q.size(1) % k.size(1) != 0: - raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") - - -def _validate_scale(scale: Optional[float]) -> None: - if scale is None: - return - if isinstance(scale, bool) or not isinstance(scale, (int, float)): - raise ValueError("scale must be a positive finite number") - if not math.isfinite(float(scale)) or float(scale) <= 0: - raise ValueError("scale must be a positive finite number") - - -def _validate_output_dtype(output_dtype: torch.dtype) -> None: - if not isinstance(output_dtype, torch.dtype): - raise ValueError("output_dtype must be a real floating-point torch dtype") - probe = torch.empty((), dtype=output_dtype) - if not torch.is_floating_point(probe) or torch.is_complex(probe): - raise ValueError("output_dtype must be a real floating-point torch dtype") - - -def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: - total = torch.tensor(0.0, device=tensors[0].device) - for tensor in tensors: - total = total + tensor.sum() - return total * 0.0 - - -def _normalize_position_offsets( - offsets: Optional[torch.Tensor], - batch: int, - device: torch.device, - *, - default: int, - name: str, -) -> torch.Tensor: - if offsets is None: - return torch.full((batch,), default, dtype=torch.long, device=device) - if offsets.ndim != 1 or offsets.numel() != batch: - raise ValueError(f"{name} must have shape [B]") - if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: - raise ValueError(f"{name} must contain integer positions") - return offsets.to(device=device, dtype=torch.long) - - -def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: - base, extra = divmod(length, parts) - bounds: list[tuple[int, int]] = [] - start = 0 - for index in range(parts): - width = base + (1 if index < extra else 0) - end = start + width - bounds.append((start, end)) - start = end - return bounds - - -def _kv_block_bounds( - length: int, - cp_world_size: int, - kv_chunk_size: Optional[int], -) -> list[tuple[int, int]]: - bounds: list[tuple[int, int]] = [] +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +@dataclass(frozen=True) +class AttentionBackwardGradients: + """Training-side gradients emitted by the CP attention backward reference.""" + + dq: torch.Tensor + dk: torch.Tensor + dv: torch.Tensor + + +@dataclass(frozen=True) +class AttentionBackwardPathResult: + """One materialized CP attention backward path.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + gradients: AttentionBackwardGradients + provenance: dict[str, object] + + +@dataclass(frozen=True) +class GradientDriftStats: + """Shape-aware absolute drift summary for backward validation reports.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, object]: + 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 AttentionBackwardRankDrift: + """Backward drift for one logical CP rank's sequence ownership.""" + + rank: int + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + + def to_dict(self) -> dict[str, object]: + return { + "rank": self.rank, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + } + + +@dataclass(frozen=True) +class AttentionBackwardPathDrift: + """Candidate-vs-reference backward drift for one CP path.""" + + candidate_name: str + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + out: GradientDriftStats + lse: GradientDriftStats + per_rank: tuple[AttentionBackwardRankDrift, ...] + provenance: dict[str, object] + + def to_dict(self) -> dict[str, object]: + return { + "candidate_name": self.candidate_name, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "per_rank": [item.to_dict() for item in self.per_rank], + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionBackwardComparisonReport: + """Structured PR8 report for CP attention gradient drift validation.""" + + reference_name: str + drifts: tuple[AttentionBackwardPathDrift, ...] + + def to_dict(self) -> dict[str, object]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + } + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + op_class = "attention" + + @staticmethod + def split_kv_execution_plans( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> list[dict[str, object]]: + """Export the actual logical Split-KV plan before execution.""" + + return split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(resolved_output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def backward_reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, + name: Optional[str] = None, + ) -> AttentionBackwardPathResult: + """Run the deterministic training-side backward validation path. + + The semantic backward input is ``dout`` plus the forward attention state + produced from the same Q/K/V, masks, position offsets, CP world, and KV + block order. The reference keeps the softmax/merge math in fp32 and + records the final-write dtype in provenance; decode backward is + intentionally out of scope for PR8. + """ + + _validate_qkv(q, k, v) + if dout.shape != q.shape: + raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") + if not torch.is_floating_point(dout) or torch.is_complex(dout): + raise ValueError("dout must be a real floating-point tensor") + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") + q_leaf = q.detach().clone().requires_grad_(True) + k_leaf = k.detach().clone().requires_grad_(True) + v_leaf = v.detach().clone().requires_grad_(True) + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=resolved_output_dtype, + ) + torch.autograd.backward(out, dout.to(dtype=out.dtype)) + if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: + raise RuntimeError("CP attention backward did not produce dq/dk/dv") + + return AttentionBackwardPathResult( + name=name + or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), + out=out.detach(), + lse=lse.detach(), + gradients=AttentionBackwardGradients( + dq=q_leaf.grad.detach(), + dk=k_leaf.grad.detach(), + dv=v_leaf.grad.detach(), + ), + provenance={ + "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", + "gradient_mode": "training_backward", + "gradient_inputs": ["q", "k", "v"], + "gradient_outputs": ["out"], + "saved_forward_state": [ + "out", + "attention_lse", + "causal_mask", + "key_padding_mask", + "query_position_offsets", + "key_position_offsets", + "global_block_index", + ], + "cp_world_size": cp_world_size, + "kv_chunk_size": kv_chunk_size, + "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_backward_reference", + ), + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(resolved_output_dtype).replace("torch.", ""), + "q_dtype": str(q.dtype).replace("torch.", ""), + "k_dtype": str(k.dtype).replace("torch.", ""), + "v_dtype": str(v.dtype).replace("torch.", ""), + "dout_dtype": str(dout.dtype).replace("torch.", ""), + "te_backward_oracle": "not_used", + "decode_backward": "not_supported", + }, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. + """ + + _validate_qkv(q, k, v) + _validate_scale(scale) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = 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) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + 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_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, 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 AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + _validate_scale(scale) + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def compare_cp_attention_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + candidate_cp_world_size: int = 2, + candidate_kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, +) -> AttentionBackwardComparisonReport: + """Compare CP=1 backward with a CP/chunked-prefill candidate. + + The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank + slices. It is a validation/reporting helper, not a separate production + backward kernel. + """ + + op = DeterministicCPAttentionReferenceOp() + reference = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=1, + kv_chunk_size=None, + output_dtype=output_dtype, + name="cp1_backward_reference", + ) + candidate = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=candidate_cp_world_size, + kv_chunk_size=candidate_kv_chunk_size, + output_dtype=output_dtype, + ) + return AttentionBackwardComparisonReport( + reference_name=reference.name, + drifts=(_compare_backward_path(candidate, reference),), + ) + + +def _compare_backward_path( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, +) -> AttentionBackwardPathDrift: + cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") + return AttentionBackwardPathDrift( + candidate_name=candidate.name, + dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), + dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), + dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), + provenance=candidate.provenance, + ) + + +def _per_rank_backward_drifts( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, + cp_world_size: int, +) -> tuple[AttentionBackwardRankDrift, ...]: + q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) + kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) + per_rank = [] + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + per_rank.append( + AttentionBackwardRankDrift( + rank=rank, + dq=_drift_stats( + candidate.gradients.dq[:, :, q_start:q_end, :], + reference.gradients.dq[:, :, q_start:q_end, :], + ), + dk=_drift_stats( + candidate.gradients.dk[:, :, kv_start:kv_end, :], + reference.gradients.dk[:, :, kv_start:kv_end, :], + ), + dv=_drift_stats( + candidate.gradients.dv[:, :, kv_start:kv_end, :], + reference.gradients.dv[:, :, kv_start:kv_end, :], + ), + ) + ) + return tuple(per_rank) + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: + 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().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return GradientDriftStats( + max_abs=float(diff.max().item()), + mean_abs=float(diff.mean().item()), + p95_abs=float(torch.quantile(diff, 0.95).item()), + p99_abs=float(torch.quantile(diff, 0.99).item()), + active_count=active_count, + ) + + +def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: + prefix = f"cp{cp_world_size}" + if kv_chunk_size is None: + return f"{prefix}_backward" + return f"{prefix}_chunked_backward" + + +def _provenance_int(provenance: dict[str, object], key: str) -> int: + value = provenance[key] + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"provenance field {key!r} must be an int") + return value + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> 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 out/lse shapes") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") + previous_end = state.block_end + + +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 the same 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) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] for start, end in _split_bounds(length, cp_world_size): if kv_chunk_size is None: - bounds.append((start, end)) - continue - cursor = start - while cursor < end: - chunk_end = min(cursor + kv_chunk_size, end) - bounds.append((cursor, chunk_end)) - cursor = chunk_end - return bounds - - -def split_kv_execution_plan_provenance( - length: int, - *, - cp_world_size: int, - kv_chunk_size: Optional[int], - backend: str, -) -> list[dict[str, object]]: - """Return the actual backend-local Split-KV plan for every CP owner.""" - - if length < 1: - raise ValueError("Split-KV sequence length must be >= 1") - if cp_world_size < 1: - raise ValueError("cp_world_size must be >= 1") - if kv_chunk_size is not None and kv_chunk_size < 1: - raise ValueError("kv_chunk_size must be >= 1 when provided") - result: list[dict[str, object]] = [] + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +def split_kv_execution_plan_provenance( + length: int, + *, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str, +) -> list[dict[str, object]]: + """Return the actual backend-local Split-KV plan for every CP owner.""" + + if length < 1: + raise ValueError("Split-KV sequence length must be >= 1") + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + result: list[dict[str, object]] = [] for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): if rank_start == rank_end: continue + boundaries: tuple[tuple[int, int], ...] if kv_chunk_size is None: - boundaries = ((rank_start, rank_end),) - mode = SplitKVMode.DISABLED - else: - boundaries = tuple( - (start, min(start + kv_chunk_size, rank_end)) - for start in range(rank_start, rank_end, kv_chunk_size) - ) - mode = SplitKVMode.FIXED - plan = SplitKVExecutionPlan( - requested_mode=mode, - requested_split_size=kv_chunk_size, - actual_mode=mode, - actual_split_size=kv_chunk_size, - boundaries=boundaries, - backend=backend, - source="reference_execution", - ) - result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) - return result - - -def build_reference_split_kv_runtime_plan_set( - total_kv_tokens: Sequence[int], - *, - tp_world_size: int, - cp_world_size: int, - kv_chunk_size: Optional[int], - backend: str = "deterministic_cp_reference", -) -> SplitKVRuntimePlanSet: - """Build complete per-batch/TP/CP/owner plans for the reference path.""" - - totals = tuple(total_kv_tokens) - if not totals or any(total < cp_world_size for total in totals): - raise ValueError("reference runtime plan sets require at least one KV token per CP owner") - if tp_world_size < 1 or cp_world_size < 1: - raise ValueError("TP and CP world sizes must be >= 1") - if kv_chunk_size is not None and kv_chunk_size < 1: - raise ValueError("kv_chunk_size must be >= 1 when provided") - - entries: list[SplitKVRuntimePlanEntry] = [] - for batch_index, total in enumerate(totals): - owner_ranges = _split_bounds(total, cp_world_size) - for tp_rank in range(tp_world_size): - for cp_rank in range(cp_world_size): - for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + boundaries = ((rank_start, rank_end),) + mode = SplitKVMode.DISABLED + else: + boundaries = tuple( + (start, min(start + kv_chunk_size, rank_end)) + for start in range(rank_start, rank_end, kv_chunk_size) + ) + mode = SplitKVMode.FIXED + plan = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) + return result + + +def build_reference_split_kv_runtime_plan_set( + total_kv_tokens: Sequence[int], + *, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str = "deterministic_cp_reference", +) -> SplitKVRuntimePlanSet: + """Build complete per-batch/TP/CP/owner plans for the reference path.""" + + totals = tuple(total_kv_tokens) + if not totals or any(total < cp_world_size for total in totals): + raise ValueError("reference runtime plan sets require at least one KV token per CP owner") + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("TP and CP world sizes must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + owner_ranges = _split_bounds(total, cp_world_size) + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + boundaries: tuple[tuple[int, int], ...] if kv_chunk_size is None: - mode = SplitKVMode.DISABLED - boundaries = ((owner_start, owner_end),) - else: - mode = SplitKVMode.FIXED - boundaries = tuple( - (start, min(start + kv_chunk_size, owner_end)) - for start in range(owner_start, owner_end, kv_chunk_size) - ) - execution = SplitKVExecutionPlan( - requested_mode=mode, - requested_split_size=kv_chunk_size, - actual_mode=mode, - actual_split_size=kv_chunk_size, - boundaries=boundaries, - backend=backend, - source="reference_execution", - ) - entries.append( - SplitKVRuntimePlanEntry( - coordinate=SplitKVRuntimeCoordinate( - batch_index=batch_index, - tp_rank=tp_rank, - cp_rank=cp_rank, - owner_cp_rank=owner_cp_rank, - ), - expected_kv_range=(owner_start, owner_end), - execution=execution, - ) - ) - return SplitKVRuntimePlanSet( - batch_size=len(totals), - tp_world_size=tp_world_size, - cp_world_size=cp_world_size, - total_kv_tokens=totals, - entries=tuple(entries), - ) - - -CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp - -__all__ = [ - "AttentionBackwardComparisonReport", - "AttentionBackwardGradients", - "AttentionBackwardPathDrift", - "AttentionBackwardPathResult", - "AttentionBackwardRankDrift", - "AttentionPartialState", - "build_reference_split_kv_runtime_plan_set", - "CPAttentionReferenceOp", - "DeterministicCPAttentionReferenceOp", - "GradientDriftStats", - "compare_cp_attention_backward", - "merge_attention_partial_states", - "split_kv_execution_plan_provenance", -] + mode = SplitKVMode.DISABLED + boundaries = ((owner_start, owner_end),) + else: + mode = SplitKVMode.FIXED + boundaries = tuple( + (start, min(start + kv_chunk_size, owner_end)) + for start in range(owner_start, owner_end, kv_chunk_size) + ) + execution = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", + "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", + "merge_attention_partial_states", + "split_kv_execution_plan_provenance", +] From da12ac11056ca58f2c40bc73027f8523269ecb89 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 23:12:21 +0800 Subject: [PATCH 09/16] docs(attention): link PR1 contract dependency --- docs/operators/attention.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index bc7f4a38..c8aac94e 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -94,7 +94,8 @@ before selecting a backend. Legacy `get_op("attention")` behavior remains unchan Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as -a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). +a silent fallback. See the +[WS2 CP-aware Attention contract in PR #236](https://github.com/RL-Align/RL-Kernel/blob/feat/issue-235-attention-cp-contract/docs/design/ws2-cp-attention-contract.md). Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow `disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block From c1cdfe5321dc36831acb823ffbdcfd8eae0e4f21 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:48:04 +0800 Subject: [PATCH 10/16] docs(attention): record owner-local reference boundary --- .../ops/pytorch/attention/cp_attention.py | 40 +++++++++++++++++-- tests/test_cp_attention.py | 21 ++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 2394b98d..840ec655 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -209,7 +209,7 @@ class DeterministicCPAttentionReferenceOp: op_class = "attention" @staticmethod - def split_kv_execution_plans( + def split_kv_execution_plans( total_kv_tokens: int, *, cp_world_size: int = 1, @@ -217,12 +217,44 @@ def split_kv_execution_plans( ) -> list[dict[str, object]]: """Export the actual logical Split-KV plan before execution.""" - return split_kv_execution_plan_provenance( + return split_kv_execution_plan_provenance( total_kv_tokens, cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, - backend="deterministic_cp_reference", - ) + backend="deterministic_cp_reference", + ) + + @staticmethod + def execution_provenance( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> dict[str, object]: + """Describe the reference boundary without claiming production communication.""" + + plans = split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + return { + "execution_scope": "logical_single_process_cp_reference", + "runtime_verified": False, + "input_boundary": "projected_post_qk_norm_post_rope_qkv", + "query_scope": "logical_global_query_reference", + "kv_scope": "logical_owner_local_cp_shards", + "production_cp_protocol": "ag_query_local_kv_rs_out_lse", + "communication_executed": "none", + "partial_state": "fp32_out_attention_lse", + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "requested_split_kv_policy": "disabled" if kv_chunk_size is None else "fixed", + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": plans, + } def __call__( self, diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index cc93874b..db42337d 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -626,6 +626,27 @@ def test_invalid_scale_fails_before_attention_math(scale): op.forward_fp32_with_lse(q, k, v, scale=scale) +def test_forward_reference_provenance_does_not_claim_production_communication(): + provenance = DeterministicCPAttentionReferenceOp.execution_provenance( + 8, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert provenance["execution_scope"] == "logical_single_process_cp_reference" + assert provenance["query_scope"] == "logical_global_query_reference" + assert provenance["kv_scope"] == "logical_owner_local_cp_shards" + assert provenance["production_cp_protocol"] == "ag_query_local_kv_rs_out_lse" + assert provenance["communication_executed"] == "none" + assert provenance["merge_order"] == "global_block_index" + plans = provenance["actual_split_kv_plans"] + assert [plan["owner_cp_rank"] for plan in plans] == [0, 1] + assert [plan["actual_split_boundaries"] for plan in plans] == [ + [[0, 2], [2, 4]], + [[4, 6], [6, 8]], + ] + + def test_qkv_dtype_and_floating_contract_fails_closed(): op = DeterministicCPAttentionReferenceOp() q, k, v = _qkv(1, 4, 4, seed=25, heads=4, kv_heads=2, dim=8) From afa00dc077711c251a95bcbf13cc0c600aeb8a28 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:49:12 +0800 Subject: [PATCH 11/16] feat(attention): align projection contract --- rl_engine/kernels/attention_contract.py | 144 ++++++++++++++++++++++++ 1 file changed, 144 insertions(+) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 1750476d..49aae668 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -61,6 +61,13 @@ class SplitKVMode(str, Enum): AUTO = "auto" +class ProjectionCollective(str, Enum): + NONE = "none" + ALL_REDUCE = "all_reduce" + ALL_GATHER = "all_gather" + REDUCE_SCATTER = "reduce_scatter" + + class RoPEState(str, Enum): PRE_ROPE = "pre_rope" POST_ROPE = "post_rope" @@ -134,12 +141,16 @@ class ShardingSpec: global_block_token_starts: tuple[int, ...] local_block_offsets: tuple[int, ...] packed_sequence_offsets: tuple[int, ...] | None = None + sp_rank: int = 0 + sp_world_size: int = 1 def __post_init__(self) -> None: tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + sp_world_size = _positive_int(self.sp_world_size, "sp_world_size") tp_rank = _non_negative_int(self.tp_rank, "tp_rank") cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + sp_rank = _non_negative_int(self.sp_rank, "sp_rank") if tp_rank >= tp_world_size: raise AttentionContractError( f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" @@ -148,6 +159,10 @@ def __post_init__(self) -> None: raise AttentionContractError( f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" ) + if sp_rank >= sp_world_size: + raise AttentionContractError( + f"sp_rank={sp_rank} must be smaller than sp_world_size={sp_world_size}" + ) global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") @@ -1076,6 +1091,98 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class AttentionProjectionSpec: + """Deterministic QKV or output-projection execution contract.""" + + name: str + input_dtype: AttentionDType = AttentionDType.BF16 + output_dtype: AttentionDType = AttentionDType.BF16 + acc_dtype: AttentionDType = AttentionDType.FP32 + split_kv: SplitKVMode = SplitKVMode.DISABLED + k_order: str = "ascending" + backend_policy: str = "native_verified_then_common_deterministic" + deterministic_backend: str = "rlkernel.cuda.det_gemm" + tp_forward_collective: ProjectionCollective = ProjectionCollective.NONE + tp_backward_dgrad_collective: ProjectionCollective = ProjectionCollective.NONE + sp_forward_collective: ProjectionCollective = ProjectionCollective.NONE + sp_backward_collective: ProjectionCollective = ProjectionCollective.NONE + qkv_split_order: tuple[str, ...] = () + require_runtime_readback: bool = True + + @classmethod + def qkv(cls) -> "AttentionProjectionSpec": + return cls( + name="qkv", + tp_backward_dgrad_collective=ProjectionCollective.ALL_REDUCE, + sp_forward_collective=ProjectionCollective.ALL_GATHER, + sp_backward_collective=ProjectionCollective.REDUCE_SCATTER, + qkv_split_order=("q", "k", "v"), + ) + + @classmethod + def output(cls) -> "AttentionProjectionSpec": + return cls( + name="o_proj", + tp_forward_collective=ProjectionCollective.ALL_REDUCE, + sp_forward_collective=ProjectionCollective.REDUCE_SCATTER, + sp_backward_collective=ProjectionCollective.ALL_GATHER, + ) + + def __post_init__(self) -> None: + if self.name not in {"qkv", "o_proj"}: + raise AttentionContractError("projection name must be qkv or o_proj") + for field_name in ("input_dtype", "output_dtype", "acc_dtype"): + object.__setattr__( + self, + field_name, + _enum_value(AttentionDType, getattr(self, field_name), field_name), + ) + object.__setattr__( + self, + "split_kv", + _enum_value(SplitKVMode, self.split_kv, "projection.split_kv"), + ) + for field_name in ( + "tp_forward_collective", + "tp_backward_dgrad_collective", + "sp_forward_collective", + "sp_backward_collective", + ): + object.__setattr__( + self, + field_name, + _enum_value( + ProjectionCollective, + getattr(self, field_name), + field_name, + ), + ) + if self.input_dtype is not AttentionDType.BF16: + raise AttentionContractError("Attention projections require BF16 input") + if self.output_dtype is not AttentionDType.BF16: + raise AttentionContractError("Attention projections require BF16 output") + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Attention projections require FP32 accumulation") + if self.split_kv is not SplitKVMode.DISABLED: + raise AttentionContractError("Attention projection GEMMs must disable Split-K") + if self.k_order != "ascending": + raise AttentionContractError("Attention projection GEMMs require ascending K order") + for field_name in ("backend_policy", "deterministic_backend"): + value = getattr(self, field_name) + if not isinstance(value, str) or not value.strip(): + raise AttentionContractError(f"{field_name} must be a non-empty string") + if not isinstance(self.require_runtime_readback, bool) or not self.require_runtime_readback: + raise AttentionContractError("projection runtime readback must be required") + split_order = tuple(self.qkv_split_order) + if self.name == "qkv": + if split_order != ("q", "k", "v"): + raise AttentionContractError("QKV projection must split in Q, K, V order") + elif split_order: + raise AttentionContractError("o_proj must not declare a QKV split order") + object.__setattr__(self, "qkv_split_order", split_order) + + @dataclass(frozen=True) class AttentionContract: """Complete semantic request consumed by contract-aware dispatch.""" @@ -1094,6 +1201,10 @@ class AttentionContract: kv_cache: KVCacheSpec | None = None rope: RoPESpec | None = None export_lse: bool = True + qkv_projection: AttentionProjectionSpec = field(default_factory=AttentionProjectionSpec.qkv) + output_projection: AttentionProjectionSpec = field( + default_factory=AttentionProjectionSpec.output + ) def __post_init__(self) -> None: object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) @@ -1108,6 +1219,16 @@ def __post_init__(self) -> None: raise AttentionContractError("reduction must be a ReductionSpec") if not isinstance(self.split_kv, SplitKVSpec): raise AttentionContractError("split_kv must be a SplitKVSpec") + if not isinstance(self.qkv_projection, AttentionProjectionSpec): + raise AttentionContractError("qkv_projection must be an AttentionProjectionSpec") + if self.qkv_projection.name != "qkv": + raise AttentionContractError("qkv_projection must use name='qkv'") + if not isinstance(self.output_projection, AttentionProjectionSpec): + raise AttentionContractError( + "output_projection must be an AttentionProjectionSpec" + ) + if self.output_projection.name != "o_proj": + raise AttentionContractError("output_projection must use name='o_proj'") if ( self.mode is AttentionMode.PREFILL and query_sequence_length != self.sharding.local_sequence_length @@ -1192,6 +1313,8 @@ def to_dict(self) -> dict[str, Any]: "tp_world_size": self.sharding.tp_world_size, "cp_rank": self.sharding.cp_rank, "cp_world_size": self.sharding.cp_world_size, + "sp_rank": self.sharding.sp_rank, + "sp_world_size": self.sharding.sp_world_size, "global_q_heads": self.sharding.global_q_heads, "global_kv_heads": self.sharding.global_kv_heads, "local_q_head_start": self.sharding.local_q_head_start, @@ -1254,6 +1377,24 @@ def to_dict(self) -> dict[str, Any]: "output_dtype": self.rope.output_dtype.value, "fusion_boundary": self.rope.fusion_boundary.value, } + projections = { + spec.name: { + "input_dtype": spec.input_dtype.value, + "output_dtype": spec.output_dtype.value, + "acc_dtype": spec.acc_dtype.value, + "split_kv": spec.split_kv.value, + "k_order": spec.k_order, + "backend_policy": spec.backend_policy, + "deterministic_backend": spec.deterministic_backend, + "tp_forward_collective": spec.tp_forward_collective.value, + "tp_backward_dgrad_collective": spec.tp_backward_dgrad_collective.value, + "sp_forward_collective": spec.sp_forward_collective.value, + "sp_backward_collective": spec.sp_backward_collective.value, + "qkv_split_order": list(spec.qkv_split_order), + "require_runtime_readback": spec.require_runtime_readback, + } + for spec in (self.qkv_projection, self.output_projection) + } return { "semantic_operator": "standard_softmax_attention", "role": self.role.value, @@ -1273,6 +1414,7 @@ def to_dict(self) -> dict[str, Any]: "split_kv": self.split_kv.to_dict(), "kv_cache": kv_cache, "rope": rope, + "projections": projections, } @@ -1428,6 +1570,7 @@ class AttentionDispatchResult: "AttentionBackendCapability", "AttentionDispatchResult", "AttentionDType", + "AttentionProjectionSpec", "AttentionMerge", "AttentionMode", "AttentionRole", @@ -1436,6 +1579,7 @@ class AttentionDispatchResult: "ReductionEngine", "ReductionOrder", "ReductionSpec", + "ProjectionCollective", "RoPECastPoint", "RoPEFusionBoundary", "RoPESpec", From 416642ae4c03f1a0ec724bda3b065ec2a65f4e88 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Mon, 17 Aug 2026 18:07:28 +0800 Subject: [PATCH 12/16] feat(attention): expose deterministic core reference contract --- rl_engine/kernels/attention_contract.py | 4 + .../ops/pytorch/attention/cp_attention.py | 115 +++++++++++++++--- tests/test_cp_attention.py | 23 ++++ 3 files changed, 123 insertions(+), 19 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 49aae668..cc532f93 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -18,6 +18,9 @@ _EnumT = TypeVar("_EnumT", bound=Enum) +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" + + class AttentionContractError(ValueError): """Raised when attention metadata does not describe a valid invocation.""" @@ -1591,6 +1594,7 @@ class AttentionDispatchResult: "SplitKVRuntimePlanEntry", "SplitKVRuntimePlanSet", "SplitKVSpec", + "STRICT_ATTENTION_CORE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 840ec655..f3ad3bbe 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -16,18 +16,20 @@ import torch -from rl_engine.kernels.attention_contract import ( - SplitKVExecutionPlan, - SplitKVMode, - SplitKVRuntimeCoordinate, - SplitKVRuntimePlanEntry, - SplitKVRuntimePlanSet, -) -from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp - - -@dataclass(frozen=True) -class AttentionPartialState: +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + STRICT_ATTENTION_CORE_ID, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: """One KV block's attention state before deterministic LSE merge. ``out`` is already normalized within the local KV block and has shape @@ -52,8 +54,80 @@ def __post_init__(self) -> None: raise ValueError("partial attention out/lse must remain FP32 before merge") if self.block_start < 0: raise ValueError("block_start must be non-negative") - if self.block_end < self.block_start: - raise ValueError("block_end must be >= block_start") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +@dataclass(frozen=True) +class DeterministicAttentionCoreResult: + """Strict-core output and the exact arithmetic plan used for it.""" + + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, object] + + +class DeterministicAttentionCore: + """Common FP32 Attention arithmetic used by both train and rollout.""" + + core_id = STRICT_ATTENTION_CORE_ID + backend_id = "rlkernel.attention.cp_reference" + merge_order = "global_block_index" + accum_dtype = "fp32" + downcast_at = "final_write" + + def __init__(self, *, split_kv: SplitKVSpec | None = None) -> None: + self.split_kv = SplitKVSpec.disabled() if split_kv is None else split_kv + if not isinstance(self.split_kv, SplitKVSpec): + raise TypeError("split_kv must be a SplitKVSpec") + if self.split_kv.mode is SplitKVMode.AUTO: + raise ValueError("strict deterministic Attention core does not allow auto Split-KV") + self._reference = DeterministicCPAttentionReferenceOp() + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: float | None = None, + key_padding_mask: torch.Tensor | None = None, + query_position_offsets: torch.Tensor | None = None, + key_position_offsets: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + ) -> DeterministicAttentionCoreResult: + out, lse = self._reference.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=1, + kv_chunk_size=self.split_kv.fixed_split_size, + ) + resolved_dtype = q.dtype if output_dtype is None else output_dtype + if resolved_dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict Attention output_dtype must be FP16 or BF16") + plan = self.split_kv.resolve(k.size(2), backend=self.backend_id) + return DeterministicAttentionCoreResult( + out=out.to(dtype=resolved_dtype), + lse=lse, + provenance={ + "strict_core_id": self.core_id, + "attention_backend": self.backend_id, + "split_kv": plan.to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "fallback": False, + "fallback_reason": None, + "native_attention_arithmetic": False, + }, + ) @dataclass(frozen=True) @@ -1100,11 +1174,14 @@ def build_reference_split_kv_runtime_plan_set( "AttentionBackwardPathDrift", "AttentionBackwardPathResult", "AttentionBackwardRankDrift", - "AttentionPartialState", - "build_reference_split_kv_runtime_plan_set", - "CPAttentionReferenceOp", - "DeterministicCPAttentionReferenceOp", - "GradientDriftStats", + "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", + "DeterministicAttentionCore", + "DeterministicAttentionCoreResult", + "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "STRICT_ATTENTION_CORE_ID", "compare_cp_attention_backward", "merge_attention_partial_states", "split_kv_execution_plan_provenance", diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index db42337d..e057b2fb 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -14,9 +14,12 @@ import pytest import torch +from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, + DeterministicAttentionCore, DeterministicCPAttentionReferenceOp, + STRICT_ATTENTION_CORE_ID, compare_cp_attention_backward, merge_attention_partial_states, split_kv_execution_plan_provenance, @@ -83,6 +86,26 @@ def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): return torch.logsumexp(scores, dim=-1) +def test_strict_attention_core_freezes_plan_and_final_write(): + q, k, v = _qkv(1, 3, 4, seed=17, dtype=torch.bfloat16) + core = DeterministicAttentionCore() + result = core.forward_with_lse(q, k, v, output_dtype=torch.bfloat16) + + assert result.out.dtype is torch.bfloat16 + assert result.lse.dtype is torch.float32 + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["merge_order"] == "global_block_index" + assert result.provenance["accum_dtype"] == "fp32" + assert result.provenance["downcast_at"] == "final_write" + assert result.provenance["fallback"] is False + assert result.provenance["native_attention_arithmetic"] is False + + +def test_strict_attention_core_rejects_auto_split_kv(): + with pytest.raises(ValueError, match="does not allow auto Split-KV"): + DeterministicAttentionCore(split_kv=SplitKVSpec.auto()) + + def test_cp1_matches_native_attention_and_exports_lse(): op = DeterministicCPAttentionReferenceOp() native = NativeAttentionOp() From 05a22e67e58951401e8abb533261ab74b5085248 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Mon, 17 Aug 2026 18:48:32 +0800 Subject: [PATCH 13/16] fix(attention): keep strict reference on no-split core Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 8 +++++--- tests/test_cp_attention.py | 7 ++++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index f3ad3bbe..f20ab9f5 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -80,8 +80,10 @@ def __init__(self, *, split_kv: SplitKVSpec | None = None) -> None: self.split_kv = SplitKVSpec.disabled() if split_kv is None else split_kv if not isinstance(self.split_kv, SplitKVSpec): raise TypeError("split_kv must be a SplitKVSpec") - if self.split_kv.mode is SplitKVMode.AUTO: - raise ValueError("strict deterministic Attention core does not allow auto Split-KV") + if self.split_kv.mode is not SplitKVMode.DISABLED: + raise ValueError( + "strict deterministic Attention core requires Split-KV to be disabled" + ) self._reference = DeterministicCPAttentionReferenceOp() def forward_with_lse( @@ -107,7 +109,7 @@ def forward_with_lse( query_position_offsets=query_position_offsets, key_position_offsets=key_position_offsets, cp_world_size=1, - kv_chunk_size=self.split_kv.fixed_split_size, + kv_chunk_size=None, ) resolved_dtype = q.dtype if output_dtype is None else output_dtype if resolved_dtype not in (torch.float16, torch.bfloat16): diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index e057b2fb..5d60c959 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -101,9 +101,10 @@ def test_strict_attention_core_freezes_plan_and_final_write(): assert result.provenance["native_attention_arithmetic"] is False -def test_strict_attention_core_rejects_auto_split_kv(): - with pytest.raises(ValueError, match="does not allow auto Split-KV"): - DeterministicAttentionCore(split_kv=SplitKVSpec.auto()) +@pytest.mark.parametrize("split_kv", [SplitKVSpec.fixed(2), SplitKVSpec.auto()]) +def test_strict_attention_core_rejects_split_kv(split_kv): + with pytest.raises(ValueError, match="requires Split-KV to be disabled"): + DeterministicAttentionCore(split_kv=split_kv) def test_cp1_matches_native_attention_and_exports_lse(): From 321b6ed2f72908f10b8da19bd80a52bebd43ce28 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 21:25:19 +0800 Subject: [PATCH 14/16] feat(attention): add canonical bitwise reference schedule --- .../ops/pytorch/attention/cp_attention.py | 170 ++++++++++++++++-- tests/test_cp_attention.py | 49 +++++ 2 files changed, 201 insertions(+), 18 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index f20ab9f5..bd4d0cb5 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -84,7 +84,7 @@ def __init__(self, *, split_kv: SplitKVSpec | None = None) -> None: raise ValueError( "strict deterministic Attention core requires Split-KV to be disabled" ) - self._reference = DeterministicCPAttentionReferenceOp() + self._reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) def forward_with_lse( self, @@ -128,6 +128,7 @@ def forward_with_lse( "fallback": False, "fallback_reason": None, "native_attention_arithmetic": False, + "strict_schedule": "single_batch_single_query_global_kv_blocks", }, ) @@ -266,7 +267,7 @@ def merge_attention_partial_states( ) -class DeterministicCPAttentionReferenceOp: +class DeterministicCPAttentionReferenceOp: """Correctness-first CP attention reference for prefill and chunked prefill. The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K @@ -282,7 +283,20 @@ class DeterministicCPAttentionReferenceOp: ``forward_fp32`` keeps the fp32 merged output. """ - op_class = "attention" + op_class = "attention" + + def __init__(self, *, strict_bitwise: bool = False) -> None: + """Create the reference op. + + ``strict_bitwise`` uses a canonical schedule that is independent of + batch size and CP ownership. The regular path remains vectorized for + drift/performance experiments; the strict path is intentionally slower + because it is the shared arithmetic reference for train and rollout. + """ + + if not isinstance(strict_bitwise, bool): + raise TypeError("strict_bitwise must be a bool") + self.strict_bitwise = strict_bitwise @staticmethod def split_kv_execution_plans( @@ -444,18 +458,31 @@ def forward_with_lse( resolved_output_dtype = q.dtype if output_dtype is None else output_dtype _validate_output_dtype(resolved_output_dtype) - out, lse = self._forward_impl( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - ) + if self.strict_bitwise: + out, lse = self._forward_strict_bitwise( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + kv_chunk_size=kv_chunk_size, + ) + else: + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) out = out.to(resolved_output_dtype) return out, lse @@ -475,7 +502,7 @@ def forward_fp32_with_lse( ) -> tuple[torch.Tensor, torch.Tensor]: """Return fp32 ``(out, lse)`` for the CP reference path.""" - return self.forward_with_lse( + return self.forward_with_lse( q, k, v, @@ -486,8 +513,115 @@ def forward_fp32_with_lse( key_position_offsets=key_position_offsets, cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, - output_dtype=torch.float32, - ) + output_dtype=torch.float32, + ) + + def _forward_strict_bitwise( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run one canonical arithmetic schedule for every caller. + + The regular implementation changes GEMM shapes when batch/CP changes; + that is numerically valid but cannot be bitwise invariant. Strict mode + fixes every matmul to ``[1, H, 1, D]`` by processing one batch row and + one query position at a time. KV blocks are global and independent of + ``cp_world_size``; communication only determines ownership outside this + reference. Both training and rollout therefore execute the same + score, softmax, value and LSE-merge operations in the same order. + """ + + _validate_qkv(q, k, v) + _validate_scale(scale) + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + # The canonical schedule is global. CP ownership and arrival order + # must not change which partial states are generated or merged. + kv_bounds = _kv_block_bounds(skv, 1, kv_chunk_size) + out_rows: list[torch.Tensor] = [] + lse_rows: list[torch.Tensor] = [] + for batch_index in range(batch): + q_batch = q[batch_index : batch_index + 1].contiguous() + k_batch = k[batch_index : batch_index + 1].contiguous() + v_batch = v[batch_index : batch_index + 1].contiguous() + pad_batch = ( + None + if key_padding_mask is None + else key_padding_mask[batch_index : batch_index + 1].contiguous() + ) + query_offset = query_offsets[batch_index : batch_index + 1] + key_offset = key_offsets[batch_index : batch_index + 1] + query_rows: list[torch.Tensor] = [] + lse_query_rows: list[torch.Tensor] = [] + for query_index in range(sq): + q_row = q_batch[:, :, query_index : query_index + 1, :].contiguous() + states = [ + self.local_partial_state( + q_row, + k_batch[:, :, key_start:key_end, :].contiguous(), + v_batch[:, :, key_start:key_end, :].contiguous(), + q_start=query_index, + k_start=key_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None + if pad_batch is None + else pad_batch[:, key_start:key_end].contiguous() + ), + query_position_offsets=query_offset, + key_position_offsets=key_offset, + ) + for key_start, key_end in kv_bounds + if key_start != key_end + ] + if not states: + zero_dep = _zero_dependency(q_row.float(), k_batch.float(), v_batch.float()) + merged_out = torch.zeros( + 1, hq, 1, dim, device=q.device, dtype=torch.float32 + ) + zero_dep + merged_lse = torch.full( + (1, hq, 1), float("-inf"), device=q.device, dtype=torch.float32 + ) + zero_dep + else: + merged = merge_attention_partial_states(states) + merged_out, merged_lse = merged.out, merged.lse + query_rows.append(merged_out) + lse_query_rows.append(merged_lse) + out_rows.append(torch.cat(query_rows, dim=2)) + lse_rows.append(torch.cat(lse_query_rows, dim=2)) + return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) def backward_reference( self, diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 5d60c959..311f2c0d 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -733,3 +733,52 @@ def test_gapped_partial_ranges_raise(): def test_registry_dispatches_cp_attention_reference(): assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype): + """The strict candidate must not change arithmetic with batch/CP shape.""" + + q, k, v = _qkv(2, 5, 9, seed=41, dtype=dtype, heads=4, kv_heads=2, dim=8) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + cp1_out, cp1_lse = op.forward_with_lse( + q, + k, + v, + cp_world_size=1, + kv_chunk_size=3, + ) + cp2_out, cp2_lse = op.forward_with_lse( + q, + k, + v, + cp_world_size=2, + kv_chunk_size=3, + ) + assert torch.equal(cp1_out, cp2_out) + assert torch.equal(cp1_lse, cp2_lse) + + single_out, single_lse = op.forward_with_lse( + q[:1], + k[:1], + v[:1], + cp_world_size=1, + kv_chunk_size=3, + ) + assert torch.equal(single_out, cp1_out[:1]) + assert torch.equal(single_lse, cp1_lse[:1]) + + +def test_strict_core_backward_is_bitwise_invariant_to_cp_schedule(): + q, k, v = _qkv(1, 4, 8, seed=42, dtype=torch.float32, heads=4, kv_heads=2, dim=8) + dout = torch.randn_like(q) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + cp1 = op.backward_reference(q, k, v, dout, cp_world_size=1, kv_chunk_size=3) + cp2 = op.backward_reference(q, k, v, dout, cp_world_size=2, kv_chunk_size=3) + assert torch.equal(cp1.out, cp2.out) + assert torch.equal(cp1.lse, cp2.lse) + assert torch.equal(cp1.gradients.dq, cp2.gradients.dq) + assert torch.equal(cp1.gradients.dk, cp2.gradients.dk) + assert torch.equal(cp1.gradients.dv, cp2.gradients.dv) From 5d9bc13bed085e95f3ba5790c8f2df386fe7bd43 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 21:49:31 +0800 Subject: [PATCH 15/16] refactor(attention): share strict schedule identity --- rl_engine/kernels/attention_contract.py | 2 ++ rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 4 +++- tests/test_cp_attention.py | 2 ++ 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index cc532f93..fb1b572a 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -19,6 +19,7 @@ STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" class AttentionContractError(ValueError): @@ -1595,6 +1596,7 @@ class AttentionDispatchResult: "SplitKVRuntimePlanSet", "SplitKVSpec", "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_SCHEDULE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index bd4d0cb5..e3d6a343 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -24,6 +24,7 @@ SplitKVRuntimePlanSet, SplitKVSpec, STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp @@ -128,7 +129,7 @@ def forward_with_lse( "fallback": False, "fallback_reason": None, "native_attention_arithmetic": False, - "strict_schedule": "single_batch_single_query_global_kv_blocks", + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, }, ) @@ -1318,6 +1319,7 @@ def build_reference_split_kv_runtime_plan_set( "DeterministicCPAttentionReferenceOp", "GradientDriftStats", "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_SCHEDULE_ID", "compare_cp_attention_backward", "merge_attention_partial_states", "split_kv_execution_plan_provenance", diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 311f2c0d..50f944f5 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -20,6 +20,7 @@ DeterministicAttentionCore, DeterministicCPAttentionReferenceOp, STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, compare_cp_attention_backward, merge_attention_partial_states, split_kv_execution_plan_provenance, @@ -94,6 +95,7 @@ def test_strict_attention_core_freezes_plan_and_final_write(): assert result.out.dtype is torch.bfloat16 assert result.lse.dtype is torch.float32 assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID assert result.provenance["merge_order"] == "global_block_index" assert result.provenance["accum_dtype"] == "fp32" assert result.provenance["downcast_at"] == "final_write" From 988d9548e196397fec1f6cda92caea2ac67eeb4f Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 23:25:28 +0800 Subject: [PATCH 16/16] fix(attention): keep strict CP reference on canonical full KV schedule --- .../ops/pytorch/attention/cp_attention.py | 52 ++++++------------- tests/test_cp_attention.py | 6 +-- 2 files changed, 19 insertions(+), 39 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index e3d6a343..945cb67c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -567,7 +567,6 @@ def _forward_strict_bitwise( ) # The canonical schedule is global. CP ownership and arrival order # must not change which partial states are generated or merged. - kv_bounds = _kv_block_bounds(skv, 1, kv_chunk_size) out_rows: list[torch.Tensor] = [] lse_rows: list[torch.Tensor] = [] for batch_index in range(batch): @@ -585,41 +584,22 @@ def _forward_strict_bitwise( lse_query_rows: list[torch.Tensor] = [] for query_index in range(sq): q_row = q_batch[:, :, query_index : query_index + 1, :].contiguous() - states = [ - self.local_partial_state( - q_row, - k_batch[:, :, key_start:key_end, :].contiguous(), - v_batch[:, :, key_start:key_end, :].contiguous(), - q_start=query_index, - k_start=key_start, - total_kv_len=skv, - total_query_len=sq, - causal=causal, - scale=scale, - key_padding_mask=( - None - if pad_batch is None - else pad_batch[:, key_start:key_end].contiguous() - ), - query_position_offsets=query_offset, - key_position_offsets=key_offset, - ) - for key_start, key_end in kv_bounds - if key_start != key_end - ] - if not states: - zero_dep = _zero_dependency(q_row.float(), k_batch.float(), v_batch.float()) - merged_out = torch.zeros( - 1, hq, 1, dim, device=q.device, dtype=torch.float32 - ) + zero_dep - merged_lse = torch.full( - (1, hq, 1), float("-inf"), device=q.device, dtype=torch.float32 - ) + zero_dep - else: - merged = merge_attention_partial_states(states) - merged_out, merged_lse = merged.out, merged.lse - query_rows.append(merged_out) - lse_query_rows.append(merged_lse) + state = self.local_partial_state( + q_row, + k_batch, + v_batch, + q_start=query_index, + k_start=0, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=pad_batch, + query_position_offsets=query_offset, + key_position_offsets=key_offset, + ) + query_rows.append(state.out) + lse_query_rows.append(state.lse) out_rows.append(torch.cat(query_rows, dim=2)) lse_rows.append(torch.cat(lse_query_rows, dim=2)) return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 50f944f5..827ce865 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -749,7 +749,7 @@ def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype): k, v, cp_world_size=1, - kv_chunk_size=3, + kv_chunk_size=None, ) cp2_out, cp2_lse = op.forward_with_lse( q, @@ -766,7 +766,7 @@ def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype): k[:1], v[:1], cp_world_size=1, - kv_chunk_size=3, + kv_chunk_size=1, ) assert torch.equal(single_out, cp1_out[:1]) assert torch.equal(single_lse, cp1_lse[:1]) @@ -777,7 +777,7 @@ def test_strict_core_backward_is_bitwise_invariant_to_cp_schedule(): dout = torch.randn_like(q) op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) - cp1 = op.backward_reference(q, k, v, dout, cp_world_size=1, kv_chunk_size=3) + cp1 = op.backward_reference(q, k, v, dout, cp_world_size=1, kv_chunk_size=None) cp2 = op.backward_reference(q, k, v, dout, cp_world_size=2, kv_chunk_size=3) assert torch.equal(cp1.out, cp2.out) assert torch.equal(cp1.lse, cp2.lse)