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 1/6] 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 2/6] 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 82aa59932876d3c04288b6ce6040230236e26d5e Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Fri, 7 Aug 2026 22:55:21 +0800 Subject: [PATCH 3/6] feat(attention): add CP backward drift validation Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 37 +- rl_engine/kernels/gtest/operator_specs.py | 16 + .../ops/pytorch/attention/cp_attention.py | 334 ++++++++++++++++++ tests/test_cp_attention.py | 149 ++++++++ tests/test_operator_inputs.py | 18 + 5 files changed, 552 insertions(+), 2 deletions(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 988fad8c..92f646a4 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -100,6 +100,30 @@ 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. +PR8 extends the CP reference with training-side backward validation: + +```python +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + compare_cp_attention_backward, +) + +report = compare_cp_attention_backward( + q, + k, + v, + dout, + candidate_cp_world_size=2, + candidate_kv_chunk_size=512, +) +``` + +The report compares CP=1 against the CP/chunked candidate for `dq`, `dk`, `dv`, +`out`, and attention-domain `lse`. It also includes per-logical-CP-rank gradient +drift slices, fixed `global_block_index` merge provenance, final-write downcast +metadata, and the saved forward state required by training backward validation. +Decode backward remains out of scope. Transformer Engine backward is not claimed +in this reference path because compatible saved forward state is not exposed here. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -167,8 +191,9 @@ GPU-only LARGE Qwen3-8B real-shape smoke test. 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. +padding/all-masked stability, BF16 final-write behavior, backward drift reports +for `dq/dk/dv`, Qwen3-8B local TP=2/CP=2 BF16 backward smoke coverage, 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 @@ -227,6 +252,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 @@ -255,6 +284,10 @@ for measured peak memory at representative shapes. 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. +- PR8 backward validation is for training prefill/chunked-prefill only. Decode + replay remains forward-only unless a future issue scopes differentiable decode. +- Transformer Engine backward is not used or claimed by PR8 unless a later backend + exposes compatible saved forward state for `dq/dk/dv` comparison. - `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/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index c3d27848..76ba3edc 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -63,6 +63,22 @@ def _load_object(path: str) -> Any: }, grad_input_names=("q", "k", "v"), ), + "cp_attention": OperatorSpec( + name="cp_attention", + op_class="attention", + gold_path=( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + }, + grad_input_names=("q", "k", "v"), + ), "logp": OperatorSpec( name="logp", op_class="logprob", diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 302ad73f..d34dd729 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -45,6 +45,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: @@ -252,6 +350,98 @@ 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, + "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 +646,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, @@ -547,7 +874,14 @@ def _kv_block_bounds( __all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", "AttentionPartialState", "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", "merge_attention_partial_states", ] diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 03c861dd..e486a24b 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,6 +17,7 @@ from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, merge_attention_partial_states, ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp @@ -26,6 +28,7 @@ _N_KV = 8 _HEAD_DIM = 128 _ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 @contextlib.contextmanager @@ -350,6 +353,152 @@ 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 + + +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) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 66395e8c..fee1cf94 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -9,6 +9,11 @@ import torch from rl_engine.kernels.gtest.operator_inputs import make_operator_inputs, operator_shape_name +from rl_engine.kernels.gtest.operator_specs import ( + make_candidate, + make_operator_case, + operator_names, +) def _args(**overrides): @@ -25,6 +30,7 @@ def _args(**overrides): "n_dim": 32, "theta": 1.0e6, "eps": 1.0e-6, + "arch_key": None, } values.update(overrides) return argparse.Namespace(**values) @@ -81,6 +87,18 @@ def test_random_logp_inputs_are_seeded(): assert torch.equal(first["token_ids"], second["token_ids"]) +def test_cp_attention_operator_spec_registers_backward_grad_inputs(): + args = _args(op="cp_attention", input_mode="constant", batch=1, seq=2) + + assert "cp_attention" in operator_names() + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(argparse.Namespace(**{**vars(args), "candidate": "pytorch"})) + + assert case.op_class == "attention" + assert case.grad_input_names == ("q", "k", "v") + assert candidate.name == "pytorch-cp_attention" + + def test_constant_linear_logp_inputs_match_operator_contract(): args = _args(input_mode="constant", constant_value=0.5, token_value=3) inputs = make_operator_inputs("linear_logp", args, torch.float32, torch.device("cpu")) From 561dce6beec6cb96b420f0851c73ef19690fc5ca Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sat, 8 Aug 2026 15:49:30 +0800 Subject: [PATCH 4/6] feat(attention): add CP drift benchmark artifacts Signed-off-by: inaniloquentee <3051000145@qq.com> --- .../benchmark_ws2_cp_attention_drift.py | 957 ++++++++++++++++++ ...tention-pr5-distributed-drift-benchmark.md | 112 ++ docs/operators/attention.md | 24 + .../test_ws2_cp_attention_drift_benchmark.py | 133 +++ 4 files changed, 1226 insertions(+) create mode 100644 benchmarks/benchmark_ws2_cp_attention_drift.py create mode 100644 docs/design/ws2-attention-pr5-distributed-drift-benchmark.md create mode 100644 tests/test_ws2_cp_attention_drift_benchmark.py diff --git a/benchmarks/benchmark_ws2_cp_attention_drift.py b/benchmarks/benchmark_ws2_cp_attention_drift.py new file mode 100644 index 00000000..9f26d41c --- /dev/null +++ b/benchmarks/benchmark_ws2_cp_attention_drift.py @@ -0,0 +1,957 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""WS2 CP attention drift benchmark and report artifact generator. + +This is the PR5 artifact path for issue #235. It is intentionally rank-aware +and torchrun-friendly, but the correctness surface remains the deterministic +PyTorch CP reference. The benchmark can run as a CPU smoke test on one process +or under torchrun; rank 0 writes the shared JSON report. +""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime as _datetime +import hashlib +import importlib +import importlib.metadata +import json +import os +import platform +import shlex +import sys +from pathlib import Path +from typing import Any, Iterator, Sequence + +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, + merge_attention_partial_states, +) +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp + +SCHEMA_VERSION = "ws2_cp_attention_drift/v1" +ISSUE = 235 +PR = 5 +DEFAULT_SEQ_LEN = 16 +QWEN3_8B_HEADS = 32 +QWEN3_8B_KV_HEADS = 8 +QWEN3_8B_HEAD_DIM = 128 +QWEN3_8B_ROPE_THETA = 1_000_000.0 +TE_CONTEXT_PARALLEL_MODULE = ( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" +) +TE_SYMBOLS = ( + "flash_attn_fwd_softmax_lse_correction", + "flash_attn_fwd_out_correction_init", + "flash_attn_fwd_out_correction", +) + + +class TEContextParallelMergeAdapter: + """Optional Transformer Engine CP merge oracle used only by PR5 reports.""" + + def __init__(self, module: Any, *, version: str) -> None: + self._module = module + self.version = version + + @classmethod + def probe(cls) -> tuple["TEContextParallelMergeAdapter | None", dict[str, object]]: + status: dict[str, object] = { + "te_available": False, + "te_version": None, + "te_module": TE_CONTEXT_PARALLEL_MODULE, + "te_symbols": list(TE_SYMBOLS), + "te_capability_probe": "unavailable", + "te_signature_checked": False, + "te_numeric_selftest": "not_run", + "fallback": True, + "fallback_reason": None, + } + try: + module = importlib.import_module(TE_CONTEXT_PARALLEL_MODULE) + version = _transformer_engine_version() + missing = [name for name in TE_SYMBOLS if not hasattr(module, name)] + if missing: + status.update( + { + "te_version": version, + "te_capability_probe": "missing_symbols", + "fallback_reason": f"missing symbols: {', '.join(missing)}", + } + ) + return None, status + adapter = cls(module, version=version) + status.update( + { + "te_available": True, + "te_version": version, + "te_signature_checked": True, + } + ) + adapter._numeric_selftest() + except ( + ImportError, + OSError, + RuntimeError, + AttributeError, + TypeError, + AssertionError, + ) as exc: + status.update( + { + "te_capability_probe": "failed", + "te_numeric_selftest": "failed", + "fallback_reason": str(exc), + } + ) + return None, status + + status.update( + { + "te_capability_probe": "passed", + "te_numeric_selftest": "passed", + "fallback": False, + "fallback_reason": None, + } + ) + return adapter, status + + def merge(self, states: Sequence[AttentionPartialState]) -> AttentionPartialState: + if not states: + raise ValueError("at least one partial state is required") + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + if len(ordered) == 1: + state = ordered[0] + return AttentionPartialState( + out=state.out.float().clone(), + lse=state.lse.float().clone(), + block_start=state.block_start, + block_end=state.block_end, + ) + + merged_lse = ordered[0].lse.float().clone() + merged_out = ordered[0].out.float().clone() + for state in ordered[1:]: + next_lse = state.lse.float() + previous_lse = merged_lse.clone() + self._module.flash_attn_fwd_softmax_lse_correction(merged_lse, next_lse) + merged_out = self._module.flash_attn_fwd_out_correction_init( + merged_out, + merged_lse, + previous_lse, + seq_dim=2, + ) + self._module.flash_attn_fwd_out_correction( + merged_out, + state.out.float(), + merged_lse, + next_lse, + seq_dim=2, + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + def _numeric_selftest(self) -> None: + gen = torch.Generator().manual_seed(5) + states = [ + AttentionPartialState( + out=torch.randn(1, 2, 3, 4, generator=gen), + lse=torch.randn(1, 2, 3, generator=gen), + block_start=0, + block_end=2, + ), + AttentionPartialState( + out=torch.randn(1, 2, 3, 4, generator=gen), + lse=torch.randn(1, 2, 3, generator=gen), + block_start=2, + block_end=5, + ), + ] + ours = merge_attention_partial_states(states) + te = self.merge(states) + torch.testing.assert_close(te.lse, ours.lse, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(te.out, ours.out, atol=1.0e-6, rtol=0.0) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the WS2 CP attention drift benchmark for issue #235 PR5." + ) + parser.add_argument("--model", default="qwen3-8b", choices=["qwen3-8b"]) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=DEFAULT_SEQ_LEN) + parser.add_argument("--seed", type=int, default=2355) + parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="cpu") + parser.add_argument("--dtype", choices=["bf16", "fp32"], default="bf16") + parser.add_argument("--tp-world-sizes", default="1,2") + parser.add_argument("--cp-world-sizes", default="1,2") + parser.add_argument( + "--kv-chunk-sizes", + default="none,4", + help="Comma list such as 'none,4'. 'none' means full prefill.", + ) + parser.add_argument("--smoke", action="store_true", help="Use a tiny CPU-friendly shape.") + parser.add_argument( + "--include-backward", + action="store_true", + help="Include optional PR8 dq/dk/dv drift fields.", + ) + parser.add_argument( + "--no-rope", + action="store_false", + dest="compose_rope", + help="Disable the pre-attention RoPE composition step.", + ) + parser.set_defaults(compose_rope=True) + parser.add_argument("--num-threads", type=int, default=1) + parser.add_argument( + "--init-process-group", + action="store_true", + help="Initialize torch.distributed from torchrun env vars before benchmarking.", + ) + parser.add_argument("--output", type=Path, help="Optional JSON artifact path.") + parser.add_argument("--json", action="store_true", help="Print the JSON report on rank 0.") + return parser.parse_args(argv) + + +def run_benchmark(args: argparse.Namespace) -> dict[str, object]: + rank_env = _rank_env() + device = _resolve_device(args.device, rank_env) + _validate_args(args) + distributed = _maybe_init_process_group(args, device, rank_env) + te_adapter, te_status = TEContextParallelMergeAdapter.probe() + seq_len = 4 if args.smoke and args.seq_len == DEFAULT_SEQ_LEN else args.seq_len + kv_chunk_sizes = _parse_kv_chunk_sizes(args.kv_chunk_sizes) + if args.smoke and args.kv_chunk_sizes == "none,4": + kv_chunk_sizes = (None, 1) + + cases: list[dict[str, object]] = [] + report: dict[str, object] = { + "schema_version": SCHEMA_VERSION, + "report_family": "ws2_cross_config_drift_report", + "tolerance_source": "#108", + "issue": ISSUE, + "pr": PR, + "created_at_utc": _datetime.datetime.now(_datetime.UTC).isoformat(), + "launch": _launch_metadata(rank_env), + "runtime": _runtime_metadata(device, distributed, rank_env), + "target": { + "model": args.model, + "global_num_query_heads": QWEN3_8B_HEADS, + "global_num_kv_heads": QWEN3_8B_KV_HEADS, + "head_dim": QWEN3_8B_HEAD_DIM, + "dtype": args.dtype, + "batch": args.batch, + "seq_len": seq_len, + "causal": True, + }, + "te_context_parallel_merge": te_status, + "dlogp": { + "status": "not_available", + "reason": "selected-logprob chain integration is outside PR5 benchmark scope", + }, + "cases": cases, + } + + try: + with _thread_limit(args.num_threads): + for tp_world_size in _parse_int_csv(args.tp_world_sizes, name="tp_world_sizes"): + for cp_world_size in _parse_int_csv(args.cp_world_sizes, name="cp_world_sizes"): + for kv_chunk_size in kv_chunk_sizes: + cases.append( + _run_case( + args, + device=device, + seq_len=seq_len, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + te_adapter=te_adapter, + ) + ) + finally: + if distributed["initialized"]: + import torch.distributed as dist + + if sys.exc_info()[0] is None: + dist.barrier() + dist.destroy_process_group() + return report + + +def write_report(report: dict[str, object], output: Path) -> None: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + report = run_benchmark(args) + rank = int(report["launch"]["rank"]) + if rank == 0: + if args.output is not None: + write_report(report, args.output) + if args.json: + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + + +def _run_case( + args: argparse.Namespace, + *, + device: torch.device, + seq_len: int, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: int | None, + te_adapter: TEContextParallelMergeAdapter | None, +) -> dict[str, object]: + _validate_topology(tp_world_size, cp_world_size) + dtype = _dtype_from_name(args.dtype) + local_hq = QWEN3_8B_HEADS // tp_world_size + local_hkv = QWEN3_8B_KV_HEADS // tp_world_size + case_seed = _case_seed(args.seed, tp_world_size, cp_world_size, kv_chunk_size) + q, k, v, rope_report = _make_qkv( + batch=args.batch, + local_hq=local_hq, + local_hkv=local_hkv, + seq_len=seq_len, + dtype=dtype, + device=device, + seed=case_seed, + compose_rope=args.compose_rope, + ) + dout = _make_dout( + batch=args.batch, + local_hq=local_hq, + seq_len=seq_len, + dtype=dtype, + device=device, + seed=case_seed + 17, + ) + attention = DeterministicCPAttentionReferenceOp() + reference_out, reference_lse = attention.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=1, + kv_chunk_size=None, + ) + candidate_fp32_out, candidate_fp32_lse = attention.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + candidate_dtype_out, candidate_dtype_lse = attention.forward_with_lse( + q, + k, + v, + causal=True, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=dtype, + ) + + split_kv_policy = "none" if kv_chunk_size is None else "fixed_kv_chunk" + attention_mode = "prefill" if kv_chunk_size is None else "chunked_prefill" + q_bounds = _split_bounds(seq_len, cp_world_size) + kv_bounds = _kv_block_bounds(seq_len, cp_world_size, kv_chunk_size) + case: dict[str, object] = { + "case_name": _case_name(tp_world_size, cp_world_size, kv_chunk_size, args.dtype), + "attention_mode": attention_mode, + "model": args.model, + "topology": { + "tp_world_size": tp_world_size, + "cp_world_size": cp_world_size, + "tp_rank": 0, + "logical_cp_ranks": cp_world_size, + "local_num_query_heads": local_hq, + "local_num_kv_heads": local_hkv, + "local_query_head_range": [0, local_hq], + "local_kv_head_range": [0, local_hkv], + "head_dim": QWEN3_8B_HEAD_DIM, + "q_sequence_bounds": [list(item) for item in q_bounds], + "kv_block_bounds": [list(item) for item in kv_bounds], + }, + "provenance": { + "backend": "deterministic_cp_reference", + "reference_backend": "cp1_fp32_prefill", + "candidate_backend": "cp_reference", + "dtype": args.dtype, + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_domain": "attention", + "merge_order": "global_block_index", + "split_kv_policy": split_kv_policy, + "kv_chunk_size": kv_chunk_size, + "block_metadata_hash": _block_metadata_hash(kv_bounds), + "scale_placement": "scores_after_qk_matmul", + "mask_application_order": ["scale", "causal_mask", "key_padding_mask"], + "dropout_policy": "disabled", + "deterministic_controls": { + "reference": "strict_fp32_math_inside_cp_attention", + "num_threads": args.num_threads, + }, + "rope": rope_report["provenance"], + }, + "drift": { + "cp_merge_fp32": { + "out": _drift_stats(candidate_fp32_out, reference_out), + "lse": _drift_stats(candidate_fp32_lse, reference_lse), + "source_class": "reduction_and_collective_drift", + }, + "dtype_path_vs_fp32": { + "out": _drift_stats(candidate_dtype_out, reference_out), + "lse": _drift_stats(candidate_dtype_lse, reference_lse), + "source_class": "arithmetic_schedule_drift", + }, + "rope": rope_report["drift"], + }, + "merge_order_probe": _merge_order_probe( + attention, + q, + k, + v, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ), + "te_merge_oracle": _te_merge_oracle_probe( + attention, + q, + k, + v, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + te_adapter=te_adapter, + ), + "per_rank": _per_rank_forward_drifts( + candidate_fp32_out, + candidate_fp32_lse, + reference_out, + reference_lse, + cp_world_size, + ), + "backward": {"status": "not_requested"}, + } + if args.include_backward: + backward = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=cp_world_size, + candidate_kv_chunk_size=kv_chunk_size, + output_dtype=dtype, + ) + case["backward"] = { + "status": "available", + "report": backward.to_dict(), + } + return case + + +def _make_qkv( + *, + batch: int, + local_hq: int, + local_hkv: int, + seq_len: int, + dtype: torch.dtype, + device: torch.device, + seed: int, + compose_rope: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, object]]: + gen = torch.Generator(device="cpu").manual_seed(seed) + q_pre = torch.randn(batch, local_hq, seq_len, QWEN3_8B_HEAD_DIM, generator=gen) + k_pre = torch.randn(batch, local_hkv, seq_len, QWEN3_8B_HEAD_DIM, generator=gen) + v = torch.randn(batch, local_hkv, seq_len, QWEN3_8B_HEAD_DIM, generator=gen) + q_pre = q_pre.to(device=device, dtype=dtype) + k_pre = k_pre.to(device=device, dtype=dtype) + v = v.to(device=device, dtype=dtype) + + positions = ( + torch.arange(seq_len, dtype=torch.long, device=device) + .unsqueeze(0) + .expand( + batch, + -1, + ) + ) + if not compose_rope: + return q_pre, k_pre, v, _rope_report_disabled() + + rope = NativeRoPEOp() + q_rope_dtype = rope.forward(q_pre, positions, theta=QWEN3_8B_ROPE_THETA) + k_rope_dtype = rope.forward(k_pre, positions, theta=QWEN3_8B_ROPE_THETA) + q_rope_fp32 = rope.forward_fp32(q_pre, positions, theta=QWEN3_8B_ROPE_THETA) + k_rope_fp32 = rope.forward_fp32(k_pre, positions, theta=QWEN3_8B_ROPE_THETA) + return ( + q_rope_dtype, + k_rope_dtype, + v, + { + "provenance": { + "rope_state": "post_rope", + "rope_theta": QWEN3_8B_ROPE_THETA, + "rope_scaling": None, + "rotary_dim": QWEN3_8B_HEAD_DIM, + "position_ids": "arange(seq_len)", + "cache_position": "same_as_position_ids", + "query_position_offsets": [0 for _ in range(batch)], + "key_position_offsets": [0 for _ in range(batch)], + "k_cache_rope_state": "post_rope", + "rope_cast_at": "rope_output", + "rope_output_dtype": _dtype_name(dtype), + "fusion_boundary": "unfused_rope_attention_reference", + }, + "drift": { + "status": "available", + "q": _drift_stats(q_rope_dtype, q_rope_fp32), + "k": _drift_stats(k_rope_dtype, k_rope_fp32), + }, + }, + ) + + +def _make_dout( + *, + batch: int, + local_hq: int, + seq_len: int, + dtype: torch.dtype, + device: torch.device, + seed: int, +) -> torch.Tensor: + gen = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn( + batch, + local_hq, + seq_len, + QWEN3_8B_HEAD_DIM, + generator=gen, + dtype=dtype, + ).to(device=device) + + +def _rope_report_disabled() -> dict[str, object]: + return { + "provenance": { + "rope_state": "not_composed", + "fusion_boundary": "attention_only", + }, + "drift": { + "status": "not_composed", + "q": None, + "k": None, + }, + } + + +def _merge_order_probe( + attention: DeterministicCPAttentionReferenceOp, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cp_world_size: int, + kv_chunk_size: int | None, +) -> dict[str, object]: + reversed_out: list[torch.Tensor] = [] + reversed_lse: list[torch.Tensor] = [] + sorted_out: list[torch.Tensor] = [] + sorted_lse: list[torch.Tensor] = [] + kv_bounds = _kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size) + for q_start, q_end in _split_bounds(q.size(2), cp_world_size): + if q_start == q_end: + continue + states = _partial_states_for_query_block( + attention, + q, + k, + v, + q_start=q_start, + q_end=q_end, + kv_bounds=kv_bounds, + ) + sorted_merge = merge_attention_partial_states(states) + reversed_merge = merge_attention_partial_states(list(reversed(states))) + sorted_out.append(sorted_merge.out) + sorted_lse.append(sorted_merge.lse) + reversed_out.append(reversed_merge.out) + reversed_lse.append(reversed_merge.lse) + if not sorted_out: + return { + "status": "empty_query", + "arrival_order_policy": "ignored_then_sorted_by_global_block_index", + } + return { + "status": "available", + "arrival_order_policy": "ignored_then_sorted_by_global_block_index", + "out": _drift_stats(torch.cat(reversed_out, dim=2), torch.cat(sorted_out, dim=2)), + "lse": _drift_stats(torch.cat(reversed_lse, dim=2), torch.cat(sorted_lse, dim=2)), + } + + +def _te_merge_oracle_probe( + attention: DeterministicCPAttentionReferenceOp, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cp_world_size: int, + kv_chunk_size: int | None, + te_adapter: TEContextParallelMergeAdapter | None, +) -> dict[str, object]: + if te_adapter is None: + return { + "status": "unavailable", + "te_module": TE_CONTEXT_PARALLEL_MODULE, + "fallback": "deterministic_cp_reference", + } + + ours_out: list[torch.Tensor] = [] + ours_lse: list[torch.Tensor] = [] + te_out: list[torch.Tensor] = [] + te_lse: list[torch.Tensor] = [] + kv_bounds = _kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size) + for q_start, q_end in _split_bounds(q.size(2), cp_world_size): + if q_start == q_end: + continue + states = _partial_states_for_query_block( + attention, + q, + k, + v, + q_start=q_start, + q_end=q_end, + kv_bounds=kv_bounds, + ) + ours = merge_attention_partial_states(states) + te = te_adapter.merge(states) + ours_out.append(ours.out) + ours_lse.append(ours.lse) + te_out.append(te.out) + te_lse.append(te.lse) + if not ours_out: + return {"status": "empty_query", "te_version": te_adapter.version} + return { + "status": "available", + "te_version": te_adapter.version, + "te_module": TE_CONTEXT_PARALLEL_MODULE, + "te_symbols": list(TE_SYMBOLS), + "out": _drift_stats(torch.cat(te_out, dim=2), torch.cat(ours_out, dim=2)), + "lse": _drift_stats(torch.cat(te_lse, dim=2), torch.cat(ours_lse, dim=2)), + } + + +def _partial_states_for_query_block( + attention: DeterministicCPAttentionReferenceOp, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + q_end: int, + kv_bounds: Sequence[tuple[int, int]], +) -> list[AttentionPartialState]: + return [ + attention.local_partial_state( + q[:, :, q_start:q_end, :], + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=k.size(2), + total_query_len=q.size(2), + causal=True, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + + +def _per_rank_forward_drifts( + candidate_out: torch.Tensor, + candidate_lse: torch.Tensor, + reference_out: torch.Tensor, + reference_lse: torch.Tensor, + cp_world_size: int, +) -> list[dict[str, object]]: + per_rank = [] + for rank, (q_start, q_end) in enumerate(_split_bounds(candidate_out.size(2), cp_world_size)): + per_rank.append( + { + "rank": rank, + "query_start": q_start, + "query_end": q_end, + "out": _drift_stats( + candidate_out[:, :, q_start:q_end, :], + reference_out[:, :, q_start:q_end, :], + ), + "lse": _drift_stats( + candidate_lse[:, :, q_start:q_end], + reference_lse[:, :, q_start:q_end], + ), + } + ) + return per_rank + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict[str, object]: + 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 { + "max_abs": 0.0, + "mean_abs": 0.0, + "p95_abs": 0.0, + "p99_abs": 0.0, + "active_count": 0, + } + return { + "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 _rank_env() -> dict[str, int | bool]: + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + return { + "rank": rank, + "world_size": world_size, + "local_rank": local_rank, + "torchrun": "RANK" in os.environ or "WORLD_SIZE" in os.environ, + } + + +def _resolve_device(device_arg: str, rank_env: dict[str, int | bool]) -> torch.device: + if device_arg == "auto": + device_arg = "cuda" if torch.cuda.is_available() else "cpu" + if device_arg == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available") + local_rank = int(rank_env["local_rank"]) + if torch.cuda.device_count() > 0: + torch.cuda.set_device(local_rank % torch.cuda.device_count()) + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + +def _maybe_init_process_group( + args: argparse.Namespace, + device: torch.device, + rank_env: dict[str, int | bool], +) -> dict[str, object]: + initialized = False + backend = None + if args.init_process_group and int(rank_env["world_size"]) > 1: + import torch.distributed as dist + + backend = "nccl" if device.type == "cuda" else "gloo" + dist.init_process_group(backend=backend, init_method="env://") + initialized = True + return { + "initialized": initialized, + "backend": backend, + "transport": "torchrun_env_rank_aware", + } + + +def _runtime_metadata( + device: torch.device, + distributed: dict[str, object], + rank_env: dict[str, int | bool], +) -> dict[str, object]: + return { + "python": sys.version.split()[0], + "platform": platform.platform(), + "torch_version": torch.__version__, + "cuda_available": torch.cuda.is_available(), + "device": str(device), + "distributed": distributed, + "rank_env": rank_env, + } + + +def _launch_metadata(rank_env: dict[str, int | bool]) -> dict[str, object]: + return { + "command": _shell_join(sys.argv), + "rank": int(rank_env["rank"]), + "world_size": int(rank_env["world_size"]), + "local_rank": int(rank_env["local_rank"]), + "torchrun": bool(rank_env["torchrun"]), + } + + +def _validate_args(args: argparse.Namespace) -> None: + if args.batch < 1: + raise ValueError("batch must be >= 1") + if args.seq_len < 1: + raise ValueError("seq_len must be >= 1") + if args.num_threads < 1: + raise ValueError("num_threads must be >= 1") + for tp_world_size in _parse_int_csv(args.tp_world_sizes, name="tp_world_sizes"): + _validate_topology(tp_world_size, 1) + for cp_world_size in _parse_int_csv(args.cp_world_sizes, name="cp_world_sizes"): + _validate_topology(1, cp_world_size) + _parse_kv_chunk_sizes(args.kv_chunk_sizes) + + +def _validate_topology(tp_world_size: int, cp_world_size: int) -> None: + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("tp_world_size and cp_world_size must be >= 1") + if QWEN3_8B_HEADS % tp_world_size != 0: + raise ValueError("Qwen3 query heads must be divisible by tp_world_size") + if QWEN3_8B_KV_HEADS % tp_world_size != 0: + raise ValueError("Qwen3 KV heads must be divisible by tp_world_size") + + +def _parse_int_csv(value: str, *, name: str) -> tuple[int, ...]: + parsed: list[int] = [] + for raw in value.split(","): + item = raw.strip() + if not item: + continue + try: + parsed.append(int(item)) + except ValueError as exc: + raise ValueError(f"{name} must be a comma-separated integer list") from exc + if not parsed: + raise ValueError(f"{name} must contain at least one integer") + return tuple(parsed) + + +def _parse_kv_chunk_sizes(value: str) -> tuple[int | None, ...]: + parsed: list[int | None] = [] + for raw in value.split(","): + item = raw.strip().lower() + if not item: + continue + if item in {"none", "full", "no_split"}: + parsed.append(None) + continue + try: + size = int(item) + except ValueError as exc: + raise ValueError("kv_chunk_sizes must contain integers or 'none'") from exc + if size < 1: + raise ValueError("kv chunk sizes must be >= 1") + parsed.append(size) + if not parsed: + raise ValueError("kv_chunk_sizes must contain at least one entry") + return tuple(parsed) + + +def _dtype_from_name(name: str) -> torch.dtype: + if name == "bf16": + return torch.bfloat16 + if name == "fp32": + return torch.float32 + raise ValueError(f"unsupported dtype: {name}") + + +def _dtype_name(dtype: torch.dtype) -> str: + return str(dtype).replace("torch.", "") + + +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: int | None, +) -> 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 _block_metadata_hash(bounds: Sequence[tuple[int, int]]) -> str: + payload = json.dumps([list(item) for item in bounds], separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest()[:16] + + +def _case_seed(seed: int, tp_world_size: int, cp_world_size: int, kv_chunk_size: int | None) -> int: + return seed + tp_world_size * 101 + cp_world_size * 17 + (kv_chunk_size or 0) + + +def _case_name( + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: int | None, + dtype: str, +) -> str: + mode = "prefill" if kv_chunk_size is None else f"chunk{kv_chunk_size}" + return f"qwen3_8b_tp{tp_world_size}_cp{cp_world_size}_{mode}_{dtype}" + + +def _transformer_engine_version() -> str: + for package in ("transformer-engine", "transformer_engine"): + try: + return importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + continue + return "unknown" + + +def _shell_join(argv: Sequence[str]) -> str: + if os.name == "nt": + return " ".join(argv) + return shlex.join(argv) + + +@contextlib.contextmanager +def _thread_limit(num_threads: int) -> Iterator[None]: + previous = torch.get_num_threads() + torch.set_num_threads(num_threads) + try: + yield + finally: + torch.set_num_threads(previous) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md new file mode 100644 index 00000000..e5f061e4 --- /dev/null +++ b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md @@ -0,0 +1,112 @@ +# WS2 Attention PR5 Drift Benchmark + +PR5 adds the report artifact path for issue #235. It does not introduce a +production communication kernel. The benchmark is a rank-aware, torchrun-style +driver around the deterministic CP attention reference so the same artifact can +be produced locally or under a multi-process launcher. + +## Scope + +The benchmark covers the Qwen3-8B Attention target: + +- global heads: `Hq=32`, `Hkv=8`, `D=128` +- TP sweep: `TP=1/2`; TP only changes the local head shard shape +- CP sweep: `CP=1/2` +- modes: full prefill and chunked-prefill replay +- dtype path: BF16 candidate path compared with FP32 reference +- optional backward: `dq`, `dk`, `dv` drift from the PR8 reference +- optional RoPE composition before Attention, while CP Attention still consumes + post-RoPE Q/K + +The report separates two drift classes: + +| Field | Meaning | +| --- | --- | +| `drift.cp_merge_fp32` | CP/chunked candidate with FP32 output vs CP=1 FP32 prefill. This isolates CP merge and split-KV order. | +| `drift.dtype_path_vs_fp32` | BF16 candidate path vs FP32 reference. This exposes arithmetic/final-write drift. | +| `merge_order_probe` | Reversed-arrival partial states vs canonical sorted merge. This verifies that arrival order is ignored. | +| `te_merge_oracle` | Optional Transformer Engine merge-oracle drift when TE is installed and passes capability probes. | +| `backward` | Optional PR8 `dq/dk/dv` drift report when `--include-backward` is used. | + +Selected-logprob `dlogp` remains `not_available` here because the full logprob +chain integration is outside PR5. PR4/WS2 runtime integration should fill that +field once Attention is wired into the chain. + +## Commands + +Local smoke: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +``` + +Qwen3 TP=2 / CP=2 with backward drift and a JSON artifact: + +```bash +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 +``` + +Torchrun-style launch: + +```bash +torchrun --nproc-per-node=2 benchmarks/benchmark_ws2_cp_attention_drift.py \ + --smoke \ + --json +``` + +Rank 0 prints or writes the shared report. Other ranks can run the same +rank-aware benchmark without changing the numerical reducer. + +## Transformer Engine Reuse + +PR5 reuses Transformer Engine only as an optional merge oracle, not as the +source of truth. The adapter imports: + +```text +transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +``` + +and uses these APIs when available: + +```text +flash_attn_fwd_softmax_lse_correction +flash_attn_fwd_out_correction_init +flash_attn_fwd_out_correction +``` + +The benchmark first builds RL-Kernel partial states: + +```text +state_i = (out_i, lse_i, global_block_index_i) +``` + +then sorts them by `global_block_index`. TE is allowed to perform only the +online-softmax correction arithmetic for those already-sorted states. If TE is +missing, incompatible, or fails the numeric self-test, the report records a +provenance fallback and continues with the deterministic RL-Kernel merge. + +## Report Contract + +The JSON root contains: + +```text +schema_version +issue / pr +launch.command +runtime.rank_env +target +te_context_parallel_merge +dlogp +cases[] +``` + +Each case records topology, RoPE/cache provenance, split-KV policy, block +metadata hash, drift summaries, per-logical-CP-rank metrics, and optional +backward drift. The merge order is always `global_block_index`, and +`downcast_at` is always `final_write`. diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 92f646a4..8ebd9444 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -124,6 +124,30 @@ metadata, and the saved forward state required by training backward validation. Decode backward remains out of scope. Transformer Engine backward is not claimed in this reference path because compatible saved forward state is not exposed here. +PR5 adds a rank-aware drift artifact benchmark around the same reference path: + +```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 +``` + +The report covers Qwen3-8B-style TP-local head shards, CP=1/2, full prefill, +chunked-prefill, BF16-vs-FP32 drift, per-logical-CP-rank metrics, optional +PR8 backward drift, RoPE provenance, and optional Transformer Engine +context-parallel merge-oracle drift. TE reuse is limited to +`transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py` +helpers: +`flash_attn_fwd_softmax_lse_correction`, +`flash_attn_fwd_out_correction_init`, and +`flash_attn_fwd_out_correction`. +See `docs/design/ws2-attention-pr5-distributed-drift-benchmark.md`. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): diff --git a/tests/test_ws2_cp_attention_drift_benchmark.py b/tests/test_ws2_cp_attention_drift_benchmark.py new file mode 100644 index 00000000..949aca6a --- /dev/null +++ b/tests/test_ws2_cp_attention_drift_benchmark.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 CP attention PR5 drift benchmark artifact.""" + +from __future__ import annotations + +import json + +import pytest + +from benchmarks.benchmark_ws2_cp_attention_drift import ( + SCHEMA_VERSION, + parse_args, + run_benchmark, + write_report, +) + + +def test_smoke_report_has_pr5_schema_and_qwen3_tp2_cp2_case(): + report = run_benchmark( + parse_args( + [ + "--smoke", + "--tp-world-sizes", + "2", + "--cp-world-sizes", + "2", + "--kv-chunk-sizes", + "none,1", + ] + ) + ) + + assert report["schema_version"] == SCHEMA_VERSION + assert report["report_family"] == "ws2_cross_config_drift_report" + assert report["tolerance_source"] == "#108" + assert report["issue"] == 235 + assert report["pr"] == 5 + assert report["target"]["model"] == "qwen3-8b" + assert report["target"]["global_num_query_heads"] == 32 + assert report["target"]["global_num_kv_heads"] == 8 + assert report["te_context_parallel_merge"]["te_module"].endswith("context_parallel") + assert report["dlogp"]["status"] == "not_available" + assert len(report["cases"]) == 2 + + names = {case["case_name"] for case in report["cases"]} + assert "qwen3_8b_tp2_cp2_prefill_bf16" in names + assert "qwen3_8b_tp2_cp2_chunk1_bf16" in names + + chunked = next(case for case in report["cases"] if case["attention_mode"] == "chunked_prefill") + assert chunked["topology"]["local_num_query_heads"] == 16 + assert chunked["topology"]["local_num_kv_heads"] == 4 + assert chunked["topology"]["local_query_head_range"] == [0, 16] + assert chunked["topology"]["local_kv_head_range"] == [0, 4] + assert chunked["provenance"]["merge_order"] == "global_block_index" + assert chunked["provenance"]["split_kv_policy"] == "fixed_kv_chunk" + assert chunked["provenance"]["block_metadata_hash"] + assert chunked["provenance"]["rope"]["rope_state"] == "post_rope" + assert chunked["drift"]["rope"]["status"] == "available" + assert chunked["drift"]["cp_merge_fp32"]["out"]["max_abs"] <= 1.0e-5 + assert chunked["drift"]["cp_merge_fp32"]["lse"]["max_abs"] <= 1.0e-5 + assert chunked["merge_order_probe"]["out"]["max_abs"] == 0.0 + assert len(chunked["per_rank"]) == 2 + assert chunked["per_rank"][0]["out"]["active_count"] > 0 + + +def test_report_writes_reproducible_json_artifact(tmp_path): + output = tmp_path / "ws2-cp-attention-drift.json" + report = run_benchmark( + parse_args( + [ + "--smoke", + "--no-rope", + "--tp-world-sizes", + "1", + "--cp-world-sizes", + "1", + "--kv-chunk-sizes", + "none", + ] + ) + ) + + write_report(report, output) + loaded = json.loads(output.read_text(encoding="utf-8")) + + assert loaded["schema_version"] == SCHEMA_VERSION + assert loaded["cases"][0]["provenance"]["rope"]["rope_state"] == "not_composed" + assert loaded["cases"][0]["attention_mode"] == "prefill" + + +def test_include_backward_adds_pr8_gradient_drift_report(): + report = run_benchmark( + parse_args( + [ + "--smoke", + "--include-backward", + "--tp-world-sizes", + "2", + "--cp-world-sizes", + "2", + "--kv-chunk-sizes", + "1", + ] + ) + ) + + backward = report["cases"][0]["backward"] + assert backward["status"] == "available" + drift = backward["report"]["drifts"][0] + assert drift["candidate_name"] == "cp2_chunked_backward" + assert drift["provenance"]["attention_mode"] == "chunked_prefill" + 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 + assert len(drift["per_rank"]) == 2 + + +def test_invalid_qwen3_tp_topology_is_rejected(): + with pytest.raises(ValueError, match="query heads"): + run_benchmark( + parse_args( + [ + "--tp-world-sizes", + "3", + "--cp-world-sizes", + "1", + "--kv-chunk-sizes", + "none", + ] + ) + ) 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 5/6] 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 a013e1c5cb0cbf5156f887bba901f15de95b944b Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:34:09 +0800 Subject: [PATCH 6/6] feat(attention): add P2P NCCL CP drift benchmark --- .../benchmark_ws2_cp_attention_drift.py | 178 +++++++++++++++- ...tention-pr5-distributed-drift-benchmark.md | 27 ++- .../ws2_p2p_nccl_attention_reference_check.py | 195 ++++++++++++++++++ .../test_ws2_cp_attention_drift_benchmark.py | 16 +- 4 files changed, 409 insertions(+), 7 deletions(-) create mode 100644 scripts/ws2_p2p_nccl_attention_reference_check.py diff --git a/benchmarks/benchmark_ws2_cp_attention_drift.py b/benchmarks/benchmark_ws2_cp_attention_drift.py index 9f26d41c..48cf9be0 100644 --- a/benchmarks/benchmark_ws2_cp_attention_drift.py +++ b/benchmarks/benchmark_ws2_cp_attention_drift.py @@ -26,11 +26,25 @@ import torch +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, DeterministicCPAttentionReferenceOp, + build_reference_split_kv_runtime_plan_set, compare_cp_attention_backward, merge_attention_partial_states, + split_kv_execution_plan_provenance, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + P2PNCCLAttentionCPCommunication, ) from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp @@ -366,10 +380,17 @@ def _run_case( output_dtype=dtype, ) - split_kv_policy = "none" if kv_chunk_size is None else "fixed_kv_chunk" + split_kv_policy = "disabled" if kv_chunk_size is None else "fixed" attention_mode = "prefill" if kv_chunk_size is None else "chunked_prefill" q_bounds = _split_bounds(seq_len, cp_world_size) kv_bounds = _kv_block_bounds(seq_len, cp_world_size, kv_chunk_size) + runtime_plan_set = build_reference_split_kv_runtime_plan_set( + (seq_len,) * args.batch, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) case: dict[str, object] = { "case_name": _case_name(tp_world_size, cp_world_size, kv_chunk_size, args.dtype), "attention_mode": attention_mode, @@ -397,6 +418,15 @@ def _run_case( "lse_domain": "attention", "merge_order": "global_block_index", "split_kv_policy": split_kv_policy, + "requested_split_kv_policy": split_kv_policy, + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + seq_len, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ), + "actual_split_kv_plan_set": runtime_plan_set.to_dict(), "kv_chunk_size": kv_chunk_size, "block_metadata_hash": _block_metadata_hash(kv_bounds), "scale_placement": "scores_after_qk_matmul", @@ -447,6 +477,19 @@ def _run_case( ), "backward": {"status": "not_requested"}, } + distributed_reference = _run_distributed_p2p_reference( + q, + k, + v, + reference_out, + reference_lse, + device=device, + seq_len=seq_len, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + case["distributed_p2p_reference"] = distributed_reference if args.include_backward: backward = compare_cp_attention_backward( q, @@ -465,6 +508,139 @@ def _run_case( return case +def _run_distributed_p2p_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + reference_out: torch.Tensor, + reference_lse: torch.Tensor, + *, + device: torch.device, + seq_len: int, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: int | None, +) -> dict[str, object]: + """Exercise the actual P2P reference when launched as a matching NCCL job.""" + + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + return {"status": "not_requested", "reason": "process_group_not_initialized"} + backend = str(dist.get_backend()).lower() + world_size = int(dist.get_world_size()) + if device.type != "cuda": + return { + "status": "skipped", + "reason": "P2P NCCL reference requires CUDA", + "backend": backend, + } + if "nccl" not in backend: + return { + "status": "skipped", + "reason": "P2P reference requires NCCL", + "backend": backend, + } + if world_size != cp_world_size: + return { + "status": "skipped", + "reason": "WORLD_SIZE must equal cp_world_size for the CP reference", + "world_size": world_size, + "cp_world_size": cp_world_size, + } + + rank = int(dist.get_rank()) + owner_ranges = _split_bounds(seq_len, cp_world_size) + block_bounds = _kv_block_bounds(seq_len, cp_world_size, kv_chunk_size) + blocks: list[AttentionCPBlockMetadata] = [] + owner_block_counts = [0] * cp_world_size + for block_index, (start, end) in enumerate(block_bounds): + owner = next( + owner_rank + for owner_rank, (owner_start, owner_end) in enumerate(owner_ranges) + if owner_start <= start < owner_end + ) + blocks.append( + AttentionCPBlockMetadata( + global_block_index=block_index, + kv_block_start=start, + kv_block_end=end, + owner_cp_rank=owner, + owner_tp_rank=0, + ) + ) + owner_block_counts[owner] += 1 + + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=tp_world_size, + tp_rank=0, + cp_world_size=cp_world_size, + cp_rank=rank, + ), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=tuple(blocks), + expected_kv_token_range=(0, seq_len), + query_token_ranges=tuple(_split_bounds(q.size(2), cp_world_size)), + ) + attention = DeterministicCPAttentionReferenceOp() + local_states: list[AttentionCPPartialState] = [] + for block in blocks: + if block.owner_cp_rank != rank: + continue + state = attention.local_partial_state( + q, + k[:, :, block.kv_block_start : block.kv_block_end, :], + v[:, :, block.kv_block_start : block.kv_block_end, :], + q_start=0, + k_start=block.kv_block_start, + total_kv_len=seq_len, + total_query_len=q.size(2), + causal=True, + ) + local_states.append( + AttentionCPPartialState( + out=state.out, + lse=state.lse, + block=block, + ) + ) + communication = P2PNCCLAttentionCPCommunication() + gathered = communication.all_gather_partial_states(tuple(local_states), plan) + merged = merge_attention_partial_states( + [ + AttentionPartialState( + out=state.out, + lse=state.lse, + block_start=state.block.kv_block_start, + block_end=state.block.kv_block_end, + ) + for state in gathered + ] + ) + local = communication.reduce_scatter_merged_state( + AttentionCPMergedState(out=merged.out, lse=merged.lse), + plan, + ) + q_start, q_end = plan.query_token_ranges[rank] + reference_local_out = reference_out[:, :, q_start:q_end, :] + reference_local_lse = reference_lse[:, :, q_start:q_end] + return { + "status": "available", + "backend": backend, + "rank": rank, + "world_size": world_size, + "transport": "p2p_nccl_reference", + "manifest_block_count": len(blocks), + "owner_block_counts": owner_block_counts, + "gathered_block_indices": [state.block.global_block_index for state in gathered], + "query_range": [q_start, q_end], + "out": _drift_stats(local.out, reference_local_out), + "lse": _drift_stats(local.lse, reference_local_lse), + } + + def _make_qkv( *, batch: int, diff --git a/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md index e5f061e4..88fbd199 100644 --- a/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md +++ b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md @@ -2,8 +2,9 @@ PR5 adds the report artifact path for issue #235. It does not introduce a production communication kernel. The benchmark is a rank-aware, torchrun-style -driver around the deterministic CP attention reference so the same artifact can -be produced locally or under a multi-process launcher. +driver around the deterministic CP attention reference. Under a matching +two-rank CUDA/NCCL launch it executes the P2P reference transport; CPU/Gloo +remains a report-generation smoke path. ## Scope @@ -27,6 +28,7 @@ The report separates two drift classes: | `merge_order_probe` | Reversed-arrival partial states vs canonical sorted merge. This verifies that arrival order is ignored. | | `te_merge_oracle` | Optional Transformer Engine merge-oracle drift when TE is installed and passes capability probes. | | `backward` | Optional PR8 `dq/dk/dv` drift report when `--include-backward` is used. | +| `distributed_p2p_reference` | Real NCCL P2P partial-state gather, FP32 merge, and query scatter drift. | Selected-logprob `dlogp` remains `not_available` here because the full logprob chain integration is outside PR5. PR4/WS2 runtime integration should fill that @@ -52,16 +54,31 @@ python benchmarks/benchmark_ws2_cp_attention_drift.py \ --output artifacts/ws2-cp-attention-drift.json ``` -Torchrun-style launch: +Two-GPU NCCL transport check: ```bash -torchrun --nproc-per-node=2 benchmarks/benchmark_ws2_cp_attention_drift.py \ +torchrun --standalone --nproc-per-node=2 \ + scripts/ws2_p2p_nccl_attention_reference_check.py +``` + +Two-GPU benchmark report with real P2P transport: + +```bash +torchrun --standalone --nproc-per-node=2 \ + benchmarks/benchmark_ws2_cp_attention_drift.py \ --smoke \ + --device cuda \ + --init-process-group \ + --tp-world-sizes 2 \ + --cp-world-sizes 2 \ --json ``` Rank 0 prints or writes the shared report. Other ranks can run the same -rank-aware benchmark without changing the numerical reducer. +rank-aware benchmark without changing the numerical reducer. The recommended +container is the repository CUDA image built from `docker/Dockerfile.cuda` +(`ghcr.io/rl-align/rl-kernel/rl-kernel-ci:cuda` when using the repository image +workflow). It is based on PyTorch 2.4 / CUDA 12.4 and includes NCCL support. ## Transformer Engine Reuse diff --git a/scripts/ws2_p2p_nccl_attention_reference_check.py b/scripts/ws2_p2p_nccl_attention_reference_check.py new file mode 100644 index 00000000..8d95c99d --- /dev/null +++ b/scripts/ws2_p2p_nccl_attention_reference_check.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Two-GPU P2P NCCL reference check for issue #235. + +Run with: + + torchrun --standalone --nproc-per-node=2 \ + scripts/ws2_p2p_nccl_attention_reference_check.py +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Sequence + +import torch +import torch.distributed as dist + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + P2PNCCLAttentionCPCommunication, +) +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + merge_attention_partial_states, +) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq-len", type=int, default=16) + parser.add_argument("--q-heads", type=int, default=16) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--chunk-size", type=int, default=4) + parser.add_argument("--seed", type=int, default=2357) + parser.add_argument("--atol", type=float, default=2.0e-4) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + raise RuntimeError("this check requires at least two visible CUDA devices") + dist.init_process_group("nccl", init_method="env://") + try: + world_size = dist.get_world_size() + rank = dist.get_rank() + if world_size != 2: + raise RuntimeError("this reference check requires exactly two NCCL ranks") + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + result = run_check(args, rank=rank, device=device) + failures = torch.tensor( + [0 if result["passed"] else 1], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(failures, op=dist.ReduceOp.SUM) + result["global_failure_count"] = int(failures.item()) + reports: list[dict[str, object] | None] = [None] * world_size + dist.all_gather_object(reports, result) + if rank == 0: + print(json.dumps({"ranks": reports}, indent=2, sort_keys=True)) + return 0 if int(failures.item()) == 0 else 1 + finally: + dist.destroy_process_group() + + +def run_check( + args: argparse.Namespace, + *, + rank: int, + device: torch.device, +) -> dict[str, object]: + if args.seq_len < 2 or args.seq_len % 2 != 0: + raise ValueError("seq_len must be positive and divisible by CP=2") + if args.chunk_size < 1: + raise ValueError("chunk_size must be positive") + if args.q_heads % args.kv_heads != 0: + raise ValueError("q_heads must be divisible by kv_heads") + + generator = torch.Generator(device="cpu").manual_seed(args.seed) + shape_q = (args.batch, args.q_heads, args.seq_len, args.head_dim) + shape_kv = (args.batch, args.kv_heads, args.seq_len, args.head_dim) + q = torch.randn(shape_q, generator=generator, dtype=torch.bfloat16).to(device) + k = torch.randn(shape_kv, generator=generator, dtype=torch.bfloat16).to(device) + v = torch.randn(shape_kv, generator=generator, dtype=torch.bfloat16).to(device) + owner_ranges = ((0, args.seq_len // 2), (args.seq_len // 2, args.seq_len)) + query_ranges = owner_ranges + blocks: list[AttentionCPBlockMetadata] = [] + for owner, (owner_start, owner_end) in enumerate(owner_ranges): + for start in range(owner_start, owner_end, args.chunk_size): + blocks.append( + AttentionCPBlockMetadata( + global_block_index=len(blocks), + kv_block_start=start, + kv_block_end=min(start + args.chunk_size, owner_end), + owner_cp_rank=owner, + owner_tp_rank=0, + ) + ) + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=2, + tp_rank=0, + cp_world_size=2, + cp_rank=rank, + ), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=tuple(blocks), + expected_kv_token_range=(0, args.seq_len), + query_token_ranges=query_ranges, + ) + reference = DeterministicCPAttentionReferenceOp() + local_states: list[AttentionCPPartialState] = [] + for block in reversed(blocks): + if block.owner_cp_rank != rank: + continue + state = reference.local_partial_state( + q, + k[:, :, block.kv_block_start : block.kv_block_end, :], + v[:, :, block.kv_block_start : block.kv_block_end, :], + q_start=0, + k_start=block.kv_block_start, + total_kv_len=args.seq_len, + total_query_len=args.seq_len, + causal=True, + ) + local_states.append( + AttentionCPPartialState(state.out, state.lse, block) + ) + + communication = P2PNCCLAttentionCPCommunication() + gathered = communication.all_gather_partial_states(tuple(local_states), plan) + merged = merge_attention_partial_states( + [ + AttentionPartialState( + state.out, + state.lse, + state.block.kv_block_start, + state.block.kv_block_end, + ) + for state in gathered + ] + ) + local = communication.reduce_scatter_merged_state( + AttentionCPMergedState(merged.out, merged.lse), + plan, + ) + full_out, full_lse = reference.forward_fp32_with_lse(q, k, v, causal=True) + start, end = query_ranges[rank] + out_max_abs = float((local.out - full_out[:, :, start:end, :]).abs().max().item()) + lse_max_abs = float((local.lse - full_lse[:, :, start:end]).abs().max().item()) + expected_indices = list(range(len(blocks))) + gathered_indices = [state.block.global_block_index for state in gathered] + passed = ( + gathered_indices == expected_indices + and out_max_abs <= args.atol + and lse_max_abs <= args.atol + ) + return { + "rank": rank, + "device": str(device), + "dtype": "bf16", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "transport": "p2p_nccl_reference", + "query_range": [start, end], + "gathered_block_indices": gathered_indices, + "out_max_abs": out_max_abs, + "lse_max_abs": lse_max_abs, + "atol": args.atol, + "passed": passed, + } + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_ws2_cp_attention_drift_benchmark.py b/tests/test_ws2_cp_attention_drift_benchmark.py index 949aca6a..0f348505 100644 --- a/tests/test_ws2_cp_attention_drift_benchmark.py +++ b/tests/test_ws2_cp_attention_drift_benchmark.py @@ -53,7 +53,21 @@ def test_smoke_report_has_pr5_schema_and_qwen3_tp2_cp2_case(): assert chunked["topology"]["local_query_head_range"] == [0, 16] assert chunked["topology"]["local_kv_head_range"] == [0, 4] assert chunked["provenance"]["merge_order"] == "global_block_index" - assert chunked["provenance"]["split_kv_policy"] == "fixed_kv_chunk" + assert chunked["provenance"]["split_kv_policy"] == "fixed" + assert chunked["provenance"]["requested_split_kv_size"] == 1 + assert chunked["provenance"]["actual_split_kv_plans"][0][ + "actual_split_boundaries" + ] + assert chunked["provenance"]["actual_split_kv_plans"][0][ + "split_kv_accum_dtype" + ] == "fp32" + assert chunked["provenance"]["actual_split_kv_plans"][0][ + "split_kv_downcast_at" + ] == "final_write" + plan_set = chunked["provenance"]["actual_split_kv_plan_set"] + assert plan_set["coverage"] == "complete_batch_tp_cp_owner_cartesian_product" + assert len(plan_set["entries"]) == 8 + assert chunked["distributed_p2p_reference"]["status"] == "not_requested" assert chunked["provenance"]["block_metadata_hash"] assert chunked["provenance"]["rope"]["rope_state"] == "post_rope" assert chunked["drift"]["rope"]["status"] == "available"