From 0cd715f74f9a7dfce2e3778c492911b6ac26620c Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:29:31 +0800 Subject: [PATCH 1/6] fix(attention): align CP backward Split-KV provenance --- docs/operators/attention.md | 85 +- rl_engine/kernels/attention_contract.py | 1472 +++++++++++++++++ .../ops/pytorch/attention/cp_attention.py | 152 +- tests/test_cp_attention.py | 64 + tests/test_operator_inputs.py | 21 +- 5 files changed, 1729 insertions(+), 65 deletions(-) create mode 100644 rl_engine/kernels/attention_contract.py diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 92f646a4..bc7f4a38 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -85,44 +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. - -PR8 extends the CP reference with training-side backward validation: +### WS2 CP-aware dispatch -```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, -) -``` +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 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. +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 @@ -177,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 @@ -187,31 +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, 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 -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) @@ -280,14 +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. -- 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/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 d34dd729..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 @@ -195,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, @@ -429,6 +454,16 @@ def backward_reference( ], "cp_world_size": cp_world_size, "kv_chunk_size": kv_chunk_size, + "requested_split_kv_policy": ( + "disabled" if kv_chunk_size is None else "fixed" + ), + "requested_split_kv_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", @@ -803,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 @@ -873,6 +908,116 @@ 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", @@ -880,8 +1025,11 @@ def _kv_block_bounds( "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 e486a24b..196cfcaa 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -19,6 +19,7 @@ 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 @@ -409,6 +410,57 @@ def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): 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(): @@ -568,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) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index fee1cf94..8840eed3 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -41,8 +41,8 @@ def _args(**overrides): [ "rms_norm", "matmul", + "det_gemm", "attention", - "cp_attention", "logp", "linear_logp", "batch_invariant_logp", @@ -107,3 +107,22 @@ def test_constant_linear_logp_inputs_match_operator_contract(): assert torch.equal(inputs["lm_head_weight"], torch.full((17, 128), 0.51)) assert torch.equal(inputs["target_ids"], torch.full((1, 2), 3, dtype=torch.long)) assert inputs["bias"] is None + + +def test_constant_embedding_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5, token_value=3) + inputs = make_operator_inputs("embedding", args, torch.float32, torch.device("cpu")) + + assert torch.equal(inputs["token_ids"], torch.full((1, 2), 3, dtype=torch.long)) + assert torch.equal(inputs["weight"], torch.full((17, 128), 0.5)) + assert operator_shape_name("embedding", args) == "1x2x17x128" + + +def test_constant_lm_head_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5) + inputs = make_operator_inputs("lm_head", args, torch.float32, torch.device("cpu")) + + assert torch.equal(inputs["hidden"], torch.full((1, 2, 128), 0.5)) + assert torch.equal(inputs["weight"], torch.full((17, 128), 0.51)) + assert inputs["bias"] is None + assert operator_shape_name("lm_head", args) == "1x2x128x17" From 12f35e1ab57f8f3c58fbe43ab0da1ec57b5b7815 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:34:53 +0800 Subject: [PATCH 2/6] feat(attention): consume saved state in CP backward Signed-off-by: lamentropetion <3051000145@qq.com> --- docs/operators/attention.md | 8 +- .../ops/pytorch/attention/cp_attention.py | 428 +++++++++++++++++- tests/test_cp_attention.py | 232 ++++++++++ 3 files changed, 648 insertions(+), 20 deletions(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index bc7f4a38..e1aff500 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -221,8 +221,12 @@ 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. +- `save_forward_state(q, k, v, ...)` captures FP32 `out/lse`, masks, position metadata, + topology, Split-KV boundaries, and content fingerprints for the exact forward invocation. +- `backward_reference(q, k, v, dout, ..., saved_forward_state=state)` consumes that saved + state in canonical global KV-block order. It fails closed if Q/K/V, saved tensors, masks, + offsets, topology, or Split-KV metadata changed, 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. diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index c334f49a..8e766a6e 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -10,6 +10,7 @@ from __future__ import annotations +import hashlib import math from dataclasses import dataclass from typing import Optional, Sequence @@ -46,6 +47,10 @@ def __post_init__(self) -> None: raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") if self.lse.shape != self.out.shape[:3]: raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") if self.block_start < 0: raise ValueError("block_start must be non-negative") if self.block_end < self.block_start: @@ -61,6 +66,45 @@ class AttentionBackwardGradients: dv: torch.Tensor +@dataclass(frozen=True) +class AttentionSavedForwardState: + """Exact FP32 forward state consumed by the PR8 backward reference.""" + + out: torch.Tensor + lse: torch.Tensor + causal: bool + scale: float + key_padding_mask: Optional[torch.Tensor] + query_position_offsets: torch.Tensor + key_position_offsets: torch.Tensor + cp_world_size: int + kv_chunk_size: Optional[int] + query_bounds: tuple[tuple[int, int], ...] + kv_block_bounds: tuple[tuple[int, int], ...] + q_shape: tuple[int, ...] + k_shape: tuple[int, ...] + v_shape: tuple[int, ...] + q_dtype: torch.dtype + k_dtype: torch.dtype + v_dtype: torch.dtype + q_fingerprint: str + k_fingerprint: str + v_fingerprint: str + out_fingerprint: str + lse_fingerprint: str + key_padding_mask_fingerprint: Optional[str] + query_position_offsets_fingerprint: str + key_position_offsets_fingerprint: str + + def __post_init__(self) -> None: + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("saved attention out/lse must be FP32") + if self.out.ndim != 4 or self.lse.shape != self.out.shape[:3]: + raise ValueError("saved attention out/lse shapes are invalid") + if not math.isfinite(self.scale) or self.scale <= 0: + raise ValueError("saved attention scale must be positive and finite") + + @dataclass(frozen=True) class AttentionBackwardPathResult: """One materialized CP attention backward path.""" @@ -69,6 +113,7 @@ class AttentionBackwardPathResult: out: torch.Tensor lse: torch.Tensor gradients: AttentionBackwardGradients + saved_forward_state: AttentionSavedForwardState provenance: dict[str, object] @@ -330,6 +375,8 @@ def forward_with_lse( ``output_dtype`` defaults to the input dtype. """ + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) out, lse = self._forward_impl( q, k, @@ -342,7 +389,7 @@ def forward_with_lse( cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, ) - out = out.to(q.dtype if output_dtype is None else output_dtype) + out = out.to(resolved_output_dtype) return out, lse def forward_fp32_with_lse( @@ -391,6 +438,7 @@ def backward_reference( kv_chunk_size: Optional[int] = None, output_dtype: Optional[torch.dtype] = torch.float32, name: Optional[str] = None, + saved_forward_state: Optional[AttentionSavedForwardState] = None, ) -> AttentionBackwardPathResult: """Run the deterministic training-side backward validation path. @@ -406,15 +454,16 @@ def backward_reference( raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") if not torch.is_floating_point(dout) or torch.is_complex(dout): raise ValueError("dout must be a real floating-point tensor") - q_leaf = q.detach().clone().requires_grad_(True) - k_leaf = k.detach().clone().requires_grad_(True) - v_leaf = v.detach().clone().requires_grad_(True) - + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") 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, + _validate_output_dtype(resolved_output_dtype) + state = saved_forward_state or self.save_forward_state( + q, + k, + v, causal=causal, scale=scale, key_padding_mask=key_padding_mask, @@ -422,11 +471,23 @@ def backward_reference( 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") + _validate_saved_forward_state( + state, + 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, + ) + gradients = _backward_from_saved_state(q, k, v, dout, state) + out = state.out.to(resolved_output_dtype) + lse = state.lse return AttentionBackwardPathResult( name=name @@ -434,10 +495,11 @@ def backward_reference( out=out.detach(), lse=lse.detach(), gradients=AttentionBackwardGradients( - dq=q_leaf.grad.detach(), - dk=k_leaf.grad.detach(), - dv=v_leaf.grad.detach(), + dq=gradients.dq, + dk=gradients.dk, + dv=gradients.dv, ), + saved_forward_state=state, provenance={ "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", "gradient_mode": "training_backward", @@ -472,11 +534,89 @@ def backward_reference( "k_dtype": str(k.dtype).replace("torch.", ""), "v_dtype": str(v.dtype).replace("torch.", ""), "dout_dtype": str(dout.dtype).replace("torch.", ""), + "saved_forward_state_source": ( + "caller" if saved_forward_state is not None else "captured_reference" + ), + "backward_algorithm": "saved_out_lse_block_order_reference", "te_backward_oracle": "not_used", "decode_backward": "not_supported", }, ) + def save_forward_state( + 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, + ) -> AttentionSavedForwardState: + """Capture the exact state a production training backward must consume.""" + + out, lse = 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, + ) + batch, _, sq, dim = q.shape + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=k.size(2) - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + mask = None if key_padding_mask is None else key_padding_mask.detach().clone() + return AttentionSavedForwardState( + out=out.detach().clone(), + lse=lse.detach().clone(), + causal=causal, + scale=float(scale if scale is not None else 1.0 / math.sqrt(dim)), + key_padding_mask=mask, + query_position_offsets=query_offsets.detach().clone(), + key_position_offsets=key_offsets.detach().clone(), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + query_bounds=tuple(_split_bounds(sq, cp_world_size)), + kv_block_bounds=tuple(_kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size)), + q_shape=tuple(q.shape), + k_shape=tuple(k.shape), + v_shape=tuple(v.shape), + q_dtype=q.dtype, + k_dtype=k.dtype, + v_dtype=v.dtype, + q_fingerprint=_tensor_fingerprint(q), + k_fingerprint=_tensor_fingerprint(k), + v_fingerprint=_tensor_fingerprint(v), + out_fingerprint=_tensor_fingerprint(out), + lse_fingerprint=_tensor_fingerprint(lse), + key_padding_mask_fingerprint=( + None if mask is None else _tensor_fingerprint(mask) + ), + query_position_offsets_fingerprint=_tensor_fingerprint(query_offsets), + key_position_offsets_fingerprint=_tensor_fingerprint(key_offsets), + ) + def local_partial_state( self, q: torch.Tensor, @@ -503,6 +643,7 @@ def local_partial_state( """ _validate_qkv(q, k, v) + _validate_scale(scale) if q_start < 0 or k_start < 0: raise ValueError("q_start and k_start must be non-negative") if total_kv_len < k_start + k.size(2): @@ -598,9 +739,14 @@ def _forward_impl( kv_chunk_size: Optional[int], ) -> tuple[torch.Tensor, torch.Tensor]: _validate_qkv(q, k, v) - if cp_world_size < 1: + _validate_scale(scale) + if isinstance(cp_world_size, bool) or not isinstance(cp_world_size, int) or cp_world_size < 1: raise ValueError("cp_world_size must be >= 1") - if kv_chunk_size is not None and kv_chunk_size < 1: + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): raise ValueError("kv_chunk_size must be >= 1 when provided") batch, hq, sq, dim = q.shape @@ -739,6 +885,224 @@ def compare_cp_attention_backward( ) +def _backward_from_saved_state( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + state: AttentionSavedForwardState, +) -> AttentionBackwardGradients: + """Apply standard-softmax backward in canonical global KV-block order.""" + + batch, hq, _, dim = q.shape + hkv = k.size(1) + group_size = hq // hkv + with NativeAttentionOp._strict_fp32_math(q.device.type): + qf = q.float() + kf = k.float() + vf = v.float() + doutf = dout.float() + k_expanded = kf.repeat_interleave(group_size, dim=1) + v_expanded = vf.repeat_interleave(group_size, dim=1) + dq = torch.zeros_like(qf) + dk_expanded = torch.zeros( + batch, + hq, + k.size(2), + dim, + dtype=torch.float32, + device=q.device, + ) + dv_expanded = torch.zeros_like(dk_expanded) + + for q_start, q_end in state.query_bounds: + if q_start == q_end: + continue + q_block = qf[:, :, q_start:q_end, :] + dout_block = doutf[:, :, q_start:q_end, :] + out_block = state.out[:, :, q_start:q_end, :] + lse_block = state.lse[:, :, q_start:q_end] + dq_block = torch.zeros_like(q_block) + for k_start, k_end in state.kv_block_bounds: + if k_start == k_end: + continue + k_block = k_expanded[:, :, k_start:k_end, :] + v_block = v_expanded[:, :, k_start:k_end, :] + scores = torch.matmul(q_block, k_block.transpose(-1, -2)) * state.scale + if state.causal: + query_base = state.query_position_offsets[:, None] + q_start + key_base = state.key_position_offsets[:, None] + k_start + q_pos = torch.arange( + q_end - q_start, + device=q.device, + dtype=torch.long, + ) + query_base + k_pos = torch.arange( + k_end - k_start, + device=q.device, + dtype=torch.long, + ) + key_base + scores = scores.masked_fill( + (k_pos[:, None, :] > q_pos[:, :, None])[:, None, :, :], + float("-inf"), + ) + if state.key_padding_mask is not None: + scores = scores.masked_fill( + ~state.key_padding_mask[:, None, None, k_start:k_end], + float("-inf"), + ) + probability = torch.exp(scores - lse_block.unsqueeze(-1)) + probability = torch.where( + torch.isfinite(lse_block).unsqueeze(-1), + probability, + torch.zeros_like(probability), + ) + dv_expanded[:, :, k_start:k_end, :] += torch.matmul( + probability.transpose(-1, -2), + dout_block, + ) + dp = torch.matmul(dout_block, v_block.transpose(-1, -2)) + # The global softmax dot term is dout dot the saved global output. + ds = probability * ( + dp - (dout_block * out_block).sum(dim=-1, keepdim=True) + ) + dq_block += torch.matmul(ds, k_block) * state.scale + dk_expanded[:, :, k_start:k_end, :] += ( + torch.matmul(ds.transpose(-1, -2), q_block) * state.scale + ) + dq[:, :, q_start:q_end, :] = dq_block + + dk = dk_expanded.reshape( + batch, + hkv, + group_size, + k.size(2), + dim, + ).sum(dim=2) + dv = dv_expanded.reshape( + batch, + hkv, + group_size, + v.size(2), + dim, + ).sum(dim=2) + return AttentionBackwardGradients( + dq=dq.to(q.dtype), + dk=dk.to(k.dtype), + dv=dv.to(v.dtype), + ) + + +def _validate_saved_forward_state( + state: AttentionSavedForwardState, + 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], +) -> None: + if not isinstance(state, AttentionSavedForwardState): + raise ValueError("saved_forward_state must be an AttentionSavedForwardState") + expected_scale = float(scale if scale is not None else 1.0 / math.sqrt(q.size(-1))) + expected_query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=k.size(2) - q.size(2), + name="query_position_offsets", + ) + expected_key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + checks = { + "out_shape": (tuple(state.out.shape), tuple(q.shape)), + "lse_shape": (tuple(state.lse.shape), tuple(q.shape[:3])), + "out_device": (state.out.device, q.device), + "lse_device": (state.lse.device, q.device), + "q_shape": (state.q_shape, tuple(q.shape)), + "k_shape": (state.k_shape, tuple(k.shape)), + "v_shape": (state.v_shape, tuple(v.shape)), + "q_dtype": (state.q_dtype, q.dtype), + "k_dtype": (state.k_dtype, k.dtype), + "v_dtype": (state.v_dtype, v.dtype), + "causal": (state.causal, causal), + "scale": (state.scale, expected_scale), + "cp_world_size": (state.cp_world_size, cp_world_size), + "kv_chunk_size": (state.kv_chunk_size, kv_chunk_size), + "query_bounds": (state.query_bounds, tuple(_split_bounds(q.size(2), cp_world_size))), + "kv_block_bounds": ( + state.kv_block_bounds, + tuple(_kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size)), + ), + "q_fingerprint": (state.q_fingerprint, _tensor_fingerprint(q)), + "k_fingerprint": (state.k_fingerprint, _tensor_fingerprint(k)), + "v_fingerprint": (state.v_fingerprint, _tensor_fingerprint(v)), + "out_fingerprint": (state.out_fingerprint, _tensor_fingerprint(state.out)), + "lse_fingerprint": (state.lse_fingerprint, _tensor_fingerprint(state.lse)), + "query_position_offsets_fingerprint": ( + state.query_position_offsets_fingerprint, + _tensor_fingerprint(state.query_position_offsets), + ), + "key_position_offsets_fingerprint": ( + state.key_position_offsets_fingerprint, + _tensor_fingerprint(state.key_position_offsets), + ), + "query_position_offsets_device": ( + state.query_position_offsets.device, + q.device, + ), + "key_position_offsets_device": ( + state.key_position_offsets.device, + q.device, + ), + } + mismatches = [name for name, (actual, expected) in checks.items() if actual != expected] + if not torch.equal(state.query_position_offsets, expected_query_offsets): + mismatches.append("query_position_offsets") + if not torch.equal(state.key_position_offsets, expected_key_offsets): + mismatches.append("key_position_offsets") + masks_match = ( + state.key_padding_mask is None + and key_padding_mask is None + or state.key_padding_mask is not None + and key_padding_mask is not None + and torch.equal(state.key_padding_mask, key_padding_mask) + ) + if not masks_match: + mismatches.append("key_padding_mask") + actual_mask_fingerprint = ( + None + if state.key_padding_mask is None + else _tensor_fingerprint(state.key_padding_mask) + ) + if state.key_padding_mask_fingerprint != actual_mask_fingerprint: + mismatches.append("key_padding_mask_fingerprint") + if mismatches: + raise ValueError( + "saved_forward_state does not match the backward invocation: " + + ", ".join(mismatches) + ) + + +def _tensor_fingerprint(tensor: torch.Tensor) -> str: + digest = hashlib.sha256() + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(str(tensor.device).encode()) + digest.update(tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + def _compare_backward_path( candidate: AttentionBackwardPathResult, reference: AttentionBackwardPathResult, @@ -850,10 +1214,37 @@ def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: raise ValueError("k and v must have the same shape") if q.size(0) != k.size(0) or q.size(3) != k.size(3): raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") if q.size(1) % k.size(1) != 0: raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") +def _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: total = torch.tensor(0.0, device=tensors[0].device) for tensor in tensors: @@ -1025,6 +1416,7 @@ def build_reference_split_kv_runtime_plan_set( "AttentionBackwardPathResult", "AttentionBackwardRankDrift", "AttentionPartialState", + "AttentionSavedForwardState", "build_reference_split_kv_runtime_plan_set", "CPAttentionReferenceOp", "DeterministicCPAttentionReferenceOp", diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 196cfcaa..0670bbc5 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -16,6 +16,7 @@ from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, + AttentionSavedForwardState, DeterministicCPAttentionReferenceOp, compare_cp_attention_backward, merge_attention_partial_states, @@ -447,6 +448,183 @@ def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): ] +def test_saved_forward_backward_matches_independent_dense_autograd(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 5, 7, seed=31, heads=4, kv_heads=2, dim=8) + mask = torch.tensor( + [[True, True, True, True, True, True, False], [True] * 7], + dtype=torch.bool, + ) + query_offsets = torch.tensor([11, 23], dtype=torch.long) + key_offsets = torch.tensor([9, 21], dtype=torch.long) + dout = torch.randn(q.shape, generator=torch.Generator().manual_seed(32)) + + state = op.save_forward_state( + q, + k, + v, + causal=True, + scale=0.37, + key_padding_mask=mask, + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + result = op.backward_reference( + q, + k, + v, + dout, + causal=True, + scale=0.37, + key_padding_mask=mask, + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + cp_world_size=2, + kv_chunk_size=2, + saved_forward_state=state, + ) + + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + k_expanded = k_ref.repeat_interleave(2, dim=1) + v_expanded = v_ref.repeat_interleave(2, dim=1) + scores = torch.matmul(q_ref, k_expanded.transpose(-1, -2)) * 0.37 + q_pos = query_offsets[:, None] + torch.arange(q.size(2)) + k_pos = key_offsets[:, None] + torch.arange(k.size(2)) + scores = scores.masked_fill( + (k_pos[:, None, :] > q_pos[:, :, None])[:, None, :, :], + float("-inf"), + ) + scores = scores.masked_fill(~mask[:, None, None, :], float("-inf")) + out_ref = torch.matmul(torch.softmax(scores, dim=-1), v_expanded) + out_ref.backward(dout) + + assert isinstance(result.saved_forward_state, AttentionSavedForwardState) + assert result.saved_forward_state is state + assert result.provenance["saved_forward_state_source"] == "caller" + torch.testing.assert_close(result.out, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(result.gradients.dq, q_ref.grad, atol=_GRAD_ATOL, rtol=0.0) + torch.testing.assert_close(result.gradients.dk, k_ref.grad, atol=_GRAD_ATOL, rtol=0.0) + torch.testing.assert_close(result.gradients.dv, v_ref.grad, atol=_GRAD_ATOL, rtol=0.0) + + +@pytest.mark.parametrize( + ("tensor_name", "message"), + [("q", "q_fingerprint"), ("k", "k_fingerprint"), ("v", "v_fingerprint")], +) +def test_saved_forward_state_rejects_stale_qkv(tensor_name, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=33, heads=4, kv_heads=2, dim=8) + state = op.save_forward_state(q, k, v, cp_world_size=2, kv_chunk_size=2) + inputs = {"q": q.clone(), "k": k.clone(), "v": v.clone()} + inputs[tensor_name].flatten()[0] += 1.0 + + with pytest.raises(ValueError, match=message): + op.backward_reference( + inputs["q"], + inputs["k"], + inputs["v"], + torch.ones_like(q), + cp_world_size=2, + kv_chunk_size=2, + saved_forward_state=state, + ) + + +@pytest.mark.parametrize( + ("tensor_name", "message"), + [ + ("out", "out_fingerprint"), + ("lse", "lse_fingerprint"), + ("key_padding_mask", "key_padding_mask_fingerprint"), + ("query_position_offsets", "query_position_offsets_fingerprint"), + ("key_position_offsets", "key_position_offsets_fingerprint"), + ], +) +def test_saved_forward_state_rejects_mutated_saved_tensors(tensor_name, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=34, heads=4, kv_heads=2, dim=8) + mask = torch.ones(1, 4, dtype=torch.bool) + offsets = torch.tensor([7], dtype=torch.long) + state = op.save_forward_state( + q, + k, + v, + key_padding_mask=mask, + query_position_offsets=offsets, + key_position_offsets=offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + tensor = getattr(state, tensor_name) + if tensor.dtype == torch.bool: + tensor.flatten()[0].logical_not_() + else: + tensor.flatten()[0].add_(1) + + with pytest.raises(ValueError, match=message): + op.backward_reference( + q, + k, + v, + torch.ones_like(q), + key_padding_mask=mask, + query_position_offsets=offsets, + key_position_offsets=offsets, + cp_world_size=2, + kv_chunk_size=2, + saved_forward_state=state, + ) + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"causal": False}, "causal"), + ({"scale": 0.5}, "scale"), + ({"cp_world_size": 1}, "cp_world_size"), + ({"kv_chunk_size": None}, "kv_chunk_size"), + ({"query_position_offsets": torch.tensor([8])}, "query_position_offsets"), + ({"key_position_offsets": torch.tensor([8])}, "key_position_offsets"), + ], +) +def test_saved_forward_state_rejects_execution_metadata_mismatch(override, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=35, heads=4, kv_heads=2, dim=8) + offsets = torch.tensor([7], dtype=torch.long) + state = op.save_forward_state( + q, + k, + v, + query_position_offsets=offsets, + key_position_offsets=offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + kwargs = { + "causal": True, + "scale": None, + "query_position_offsets": offsets, + "key_position_offsets": offsets, + "cp_world_size": 2, + "kv_chunk_size": 2, + } + kwargs.update(override) + + with pytest.raises(ValueError, match=message): + op.backward_reference( + q, + k, + v, + torch.ones_like(q), + saved_forward_state=state, + **kwargs, + ) + + def test_split_kv_plan_never_crosses_cp_owner_boundaries(): plans = split_kv_execution_plan_provenance( 10, @@ -550,6 +728,15 @@ def test_backward_report_validates_dout_shape_and_dtype(): cp_world_size=2, ) + with pytest.raises(ValueError, match="dout must have the same dtype"): + op.backward_reference( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + torch.ones_like(q), + cp_world_size=2, + ) + def test_inputs_are_not_mutated(): op = DeterministicCPAttentionReferenceOp() @@ -608,6 +795,51 @@ def test_invalid_gqa_and_mask_shapes_raise(): ) +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf"), True, "bad"]) +def test_invalid_scale_fails_before_attention_math(scale): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=24, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="scale must be a positive finite number"): + op.forward_fp32_with_lse(q, k, v, scale=scale) + + +def test_qkv_dtype_and_floating_contract_fails_closed(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=25, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="same dtype"): + op.forward_fp32_with_lse(q, k.to(torch.bfloat16), v) + with pytest.raises(ValueError, match="real floating-point"): + op.forward_fp32_with_lse(q.to(torch.long), k.to(torch.long), v.to(torch.long)) + + +@pytest.mark.parametrize("kwargs", [{"cp_world_size": True}, {"kv_chunk_size": True}]) +def test_boolean_parallelism_arguments_fail_closed(kwargs): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=26, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError): + op.forward_fp32_with_lse(q, k, v, **kwargs) + + +@pytest.mark.parametrize("output_dtype", [torch.long, torch.complex64, "fp32"]) +def test_nonfloating_output_dtype_fails_closed(output_dtype): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=27, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="output_dtype must be a real floating-point"): + op.forward_with_lse(q, k, v, output_dtype=output_dtype) + + +def test_partial_states_must_remain_fp32_and_colocated(): + out = torch.zeros(1, 1, 1, 1, dtype=torch.bfloat16) + lse = torch.zeros(1, 1, 1, dtype=torch.float32) + + with pytest.raises(ValueError, match="must remain FP32"): + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=1) + + def test_overlapping_partial_ranges_raise(): out = torch.zeros(1, 1, 1, 1) lse = torch.zeros(1, 1, 1) From 616eea711c8bf824b1735808dba1d894606856e2 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 21:03:00 +0800 Subject: [PATCH 3/6] fix(attention): align operator input shape metadata Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/gtest/operator_inputs.py | 12 +++++++----- tests/test_operator_inputs.py | 6 ++++++ 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index c9777cf7..0156907b 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -61,8 +61,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", - "embedding": f"{batch}x{seq}x{vocab}x{DEFAULT_HIDDEN}", - "lm_head": f"{batch}x{seq}x{vocab}", + "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", + "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", } try: @@ -234,9 +234,10 @@ def _make_embedding_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { "token_ids": _token_ids((batch, seq), vocab, args, device), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 0), } @@ -245,9 +246,10 @@ def _make_lm_head_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { - "hidden": _floating_tensor((batch, seq, DEFAULT_HIDDEN), args, dtype, device, 0), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 1), + "hidden": _floating_tensor((batch, seq, hidden_dim), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 1), "bias": None, } diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 8840eed3..4b9e68ec 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -114,6 +114,8 @@ def test_constant_embedding_inputs_match_operator_contract(): inputs = make_operator_inputs("embedding", args, torch.float32, torch.device("cpu")) assert torch.equal(inputs["token_ids"], torch.full((1, 2), 3, dtype=torch.long)) + assert inputs["weight"].shape == (17, 128) + assert inputs["weight"].dtype is torch.float32 assert torch.equal(inputs["weight"], torch.full((17, 128), 0.5)) assert operator_shape_name("embedding", args) == "1x2x17x128" @@ -122,6 +124,10 @@ def test_constant_lm_head_inputs_match_operator_contract(): args = _args(input_mode="constant", constant_value=0.5) inputs = make_operator_inputs("lm_head", args, torch.float32, torch.device("cpu")) + assert inputs["hidden"].shape == (1, 2, 128) + assert inputs["weight"].shape == (17, 128) + assert inputs["hidden"].dtype is torch.float32 + assert inputs["weight"].dtype is torch.float32 assert torch.equal(inputs["hidden"], torch.full((1, 2, 128), 0.5)) assert torch.equal(inputs["weight"], torch.full((17, 128), 0.51)) assert inputs["bias"] is None From 26c62fa78dd20d234d18d04597a4d8a89a3c3f0c Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:48:04 +0800 Subject: [PATCH 4/6] test(attention): record projection collective contract --- rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 7 +++++++ tests/test_cp_attention.py | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 36e4ce30..6d4ddad9 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -538,6 +538,13 @@ def backward_reference( "backward_algorithm": "saved_out_lse_block_order_reference", "te_backward_oracle": "not_used", "decode_backward": "not_supported", + "projection_scope": "attention_core_only", + "qkv_projection_backward_dgrad_collective": "all_reduce", + "qkv_projection_sp_backward_collective": "reduce_scatter", + "o_proj_backward_dgrad_collective": "none", + "o_proj_sp_backward_collective": "all_gather", + "projection_collectives_executed": False, + "projection_collectives_source": "attention_contract_runtime_adapter", }, ) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 0670bbc5..e0d0bde4 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -385,6 +385,12 @@ def test_backward_report_cp2_prefill_matches_cp1_reference(): assert drift.provenance["merge_order"] == "global_block_index" assert drift.provenance["te_backward_oracle"] == "not_used" assert drift.provenance["decode_backward"] == "not_supported" + assert drift.provenance["projection_scope"] == "attention_core_only" + assert drift.provenance["qkv_projection_backward_dgrad_collective"] == "all_reduce" + assert drift.provenance["qkv_projection_sp_backward_collective"] == "reduce_scatter" + assert drift.provenance["o_proj_backward_dgrad_collective"] == "none" + assert drift.provenance["o_proj_sp_backward_collective"] == "all_gather" + assert drift.provenance["projection_collectives_executed"] is False json.dumps(report.to_dict()) From 9da6c20cec85fb5cc81d5ffe7cade8c7a12cfdcd Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:49:12 +0800 Subject: [PATCH 5/6] feat(attention): align projection contract --- rl_engine/kernels/attention_contract.py | 142 ++++++++++++++++++++++++ 1 file changed, 142 insertions(+) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 1750476d..2e359d29 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -61,6 +61,13 @@ class SplitKVMode(str, Enum): AUTO = "auto" +class ProjectionCollective(str, Enum): + NONE = "none" + ALL_REDUCE = "all_reduce" + ALL_GATHER = "all_gather" + REDUCE_SCATTER = "reduce_scatter" + + class RoPEState(str, Enum): PRE_ROPE = "pre_rope" POST_ROPE = "post_rope" @@ -134,12 +141,16 @@ class ShardingSpec: global_block_token_starts: tuple[int, ...] local_block_offsets: tuple[int, ...] packed_sequence_offsets: tuple[int, ...] | None = None + sp_rank: int = 0 + sp_world_size: int = 1 def __post_init__(self) -> None: tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + sp_world_size = _positive_int(self.sp_world_size, "sp_world_size") tp_rank = _non_negative_int(self.tp_rank, "tp_rank") cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + sp_rank = _non_negative_int(self.sp_rank, "sp_rank") if tp_rank >= tp_world_size: raise AttentionContractError( f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" @@ -148,6 +159,10 @@ def __post_init__(self) -> None: raise AttentionContractError( f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" ) + if sp_rank >= sp_world_size: + raise AttentionContractError( + f"sp_rank={sp_rank} must be smaller than sp_world_size={sp_world_size}" + ) global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") @@ -1076,6 +1091,98 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class AttentionProjectionSpec: + """Deterministic QKV or output-projection execution contract.""" + + name: str + input_dtype: AttentionDType = AttentionDType.BF16 + output_dtype: AttentionDType = AttentionDType.BF16 + acc_dtype: AttentionDType = AttentionDType.FP32 + split_kv: SplitKVMode = SplitKVMode.DISABLED + k_order: str = "ascending" + backend_policy: str = "native_verified_then_common_deterministic" + deterministic_backend: str = "rlkernel.cuda.det_gemm" + tp_forward_collective: ProjectionCollective = ProjectionCollective.NONE + tp_backward_dgrad_collective: ProjectionCollective = ProjectionCollective.NONE + sp_forward_collective: ProjectionCollective = ProjectionCollective.NONE + sp_backward_collective: ProjectionCollective = ProjectionCollective.NONE + qkv_split_order: tuple[str, ...] = () + require_runtime_readback: bool = True + + @classmethod + def qkv(cls) -> "AttentionProjectionSpec": + return cls( + name="qkv", + tp_backward_dgrad_collective=ProjectionCollective.ALL_REDUCE, + sp_forward_collective=ProjectionCollective.ALL_GATHER, + sp_backward_collective=ProjectionCollective.REDUCE_SCATTER, + qkv_split_order=("q", "k", "v"), + ) + + @classmethod + def output(cls) -> "AttentionProjectionSpec": + return cls( + name="o_proj", + tp_forward_collective=ProjectionCollective.ALL_REDUCE, + sp_forward_collective=ProjectionCollective.REDUCE_SCATTER, + sp_backward_collective=ProjectionCollective.ALL_GATHER, + ) + + def __post_init__(self) -> None: + if self.name not in {"qkv", "o_proj"}: + raise AttentionContractError("projection name must be qkv or o_proj") + for field_name in ("input_dtype", "output_dtype", "acc_dtype"): + object.__setattr__( + self, + field_name, + _enum_value(AttentionDType, getattr(self, field_name), field_name), + ) + object.__setattr__( + self, + "split_kv", + _enum_value(SplitKVMode, self.split_kv, "projection.split_kv"), + ) + for field_name in ( + "tp_forward_collective", + "tp_backward_dgrad_collective", + "sp_forward_collective", + "sp_backward_collective", + ): + object.__setattr__( + self, + field_name, + _enum_value( + ProjectionCollective, + getattr(self, field_name), + field_name, + ), + ) + if self.input_dtype is not AttentionDType.BF16: + raise AttentionContractError("Attention projections require BF16 input") + if self.output_dtype is not AttentionDType.BF16: + raise AttentionContractError("Attention projections require BF16 output") + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Attention projections require FP32 accumulation") + if self.split_kv is not SplitKVMode.DISABLED: + raise AttentionContractError("Attention projection GEMMs must disable Split-K") + if self.k_order != "ascending": + raise AttentionContractError("Attention projection GEMMs require ascending K order") + for field_name in ("backend_policy", "deterministic_backend"): + value = getattr(self, field_name) + if not isinstance(value, str) or not value.strip(): + raise AttentionContractError(f"{field_name} must be a non-empty string") + if not isinstance(self.require_runtime_readback, bool) or not self.require_runtime_readback: + raise AttentionContractError("projection runtime readback must be required") + split_order = tuple(self.qkv_split_order) + if self.name == "qkv": + if split_order != ("q", "k", "v"): + raise AttentionContractError("QKV projection must split in Q, K, V order") + elif split_order: + raise AttentionContractError("o_proj must not declare a QKV split order") + object.__setattr__(self, "qkv_split_order", split_order) + + @dataclass(frozen=True) class AttentionContract: """Complete semantic request consumed by contract-aware dispatch.""" @@ -1094,6 +1201,10 @@ class AttentionContract: kv_cache: KVCacheSpec | None = None rope: RoPESpec | None = None export_lse: bool = True + qkv_projection: AttentionProjectionSpec = field(default_factory=AttentionProjectionSpec.qkv) + output_projection: AttentionProjectionSpec = field( + default_factory=AttentionProjectionSpec.output + ) def __post_init__(self) -> None: object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) @@ -1108,6 +1219,14 @@ def __post_init__(self) -> None: raise AttentionContractError("reduction must be a ReductionSpec") if not isinstance(self.split_kv, SplitKVSpec): raise AttentionContractError("split_kv must be a SplitKVSpec") + if not isinstance(self.qkv_projection, AttentionProjectionSpec): + raise AttentionContractError("qkv_projection must be an AttentionProjectionSpec") + if self.qkv_projection.name != "qkv": + raise AttentionContractError("qkv_projection must use name='qkv'") + if not isinstance(self.output_projection, AttentionProjectionSpec): + raise AttentionContractError("output_projection must be an AttentionProjectionSpec") + if self.output_projection.name != "o_proj": + raise AttentionContractError("output_projection must use name='o_proj'") if ( self.mode is AttentionMode.PREFILL and query_sequence_length != self.sharding.local_sequence_length @@ -1192,6 +1311,8 @@ def to_dict(self) -> dict[str, Any]: "tp_world_size": self.sharding.tp_world_size, "cp_rank": self.sharding.cp_rank, "cp_world_size": self.sharding.cp_world_size, + "sp_rank": self.sharding.sp_rank, + "sp_world_size": self.sharding.sp_world_size, "global_q_heads": self.sharding.global_q_heads, "global_kv_heads": self.sharding.global_kv_heads, "local_q_head_start": self.sharding.local_q_head_start, @@ -1254,6 +1375,24 @@ def to_dict(self) -> dict[str, Any]: "output_dtype": self.rope.output_dtype.value, "fusion_boundary": self.rope.fusion_boundary.value, } + projections = { + spec.name: { + "input_dtype": spec.input_dtype.value, + "output_dtype": spec.output_dtype.value, + "acc_dtype": spec.acc_dtype.value, + "split_kv": spec.split_kv.value, + "k_order": spec.k_order, + "backend_policy": spec.backend_policy, + "deterministic_backend": spec.deterministic_backend, + "tp_forward_collective": spec.tp_forward_collective.value, + "tp_backward_dgrad_collective": spec.tp_backward_dgrad_collective.value, + "sp_forward_collective": spec.sp_forward_collective.value, + "sp_backward_collective": spec.sp_backward_collective.value, + "qkv_split_order": list(spec.qkv_split_order), + "require_runtime_readback": spec.require_runtime_readback, + } + for spec in (self.qkv_projection, self.output_projection) + } return { "semantic_operator": "standard_softmax_attention", "role": self.role.value, @@ -1273,6 +1412,7 @@ def to_dict(self) -> dict[str, Any]: "split_kv": self.split_kv.to_dict(), "kv_cache": kv_cache, "rope": rope, + "projections": projections, } @@ -1428,6 +1568,7 @@ class AttentionDispatchResult: "AttentionBackendCapability", "AttentionDispatchResult", "AttentionDType", + "AttentionProjectionSpec", "AttentionMerge", "AttentionMode", "AttentionRole", @@ -1436,6 +1577,7 @@ class AttentionDispatchResult: "ReductionEngine", "ReductionOrder", "ReductionSpec", + "ProjectionCollective", "RoPECastPoint", "RoPEFusionBoundary", "RoPESpec", From 9e5c40036c3e445a78ca9dfe2e27b10f3c40424b Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 23:25:31 +0800 Subject: [PATCH 6/6] fix(attention): canonicalize strict backward schedule --- rl_engine/kernels/attention_contract.py | 5 + .../ops/pytorch/attention/cp_attention.py | 323 ++++++++++++++++-- tests/test_cp_attention.py | 91 +++++ 3 files changed, 400 insertions(+), 19 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 2e359d29..6f71b14c 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -17,6 +17,9 @@ _EnumT = TypeVar("_EnumT", bound=Enum) +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" + class AttentionContractError(ValueError): """Raised when attention metadata does not describe a valid invocation.""" @@ -1589,6 +1592,8 @@ class AttentionDispatchResult: "SplitKVRuntimePlanEntry", "SplitKVRuntimePlanSet", "SplitKVSpec", + "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_SCHEDULE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 6d4ddad9..544ecaa2 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -23,6 +23,8 @@ SplitKVRuntimeCoordinate, SplitKVRuntimePlanEntry, SplitKVRuntimePlanSet, + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp @@ -95,6 +97,8 @@ class AttentionSavedForwardState: key_padding_mask_fingerprint: Optional[str] query_position_offsets_fingerprint: str key_position_offsets_fingerprint: str + strict_bitwise: bool + strict_schedule: Optional[str] def __post_init__(self) -> None: if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: @@ -103,6 +107,9 @@ def __post_init__(self) -> None: raise ValueError("saved attention out/lse shapes are invalid") if not math.isfinite(self.scale) or self.scale <= 0: raise ValueError("saved attention scale must be positive and finite") + expected_schedule = STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None + if self.strict_schedule != expected_schedule: + raise ValueError("saved attention strict schedule does not match strict_bitwise") @dataclass(frozen=True) @@ -249,6 +256,13 @@ class DeterministicCPAttentionReferenceOp: op_class = "attention" + def __init__(self, *, strict_bitwise: bool = False) -> None: + """Create either the diagnostic reference or the canonical strict path.""" + + if not isinstance(strict_bitwise, bool): + raise TypeError("strict_bitwise must be a bool") + self.strict_bitwise = strict_bitwise + @staticmethod def split_kv_execution_plans( total_kv_tokens: int, @@ -377,18 +391,32 @@ def forward_with_lse( resolved_output_dtype = q.dtype if output_dtype is None else output_dtype _validate_output_dtype(resolved_output_dtype) - out, lse = self._forward_impl( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - ) + if self.strict_bitwise: + out, lse = self._forward_strict_bitwise( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + else: + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) out = out.to(resolved_output_dtype) return out, lse @@ -422,6 +450,97 @@ def forward_fp32_with_lse( output_dtype=torch.float32, ) + def _forward_strict_bitwise( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Execute the same row-local arithmetic for every batch/CP layout. + + CP ownership and a caller's chunk request remain metadata only. The + strict arithmetic always consumes the complete logical KV row once, + which is the no-Split-KV schedule shared by train and rollout. + """ + + _validate_qkv(q, k, v) + _validate_scale(scale) + _validate_partition_args(cp_world_size, kv_chunk_size) + 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", + ) + if sq == 0: + 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, + ) + + out_rows: list[torch.Tensor] = [] + lse_rows: list[torch.Tensor] = [] + for batch_index in range(batch): + q_batch = q[batch_index : batch_index + 1].contiguous() + k_batch = k[batch_index : batch_index + 1].contiguous() + v_batch = v[batch_index : batch_index + 1].contiguous() + pad_batch = ( + None + if key_padding_mask is None + else key_padding_mask[batch_index : batch_index + 1].contiguous() + ) + query_offset = query_offsets[batch_index : batch_index + 1] + key_offset = key_offsets[batch_index : batch_index + 1] + query_rows: list[torch.Tensor] = [] + lse_query_rows: list[torch.Tensor] = [] + for query_index in range(sq): + q_row = q_batch[:, :, query_index : query_index + 1, :].contiguous() + state = self.local_partial_state( + q_row, + k_batch, + v_batch, + q_start=query_index, + k_start=0, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=pad_batch, + query_position_offsets=query_offset, + key_position_offsets=key_offset, + ) + query_rows.append(state.out) + lse_query_rows.append(state.lse) + out_rows.append(torch.cat(query_rows, dim=2)) + lse_rows.append(torch.cat(lse_query_rows, dim=2)) + return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) + def backward_reference( self, q: torch.Tensor, @@ -484,8 +603,16 @@ def backward_reference( key_position_offsets=key_position_offsets, cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, + strict_bitwise=self.strict_bitwise, + ) + gradients = _backward_from_saved_state( + q, + k, + v, + dout, + state, + strict_bitwise=self.strict_bitwise, ) - gradients = _backward_from_saved_state(q, k, v, dout, state) out = state.out.to(resolved_output_dtype) lse = state.lse @@ -518,15 +645,35 @@ def backward_reference( "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", + "actual_split_kv_plans": ( + _strict_no_split_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + backend="deterministic_cp_backward_strict_reference", + ) + if self.strict_bitwise + else 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", + "strict_bitwise": self.strict_bitwise, + "strict_core_id": ( + STRICT_ATTENTION_CORE_ID if self.strict_bitwise else None + ), + "strict_schedule": ( + STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None + ), + "actual_split_kv_policy": ( + "disabled" + if self.strict_bitwise + else ("disabled" if kv_chunk_size is None else "fixed") + ), "output_dtype": str(resolved_output_dtype).replace("torch.", ""), "q_dtype": str(q.dtype).replace("torch.", ""), "k_dtype": str(k.dtype).replace("torch.", ""), @@ -535,7 +682,11 @@ def backward_reference( "saved_forward_state_source": ( "caller" if saved_forward_state is not None else "captured_reference" ), - "backward_algorithm": "saved_out_lse_block_order_reference", + "backward_algorithm": ( + "saved_out_lse_canonical_row_reference" + if self.strict_bitwise + else "saved_out_lse_block_order_reference" + ), "te_backward_oracle": "not_used", "decode_backward": "not_supported", "projection_scope": "attention_core_only", @@ -618,6 +769,10 @@ def save_forward_state( key_padding_mask_fingerprint=(None if mask is None else _tensor_fingerprint(mask)), query_position_offsets_fingerprint=_tensor_fingerprint(query_offsets), key_position_offsets_fingerprint=_tensor_fingerprint(key_offsets), + strict_bitwise=self.strict_bitwise, + strict_schedule=( + STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None + ), ) def local_partial_state( @@ -898,9 +1053,14 @@ def _backward_from_saved_state( v: torch.Tensor, dout: torch.Tensor, state: AttentionSavedForwardState, + *, + strict_bitwise: bool, ) -> AttentionBackwardGradients: """Apply standard-softmax backward in canonical global KV-block order.""" + if strict_bitwise: + return _backward_strict_from_saved_state(q, k, v, dout, state) + batch, hq, _, dim = q.shape hkv = k.size(1) group_size = hq // hkv @@ -1004,6 +1164,81 @@ def _backward_from_saved_state( ) +def _backward_strict_from_saved_state( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + state: AttentionSavedForwardState, +) -> AttentionBackwardGradients: + """Run one batch row and the complete logical Q/KV domain at a time.""" + + batch, hq, sq, dim = q.shape + hkv = k.size(1) + skv = k.size(2) + group_size = hq // hkv + if sq == 0: + return AttentionBackwardGradients( + dq=torch.zeros_like(q), + dk=torch.zeros_like(k), + dv=torch.zeros_like(v), + ) + + dq_rows: list[torch.Tensor] = [] + dk_rows: list[torch.Tensor] = [] + dv_rows: list[torch.Tensor] = [] + with NativeAttentionOp._strict_fp32_math(q.device.type): + for batch_index in range(batch): + qf = q[batch_index : batch_index + 1].float().contiguous() + kf = k[batch_index : batch_index + 1].float().contiguous() + vf = v[batch_index : batch_index + 1].float().contiguous() + doutf = dout[batch_index : batch_index + 1].float().contiguous() + k_expanded = kf.repeat_interleave(group_size, dim=1) + v_expanded = vf.repeat_interleave(group_size, dim=1) + scores = torch.matmul(qf, k_expanded.transpose(-1, -2)) * state.scale + if state.causal: + q_pos = state.query_position_offsets[batch_index : batch_index + 1, None] + q_pos = q_pos + torch.arange(sq, device=q.device, dtype=torch.long) + k_pos = state.key_position_offsets[batch_index : batch_index + 1, None] + k_pos = k_pos + torch.arange(skv, device=q.device, dtype=torch.long) + scores = scores.masked_fill( + (k_pos[:, None, :] > q_pos[:, :, None])[:, None, :, :], + float("-inf"), + ) + if state.key_padding_mask is not None: + scores = scores.masked_fill( + ~state.key_padding_mask[ + batch_index : batch_index + 1, None, None, : + ], + float("-inf"), + ) + lse = state.lse[batch_index : batch_index + 1] + probability = torch.exp(scores - lse.unsqueeze(-1)) + probability = torch.where( + torch.isfinite(lse).unsqueeze(-1), + probability, + torch.zeros_like(probability), + ) + dp = torch.matmul(doutf, v_expanded.transpose(-1, -2)) + out = state.out[batch_index : batch_index + 1] + delta = (doutf * out).sum(dim=-1, keepdim=True) + ds = probability * (dp - delta) + dq_rows.append(torch.matmul(ds, k_expanded) * state.scale) + dk_expanded = torch.matmul(ds.transpose(-1, -2), qf) * state.scale + dv_expanded = torch.matmul(probability.transpose(-1, -2), doutf) + dk_rows.append( + dk_expanded.reshape(1, hkv, group_size, skv, dim).sum(dim=2) + ) + dv_rows.append( + dv_expanded.reshape(1, hkv, group_size, skv, dim).sum(dim=2) + ) + return AttentionBackwardGradients( + dq=torch.cat(dq_rows, dim=0).to(q.dtype), + dk=torch.cat(dk_rows, dim=0).to(k.dtype), + dv=torch.cat(dv_rows, dim=0).to(v.dtype), + ) + + def _validate_saved_forward_state( state: AttentionSavedForwardState, q: torch.Tensor, @@ -1017,6 +1252,7 @@ def _validate_saved_forward_state( key_position_offsets: Optional[torch.Tensor], cp_world_size: int, kv_chunk_size: Optional[int], + strict_bitwise: bool, ) -> None: if not isinstance(state, AttentionSavedForwardState): raise ValueError("saved_forward_state must be an AttentionSavedForwardState") @@ -1050,6 +1286,11 @@ def _validate_saved_forward_state( "scale": (state.scale, expected_scale), "cp_world_size": (state.cp_world_size, cp_world_size), "kv_chunk_size": (state.kv_chunk_size, kv_chunk_size), + "strict_bitwise": (state.strict_bitwise, strict_bitwise), + "strict_schedule": ( + state.strict_schedule, + STRICT_ATTENTION_SCHEDULE_ID if strict_bitwise else None, + ), "query_bounds": (state.query_bounds, tuple(_split_bounds(q.size(2), cp_world_size))), "kv_block_bounds": ( state.kv_block_bounds, @@ -1253,6 +1494,24 @@ def _validate_output_dtype(output_dtype: torch.dtype) -> None: raise ValueError("output_dtype must be a real floating-point torch dtype") +def _validate_partition_args( + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> None: + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: total = torch.tensor(0.0, device=tensors[0].device) for tensor in tensors: @@ -1348,6 +1607,32 @@ def split_kv_execution_plan_provenance( return result +def _strict_no_split_plan_provenance( + length: int, + *, + cp_world_size: int, + backend: str, +) -> list[dict[str, object]]: + """Describe the full logical KV row consumed by each strict CP executor.""" + + if length < 1: + raise ValueError("strict no-Split-KV sequence length must be >= 1") + _validate_partition_args(cp_world_size, None) + plan = SplitKVExecutionPlan( + requested_mode=SplitKVMode.DISABLED, + requested_split_size=None, + actual_mode=SplitKVMode.DISABLED, + actual_split_size=None, + boundaries=((0, length),), + backend=backend, + source="canonical_strict_execution", + ).to_dict() + return [ + {"owner_cp_rank": cp_rank, **plan} + for cp_rank in range(cp_world_size) + ] + + def build_reference_split_kv_runtime_plan_set( total_kv_tokens: Sequence[int], *, diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index e0d0bde4..266ac3e8 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -14,6 +14,10 @@ import pytest import torch +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, +) from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, AttentionSavedForwardState, @@ -872,3 +876,90 @@ def test_gapped_partial_ranges_raise(): def test_registry_dispatches_cp_attention_reference(): assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_strict_forward_is_bitwise_invariant_to_batch_cp_and_chunk(dtype): + q, k, v = _qkv(2, 5, 9, seed=41, dtype=dtype, heads=4, kv_heads=2, dim=8) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + full_out, full_lse = op.forward_with_lse(q, k, v, cp_world_size=1) + chunked_out, chunked_lse = op.forward_with_lse( + q, + k, + v, + cp_world_size=2, + kv_chunk_size=3, + ) + single_out, single_lse = op.forward_with_lse( + q[:1], + k[:1], + v[:1], + cp_world_size=4, + kv_chunk_size=1, + ) + + assert torch.equal(full_out, chunked_out) + assert torch.equal(full_lse, chunked_lse) + assert torch.equal(full_out[:1], single_out) + assert torch.equal(full_lse[:1], single_lse) + + +def test_strict_backward_is_bitwise_invariant_to_batch_cp_and_chunk(): + q, k, v = _qkv(2, 5, 9, seed=42, heads=4, kv_heads=2, dim=8) + dout = torch.randn(q.shape, generator=torch.Generator().manual_seed(43)) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + full = op.backward_reference(q, k, v, dout, cp_world_size=1) + chunked = op.backward_reference( + q, + k, + v, + dout, + cp_world_size=2, + kv_chunk_size=3, + ) + single = op.backward_reference( + q[:1], + k[:1], + v[:1], + dout[:1], + cp_world_size=4, + kv_chunk_size=1, + ) + + for full_tensor, chunked_tensor, single_tensor in ( + (full.out, chunked.out, single.out), + (full.lse, chunked.lse, single.lse), + (full.gradients.dq, chunked.gradients.dq, single.gradients.dq), + (full.gradients.dk, chunked.gradients.dk, single.gradients.dk), + (full.gradients.dv, chunked.gradients.dv, single.gradients.dv), + ): + assert torch.equal(full_tensor, chunked_tensor) + assert torch.equal(full_tensor[:1], single_tensor) + + assert chunked.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert chunked.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID + assert chunked.provenance["actual_split_kv_policy"] == "disabled" + assert chunked.provenance["backward_algorithm"] == ( + "saved_out_lse_canonical_row_reference" + ) + assert all( + plan["actual_split_kv_policy"] == "disabled" + and plan["actual_split_boundaries"] == [[0, k.size(2)]] + for plan in chunked.provenance["actual_split_kv_plans"] + ) + + +def test_strict_backward_rejects_non_strict_saved_forward_state(): + q, k, v = _qkv(1, 4, 6, seed=44, heads=4, kv_heads=2, dim=8) + state = DeterministicCPAttentionReferenceOp().save_forward_state(q, k, v) + + with pytest.raises(ValueError, match="strict_bitwise"): + DeterministicCPAttentionReferenceOp(strict_bitwise=True).backward_reference( + q, + k, + v, + torch.ones_like(q), + saved_forward_state=state, + )