From f5dd5cabfba9c997e5ae26f5c609c4a8fafeb566 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:26:07 +0800 Subject: [PATCH 1/5] feat(attention): harden single-gpu split-kv harness --- rl_engine/kernels/attention_contract.py | 1472 +++++++++++++++++++++ rl_engine/testing/__init__.py | 12 + rl_engine/testing/attention_comparison.py | 652 ++++++++- tests/test_attention_comparison.py | 444 +++++++ 4 files changed, 2578 insertions(+), 2 deletions(-) create mode 100644 rl_engine/kernels/attention_contract.py diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..eb4994b3 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1472 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class SplitKVMode(str, Enum): + DISABLED = "disabled" + FIXED = "fixed" + AUTO = "auto" + + +class RoPEState(str, Enum): + PRE_ROPE = "pre_rope" + POST_ROPE = "post_rope" + + +class RoPECastPoint(str, Enum): + NONE = "none" + BEFORE_ROPE = "before_rope" + AFTER_ROPE = "after_rope" + + +class RoPEFusionBoundary(str, Enum): + UNFUSED_ROPE_ATTENTION = "unfused_rope_attention" + FUSED_ROPE_ATTENTION = "fused_rope_attention" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise AttentionContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _integer_tuple(values: Iterable[int], field: str) -> tuple[int, ...]: + try: + result = tuple(values) + except TypeError as exc: + raise AttentionContractError(f"{field} must be an iterable of integers") from exc + for index, value in enumerate(result): + if isinstance(value, bool) or not isinstance(value, int): + raise AttentionContractError(f"{field}[{index}] must be an integer; got {value!r}") + return result + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical TP/CP ownership for one attention invocation. + + TP head shards are currently required to be equal and contiguous. CP + sequence ownership may be uneven, but every local block must carry a stable + logical global index so a later implementation can merge by logical order + instead of collective arrival order. + """ + + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + global_q_heads: int + global_kv_heads: int + local_q_head_start: int + local_q_heads: int + local_kv_head_start: int + local_kv_heads: int + global_sequence_length: int + local_sequence_length: int + global_block_indices: tuple[int, ...] + global_block_token_starts: tuple[int, ...] + local_block_offsets: tuple[int, ...] + packed_sequence_offsets: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if tp_rank >= tp_world_size: + raise AttentionContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + if cp_rank >= cp_world_size: + raise AttentionContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") + global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") + if global_q_heads % global_kv_heads != 0: + raise AttentionContractError( + f"global_q_heads={global_q_heads} must be divisible by " + f"global_kv_heads={global_kv_heads} for GQA" + ) + if global_q_heads % tp_world_size != 0 or global_kv_heads % tp_world_size != 0: + raise AttentionContractError( + "global Q/KV heads must be evenly divisible by tp_world_size; " + f"got Hq={global_q_heads}, Hkv={global_kv_heads}, TP={tp_world_size}" + ) + + local_q_heads = _positive_int(self.local_q_heads, "local_q_heads") + local_kv_heads = _positive_int(self.local_kv_heads, "local_kv_heads") + expected_q_heads = global_q_heads // tp_world_size + expected_kv_heads = global_kv_heads // tp_world_size + if local_q_heads != expected_q_heads or local_kv_heads != expected_kv_heads: + raise AttentionContractError( + "local TP head counts do not preserve the global Q/KV mapping; " + f"expected ({expected_q_heads}, {expected_kv_heads}), got " + f"({local_q_heads}, {local_kv_heads})" + ) + + local_q_head_start = _non_negative_int(self.local_q_head_start, "local_q_head_start") + local_kv_head_start = _non_negative_int(self.local_kv_head_start, "local_kv_head_start") + expected_q_start = tp_rank * expected_q_heads + expected_kv_start = tp_rank * expected_kv_heads + if local_q_head_start != expected_q_start or local_kv_head_start != expected_kv_start: + raise AttentionContractError( + "local TP head starts do not match contiguous rank ownership; " + f"expected ({expected_q_start}, {expected_kv_start}), got " + f"({local_q_head_start}, {local_kv_head_start})" + ) + + global_sequence_length = _positive_int( + self.global_sequence_length, "global_sequence_length" + ) + local_sequence_length = _positive_int(self.local_sequence_length, "local_sequence_length") + + block_indices = _integer_tuple(self.global_block_indices, "global_block_indices") + if not block_indices: + raise AttentionContractError("global_block_indices must not be empty") + if any(index < 0 for index in block_indices): + raise AttentionContractError("global_block_indices must be non-negative") + if any( + left >= right for left, right in zip(block_indices, block_indices[1:], strict=False) + ): + raise AttentionContractError( + "global_block_indices must be unique and strictly increasing" + ) + object.__setattr__(self, "global_block_indices", block_indices) + + block_token_starts = _integer_tuple( + self.global_block_token_starts, "global_block_token_starts" + ) + local_block_offsets = _integer_tuple(self.local_block_offsets, "local_block_offsets") + if len(block_token_starts) != len(block_indices): + raise AttentionContractError( + "global_block_token_starts must contain one entry per global_block_indices entry" + ) + if any(start < 0 for start in block_token_starts): + raise AttentionContractError("global_block_token_starts must be non-negative") + if len(local_block_offsets) != len(block_indices) + 1: + raise AttentionContractError( + "local_block_offsets must contain one boundary more than global_block_indices" + ) + if local_block_offsets[0] != 0 or local_block_offsets[-1] != local_sequence_length: + raise AttentionContractError( + "local_block_offsets must start at 0 and end at local_sequence_length" + ) + if any( + left >= right + for left, right in zip(local_block_offsets, local_block_offsets[1:], strict=False) + ): + raise AttentionContractError("local_block_offsets must be strictly increasing") + + previous_global_end = 0 + for index, (global_start, local_start, local_end) in enumerate( + zip( + block_token_starts, + local_block_offsets[:-1], + local_block_offsets[1:], + strict=True, + ) + ): + global_end = global_start + (local_end - local_start) + if global_end > global_sequence_length: + raise AttentionContractError( + f"global block {block_indices[index]} exceeds global_sequence_length" + ) + if index > 0 and global_start < previous_global_end: + raise AttentionContractError( + "global block token ranges must be non-overlapping and ordered" + ) + previous_global_end = global_end + object.__setattr__(self, "global_block_token_starts", block_token_starts) + object.__setattr__(self, "local_block_offsets", local_block_offsets) + + if self.packed_sequence_offsets is not None: + offsets = _integer_tuple(self.packed_sequence_offsets, "packed_sequence_offsets") + if len(offsets) < 2 or offsets[0] != 0: + raise AttentionContractError( + "packed_sequence_offsets must start at 0 and contain an end offset" + ) + if any(left >= right for left, right in zip(offsets, offsets[1:], strict=False)): + raise AttentionContractError("packed_sequence_offsets must be strictly increasing") + if offsets[-1] != local_sequence_length: + raise AttentionContractError( + "the final packed_sequence_offsets value must equal local_sequence_length; " + f"got {offsets[-1]} and {local_sequence_length}" + ) + object.__setattr__(self, "packed_sequence_offsets", offsets) + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic CP ``(out, lse)`` merge semantics.""" + + merge: AttentionMerge = AttentionMerge.ONLINE_SOFTMAX_LSE + acc_dtype: AttentionDType = AttentionDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(AttentionMerge, self.merge, "merge")) + object.__setattr__( + self, "acc_dtype", _enum_value(AttentionDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError( + f"CP attention accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class SplitKVExecutionPlan: + """Actual backend-local Split-KV schedule emitted by a runtime. + + CP ownership is intentionally separate from this plan. ``boundaries`` + describe the canonical logical KV ranges reduced by one backend invocation; + CP may transport those partial states between ranks but must not change the + FP32 merge contract recorded here. + """ + + requested_mode: SplitKVMode + requested_split_size: int | None + actual_mode: SplitKVMode | None + actual_split_size: int | None + boundaries: tuple[tuple[int, int], ...] + merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + acc_dtype: AttentionDType = AttentionDType.FP32 + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + backend: str = "reference" + source: str = "contract_exact" + fallback: bool = False + fallback_reason: str | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "requested_mode", + _enum_value(SplitKVMode, self.requested_mode, "requested_mode"), + ) + if self.actual_mode is not None: + object.__setattr__( + self, + "actual_mode", + _enum_value(SplitKVMode, self.actual_mode, "actual_mode"), + ) + for field_name in ("requested_split_size", "actual_split_size"): + value = getattr(self, field_name) + if value is not None: + _positive_int(value, field_name) + if self.requested_mode is SplitKVMode.FIXED: + if self.requested_split_size is None: + raise AttentionContractError( + "requested fixed Split-KV mode requires requested_split_size" + ) + elif self.requested_split_size is not None: + raise AttentionContractError( + "requested_split_size is only valid for requested fixed Split-KV mode" + ) + try: + boundaries = tuple(tuple(boundary) for boundary in self.boundaries) + except TypeError as exc: + raise AttentionContractError( + "Split-KV boundaries must be an iterable of (start, end) pairs" + ) from exc + previous_end = 0 + for index, boundary in enumerate(boundaries): + if len(boundary) != 2: + raise AttentionContractError( + f"Split-KV boundary {index} must contain exactly start and end" + ) + start, end = boundary + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError( + "Split-KV boundaries must satisfy 0 <= start < end" + ) + if index > 0 and start != previous_end: + raise AttentionContractError( + "Split-KV boundaries must be contiguous and in logical KV order" + ) + previous_end = end + if self.actual_mode is None and boundaries: + raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries") + if self.actual_mode is not None and not boundaries: + raise AttentionContractError("actual Split-KV plan must declare logical boundaries") + if self.actual_mode is SplitKVMode.FIXED: + if self.actual_split_size is None: + raise AttentionContractError( + "actual fixed Split-KV mode requires actual_split_size" + ) + widths = tuple(end - start for start, end in boundaries) + if any(width != self.actual_split_size for width in widths[:-1]) or ( + widths and widths[-1] > self.actual_split_size + ): + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_size except " + "for a shorter final split" + ) + elif self.actual_split_size is not None: + raise AttentionContractError( + "actual_split_size is only valid for actual fixed Split-KV mode" + ) + if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1: + raise AttentionContractError( + "disabled Split-KV execution must contain exactly one boundary" + ) + object.__setattr__(self, "boundaries", boundaries) + object.__setattr__( + self, + "merge_order", + _enum_value(ReductionOrder, self.merge_order, "merge_order"), + ) + object.__setattr__( + self, + "acc_dtype", + _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"), + ) + object.__setattr__( + self, + "downcast_at", + _enum_value(DowncastPoint, self.downcast_at, "downcast_at"), + ) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Split-KV partial states must be merged in fp32") + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Split-KV backend must be a non-empty string") + if not isinstance(self.source, str) or not self.source.strip(): + raise AttentionContractError("Split-KV plan source must be a non-empty string") + if not isinstance(self.fallback, bool): + raise AttentionContractError("Split-KV fallback must be a bool") + if self.fallback and not self.fallback_reason: + raise AttentionContractError("Split-KV fallback_reason is required for a fallback") + if not self.fallback and self.fallback_reason is not None: + raise AttentionContractError( + "Split-KV fallback_reason must be None when fallback=False" + ) + if not self.fallback and self.actual_mode is not None: + if self.actual_mode is not self.requested_mode: + raise AttentionContractError( + "actual Split-KV mode may differ from requested mode only for a fallback" + ) + if self.actual_split_size != self.requested_split_size: + raise AttentionContractError( + "actual Split-KV size may differ from requested size only for a fallback" + ) + + @property + def actual_split_count(self) -> int | None: + return None if self.actual_mode is None else len(self.boundaries) + + def to_dict(self) -> dict[str, Any]: + return { + "requested_split_kv_policy": self.requested_mode.value, + "requested_split_kv_size": self.requested_split_size, + "actual_split_kv_policy": ( + None if self.actual_mode is None else self.actual_mode.value + ), + "actual_split_kv_size": self.actual_split_size, + "actual_split_kv_count": self.actual_split_count, + "actual_split_boundaries": [list(boundary) for boundary in self.boundaries], + "split_kv_merge_order": self.merge_order.value, + "split_kv_accum_dtype": self.acc_dtype.value, + "split_kv_downcast_at": self.downcast_at.value, + "split_kv_backend": self.backend, + "split_kv_plan_source": self.source, + "split_kv_fallback": self.fallback, + "split_kv_fallback_reason": self.fallback_reason, + } + + +@dataclass(frozen=True) +class SplitKVSpec: + """Requested Split-KV policy shared by training and rollout paths.""" + + mode: SplitKVMode = SplitKVMode.DISABLED + fixed_split_size: int | None = None + strict_consistency: bool = True + + @classmethod + def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency) + + @classmethod + def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls( + mode=SplitKVMode.FIXED, + fixed_split_size=split_size, + strict_consistency=strict_consistency, + ) + + @classmethod + def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec": + return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency) + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode")) + if not isinstance(self.strict_consistency, bool): + raise AttentionContractError("split_kv.strict_consistency must be a bool") + if self.mode is SplitKVMode.FIXED: + if self.fixed_split_size is None: + raise AttentionContractError( + "fixed Split-KV policy requires fixed_split_size" + ) + _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") + elif self.fixed_split_size is not None: + raise AttentionContractError( + "fixed_split_size is only valid for fixed Split-KV policy" + ) + if self.strict_consistency and self.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" + ) + + def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan: + """Resolve policies whose logical schedule is fully known by contract.""" + + total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens") + if self.mode is SplitKVMode.AUTO: + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=None, + actual_mode=None, + actual_split_size=None, + boundaries=(), + backend=backend, + source="runtime_required", + ) + split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size + assert split_size is not None + boundaries = tuple( + (start, min(start + split_size, total_kv_tokens)) + for start in range(0, total_kv_tokens, split_size) + ) + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=self.fixed_split_size, + actual_mode=self.mode, + actual_split_size=self.fixed_split_size, + boundaries=boundaries, + backend=backend, + source="contract_exact", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "fixed_split_size": self.fixed_split_size, + "strict_consistency": self.strict_consistency, + } + + +def validate_split_kv_alignment( + training: SplitKVExecutionPlan, + rollout: SplitKVExecutionPlan, +) -> None: + """Fail closed unless train and rollout executed the same logical plan.""" + + if training.actual_mode is None or rollout.actual_mode is None: + raise AttentionContractError( + "strict Split-KV alignment requires actual runtime plans from both sides" + ) + fields = ( + "requested_mode", + "requested_split_size", + "actual_mode", + "actual_split_size", + "boundaries", + "merge_order", + "acc_dtype", + "downcast_at", + "fallback", + ) + mismatches = [ + field_name + for field_name in fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if mismatches: + raise AttentionContractError( + "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches) + ) + + +@dataclass(frozen=True, order=True) +class SplitKVRuntimeCoordinate: + """Identity of one batch/rank/owner Split-KV runtime plan.""" + + batch_index: int + tp_rank: int + cp_rank: int + owner_cp_rank: int + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + ) -> None: + batch_index = _non_negative_int(self.batch_index, "batch_index") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank") + if batch_index >= batch_size: + raise AttentionContractError("Split-KV batch_index is outside batch_size") + if tp_rank >= tp_world_size: + raise AttentionContractError("Split-KV tp_rank is outside tp_world_size") + if cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV cp_rank is outside cp_world_size") + if owner_cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size") + + def to_dict(self) -> dict[str, int]: + return { + "batch_index": self.batch_index, + "tp_rank": self.tp_rank, + "cp_rank": self.cp_rank, + "owner_cp_rank": self.owner_cp_rank, + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanEntry: + """Actual plan for one batch/TP/CP consumer and logical KV owner.""" + + coordinate: SplitKVRuntimeCoordinate + expected_kv_range: tuple[int, int] + execution: SplitKVExecutionPlan + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + total_kv_tokens: int, + ) -> None: + self.coordinate.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + try: + start, end = self.expected_kv_range + except (TypeError, ValueError) as exc: + raise AttentionContractError( + "expected_kv_range must contain exactly (start, end)" + ) from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + or end > total_kv_tokens + ): + raise AttentionContractError( + "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" + ) + if self.execution.actual_mode is None: + raise AttentionContractError( + "complete Split-KV plan sets require actual runtime plans" + ) + if ( + self.execution.boundaries[0][0] != start + or self.execution.boundaries[-1][1] != end + ): + raise AttentionContractError( + "Split-KV execution boundaries must exactly cover expected_kv_range" + ) + if any( + boundary_start < start or boundary_end > end + for boundary_start, boundary_end in self.execution.boundaries + ): + raise AttentionContractError( + "Split-KV execution boundary escapes expected_kv_range" + ) + + def to_dict(self) -> dict[str, Any]: + return { + **self.coordinate.to_dict(), + "expected_kv_range": list(self.expected_kv_range), + **self.execution.to_dict(), + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanSet: + """Complete actual Split-KV plans across batch, TP, CP, and KV owners. + + Every CP consumer must report the plan used for every CP-owned KV range. + This duplicates owner plans across consumers intentionally: it detects one + rank silently choosing a different Split-K schedule or merge policy. + """ + + batch_size: int + tp_world_size: int + cp_world_size: int + total_kv_tokens: tuple[int, ...] + entries: tuple[SplitKVRuntimePlanEntry, ...] + + def __post_init__(self) -> None: + batch_size = _positive_int(self.batch_size, "batch_size") + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens") + if len(totals) != batch_size or any(total <= 0 for total in totals): + raise AttentionContractError( + "total_kv_tokens must contain one positive length per batch item" + ) + object.__setattr__(self, "total_kv_tokens", totals) + entries = tuple(self.entries) + object.__setattr__(self, "entries", entries) + expected_coordinates = { + SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(batch_size) + for tp_rank in range(tp_world_size) + for cp_rank in range(cp_world_size) + for owner_cp_rank in range(cp_world_size) + } + actual_coordinates = [entry.coordinate for entry in entries] + if len(set(actual_coordinates)) != len(actual_coordinates): + raise AttentionContractError( + "Split-KV runtime plan set contains duplicate coordinates" + ) + missing = expected_coordinates.difference(actual_coordinates) + extra = set(actual_coordinates).difference(expected_coordinates) + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan set coordinate coverage is incomplete; " + f"missing={_format_split_kv_coordinates(missing)}, " + f"extra={_format_split_kv_coordinates(extra)}" + ) + for entry in entries: + entry.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals[entry.coordinate.batch_index], + ) + self._validate_owner_coverage() + self._validate_rank_invariance() + + def _validate_owner_coverage(self) -> None: + for batch_index, total in enumerate(self.total_kv_tokens): + for tp_rank in range(self.tp_world_size): + ranges = [] + for owner_cp_rank in range(self.cp_world_size): + matches = [ + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.tp_rank == tp_rank + and entry.coordinate.cp_rank == 0 + and entry.coordinate.owner_cp_rank == owner_cp_rank + ] + ranges.append(matches[0].expected_kv_range) + previous_end = 0 + for start, end in ranges: + if start != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = end + if previous_end != total: + raise AttentionContractError( + "Split-KV owner ranges do not cover total_kv_tokens" + ) + + def _validate_rank_invariance(self) -> None: + for batch_index in range(self.batch_size): + for owner_cp_rank in range(self.cp_world_size): + entries = sorted( + ( + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.owner_cp_rank == owner_cp_rank + ), + key=lambda entry: ( + entry.coordinate.tp_rank, + entry.coordinate.cp_rank, + ), + ) + reference = entries[0] + for entry in entries[1:]: + if entry.expected_kv_range != reference.expected_kv_range: + raise AttentionContractError( + "Split-KV owner range differs across TP/CP consumers" + ) + try: + validate_split_kv_alignment( + reference.execution, + entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return { + "batch_size": self.batch_size, + "tp_world_size": self.tp_world_size, + "cp_world_size": self.cp_world_size, + "total_kv_tokens": list(self.total_kv_tokens), + "entries": [ + entry.to_dict() + for entry in sorted(self.entries, key=lambda entry: entry.coordinate) + ], + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + } + + +def validate_split_kv_plan_set_alignment( + training: SplitKVRuntimePlanSet, + rollout: SplitKVRuntimePlanSet, +) -> None: + """Fail closed unless complete train/rollout runtime plan sets align.""" + + topology_fields = ( + "batch_size", + "tp_world_size", + "cp_world_size", + "total_kv_tokens", + ) + topology_mismatches = [ + field_name + for field_name in topology_fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if topology_mismatches: + raise AttentionContractError( + "training/rollout Split-KV plan-set topology differs: " + + ", ".join(topology_mismatches) + ) + training_by_coordinate = { + entry.coordinate: entry for entry in training.entries + } + rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} + if training_by_coordinate.keys() != rollout_by_coordinate.keys(): + raise AttentionContractError( + "training/rollout Split-KV plan-set coordinates differ" + ) + for coordinate in sorted(training_by_coordinate): + train_entry = training_by_coordinate[coordinate] + rollout_entry = rollout_by_coordinate[coordinate] + if train_entry.expected_kv_range != rollout_entry.expected_kv_range: + raise AttentionContractError( + f"training/rollout expected KV range differs at {coordinate}" + ) + try: + validate_split_kv_alignment( + train_entry.execution, + rollout_entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + f"training/rollout Split-KV plan differs at {coordinate}: {exc}" + ) from exc + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, field), bool): + raise AttentionContractError(f"{field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append( + f"Split-KV policy={contract.split_kv.mode.value} is unsupported" + ) + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 792b36d0..3897e009 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -8,11 +8,17 @@ AttentionComparisonReport, AttentionPathDrift, AttentionPathResult, + DecodeAttentionInputs, + DecodeKVCacheMetadata, DriftStats, TransformerEngineUnavailable, + compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + decode_prefix_cache_fingerprint, run_chunked_query_attention, + run_decode_full_prefill_reference, + run_decode_kv_replay, run_full_attention, run_fused_like_rope_attention, run_paged_kv_attention, @@ -35,18 +41,24 @@ "AttentionComparisonReport", "AttentionPathDrift", "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", "DriftStats", "SyntheticRLKernelBatch", "TransformerEngineUnavailable", "active_token_count", "compare_single_gpu_rope_attention", "compare_single_gpu_attention", + "compare_decode_kv_replay", + "decode_prefix_cache_fingerprint", "compute_policy_ratio", "compute_reference_kl", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", "run_chunked_query_attention", + "run_decode_full_prefill_reference", + "run_decode_kv_replay", "run_fused_like_rope_attention", "run_full_attention", "run_paged_kv_attention", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index dbee69af..41d5abfc 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -11,6 +11,7 @@ from __future__ import annotations +import hashlib import importlib import importlib.metadata as importlib_metadata import inspect @@ -20,10 +21,16 @@ import torch +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVSpec, +) from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.testing.reference_ops import selected_logprobs_reference MergeBackend = Literal["rl_kernel", "transformer_engine"] +RoPEState = Literal["pre_rope", "post_rope"] _TE_CONTEXT_PARALLEL_MODULE = ( "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" @@ -71,6 +78,56 @@ class AttentionComparisonInputs: rope_output_dtype: torch.dtype | None = None +@dataclass(frozen=True) +class DecodeKVCacheMetadata: + """Logical identity and physical layout for decode-stage cached KV. + + ``block_table`` maps logical KV blocks to physical cache pages. Positions + are stored per physical cache slot; unused slots must contain ``-1``. + Keeping both mappings explicit lets the harness distinguish layout changes + from changes to the logical token sequence. + """ + + cache_position: torch.Tensor + kv_seq_lens: torch.Tensor + block_table: torch.Tensor + global_token_positions: torch.Tensor + query_position_ids: torch.Tensor + key_position_ids: torch.Tensor + page_size: int + prefix_cache_key: str | None = None + prefix_cache_enabled: bool = False + prefix_length: int = 0 + prefix_cache_fingerprint: str | None = None + q_rope_state: RoPEState = "post_rope" + k_cache_rope_state: RoPEState = "post_rope" + cp_block_owners: torch.Tensor | None = None + cp_world_size: int = 1 + + +@dataclass(frozen=True) +class DecodeAttentionInputs: + """Decode queries and physically paged KV cache used by the PR6 harness.""" + + q: torch.Tensor + k_cache: torch.Tensor + v_cache: torch.Tensor + metadata: DecodeKVCacheMetadata + scale: float | None = None + output_dtype: torch.dtype = torch.float32 + rope_theta: float = 1_000_000.0 + rope_rotary_dim: int | None = None + rope_cast_at: str = "after_rope" + q_rope_output_dtype: torch.dtype | None = None + k_cache_rope_output_dtype: torch.dtype | None = None + lm_head_weight: torch.Tensor | None = None + target_ids: torch.Tensor | None = None + active_token_mask: torch.Tensor | None = None + k_new: torch.Tensor | None = None + v_new: torch.Tensor | None = None + split_kv: SplitKVSpec | None = None + + @dataclass(frozen=True) class AttentionPathResult: """One materialized attention path result.""" @@ -210,6 +267,218 @@ def compare_single_gpu_rope_attention( return AttentionComparisonReport(reference_name=reference.name, drifts=drifts) +def compare_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + include_transformer_engine: bool = False, +) -> AttentionComparisonReport: + """Compare paged decode replay with a logical full-KV teacher-forcing view.""" + + _validate_decode_inputs(inputs) + reference = _run_decode_full_prefill_reference(inputs) + candidates = [_run_decode_kv_replay(inputs, merge_backend="rl_kernel")] + unavailable: list[str] = [] + if include_transformer_engine: + try: + candidates.append(_run_decode_kv_replay(inputs, merge_backend="transformer_engine")) + except TransformerEngineUnavailable as exc: + unavailable.append(f"transformer_engine_decode_kv_replay: {exc}") + drifts = tuple(_compare_decode_path(candidate, reference, inputs) for candidate in candidates) + return AttentionComparisonReport( + reference_name=reference.name, + drifts=drifts, + unavailable=tuple(unavailable), + ) + + +def run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: + """Materialize the full logical KV sequence for each decode query. + + This is the teacher-forcing side of the PR6 comparison. It deliberately + ignores physical page boundaries after restoring logical token order. + """ + + _validate_decode_inputs(inputs) + return _run_decode_full_prefill_reference(inputs) + + +def _run_decode_full_prefill_reference(inputs: DecodeAttentionInputs) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + visible = logical_positions <= query_position + out, lse = _attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, visible, :], + v[:, :, visible, :], + causal=False, + scale=inputs.scale, + key_padding_mask=None, + q_start=0, + k_start=0, + total_query_len=1, + total_kv_len=int(visible.sum().item()), + output_dtype=inputs.output_dtype, + ) + batch_out.append(out) + batch_lse.append(lse) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + return AttentionPathResult( + name="full_prefill_decode_reference", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance={ + "attention_mode": "decode", + "materialization": "full_logical_kv", + "lse_domain": "attention", + "accum_dtype": "fp32", + }, + ) + + +def run_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend = "rl_kernel", +) -> AttentionPathResult: + """Replay decode over physical KV pages and merge by logical block index.""" + + _validate_decode_inputs(inputs) + return _run_decode_kv_replay(inputs, merge_backend=merge_backend) + + +def _run_decode_kv_replay( + inputs: DecodeAttentionInputs, + *, + merge_backend: MergeBackend, +) -> AttentionPathResult: + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + merge_orders: list[list[list[int]]] = [] + actual_split_plans: list[list[dict[str, Any]]] = [] + cp_block_owners: list[list[int]] = [] + split_kv = _resolved_decode_split_kv(inputs) + for batch_index in range(inputs.q.size(0)): + q, k, v, logical_positions = _decode_logical_qkv(inputs, batch_index) + owners = _logical_block_owners(inputs, batch_index) + cp_block_owners.append(owners) + batch_out: list[torch.Tensor] = [] + batch_lse: list[torch.Tensor] = [] + batch_orders: list[list[int]] = [] + batch_split_plans: list[dict[str, Any]] = [] + for query_index in range(q.size(2)): + query_position = int(inputs.metadata.cache_position[batch_index, query_index].item()) + states: list[_PartialAttentionState] = [] + order: list[int] = [] + visible_count = int((logical_positions <= query_position).sum().item()) + split_bounds = _decode_split_bounds(visible_count, split_kv) + for block_index, (block_start, block_end) in enumerate( + split_bounds + ): + block_positions = logical_positions[block_start:block_end] + visible = block_positions <= query_position + if not bool(visible.any()): + continue + visible_end = block_start + int(visible.sum().item()) + out, lse = _attention_with_lse( + q[:, :, query_index : query_index + 1, :], + k[:, :, block_start:visible_end, :], + v[:, :, block_start:visible_end, :], + causal=False, + scale=inputs.scale, + key_padding_mask=None, + q_start=0, + k_start=block_start, + total_query_len=1, + total_kv_len=visible_end, + output_dtype=torch.float32, + ) + states.append( + _PartialAttentionState( + out=out, + lse=lse, + block_start=block_index, + block_end=block_index + 1, + ) + ) + order.append(block_index) + if not states: + raise ValueError("each decode query must have at least one visible cached KV token") + out, lse = _merge_partial_states(states, backend=merge_backend) + batch_out.append(out.to(inputs.output_dtype)) + batch_lse.append(lse) + batch_orders.append(order) + plan = SplitKVExecutionPlan( + requested_mode=split_kv.mode, + requested_split_size=split_kv.fixed_split_size, + actual_mode=split_kv.mode, + actual_split_size=split_kv.fixed_split_size, + boundaries=tuple(split_bounds), + backend=f"{merge_backend}_decode_kv_replay", + source="reference_execution", + ) + batch_split_plans.append(plan.to_dict()) + merge_orders.append(batch_orders) + actual_split_plans.append(batch_split_plans) + outs.append(torch.cat(batch_out, dim=2)) + lses.append(torch.cat(batch_lse, dim=2)) + + provenance: dict[str, Any] = { + "attention_mode": "decode", + "decode_semantics": ( + "past_kv_plus_new_kv_append" if inputs.k_new is not None else "cache_replay" + ), + "past_kv_lengths": inputs.metadata.kv_seq_lens.tolist(), + "new_kv_length": (0 if inputs.k_new is None else inputs.k_new.size(2)), + "materialization": "paged_kv_replay", + "sq": inputs.q.size(2), + "page_size": inputs.metadata.page_size, + "cache_position": inputs.metadata.cache_position.tolist(), + "kv_seq_lens": inputs.metadata.kv_seq_lens.tolist(), + "block_table": inputs.metadata.block_table.tolist(), + "global_token_positions": inputs.metadata.global_token_positions.tolist(), + "query_position_ids": inputs.metadata.query_position_ids.tolist(), + "key_position_ids": inputs.metadata.key_position_ids.tolist(), + "prefix_cache_enabled": inputs.metadata.prefix_cache_enabled, + "prefix_cache_key": inputs.metadata.prefix_cache_key, + "prefix_length": inputs.metadata.prefix_length, + "prefix_cache_fingerprint": inputs.metadata.prefix_cache_fingerprint, + "q_rope_state": inputs.metadata.q_rope_state, + "k_cache_rope_state": inputs.metadata.k_cache_rope_state, + "rope_theta": float(inputs.rope_theta), + "rotary_dim": _decode_rope_rotary_dim(inputs), + "rope_cast_at": inputs.rope_cast_at, + "q_rope_output_dtype": str(_decode_q_rope_output_dtype(inputs)).replace("torch.", ""), + "k_cache_rope_output_dtype": str(_decode_k_rope_output_dtype(inputs)).replace("torch.", ""), + "cp_block_owners": cp_block_owners, + "cp_world_size": inputs.metadata.cp_world_size, + "requested_split_kv_policy": split_kv.mode.value, + "requested_split_kv_size": split_kv.fixed_split_size, + "actual_split_kv_plans": actual_split_plans, + "merge_order": "global_block_index", + "logical_merge_orders": merge_orders, + "merge_backend": merge_backend, + "lse_domain": "attention", + "lse_exported": True, + "accum_dtype": "fp32", + "downcast_at": "final_write", + } + if merge_backend == "transformer_engine": + provenance.update(_te_context_parallel_provenance()) + return AttentionPathResult( + name=f"{merge_backend}_decode_kv_replay", + out=torch.cat(outs, dim=0), + lse=torch.cat(lses, dim=0), + provenance=provenance, + ) + + def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: """Training-style full-sequence attention with exported attention-domain LSE.""" @@ -424,6 +693,39 @@ def transformer_engine_context_parallel_available() -> bool: return True +def decode_prefix_cache_fingerprint( + inputs: DecodeAttentionInputs, + *, + prefix_length: int, +) -> str: + """Fingerprint logical prefix positions and cached K/V content. + + The fingerprint is invariant to physical page placement because cache slots + are first restored to logical token order. It intentionally includes the + cached-K RoPE state and tensor dtypes so it identifies the actual replay + boundary rather than only the token positions. + """ + + prefix_length = _positive_int(prefix_length, "prefix_length") + if bool((inputs.metadata.kv_seq_lens < prefix_length).any()): + raise ValueError("prefix_length must not exceed any kv_seq_lens entry") + digest = hashlib.sha256() + digest.update(f"k_rope_state={inputs.metadata.k_cache_rope_state}\n".encode()) + digest.update(f"k_dtype={inputs.k_cache.dtype};v_dtype={inputs.v_cache.dtype}\n".encode()) + for batch_index in range(inputs.q.size(0)): + slots = _decode_logical_slot_index(inputs, batch_index)[:prefix_length] + for tensor in ( + inputs.metadata.global_token_positions[batch_index, slots], + inputs.metadata.key_position_ids[batch_index, slots], + inputs.k_cache[batch_index, :, slots, :], + inputs.v_cache[batch_index, :, slots, :], + ): + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + def _compare_path( candidate: AttentionPathResult, reference: AttentionPathResult, @@ -454,6 +756,25 @@ def _compare_path( ) +def _compare_decode_path( + candidate: AttentionPathResult, + reference: AttentionPathResult, + inputs: DecodeAttentionInputs, +) -> AttentionPathDrift: + dlogp = None + if inputs.lm_head_weight is not None and inputs.target_ids is not None: + candidate_logp = _selected_logps_from_decode_attention(candidate.out, inputs) + reference_logp = _selected_logps_from_decode_attention(reference.out, inputs) + dlogp = _drift_stats(candidate_logp, reference_logp, mask=inputs.active_token_mask) + return AttentionPathDrift( + candidate_name=candidate.name, + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=dlogp, + provenance=candidate.provenance, + ) + + def _apply_rope_to_qk(inputs: AttentionComparisonInputs) -> tuple[torch.Tensor, torch.Tensor]: _validate_rope_inputs(inputs) assert inputs.rope_positions is not None @@ -464,6 +785,97 @@ def _apply_rope_to_qk(inputs: AttentionComparisonInputs) -> tuple[torch.Tensor, return q, k +def _decode_logical_qkv( + inputs: DecodeAttentionInputs, + batch_index: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Restore one batch's cache to logical order and materialize RoPE state.""" + + metadata = inputs.metadata + slot_index = _decode_logical_slot_index(inputs, batch_index) + logical_position_tensor = metadata.global_token_positions[batch_index, slot_index].long() + k = inputs.k_cache[batch_index : batch_index + 1, :, slot_index, :] + v = inputs.v_cache[batch_index : batch_index + 1, :, slot_index, :] + q = inputs.q[batch_index : batch_index + 1] + + rope = NativeRoPEOp() + if metadata.q_rope_state == "pre_rope": + q = rope.forward_fp32( + q, + metadata.query_position_ids[batch_index : batch_index + 1], + theta=inputs.rope_theta, + ).to(_decode_q_rope_output_dtype(inputs)) + if metadata.k_cache_rope_state == "pre_rope": + key_positions = metadata.key_position_ids[batch_index, slot_index].unsqueeze(0) + k = rope.forward_fp32(k, key_positions, theta=inputs.rope_theta).to( + _decode_k_rope_output_dtype(inputs) + ) + if inputs.k_new is not None: + assert inputs.v_new is not None + k_new = inputs.k_new[batch_index : batch_index + 1] + if metadata.k_cache_rope_state == "pre_rope": + k_new = rope.forward_fp32( + k_new, + metadata.query_position_ids[batch_index : batch_index + 1], + theta=inputs.rope_theta, + ).to(_decode_k_rope_output_dtype(inputs)) + k = torch.cat((k, k_new), dim=2) + v = torch.cat((v, inputs.v_new[batch_index : batch_index + 1]), dim=2) + logical_position_tensor = torch.cat( + ( + logical_position_tensor, + metadata.query_position_ids[batch_index].long(), + ) + ) + return q, k, v, logical_position_tensor + + +def _decode_logical_slot_index( + inputs: DecodeAttentionInputs, + batch_index: int, +) -> torch.Tensor: + metadata = inputs.metadata + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + logical_block_count = math.ceil(sequence_length / metadata.page_size) + logical_index = torch.arange( + sequence_length, + device=inputs.k_cache.device, + dtype=torch.long, + ) + pages = metadata.block_table[batch_index, :logical_block_count].long() + return ( + pages[logical_index // metadata.page_size] * metadata.page_size + + logical_index % metadata.page_size + ) + + +def _logical_block_owners(inputs: DecodeAttentionInputs, batch_index: int) -> list[int]: + block_count = math.ceil( + int(inputs.metadata.kv_seq_lens[batch_index].item()) / inputs.metadata.page_size + ) + if inputs.metadata.cp_block_owners is None: + return [0] * block_count + return [ + int(owner) for owner in inputs.metadata.cp_block_owners[batch_index, :block_count].tolist() + ] + + +def _decode_q_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return inputs.q.dtype if inputs.q_rope_output_dtype is None else inputs.q_rope_output_dtype + + +def _decode_k_rope_output_dtype(inputs: DecodeAttentionInputs) -> torch.dtype: + return ( + inputs.k_cache.dtype + if inputs.k_cache_rope_output_dtype is None + else inputs.k_cache_rope_output_dtype + ) + + +def _decode_rope_rotary_dim(inputs: DecodeAttentionInputs) -> int: + return inputs.q.size(-1) if inputs.rope_rotary_dim is None else inputs.rope_rotary_dim + + def _rope_output_dtype(inputs: AttentionComparisonInputs) -> torch.dtype: return inputs.q.dtype if inputs.rope_output_dtype is None else inputs.rope_output_dtype @@ -743,6 +1155,28 @@ def _selected_logps_from_attention( ) +def _selected_logps_from_decode_attention( + out: torch.Tensor, + inputs: DecodeAttentionInputs, +) -> torch.Tensor: + if inputs.lm_head_weight is None or inputs.target_ids is None: + raise ValueError("lm_head_weight and target_ids are required for decode dlogp drift") + batch, heads, seq, dim = out.shape + hidden = out.transpose(1, 2).reshape(batch, seq, heads * dim) + if inputs.lm_head_weight.shape[1] != hidden.size(-1): + raise ValueError( + "lm_head_weight hidden dimension must equal Hq * D; " + f"got {inputs.lm_head_weight.shape[1]} and {hidden.size(-1)}" + ) + logits = torch.matmul(hidden.float(), inputs.lm_head_weight.float().transpose(0, 1)) + return selected_logprobs_reference( + logits, + inputs.target_ids, + mask=inputs.active_token_mask, + output_dtype=torch.float32, + ) + + def _drift_stats( candidate: torch.Tensor, reference: torch.Tensor, @@ -849,6 +1283,193 @@ def _validate_rope_inputs(inputs: AttentionComparisonInputs) -> None: raise ValueError("rope_positions must have shape [Sq] or [B, Sq]") +def _validate_decode_inputs(inputs: DecodeAttentionInputs) -> None: + _validate_qkv(inputs.q, inputs.k_cache, inputs.v_cache) + if inputs.q.device != inputs.k_cache.device or inputs.q.device != inputs.v_cache.device: + raise ValueError("q, k_cache, and v_cache must be on the same device") + metadata = inputs.metadata + batch, _, sq, head_dim = inputs.q.shape + cache_capacity = inputs.k_cache.size(2) + page_size = _positive_int(metadata.page_size, "page_size") + if cache_capacity % page_size != 0: + raise ValueError("physical KV cache capacity must be divisible by page_size") + physical_page_count = cache_capacity // page_size + if metadata.cache_position.shape != (batch, sq): + raise ValueError("cache_position must have shape [B, Sq]") + if metadata.query_position_ids.shape != (batch, sq): + raise ValueError("query_position_ids must have shape [B, Sq]") + if metadata.kv_seq_lens.shape != (batch,): + raise ValueError("kv_seq_lens must have shape [B]") + if metadata.block_table.ndim != 2 or metadata.block_table.size(0) != batch: + raise ValueError("block_table must have shape [B, max_blocks]") + expected_cache_shape = (batch, cache_capacity) + if metadata.global_token_positions.shape != expected_cache_shape: + raise ValueError("global_token_positions must have shape [B, cache_capacity]") + if metadata.key_position_ids.shape != expected_cache_shape: + raise ValueError("key_position_ids must have shape [B, cache_capacity]") + integer_tensors = { + "cache_position": metadata.cache_position, + "query_position_ids": metadata.query_position_ids, + "kv_seq_lens": metadata.kv_seq_lens, + "block_table": metadata.block_table, + "global_token_positions": metadata.global_token_positions, + "key_position_ids": metadata.key_position_ids, + } + if metadata.cp_block_owners is not None: + integer_tensors["cp_block_owners"] = metadata.cp_block_owners + for name, tensor in integer_tensors.items(): + if tensor.device != inputs.q.device: + raise ValueError(f"{name} must be on the same device as q/k/v") + if tensor.dtype not in {torch.int32, torch.int64, torch.long}: + raise ValueError(f"{name} must contain integers") + if metadata.cp_block_owners is not None: + if metadata.cp_block_owners.shape != metadata.block_table.shape: + raise ValueError("cp_block_owners must have the same shape as block_table") + if bool((metadata.cp_block_owners < 0).any()): + raise ValueError("cp_block_owners must be non-negative") + cp_world_size = _positive_int(metadata.cp_world_size, "cp_world_size") + if bool((metadata.cp_block_owners >= cp_world_size).any()): + raise ValueError("cp_block_owners must be smaller than cp_world_size") + if not torch.equal(metadata.cache_position, metadata.query_position_ids): + raise ValueError("cache_position and query_position_ids must identify the same positions") + if metadata.q_rope_state not in {"pre_rope", "post_rope"}: + raise ValueError("q_rope_state must be 'pre_rope' or 'post_rope'") + if metadata.k_cache_rope_state not in {"pre_rope", "post_rope"}: + raise ValueError("k_cache_rope_state must be 'pre_rope' or 'post_rope'") + if metadata.prefix_cache_enabled: + if not metadata.prefix_cache_key: + raise ValueError("prefix_cache_key is required when prefix cache is enabled") + _positive_int(metadata.prefix_length, "prefix_length") + if not metadata.prefix_cache_fingerprint: + raise ValueError("prefix_cache_fingerprint is required when prefix cache is enabled") + elif ( + metadata.prefix_cache_key is not None + or metadata.prefix_length != 0 + or metadata.prefix_cache_fingerprint is not None + ): + raise ValueError( + "prefix cache key/fingerprint must be None and prefix_length must be 0 " + "when prefix cache is disabled" + ) + if inputs.rope_cast_at != "after_rope": + raise ValueError("rope_cast_at must be 'after_rope' for the current fp32 RoPE reference") + if inputs.rope_rotary_dim is not None: + if inputs.rope_rotary_dim != head_dim: + raise ValueError("rope_rotary_dim must equal head_dim") + _positive_int(inputs.rope_rotary_dim, "rope_rotary_dim") + if float(inputs.rope_theta) <= 0: + raise ValueError("rope_theta must be a positive number") + if inputs.q_rope_output_dtype is not None and not isinstance( + inputs.q_rope_output_dtype, torch.dtype + ): + raise ValueError("q_rope_output_dtype must be a torch.dtype when provided") + if inputs.k_cache_rope_output_dtype is not None and not isinstance( + inputs.k_cache_rope_output_dtype, torch.dtype + ): + raise ValueError("k_cache_rope_output_dtype must be a torch.dtype when provided") + if (inputs.k_new is None) != (inputs.v_new is None): + raise ValueError("k_new and v_new must be provided together") + append_mode = inputs.k_new is not None + if append_mode: + assert inputs.k_new is not None and inputs.v_new is not None + if inputs.k_new.shape != inputs.v_new.shape: + raise ValueError("k_new and v_new must have matching shapes") + expected_new_shape = (batch, inputs.k_cache.size(1), sq, head_dim) + if inputs.k_new.shape != expected_new_shape: + raise ValueError("k_new and v_new must have shape [B, Hkv, Sq, D]") + if inputs.k_new.device != inputs.q.device or inputs.v_new.device != inputs.q.device: + raise ValueError("k_new and v_new must be on the same device as q") + if inputs.split_kv is not None and not isinstance(inputs.split_kv, SplitKVSpec): + raise ValueError("split_kv must be a SplitKVSpec when provided") + if inputs.split_kv is not None and inputs.split_kv.mode is SplitKVMode.AUTO: + raise ValueError("decode replay requires disabled or fixed Split-KV, not auto") + if ( + metadata.q_rope_state == "post_rope" + and inputs.q_rope_output_dtype is not None + and inputs.q.dtype != inputs.q_rope_output_dtype + ): + raise ValueError("post-RoPE q dtype must match q_rope_output_dtype") + if ( + metadata.k_cache_rope_state == "post_rope" + and inputs.k_cache_rope_output_dtype is not None + and inputs.k_cache.dtype != inputs.k_cache_rope_output_dtype + ): + raise ValueError("post-RoPE k_cache dtype must match k_cache_rope_output_dtype") + if (inputs.lm_head_weight is None) != (inputs.target_ids is None): + raise ValueError("lm_head_weight and target_ids must be provided together") + if inputs.target_ids is not None and inputs.target_ids.shape != (batch, sq): + raise ValueError("target_ids must have shape [B, Sq]") + if inputs.active_token_mask is not None: + if inputs.active_token_mask.shape != (batch, sq): + raise ValueError("active_token_mask must have shape [B, Sq]") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + + for batch_index in range(batch): + sequence_length = int(metadata.kv_seq_lens[batch_index].item()) + if sequence_length <= 0 or sequence_length > cache_capacity: + raise ValueError("each kv_seq_lens entry must be in [1, cache_capacity]") + block_count = math.ceil(sequence_length / page_size) + if block_count > metadata.block_table.size(1): + raise ValueError("block_table does not contain enough logical KV blocks") + pages = metadata.block_table[batch_index, :block_count] + if bool(((pages < 0) | (pages >= physical_page_count)).any()): + raise ValueError("block_table contains an out-of-range physical page") + if torch.unique(pages).numel() != block_count: + raise ValueError("active block_table entries must not contain duplicate pages") + slot_index = _decode_logical_slot_index(inputs, batch_index) + active_slot_mask = torch.zeros(cache_capacity, device=inputs.q.device, dtype=torch.bool) + active_slot_mask[slot_index] = True + if bool((metadata.global_token_positions[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused global_token_positions entries must be -1") + if bool((metadata.key_position_ids[batch_index, ~active_slot_mask] != -1).any()): + raise ValueError("unused key_position_ids entries must be -1") + global_positions = metadata.global_token_positions[batch_index, slot_index] + position_offset = int(global_positions[0].item()) + expected_positions = torch.arange( + position_offset, + position_offset + sequence_length, + device=inputs.q.device, + dtype=global_positions.dtype, + ) + if not torch.equal(global_positions, expected_positions): + raise ValueError( + "block_table/global_token_positions must reconstruct logical positions " + "as one contiguous global range" + ) + key_positions = metadata.key_position_ids[batch_index, slot_index] + if not torch.equal(key_positions, global_positions): + raise ValueError("key_position_ids must match cached global token positions") + cache_positions = metadata.cache_position[batch_index] + if append_mode: + expected_new_positions = torch.arange( + position_offset + sequence_length, + position_offset + sequence_length + sq, + device=inputs.q.device, + dtype=cache_positions.dtype, + ) + if not torch.equal(cache_positions, expected_new_positions): + raise ValueError( + "append cache_position must identify the contiguous new-token suffix" + ) + elif bool((cache_positions < position_offset).any()) or bool( + (cache_positions >= position_offset + sequence_length).any() + ): + raise ValueError("cache_position must refer to a token present in the KV cache") + if sq > 1 and bool((cache_positions[1:] <= cache_positions[:-1]).any()): + raise ValueError("few-query cache_position values must be strictly increasing") + + if metadata.prefix_cache_enabled: + actual_fingerprint = decode_prefix_cache_fingerprint( + inputs, + prefix_length=metadata.prefix_length, + ) + if actual_fingerprint != metadata.prefix_cache_fingerprint: + raise ValueError( + "prefix_cache_fingerprint does not match the logical prefix positions/content" + ) + + def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: raise ValueError("q, k, and v must have shape [B, H, S, D]") @@ -861,13 +1482,17 @@ def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: def _validate_partial_states(states: list[_PartialAttentionState]) -> None: + if not states: + raise ValueError("at least one partial attention state is required") first = states[0] + if first.block_start != 0: + raise ValueError("partial state coverage must start at logical KV token 0") previous_end = first.block_end for state in states[1:]: if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: raise ValueError("all partial states must have matching shapes") - if state.block_start < previous_end: - raise ValueError("partial state block ranges must not overlap") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") previous_end = state.block_end @@ -883,6 +1508,23 @@ def _chunk_bounds(length: int, chunk_size: int) -> list[tuple[int, int]]: return bounds +def _resolved_decode_split_kv(inputs: DecodeAttentionInputs) -> SplitKVSpec: + # Existing replay behavior used one partial state per logical cache page. + # Keep that as the explicit default while allowing disabled/fixed sweeps on + # the same physical page layout. + return inputs.split_kv or SplitKVSpec.fixed(inputs.metadata.page_size) + + +def _decode_split_bounds(length: int, split_kv: SplitKVSpec) -> list[tuple[int, int]]: + if split_kv.mode is SplitKVMode.AUTO: + raise ValueError("decode replay cannot materialize an unknown auto Split-KV plan") + chunk_size = length + if split_kv.mode is SplitKVMode.FIXED: + assert split_kv.fixed_split_size is not None + chunk_size = split_kv.fixed_split_size + return _chunk_bounds(length, chunk_size) + + def _positive_int(value: int, name: str) -> int: if isinstance(value, bool) or value <= 0: raise ValueError(f"{name} must be a positive integer") @@ -894,13 +1536,19 @@ def _positive_int(value: int, name: str) -> int: "AttentionComparisonReport", "AttentionPathDrift", "AttentionPathResult", + "DecodeAttentionInputs", + "DecodeKVCacheMetadata", "DriftStats", "TransformerEngineUnavailable", "compare_single_gpu_rope_attention", "compare_single_gpu_attention", + "compare_decode_kv_replay", + "decode_prefix_cache_fingerprint", "run_chunked_query_attention", "run_fused_like_rope_attention", "run_full_attention", + "run_decode_full_prefill_reference", + "run_decode_kv_replay", "run_paged_kv_attention", "run_unfused_rope_attention", "transformer_engine_context_parallel_available", diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 061e8edd..5ae5d5e3 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -8,16 +8,25 @@ import json import sys import types +from dataclasses import replace +from typing import Literal import pytest import torch from rl_engine.kernels.gtest import run_operator_suite from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.testing.attention_comparison import ( AttentionComparisonInputs, + DecodeAttentionInputs, + DecodeKVCacheMetadata, + compare_decode_kv_replay, compare_single_gpu_attention, compare_single_gpu_rope_attention, + decode_prefix_cache_fingerprint, + run_decode_kv_replay, run_paged_kv_attention, ) @@ -57,6 +66,67 @@ def _comparison_inputs() -> AttentionComparisonInputs: ) +def _decode_inputs( + *, + page_order: tuple[int, ...] = (0, 1, 2), + prefix_cache_enabled: bool = False, + q_rope_state: Literal["pre_rope", "post_rope"] = "post_rope", + k_cache_rope_state: Literal["pre_rope", "post_rope"] = "post_rope", +) -> DecodeAttentionInputs: + q, logical_k, logical_v = _qkv(seed=17) + q = q[:, :, 4:6, :] + page_size = 2 + physical_k = torch.empty_like(logical_k) + physical_v = torch.empty_like(logical_v) + positions = torch.full((2, 6), -1, dtype=torch.long) + for logical_page, physical_page in enumerate(page_order): + logical_slice = slice(logical_page * page_size, (logical_page + 1) * page_size) + physical_slice = slice(physical_page * page_size, (physical_page + 1) * page_size) + physical_k[:, :, physical_slice, :] = logical_k[:, :, logical_slice, :] + physical_v[:, :, physical_slice, :] = logical_v[:, :, logical_slice, :] + positions[:, physical_slice] = torch.arange( + logical_page * page_size, + (logical_page + 1) * page_size, + ) + inputs = DecodeAttentionInputs( + q=q, + k_cache=physical_k, + v_cache=physical_v, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[4, 5], [4, 5]], dtype=torch.long), + kv_seq_lens=torch.tensor([6, 6], dtype=torch.long), + block_table=torch.tensor([page_order, page_order], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor([[4, 5], [4, 5]], dtype=torch.long), + key_position_ids=positions.clone(), + page_size=page_size, + q_rope_state=q_rope_state, + k_cache_rope_state=k_cache_rope_state, + cp_block_owners=torch.tensor([[0, 1, 0], [0, 1, 0]], dtype=torch.long), + cp_world_size=2, + ), + lm_head_weight=torch.randn( + 11, q.size(1) * q.size(3), generator=torch.Generator().manual_seed(18) + ), + target_ids=torch.tensor([[1, 2], [3, 4]], dtype=torch.long), + active_token_mask=torch.tensor([[True, True], [False, True]], dtype=torch.bool), + ) + if not prefix_cache_enabled: + return inputs + prefix_length = 4 + fingerprint = decode_prefix_cache_fingerprint(inputs, prefix_length=prefix_length) + return replace( + inputs, + metadata=replace( + inputs.metadata, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + prefix_length=prefix_length, + prefix_cache_fingerprint=fingerprint, + ), + ) + + def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): report = compare_single_gpu_attention( _comparison_inputs(), @@ -155,6 +225,380 @@ def test_single_gpu_rope_attention_requires_position_metadata(): compare_single_gpu_rope_attention(AttentionComparisonInputs(q=base.q, k=base.k, v=base.v)) +def test_decode_replay_matches_full_prefill_for_single_and_few_query(): + inputs = _decode_inputs() + report = compare_decode_kv_replay(inputs) + + assert report.reference_name == "full_prefill_decode_reference" + drift = report.drifts[0] + assert drift.candidate_name == "rl_kernel_decode_kv_replay" + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.dlogp is not None + assert drift.dlogp.max_abs <= 3.0e-6 + assert drift.dlogp.active_count == 3 + assert drift.provenance["attention_mode"] == "decode" + assert drift.provenance["cache_position"] == [[4, 5], [4, 5]] + assert drift.provenance["cp_block_owners"] == [[0, 1, 0], [0, 1, 0]] + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["logical_merge_orders"] == [ + [[0, 1, 2], [0, 1, 2]], + [[0, 1, 2], [0, 1, 2]], + ] + + single_query = DecodeAttentionInputs( + q=inputs.q[:, :, -1:, :], + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position[:, -1:], + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=inputs.metadata.global_token_positions, + query_position_ids=inputs.metadata.query_position_ids[:, -1:], + key_position_ids=inputs.metadata.key_position_ids, + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + cp_world_size=inputs.metadata.cp_world_size, + ), + ) + single_report = compare_decode_kv_replay(single_query) + assert single_report.drifts[0].out.max_abs <= 1.0e-6 + assert single_report.drifts[0].lse.max_abs <= 1.0e-6 + + +def test_decode_replay_is_invariant_to_physical_page_and_prefix_layout(): + contiguous = run_decode_kv_replay(_decode_inputs()) + permuted = run_decode_kv_replay(_decode_inputs(page_order=(2, 0, 1), prefix_cache_enabled=True)) + + torch.testing.assert_close(permuted.out, contiguous.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(permuted.lse, contiguous.lse, atol=1.0e-6, rtol=0.0) + assert permuted.provenance["prefix_cache_enabled"] is True + assert permuted.provenance["prefix_cache_key"] == "shared-prefix" + assert permuted.provenance["prefix_length"] == 4 + assert permuted.provenance["prefix_cache_fingerprint"] == decode_prefix_cache_fingerprint( + _decode_inputs(), prefix_length=4 + ) + + +def test_decode_replay_is_invariant_to_equivalent_cp_block_ownership(): + cp2_inputs = _decode_inputs() + cp1_inputs = replace( + cp2_inputs, + metadata=replace( + cp2_inputs.metadata, + cp_block_owners=torch.zeros_like(cp2_inputs.metadata.cp_block_owners), + ), + ) + + cp1 = run_decode_kv_replay(cp1_inputs) + cp2 = run_decode_kv_replay(cp2_inputs) + torch.testing.assert_close(cp2.out, cp1.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(cp2.lse, cp1.lse, atol=1.0e-6, rtol=0.0) + assert cp1.provenance["cp_block_owners"] == [[0, 0, 0], [0, 0, 0]] + assert cp2.provenance["cp_block_owners"] == [[0, 1, 0], [0, 1, 0]] + + +def test_decode_replay_pre_rope_cache_matches_equivalent_post_rope_cache(): + pre_rope = _decode_inputs(q_rope_state="pre_rope", k_cache_rope_state="pre_rope") + rope = NativeRoPEOp() + post_q = rope.forward_fp32( + pre_rope.q, + pre_rope.metadata.query_position_ids, + theta=pre_rope.rope_theta, + ) + post_k = rope.forward_fp32( + pre_rope.k_cache, + pre_rope.metadata.key_position_ids, + theta=pre_rope.rope_theta, + ) + post_rope = replace( + pre_rope, + q=post_q, + k_cache=post_k, + metadata=replace( + pre_rope.metadata, + q_rope_state="post_rope", + k_cache_rope_state="post_rope", + ), + ) + + pre_result = run_decode_kv_replay(pre_rope) + post_result = run_decode_kv_replay(post_rope) + torch.testing.assert_close(post_result.out, pre_result.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(post_result.lse, pre_result.lse, atol=1.0e-6, rtol=0.0) + + +def test_decode_replay_preserves_separate_q_and_k_rope_output_dtypes(): + base = _decode_inputs(q_rope_state="pre_rope", k_cache_rope_state="pre_rope") + mixed = replace( + base, + k_cache=base.k_cache.to(torch.bfloat16), + v_cache=base.v_cache.to(torch.bfloat16), + ) + rope = NativeRoPEOp() + post = replace( + mixed, + q=rope.forward_fp32( + mixed.q, + mixed.metadata.query_position_ids, + theta=mixed.rope_theta, + ).to(torch.float32), + k_cache=rope.forward_fp32( + mixed.k_cache, + mixed.metadata.key_position_ids, + theta=mixed.rope_theta, + ).to(torch.bfloat16), + metadata=replace( + mixed.metadata, + q_rope_state="post_rope", + k_cache_rope_state="post_rope", + ), + ) + + pre_result = run_decode_kv_replay(mixed) + post_result = run_decode_kv_replay(post) + torch.testing.assert_close(post_result.out, pre_result.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(post_result.lse, pre_result.lse, atol=1.0e-6, rtol=0.0) + assert pre_result.provenance["q_rope_output_dtype"] == "float32" + assert pre_result.provenance["k_cache_rope_output_dtype"] == "bfloat16" + + +def test_decode_replay_rejects_stale_prefix_cache_content(): + inputs = _decode_inputs(prefix_cache_enabled=True) + stale_k = inputs.k_cache.clone() + first_prefix_slot = int(inputs.metadata.block_table[0, 0].item()) * inputs.metadata.page_size + stale_k[0, 0, first_prefix_slot, 0] += 1.0 + + with pytest.raises(ValueError, match="prefix_cache_fingerprint"): + run_decode_kv_replay(replace(inputs, k_cache=stale_k)) + + +def test_decode_replay_fails_loudly_on_position_identity_mismatch(): + inputs = _decode_inputs() + bad_query_positions = inputs.metadata.query_position_ids.clone() + bad_query_positions[0, -1] = 4 + bad_metadata = DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position, + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=inputs.metadata.global_token_positions, + query_position_ids=bad_query_positions, + key_position_ids=inputs.metadata.key_position_ids, + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + cp_world_size=inputs.metadata.cp_world_size, + ) + + with pytest.raises(ValueError, match="cache_position and query_position_ids"): + run_decode_kv_replay( + DecodeAttentionInputs( + q=inputs.q, + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=bad_metadata, + ) + ) + + +def test_decode_replay_fails_loudly_on_invalid_page_identity(): + inputs = _decode_inputs() + bad_positions = inputs.metadata.global_token_positions.clone() + bad_positions[:, 0] = 1 + bad_metadata = DecodeKVCacheMetadata( + cache_position=inputs.metadata.cache_position, + kv_seq_lens=inputs.metadata.kv_seq_lens, + block_table=inputs.metadata.block_table, + global_token_positions=bad_positions, + query_position_ids=inputs.metadata.query_position_ids, + key_position_ids=bad_positions.clone(), + page_size=inputs.metadata.page_size, + cp_block_owners=inputs.metadata.cp_block_owners, + cp_world_size=inputs.metadata.cp_world_size, + ) + + with pytest.raises(ValueError, match="reconstruct logical positions"): + compare_decode_kv_replay( + DecodeAttentionInputs( + q=inputs.q, + k_cache=inputs.k_cache, + v_cache=inputs.v_cache, + metadata=bad_metadata, + ) + ) + + +def test_decode_replay_covers_qwen3_gqa_head_layout(): + generator = torch.Generator().manual_seed(23) + q = torch.randn(1, 32, 1, 128, generator=generator, dtype=torch.bfloat16) + k = torch.randn(1, 8, 4, 128, generator=generator, dtype=torch.bfloat16) + v = torch.randn(1, 8, 4, 128, generator=generator, dtype=torch.bfloat16) + positions = torch.arange(4, dtype=torch.long).unsqueeze(0) + inputs = DecodeAttentionInputs( + q=q, + k_cache=k, + v_cache=v, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[3]], dtype=torch.long), + kv_seq_lens=torch.tensor([4], dtype=torch.long), + block_table=torch.tensor([[0, 1]], dtype=torch.long), + global_token_positions=positions, + query_position_ids=torch.tensor([[3]], dtype=torch.long), + key_position_ids=positions.clone(), + page_size=2, + cp_block_owners=torch.tensor([[0, 1]], dtype=torch.long), + cp_world_size=2, + ), + output_dtype=torch.bfloat16, + ) + + report = compare_decode_kv_replay(inputs) + assert report.drifts[0].out.max_abs <= 2 * torch.finfo(torch.bfloat16).eps + assert report.drifts[0].lse.max_abs <= 1.0e-6 + + +def test_decode_append_matches_full_prefill_suffix(): + generator = torch.Generator().manual_seed(71) + q = torch.randn(1, 4, 2, 8, generator=generator) + k_past = torch.randn(1, 2, 4, 8, generator=generator) + v_past = torch.randn(1, 2, 4, 8, generator=generator) + k_new = torch.randn(1, 2, 2, 8, generator=generator) + v_new = torch.randn(1, 2, 2, 8, generator=generator) + inputs = DecodeAttentionInputs( + q=q, + k_cache=k_past, + v_cache=v_past, + k_new=k_new, + v_new=v_new, + metadata=DecodeKVCacheMetadata( + cache_position=torch.tensor([[104, 105]], dtype=torch.long), + kv_seq_lens=torch.tensor([4], dtype=torch.long), + block_table=torch.tensor([[0, 1]], dtype=torch.long), + global_token_positions=torch.tensor([[100, 101, 102, 103]], dtype=torch.long), + query_position_ids=torch.tensor([[104, 105]], dtype=torch.long), + key_position_ids=torch.tensor([[100, 101, 102, 103]], dtype=torch.long), + page_size=2, + cp_block_owners=torch.tensor([[0, 1]], dtype=torch.long), + cp_world_size=2, + ), + split_kv=SplitKVSpec.fixed(2), + ) + + report = compare_decode_kv_replay(inputs) + drift = report.drifts[0] + assert drift.out.max_abs <= 1.0e-6 + assert drift.lse.max_abs <= 1.0e-6 + assert drift.provenance["decode_semantics"] == "past_kv_plus_new_kv_append" + assert drift.provenance["past_kv_lengths"] == [4] + assert drift.provenance["new_kv_length"] == 2 + assert drift.provenance["actual_split_kv_plans"][0][1][ + "actual_split_boundaries" + ] == [[0, 2], [2, 4], [4, 6]] + + +def test_decode_replay_supports_nonzero_global_position_offset(): + base = _decode_inputs() + offset = 4096 + active = base.metadata.global_token_positions >= 0 + positions = torch.where( + active, + base.metadata.global_token_positions + offset, + base.metadata.global_token_positions, + ) + inputs = replace( + base, + metadata=replace( + base.metadata, + cache_position=base.metadata.cache_position + offset, + query_position_ids=base.metadata.query_position_ids + offset, + global_token_positions=positions, + key_position_ids=positions.clone(), + ), + ) + + report = compare_decode_kv_replay(inputs) + assert report.drifts[0].out.max_abs <= 1.0e-6 + assert report.drifts[0].provenance["global_token_positions"][0][0] >= offset + + +def test_decode_split_k_disabled_and_fixed_share_cache_layout(): + base = _decode_inputs() + disabled = run_decode_kv_replay(replace(base, split_kv=SplitKVSpec.disabled())) + fixed = run_decode_kv_replay(replace(base, split_kv=SplitKVSpec.fixed(2))) + + torch.testing.assert_close(fixed.out, disabled.out, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(fixed.lse, disabled.lse, atol=1.0e-6, rtol=0.0) + assert disabled.provenance["requested_split_kv_policy"] == "disabled" + assert fixed.provenance["requested_split_kv_policy"] == "fixed" + assert disabled.provenance["block_table"] == fixed.provenance["block_table"] + + +def test_decode_transformer_engine_oracle_reuses_sorted_partial_states(monkeypatch): + calls = {"lse": 0, "out": 0} + + def lse_correction(softmax_lse, softmax_lse_per_step): + calls["lse"] += 1 + softmax_lse.copy_(torch.logaddexp(softmax_lse, softmax_lse_per_step)) + + def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): + scale = torch.exp(softmax_lse_init_step - softmax_lse).movedim(2, seq_dim) + return out_init_step * scale.unsqueeze(-1) + + def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): + calls["out"] += 1 + scale = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) + out.add_(out_per_step * scale.unsqueeze(-1)) + + monkeypatch.setitem( + sys.modules, + _TE_CONTEXT_PARALLEL_MODULE, + types.SimpleNamespace( + flash_attn_fwd_softmax_lse_correction=lse_correction, + flash_attn_fwd_out_correction_init=out_correction_init, + flash_attn_fwd_out_correction=out_correction, + ), + ) + + report = compare_decode_kv_replay( + _decode_inputs(page_order=(2, 0, 1)), + include_transformer_engine=True, + ) + + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert set(by_name) == { + "rl_kernel_decode_kv_replay", + "transformer_engine_decode_kv_replay", + } + assert by_name["transformer_engine_decode_kv_replay"].out.max_abs <= 1.0e-6 + assert by_name["transformer_engine_decode_kv_replay"].lse.max_abs <= 1.0e-6 + assert by_name["transformer_engine_decode_kv_replay"].provenance["logical_merge_orders"] == [ + [[0, 1, 2], [0, 1, 2]], + [[0, 1, 2], [0, 1, 2]], + ] + assert calls["lse"] > 0 + assert calls["out"] > 0 + assert report.unavailable == () + + +def test_decode_transformer_engine_unavailable_is_reported(monkeypatch): + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == _TE_CONTEXT_PARALLEL_MODULE: + raise ImportError("decode TE unavailable") + return real_import_module(name, package) + + monkeypatch.setattr(importlib, "import_module", fake_import_module) + + report = compare_decode_kv_replay( + _decode_inputs(), + include_transformer_engine=True, + ) + + assert {drift.candidate_name for drift in report.drifts} == {"rl_kernel_decode_kv_replay"} + assert report.unavailable == ("transformer_engine_decode_kv_replay: decode TE unavailable",) + + def test_transformer_engine_merge_oracle_can_be_reused_when_available(monkeypatch): calls = {"lse": 0, "out": 0} From d73fe374e0480f42c69558ea9ceaf3471a2d8036 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 18:19:56 +0800 Subject: [PATCH 2/5] style(attention): satisfy comparison harness lint Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/attention_contract.py | 68 ++++++++--------------- rl_engine/testing/attention_comparison.py | 10 +--- tests/test_attention_comparison.py | 10 ++-- 3 files changed, 32 insertions(+), 56 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index eb4994b3..1750476d 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -359,9 +359,7 @@ def __post_init__(self) -> None: or start < 0 or end <= start ): - raise AttentionContractError( - "Split-KV boundaries must satisfy 0 <= start < end" - ) + raise AttentionContractError("Split-KV boundaries must satisfy 0 <= start < end") if index > 0 and start != previous_end: raise AttentionContractError( "Split-KV boundaries must be contiguous and in logical KV order" @@ -486,14 +484,10 @@ def __post_init__(self) -> None: raise AttentionContractError("split_kv.strict_consistency must be a bool") if self.mode is SplitKVMode.FIXED: if self.fixed_split_size is None: - raise AttentionContractError( - "fixed Split-KV policy requires fixed_split_size" - ) + raise AttentionContractError("fixed Split-KV policy requires fixed_split_size") _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") elif self.fixed_split_size is not None: - raise AttentionContractError( - "fixed_split_size is only valid for fixed Split-KV policy" - ) + raise AttentionContractError("fixed_split_size is only valid for fixed Split-KV policy") if self.strict_consistency and self.mode is SplitKVMode.AUTO: raise AttentionContractError( "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" @@ -647,13 +641,8 @@ def validate( "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" ) if self.execution.actual_mode is None: - raise AttentionContractError( - "complete Split-KV plan sets require actual runtime plans" - ) - if ( - self.execution.boundaries[0][0] != start - or self.execution.boundaries[-1][1] != end - ): + raise AttentionContractError("complete Split-KV plan sets require actual runtime plans") + if self.execution.boundaries[0][0] != start or self.execution.boundaries[-1][1] != end: raise AttentionContractError( "Split-KV execution boundaries must exactly cover expected_kv_range" ) @@ -661,9 +650,7 @@ def validate( boundary_start < start or boundary_end > end for boundary_start, boundary_end in self.execution.boundaries ): - raise AttentionContractError( - "Split-KV execution boundary escapes expected_kv_range" - ) + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") def to_dict(self) -> dict[str, Any]: return { @@ -709,9 +696,7 @@ def __post_init__(self) -> None: } actual_coordinates = [entry.coordinate for entry in entries] if len(set(actual_coordinates)) != len(actual_coordinates): - raise AttentionContractError( - "Split-KV runtime plan set contains duplicate coordinates" - ) + raise AttentionContractError("Split-KV runtime plan set contains duplicate coordinates") missing = expected_coordinates.difference(actual_coordinates) extra = set(actual_coordinates).difference(expected_coordinates) if missing or extra: @@ -821,17 +806,12 @@ def validate_split_kv_plan_set_alignment( ] if topology_mismatches: raise AttentionContractError( - "training/rollout Split-KV plan-set topology differs: " - + ", ".join(topology_mismatches) + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) ) - training_by_coordinate = { - entry.coordinate: entry for entry in training.entries - } + training_by_coordinate = {entry.coordinate: entry for entry in training.entries} rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} if training_by_coordinate.keys() != rollout_by_coordinate.keys(): - raise AttentionContractError( - "training/rollout Split-KV plan-set coordinates differ" - ) + raise AttentionContractError("training/rollout Split-KV plan-set coordinates differ") for coordinate in sorted(training_by_coordinate): train_entry = training_by_coordinate[coordinate] rollout_entry = rollout_by_coordinate[coordinate] @@ -1075,14 +1055,16 @@ def __post_init__(self) -> None: not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() ): raise AttentionContractError("rope_scaling must be a non-empty string when provided") - for field in ("position_ids", "query_position_offsets", "key_position_offsets"): - values = getattr(self, field) + for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, position_field) if values is None: continue - normalized = _integer_tuple(values, field) + normalized = _integer_tuple(values, position_field) if not normalized or any(value < 0 for value in normalized): - raise AttentionContractError(f"{field} must contain non-negative positions") - object.__setattr__(self, field, normalized) + raise AttentionContractError( + f"{position_field} must contain non-negative positions" + ) + object.__setattr__(self, position_field, normalized) object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) object.__setattr__( self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") @@ -1186,11 +1168,11 @@ def __post_init__(self) -> None: "position_ids must describe the local query sequence or full local " "sequence length" ) - for field in ("query_position_offsets", "key_position_offsets"): - offsets = getattr(self.rope, field) + for position_field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, position_field) if offsets is not None and len(offsets) != batch_size: raise AttentionContractError( - f"{field} must contain one entry per logical batch entry" + f"{position_field} must contain one entry per logical batch entry" ) if self.mode is AttentionMode.DECODE and self.kv_cache is not None: if len(self.kv_cache.kv_seq_lens) != batch_size: @@ -1336,7 +1318,7 @@ def __post_init__(self) -> None: raise AttentionContractError("tp_world_sizes must contain positive values") if len(set(tp_world_sizes)) != len(tp_world_sizes): raise AttentionContractError("tp_world_sizes must not contain duplicates") - for field in ( + for capability_field in ( "exports_attention_lse", "deterministic_cp_merge", "supports_packed_varlen", @@ -1348,8 +1330,8 @@ def __post_init__(self) -> None: "supports_split_kv_auto", "reports_actual_split_kv_plan", ): - if not isinstance(getattr(self, field), bool): - raise AttentionContractError(f"{field} must be a bool") + if not isinstance(getattr(self, capability_field), bool): + raise AttentionContractError(f"{capability_field} must be a bool") if self.implementation_kind not in {"production", "reference", "deterministic"}: raise AttentionContractError( "implementation_kind must be production, reference, or deterministic" @@ -1401,9 +1383,7 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: SplitKVMode.AUTO: self.supports_split_kv_auto, } if not split_support[contract.split_kv.mode]: - reasons.append( - f"Split-KV policy={contract.split_kv.mode.value} is unsupported" - ) + reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported") if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: reasons.append("actual Split-KV execution-plan provenance is unsupported") return tuple(reasons) diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 41d5abfc..06b39f78 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -21,11 +21,7 @@ import torch -from rl_engine.kernels.attention_contract import ( - SplitKVExecutionPlan, - SplitKVMode, - SplitKVSpec, -) +from rl_engine.kernels.attention_contract import SplitKVExecutionPlan, SplitKVMode, SplitKVSpec from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.testing.reference_ops import selected_logprobs_reference @@ -378,9 +374,7 @@ def _run_decode_kv_replay( order: list[int] = [] visible_count = int((logical_positions <= query_position).sum().item()) split_bounds = _decode_split_bounds(visible_count, split_kv) - for block_index, (block_start, block_end) in enumerate( - split_bounds - ): + for block_index, (block_start, block_end) in enumerate(split_bounds): block_positions = logical_positions[block_start:block_end] visible = block_positions <= query_position if not bool(visible.any()): diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 5ae5d5e3..86936760 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -14,10 +14,10 @@ import pytest import torch +from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.kernels.gtest import run_operator_suite from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp -from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.testing.attention_comparison import ( AttentionComparisonInputs, DecodeAttentionInputs, @@ -491,9 +491,11 @@ def test_decode_append_matches_full_prefill_suffix(): assert drift.provenance["decode_semantics"] == "past_kv_plus_new_kv_append" assert drift.provenance["past_kv_lengths"] == [4] assert drift.provenance["new_kv_length"] == 2 - assert drift.provenance["actual_split_kv_plans"][0][1][ - "actual_split_boundaries" - ] == [[0, 2], [2, 4], [4, 6]] + assert drift.provenance["actual_split_kv_plans"][0][1]["actual_split_boundaries"] == [ + [0, 2], + [2, 4], + [4, 6], + ] def test_decode_replay_supports_nonzero_global_position_offset(): From 9761f7f8587a1bb2f6db36d306075ac718b873e8 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:48:04 +0800 Subject: [PATCH 3/5] feat(attention): record single-gpu reference boundaries --- rl_engine/testing/attention_comparison.py | 25 +++++++++++++++++++++++ tests/test_attention_comparison.py | 9 ++++++++ 2 files changed, 34 insertions(+) diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 06b39f78..6d8131d4 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -497,6 +497,7 @@ def run_full_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult "attention_mode": "prefill", "materialization": "full_sequence", "lse_domain": "attention", + **_single_gpu_reference_provenance(), }, ) @@ -604,6 +605,7 @@ def run_chunked_query_attention( "query_chunk_size": chunk_size, "chunk_bounds": [list(bound) for bound in chunk_bounds], "lse_domain": "attention", + **_single_gpu_reference_provenance(), }, ) @@ -666,6 +668,7 @@ def run_paged_kv_attention( "lse_exported": True, "accum_dtype": "fp32", "downcast_at": "final_write", + **_single_gpu_reference_provenance(), } if merge_backend == "transformer_engine": provenance.update(_te_context_parallel_provenance()) @@ -901,6 +904,28 @@ def _rope_attention_provenance( "rope_output_dtype": str(_rope_output_dtype(inputs)).replace("torch.", ""), "fusion_boundary": fusion_boundary, "lse_domain": "attention", + "preprocess_backends": { + "rope": "rlkernel.pytorch.rope_reference", + "qk_rmsnorm": "not_executed_projected_qk_input", + }, + **_single_gpu_reference_provenance(), + } + + +def _single_gpu_reference_provenance() -> dict[str, Any]: + return { + "execution_scope": "single_device_correctness_reference", + "runtime_verified": False, + "preprocess_policy": "reference_only_not_production", + "native_backend_executed": False, + "preprocess_fallback": True, + "preprocess_fallback_reason": ( + "single-GPU attribution harness intentionally uses the PyTorch deterministic reference" + ), + "qkv_input_boundary": "projected_qkv", + "qkv_projection_executed": False, + "output_projection_executed": False, + "communication": "none", } diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 86936760..9bd33fef 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -142,6 +142,10 @@ def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): assert drift.dlogp is not None assert drift.dlogp.active_count == 7 assert drift.dlogp.p95_abs <= 1.0e-6 + assert drift.provenance["execution_scope"] == "single_device_correctness_reference" + assert drift.provenance["native_backend_executed"] is False + assert drift.provenance["preprocess_fallback"] is True + assert drift.provenance["qkv_projection_executed"] is False payload = report.to_dict() assert payload["reference_name"] == "full_prefill" @@ -212,6 +216,11 @@ def test_single_gpu_rope_attention_harness_reports_rope_and_attention_drift(): assert drift.provenance["rotary_dim"] == base.q.size(-1) assert drift.provenance["rope_cast_at"] == "after_rope" assert drift.provenance["fusion_boundary"] == "fused_rope_attention" + assert drift.provenance["preprocess_backends"] == { + "rope": "rlkernel.pytorch.rope_reference", + "qk_rmsnorm": "not_executed_projected_qk_input", + } + assert drift.provenance["preprocess_policy"] == "reference_only_not_production" payload = report.to_dict() assert payload["drifts"][0]["post_rope_q"]["active_count"] == base.q.numel() From 115d34d3e72f5837181023822500fe2ddad04cef Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 15:49:12 +0800 Subject: [PATCH 4/5] 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 1ce2349862996ef03eb7ee2bb0f52e382962209c Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 23:25:28 +0800 Subject: [PATCH 5/5] test(attention): verify strict shared core layout invariance --- rl_engine/testing/__init__.py | 6 + rl_engine/testing/attention_comparison.py | 271 +++++++++++++++++++++- tests/test_attention_comparison.py | 26 +++ 3 files changed, 298 insertions(+), 5 deletions(-) diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 7bd12880..94b82f8e 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -22,6 +22,9 @@ run_full_attention, run_fused_like_rope_attention, run_paged_kv_attention, + run_strict_shared_core_attention, + run_strict_shared_core_chunked_attention, + run_strict_shared_core_paged_layout_attention, run_unfused_rope_attention, transformer_engine_context_parallel_available, ) @@ -90,6 +93,9 @@ "run_full_attention", "run_paged_kv_attention", "run_unfused_rope_attention", + "run_strict_shared_core_attention", + "run_strict_shared_core_chunked_attention", + "run_strict_shared_core_paged_layout_attention", "reference_payload", "restore_logical_order", "restore_logical_order_from_padded", diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index 6d8131d4..ec0e3d32 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -210,6 +210,7 @@ def compare_single_gpu_attention( query_chunk_size: int | None = None, kv_page_size: int | None = None, include_transformer_engine: bool = False, + strict_bitwise: bool = False, ) -> AttentionComparisonReport: """Compare full attention with chunked/paged single-GPU materializations. @@ -219,11 +220,20 @@ def compare_single_gpu_attention( """ _validate_comparison_inputs(inputs) - reference = run_full_attention(inputs) - candidates = [ - run_chunked_query_attention(inputs, query_chunk_size=query_chunk_size), - run_paged_kv_attention(inputs, kv_page_size=kv_page_size, merge_backend="rl_kernel"), - ] + if strict_bitwise: + reference = run_strict_shared_core_attention(inputs) + candidates = [ + run_strict_shared_core_chunked_attention( + inputs, query_chunk_size=query_chunk_size + ), + run_strict_shared_core_paged_layout_attention(inputs, kv_page_size=kv_page_size), + ] + else: + reference = run_full_attention(inputs) + candidates = [ + run_chunked_query_attention(inputs, query_chunk_size=query_chunk_size), + run_paged_kv_attention(inputs, kv_page_size=kv_page_size, merge_backend="rl_kernel"), + ] unavailable: list[str] = [] if include_transformer_engine: try: @@ -245,6 +255,129 @@ def compare_single_gpu_attention( ) +def run_strict_shared_core_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: + """Run the canonical single-row arithmetic schedule for the full path.""" + + _validate_comparison_inputs(inputs) + out, lse = _strict_attention_with_lse( + inputs.q, + inputs.k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=inputs.q.size(2), + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="strict_shared_core_full_prefill", + out=out, + lse=lse, + provenance=_strict_shared_core_provenance( + inputs, materialization="full_logical_kv_shared_core" + ), + ) + + +def run_strict_shared_core_chunked_attention( + inputs: AttentionComparisonInputs, + *, + query_chunk_size: int | None, +) -> AttentionPathResult: + """Replay query chunks through the same one-row schedule as full prefill.""" + + _validate_comparison_inputs(inputs) + chunk_size = ( + inputs.q.size(2) + if query_chunk_size is None + else _positive_int(query_chunk_size, "query_chunk_size") + ) + bounds = _chunk_bounds(inputs.q.size(2), chunk_size) + outs: list[torch.Tensor] = [] + lses: list[torch.Tensor] = [] + for q_start, q_end in bounds: + out, lse = _strict_attention_with_lse( + inputs.q[:, :, q_start:q_end, :], + inputs.k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=q_start, + k_start=0, + total_query_len=inputs.q.size(2), + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + outs.append(out) + lses.append(lse) + return AttentionPathResult( + name="strict_shared_core_chunked_prefill", + out=torch.cat(outs, dim=2), + lse=torch.cat(lses, dim=2), + provenance={ + **_strict_shared_core_provenance( + inputs, materialization="query_chunks_shared_core" + ), + "query_chunk_size": chunk_size, + "chunk_bounds": [list(bound) for bound in bounds], + }, + ) + + +def run_strict_shared_core_paged_layout_attention( + inputs: AttentionComparisonInputs, + *, + kv_page_size: int | None, +) -> AttentionPathResult: + """Restore paged KV to logical order, then call the shared core once. + + Strict mode deliberately does not create per-page partial states: page size + is a storage detail, while Split-KV is disabled in the arithmetic contract. + """ + + _validate_comparison_inputs(inputs) + page_size = ( + inputs.k.size(2) + if kv_page_size is None + else _positive_int(kv_page_size, "kv_page_size") + ) + pages = [ + (start, end) + for start, end in _chunk_bounds(inputs.k.size(2), page_size) + ] + logical_k = torch.cat([inputs.k[:, :, start:end, :] for start, end in pages], dim=2) + logical_v = torch.cat([inputs.v[:, :, start:end, :] for start, end in pages], dim=2) + out, lse = _strict_attention_with_lse( + inputs.q, + logical_k, + logical_v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=inputs.q.size(2), + total_kv_len=inputs.k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name="strict_shared_core_paged_kv", + out=out, + lse=lse, + provenance={ + **_strict_shared_core_provenance( + inputs, materialization="paged_kv_layout_shared_core" + ), + "kv_page_size": page_size, + "kv_page_bounds": [list(bound) for bound in pages], + }, + ) + + def compare_single_gpu_rope_attention( inputs: AttentionComparisonInputs, ) -> AttentionComparisonReport: @@ -465,6 +598,13 @@ def _run_decode_kv_replay( } if merge_backend == "transformer_engine": provenance.update(_te_context_parallel_provenance()) + provenance.update( + { + "actual_backend": "te_context_parallel_merge_helpers", + "communication_backend": "none", + "production_ready": False, + } + ) return AttentionPathResult( name=f"{merge_backend}_decode_kv_replay", out=torch.cat(outs, dim=0), @@ -672,6 +812,13 @@ def run_paged_kv_attention( } if merge_backend == "transformer_engine": provenance.update(_te_context_parallel_provenance()) + provenance.update( + { + "actual_backend": "te_context_parallel_merge_helpers", + "communication_backend": "none", + "production_ready": False, + } + ) return AttentionPathResult( name=f"{merge_backend}_paged_kv", out=out.to(inputs.output_dtype), @@ -916,6 +1063,9 @@ def _single_gpu_reference_provenance() -> dict[str, Any]: return { "execution_scope": "single_device_correctness_reference", "runtime_verified": False, + "actual_backend": "rlkernel.pytorch.attention_reference", + "communication_backend": "none", + "production_ready": False, "preprocess_policy": "reference_only_not_production", "native_backend_executed": False, "preprocess_fallback": True, @@ -929,6 +1079,117 @@ def _single_gpu_reference_provenance() -> dict[str, Any]: } +def _strict_shared_core_provenance( + inputs: AttentionComparisonInputs, + *, + materialization: str, +) -> dict[str, Any]: + return { + "execution_scope": "single_device_strict_shared_core", + "runtime_verified": False, + "actual_backend": "rlkernel.pytorch.strict_attention_reference", + "communication_backend": "none", + "production_ready": False, + "strict_core_id": "rlkernel.attention.deterministic_core.v1", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "strict_mode": True, + "native_backend_executed": False, + "native_attention_arithmetic": False, + "fallback": False, + "fallback_reason": None, + "materialization": materialization, + "attention_mode": "prefill", + "qkv_input_boundary": "projected_qkv", + "qkv_projection_executed": False, + "output_projection_executed": False, + "communication": "none", + "split_kv_policy": "disabled", + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_domain": "attention", + "lse_exported": True, + "dtype": str(inputs.q.dtype).replace("torch.", ""), + } + + +def _strict_attention_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: float | None, + key_padding_mask: torch.Tensor | None, + q_start: int, + k_start: int, + total_query_len: int, + total_kv_len: int, + output_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the fixed one-batch/one-query arithmetic schedule locally. + + PR2 must remain independently runnable before PR3 is merged, so this small + reference intentionally mirrors the PR3 schedule instead of importing its + implementation. The contract is the same: no Split-KV, FP32 intermediates, + logical position offsets, and one reduction per query row. + """ + + if q.shape[0] != k.shape[0] or k.shape[:1] != v.shape[:1] or k.shape[2:] != v.shape[2:]: + raise ValueError("strict Attention q/k/v batch and KV shapes must match") + hq, hkv = q.size(1), k.size(1) + if hq % hkv: + raise ValueError("strict Attention requires GQA-compatible head counts") + scale_value = 1.0 / math.sqrt(q.size(-1)) if scale is None else float(scale) + output_rows: list[torch.Tensor] = [] + lse_rows: list[torch.Tensor] = [] + for batch_index in range(q.size(0)): + q_rows: list[torch.Tensor] = [] + lse_batch: list[torch.Tensor] = [] + k_batch = k[batch_index : batch_index + 1].float() + v_batch = v[batch_index : batch_index + 1].float() + if hq != hkv: + k_batch = k_batch.repeat_interleave(hq // hkv, dim=1) + v_batch = v_batch.repeat_interleave(hq // hkv, dim=1) + for query_index in range(q.size(2)): + q_row = q[batch_index : batch_index + 1, :, query_index : query_index + 1, :].float() + scores = torch.matmul(q_row, k_batch.transpose(-1, -2)) * scale_value + if causal: + query_position = total_kv_len - total_query_len + q_start + query_index + key_positions = torch.arange( + k.size(2), device=q.device, dtype=torch.long + ) + k_start + scores = scores.masked_fill( + key_positions.view(1, 1, 1, -1) > query_position, + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill( + ~key_padding_mask[batch_index : batch_index + 1] + .view(1, 1, 1, -1), + float("-inf"), + ) + row_max = scores.amax(dim=-1, keepdim=True) + finite = torch.isfinite(row_max) + exp_scores = torch.where( + finite, + torch.exp(scores - row_max), + torch.zeros_like(scores), + ) + row_sum = exp_scores.sum(dim=-1, keepdim=True) + lse = torch.where( + row_sum > 0, + row_max + torch.log(row_sum), + torch.full_like(row_sum, float("-inf")), + ) + weights = torch.where(row_sum > 0, exp_scores / row_sum, torch.zeros_like(exp_scores)) + q_rows.append(torch.matmul(weights, v_batch).to(output_dtype)) + lse_batch.append(lse.squeeze(-1)) + output_rows.append(torch.cat(q_rows, dim=2)) + lse_rows.append(torch.cat(lse_batch, dim=2)) + return torch.cat(output_rows, dim=0), torch.cat(lse_rows, dim=0) + + def _attention_with_lse( q: torch.Tensor, k: torch.Tensor, diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index 9bd33fef..fa55f7a8 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -153,6 +153,32 @@ def test_single_gpu_attention_harness_reports_out_lse_and_dlogp_drift(): json.dumps(payload) +def test_strict_shared_core_is_bitwise_across_full_chunked_and_paged_layouts(): + report = compare_single_gpu_attention( + _comparison_inputs(), + query_chunk_size=2, + kv_page_size=3, + strict_bitwise=True, + ) + + assert report.reference_name == "strict_shared_core_full_prefill" + by_name = {drift.candidate_name: drift for drift in report.drifts} + assert set(by_name) == { + "strict_shared_core_chunked_prefill", + "strict_shared_core_paged_kv", + } + for drift in by_name.values(): + assert drift.out.max_abs == 0.0 + assert drift.lse.max_abs == 0.0 + assert drift.provenance["strict_core_id"] == ( + "rlkernel.attention.deterministic_core.v1" + ) + assert drift.provenance["strict_schedule"] == ( + "single_batch_single_query_global_kv_blocks" + ) + assert drift.provenance["split_kv_policy"] == "disabled" + + def test_single_gpu_attention_harness_preserves_key_padding_mask(): q, k, v = _qkv(seed=3) key_padding_mask = torch.tensor(