From cb359fe6c7f9d9f5900a9515d0fcf2fd811404d2 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 14:49:37 +0000 Subject: [PATCH 01/11] feat: add CP-aware attention contract --- docs/design/runtime-dispatch.md | 6 + docs/design/ws2-cp-attention-contract.md | 167 +++++++ docs/operators/attention.md | 11 + rl_engine/kernels/attention_contract.py | 593 +++++++++++++++++++++++ rl_engine/kernels/registry.py | 156 ++++++ tests/test_attention_contract.py | 286 +++++++++++ 6 files changed, 1219 insertions(+) create mode 100644 docs/design/ws2-cp-attention-contract.md create mode 100644 rl_engine/kernels/attention_contract.py create mode 100644 tests/test_attention_contract.py diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..41946a97 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,12 @@ logical type, and the registry selects the first available backend for the curre 4. Cache successfully constructed operator instances. 5. Skip backends that already failed in the current process. +WS2 Attention uses the stricter `KernelRegistry.get_attention_op(contract)` path. In addition to +platform priority, this path requires a backend capability descriptor and checks the requested +role, mode, dtype, TP/CP layout, LSE export, deterministic merge, packed varlen, and KV-cache +semantics. Incompatible candidates produce explicit rejection reasons and are never used as an +undeclared fallback. See [WS2 CP-aware Attention contract](ws2-cp-attention-contract.md). + ## LogP Priority | Platform | Priority | diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md new file mode 100644 index 00000000..dfd053ff --- /dev/null +++ b/docs/design/ws2-cp-attention-contract.md @@ -0,0 +1,167 @@ +# WS2 CP-Aware Attention Contract + +Status: PR1 contract and dispatch metadata + +Tracking and shared contracts: + +- [#235: CP-aware deterministic Attention](https://github.com/RL-Align/RL-Kernel/issues/235) +- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) +- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) +- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) +- [#207: cross-config logprob drift contract](https://github.com/RL-Align/RL-Kernel/issues/207) + +## Scope + +This contract describes the logical inputs and deterministic reduction semantics for standard +softmax Attention under tensor parallelism (TP) and context parallelism (CP). It lets runtime +dispatch reject a backend whose numerical semantics do not match the requested layout. + +This PR1 layer does not shard tensors, launch a collective, merge CP partial states, or implement +a fused kernel. The deterministic CP reference implementation and its distributed numerical tests +belong to later work in #235. + +## Contract Objects + +`rl_engine.kernels.attention_contract` defines: + +- `AttentionContract`: role, mode, dtype, causal metadata, sharding, reduction, and optional cache + identity; +- `ShardingSpec`: TP-local head ownership and CP block-to-token ownership; +- `ReductionSpec`: fixed `(out, lse)` merge semantics; +- `KVCacheSpec`: decode replay cache identity; +- `AttentionBackendCapability`: the layouts and semantics a backend explicitly supports. + +Construction performs validation immediately. A structurally valid contract means that the +request is complete and internally consistent; it does not mean that an installed backend can +materialize it. + +## Qwen3-8B TP=4 CP=4 Example + +```python +from rl_engine.kernels.attention_contract import ( + AttentionContract, + ReductionSpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=4, + cp_rank=0, + cp_world_size=4, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=8, + local_kv_head_start=0, + local_kv_heads=2, + global_sequence_length=4096, + local_sequence_length=1024, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 1024), +) + +contract = AttentionContract( + role="infer", + mode="prefill", + dtype="bf16", + batch_size=1, + query_sequence_length=1024, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), +) +``` + +The TP fields preserve the global Qwen3 GQA mapping: each rank owns 8 of 32 query heads and 2 of +8 KV heads. The CP fields map local tensor slices to stable logical global block ids. A rank that +owns non-contiguous blocks uses one global token start per block and one extra local boundary: + +```python +global_block_indices=(0, 7) +global_block_token_starts=(0, 3584) +local_block_offsets=(0, 512, 1024) +``` + +This metadata is sufficient for a later implementation to restore logical global order without +using ring arrival order. + +## Reduction Semantics + +The only PR1 reduction contract is: + +```text +partial state: (out, attention-domain lse) +merge: online_softmax_lse +acc_dtype: fp32 +order: global_block_index +downcast_at: final_write +engine: in_op_reference +``` + +CP output is not a plain sum. A backend that cannot export attention-domain LSE or cannot merge +partial states in fixed logical order is incompatible with this contract. + +The acceptable output and selected-logprob drift thresholds remain owned by #108. This contract +does not introduce another tolerance table. When connected to the rollout/training chain, the +selected-token metric remains the #207 convention: + +```text +dlogp = training-side recomputed logp - rollout-side old logp +``` + +## Mode-Specific Metadata + +All causal calls provide `causal_offsets`. Packed varlen calls provide one causal offset per +packed sequence and validated `packed_sequence_offsets`. + +Decode additionally requires `KVCacheSpec` with: + +- one cache position and KV sequence length per logical sequence; +- a block/page table; +- global token positions for every logical cached token; +- a prefix-cache key when prefix caching is enabled. + +Missing decode cache identity is an error at contract construction time. + +## Contract-Aware Dispatch + +Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: + +```python +result = kernel_registry.get_attention_op(contract) +op = result.op +provenance = result.provenance +``` + +Dispatch considers only backends with an `AttentionBackendCapability`. It checks role, attention +mode, dtype, TP/CP degree, LSE export, deterministic CP merge, packed varlen, and KV-cache support. +An undeclared or incompatible backend is skipped with an explicit rejection reason. + +The current WS1 PyTorch Attention implementations support local reference math but do not export +attention-domain LSE or materialize deterministic CP merge. Strict WS2 requests therefore fail +clearly today. A later deterministic backend becomes selectable by registering a capability that +truthfully declares those features; no grid-planner branch or silent fallback is required. + +Successful dispatch provenance records: + +- requested and actual backend ids; +- platform and fallback status; +- prior candidate rejection reasons; +- the complete requested contract; +- the selected backend capability descriptor. + +## Validation + +Contract and dispatch behavior are covered by: + +```bash +python -m pytest tests/test_attention_contract.py -q +``` + +The tests include Qwen3 TP=4/CP=4 construction, GQA ownership errors, non-contiguous CP blocks, +packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible +fallback, and JSON-compatible provenance. diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3cb5f9b..8bec7d22 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -83,6 +83,17 @@ path. When fused attention kernels land, they are prepended to the priority list op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. +### WS2 CP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed +`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity +before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. + +Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic +CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as +a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..b2269d04 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,593 @@ +# 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 +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" + + +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 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, ...] + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + + 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") + 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 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)}" + ) + + 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" + ) + for row_index, row in enumerate(block_table): + 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 not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + 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" + ) + + 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 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 + kv_cache: KVCacheSpec | 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") + _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.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 = len(self.sharding.packed_sequence_offsets) - 1 + 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.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), + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + } + 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, + "kv_cache": kv_cache, + } + + +@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 + 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", + ): + 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") + 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, + "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", + "ShardingSpec", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..cab339fd 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -6,6 +6,15 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDispatchResult, + AttentionDType, + AttentionMode, + AttentionRole, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -135,6 +144,39 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + # These descriptors report what existing WS1 implementations actually + # support. Neither implementation exports attention-domain LSE yet, so + # a strict WS2 request is rejected until the deterministic CP reference + # backend lands instead of silently selecting an incompatible fallback. + common_roles = frozenset({AttentionRole.TRAIN, AttentionRole.INFER}) + common_dtypes = frozenset({AttentionDType.BF16, AttentionDType.FP16, AttentionDType.FP32}) + self._attention_capabilities = { + OpBackend.PYTORCH_NATIVE_ATTENTION: AttentionBackendCapability( + backend_id="pytorch-native-attention-ws1", + roles=common_roles, + modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}), + dtypes=common_dtypes, + cp_world_sizes=(1,), + exports_attention_lse=False, + deterministic_cp_merge=False, + supports_packed_varlen=False, + supports_kv_cache=False, + implementation_kind="reference", + ), + OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN: AttentionBackendCapability( + backend_id="pytorch-native-kv-cache-attention-ws1", + roles=common_roles, + modes=frozenset({AttentionMode.DECODE}), + dtypes=common_dtypes, + cp_world_sizes=(1,), + exports_attention_lse=False, + deterministic_cp_merge=False, + supports_packed_varlen=False, + supports_kv_cache=False, + implementation_kind="reference", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -314,6 +356,120 @@ def get_op(self, op_type: str) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def get_attention_op( + self, + contract: AttentionContract, + *, + requested_backend: str = "deterministic", + ) -> AttentionDispatchResult: + """Resolve only a backend that explicitly supports the WS2 contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + """ + + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise AttentionContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip().lower() + + platform = self._platform() + op_type = "kv_cache_attention" if contract.mode is AttentionMode.DECODE else "attention" + candidates = self._priority_map.get(platform, {}).get(op_type, []) + rejected: list[str] = [] + + for backend in candidates: + capability = self._attention_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no AttentionBackendCapability declared") + continue + incompatibilities = list(capability.incompatibilities(contract)) + policy_mismatch = self._attention_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + incompatibilities.append(policy_mismatch) + if incompatibilities: + rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": bool(rejected), + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return AttentionDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No attention backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, mode={requested['mode']}, " + f"dtype={requested['dtype']}, TP={contract.sharding.tp_world_size}, " + f"CP={contract.sharding.cp_world_size}. Rejections: {details}" + ) + + @staticmethod + def _attention_policy_mismatch( + requested_backend: str, + capability: AttentionBackendCapability, + ) -> str | None: + if requested_backend == "auto": + return None + if requested_backend in {"production", "reference", "deterministic"}: + if capability.implementation_kind == requested_backend: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={requested_backend}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + def _platform(self) -> str: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + def _get_or_create_backend(self, backend: OpBackend) -> Any | None: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op + def _load_backend(self, backend: OpBackend) -> Optional[Type]: """Dynamic loading technique: Import modules only when needed and check environment dependencies. diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py new file mode 100644 index 00000000..907b776d --- /dev/null +++ b/tests/test_attention_contract.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 Attention CP contract and contract-aware dispatch tests (issue #235).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + KVCacheSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 4, + cp_rank: int = 0, + cp_world_size: int = 4, + global_sequence_length: int = 4096, + local_sequence_length: int = 1024, + global_block_indices: tuple[int, ...] = (0,), + global_block_token_starts: tuple[int, ...] = (0,), + local_block_offsets: tuple[int, ...] = (0, 1024), + packed_sequence_offsets: tuple[int, ...] | None = None, +) -> ShardingSpec: + local_q_heads = 32 // tp_world_size + local_kv_heads = 8 // tp_world_size + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + global_block_indices=global_block_indices, + global_block_token_starts=global_block_token_starts, + local_block_offsets=local_block_offsets, + packed_sequence_offsets=packed_sequence_offsets, + ) + + +def _contract( + *, + role: str = "infer", + mode: str = "prefill", + sharding: ShardingSpec | None = None, + kv_cache: KVCacheSpec | None = None, + causal_offsets: tuple[int, ...] = (0,), + batch_size: int = 1, +) -> AttentionContract: + resolved_sharding = sharding or _sharding() + return AttentionContract( + role=role, + mode=mode, + dtype="bf16", + batch_size=batch_size, + query_sequence_length=(1 if mode == "decode" else resolved_sharding.local_sequence_length), + head_dim=128, + causal=True, + causal_offsets=causal_offsets, + sharding=resolved_sharding, + reduction=ReductionSpec(), + kv_cache=kv_cache, + ) + + +def _declared_cp_backend() -> AttentionBackendCapability: + return AttentionBackendCapability( + backend_id="test-deterministic-cp-attention", + roles=frozenset({AttentionRole.TRAIN, AttentionRole.INFER}), + modes=frozenset( + {AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL, AttentionMode.DECODE} + ), + dtypes=frozenset({AttentionDType.BF16}), + tp_world_sizes=(4,), + cp_world_sizes=(1, 2, 4), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=True, + supports_kv_cache=True, + implementation_kind="deterministic", + ) + + +def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.local_q_heads == 8 + assert contract.sharding.local_kv_heads == 2 + assert contract.reduction.acc_dtype is AttentionDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "online_softmax_lse", + "acc_dtype": "fp32", + "order": "global_block_index", + "downcast_at": "final_write", + "engine": "in_op_reference", + } + json.dumps(contract.to_dict()) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 4, "tp_rank=4"), + ("cp_rank", 4, "cp_rank=4"), + ("global_block_indices", (), "must not be empty"), + ("global_block_indices", (1, 0), "strictly increasing"), + ], +) +def test_invalid_rank_and_cp_order_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 4, + "cp_rank": 0, + "cp_world_size": 4, + "global_q_heads": 32, + "global_kv_heads": 8, + "local_q_head_start": 0, + "local_q_heads": 8, + "local_kv_head_start": 0, + "local_kv_heads": 2, + "global_sequence_length": 4096, + "local_sequence_length": 1024, + "global_block_indices": (0,), + "global_block_token_starts": (0,), + "local_block_offsets": (0, 1024), + } + values[field] = value + + with pytest.raises(AttentionContractError, match=message): + ShardingSpec(**values) + + +def test_tp_local_heads_must_preserve_global_gqa_mapping(): + with pytest.raises(AttentionContractError, match="local TP head counts"): + replace(_sharding(), local_q_heads=7) + + with pytest.raises(AttentionContractError, match="head starts"): + replace(_sharding(tp_rank=1), local_q_head_start=0) + + +def test_sequence_range_and_packed_offsets_are_validated(): + with pytest.raises(AttentionContractError, match="exceeds global_sequence_length"): + _sharding(global_block_token_starts=(4000,)) + + with pytest.raises(AttentionContractError, match="final packed_sequence_offsets"): + _sharding(packed_sequence_offsets=(0, 512)) + + sharding = _sharding(packed_sequence_offsets=(0, 256, 1024)) + assert sharding.packed_sequence_offsets == (0, 256, 1024) + + +def test_non_contiguous_cp_blocks_have_explicit_global_and_local_offsets(): + sharding = _sharding( + global_block_indices=(0, 7), + global_block_token_starts=(0, 3584), + local_block_offsets=(0, 512, 1024), + ) + + assert sharding.global_block_indices == (0, 7) + assert sharding.global_block_token_starts == (0, 3584) + assert sharding.local_block_offsets == (0, 512, 1024) + + with pytest.raises(AttentionContractError, match="non-overlapping and ordered"): + _sharding( + global_block_indices=(0, 1), + global_block_token_starts=(0, 256), + local_block_offsets=(0, 512, 1024), + ) + + +def test_reduction_requires_fp32_accumulation(): + with pytest.raises(AttentionContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + +def test_causal_attention_requires_explicit_offset(): + contract = _contract() + with pytest.raises(AttentionContractError, match="causal_offsets are required"): + replace(contract, causal_offsets=None) + + +def test_decode_requires_complete_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="kv_cache metadata is required"): + _contract(mode="decode") + + cache = KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1, -1),), + global_token_positions=tuple(range(17)), + prefix_cache_enabled=True, + prefix_cache_key="prefix:sample-0", + ) + contract = _contract(mode="decode", kv_cache=cache) + assert contract.to_dict()["kv_cache"]["block_table"] == [[0, 1, -1]] + + +def test_prefix_cache_key_is_required_only_when_prefix_cache_is_enabled(): + with pytest.raises(AttentionContractError, match="prefix_cache_key is required"): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + prefix_cache_enabled=True, + ) + + +def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): + registry = KernelRegistry() + + with pytest.raises(RuntimeError) as exc_info: + registry.get_attention_op(_contract()) + + message = str(exc_info.value) + assert "CP=4 is unsupported" in message + assert "attention-domain LSE export is unsupported" in message + assert "deterministic CP (out, lse) merge is unsupported" in message + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["attention"] = [OpBackend.PYTORCH_ATTN] + + with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"): + registry.get_attention_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + result = registry.get_attention_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-cp-attention" + assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["cp_world_size"] == 4 + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_attention_op(_contract(), requested_backend="another-backend") + + result = registry.get_attention_op( + _contract(), requested_backend="test-deterministic-cp-attention" + ) + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + + +def test_packed_layout_requires_declared_backend_support(): + capability = replace(_declared_cp_backend(), supports_packed_varlen=False) + contract = _contract( + sharding=_sharding(packed_sequence_offsets=(0, 512, 1024)), + causal_offsets=(0, 0), + ) + + assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) From 33119ffebf6e7912fce59358b9f6ef5d9d65346d Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 15:12:26 +0000 Subject: [PATCH 02/11] fix: tighten attention contract validation --- docs/design/ws2-cp-attention-contract.md | 11 +++- rl_engine/kernels/attention_contract.py | 56 ++++++++++++++++++- tests/test_attention_contract.py | 69 ++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 3 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index dfd053ff..0bc1be02 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -35,6 +35,10 @@ Construction performs validation immediately. A structurally valid contract mean request is complete and internally consistent; it does not mean that an installed backend can materialize it. +`AttentionContract.batch_size` is the logical sequence count. For packed varlen input it must +equal `len(packed_sequence_offsets) - 1`; it is not the physical leading dimension of a flattened +token tensor. + ## Qwen3-8B TP=4 CP=4 Example ```python @@ -122,10 +126,15 @@ Decode additionally requires `KVCacheSpec` with: - one cache position and KV sequence length per logical sequence; - a block/page table; +- the physical page size; - global token positions for every logical cached token; - a prefix-cache key when prefix caching is enabled. -Missing decode cache identity is an error at contract construction time. +Within each logical sequence, global token positions must be strictly increasing. Block-table +padding must be trailing, the active page count must match `ceil(kv_seq_len / page_size)`, and a +sequence cannot repeat one physical page id. Different sequences may share physical pages for an +equivalent prefix. Missing or inconsistent decode cache identity is an error at contract +construction time. ## Contract-Aware Dispatch diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index b2269d04..2607b737 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -276,12 +276,14 @@ class KVCacheSpec: 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 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" ) @@ -289,6 +291,10 @@ def __post_init__(self) -> None: 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" @@ -298,6 +304,20 @@ def __post_init__(self) -> None: "global_token_positions must describe every logical cached token; " f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" ) + token_offset = 0 + 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" + ) + token_offset += sequence_length try: block_table = tuple(tuple(row) for row in self.block_table) @@ -309,13 +329,37 @@ def __post_init__(self) -> None: raise AttentionContractError( "block_table must contain one non-empty row per kv_seq_lens entry" ) - for row_index, row in enumerate(block_table): + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + 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" + ) + active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(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(active_blocks)}" + ) + if len(set(active_blocks)) != len(active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) if not isinstance(self.prefix_cache_enabled, bool): raise AttentionContractError("prefix_cache_enabled must be a bool") @@ -362,6 +406,13 @@ def __post_init__(self) -> None: raise AttentionContractError("sharding must be a ShardingSpec") if not isinstance(self.reduction, ReductionSpec): raise AttentionContractError("reduction must be a ReductionSpec") + 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: @@ -372,7 +423,7 @@ def __post_init__(self) -> None: 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 = len(self.sharding.packed_sequence_offsets) - 1 + expected_causal_offsets = batch_size offset_owner = "packed sequence" else: expected_causal_offsets = batch_size @@ -441,6 +492,7 @@ def to_dict(self) -> dict[str, Any]: "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, } diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 907b776d..a07274be 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -209,6 +209,7 @@ def test_decode_requires_complete_kv_cache_identity(): kv_seq_lens=(17,), block_table=((0, 1, -1),), global_token_positions=tuple(range(17)), + page_size=16, prefix_cache_enabled=True, prefix_cache_key="prefix:sample-0", ) @@ -223,10 +224,67 @@ def test_prefix_cache_key_is_required_only_when_prefix_cache_is_enabled(): kv_seq_lens=(17,), block_table=((0, 1),), global_token_positions=tuple(range(17)), + page_size=16, prefix_cache_enabled=True, ) +def test_cache_positions_must_match_kv_sequence_count(): + with pytest.raises(AttentionContractError, match="one entry per kv_seq_lens"): + KVCacheSpec( + cache_positions=(1,), + kv_seq_lens=(2, 2), + block_table=((0,), (1,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + ) + + +@pytest.mark.parametrize("positions", [(7, 6), (7, 7)]) +def test_kv_cache_positions_must_be_strictly_increasing_per_sequence(positions): + with pytest.raises(AttentionContractError, match="strictly increasing"): + KVCacheSpec( + cache_positions=(7,), + kv_seq_lens=(2,), + block_table=((0,),), + global_token_positions=positions, + page_size=2, + ) + + +@pytest.mark.parametrize( + ("block_table", "message"), + [ + ((0, -1, 1), "padding must be trailing"), + ((0, 0, -1), "duplicate active page ids"), + ((0, -1, -1), "active page count"), + ], +) +def test_kv_cache_block_table_page_mapping_is_validated(block_table, message): + with pytest.raises(AttentionContractError, match=message): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=(block_table,), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +def test_prefix_pages_may_be_shared_across_sequences(): + cache = KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + ) + + assert cache.block_table == ((3,), (3,)) + + def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): registry = KernelRegistry() @@ -281,6 +339,17 @@ def test_packed_layout_requires_declared_backend_support(): contract = _contract( sharding=_sharding(packed_sequence_offsets=(0, 512, 1024)), causal_offsets=(0, 0), + batch_size=2, ) assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) + + +def test_packed_sequence_count_must_match_logical_batch_size(): + sharding = _sharding(packed_sequence_offsets=(0, 512, 1024)) + + with pytest.raises(AttentionContractError, match="must equal logical batch_size"): + _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=1) + + contract = _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=2) + assert contract.batch_size == 2 From b7ba64b58977cd59bec17b8a4c3f37cddbbb2ce0 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 15:28:36 +0000 Subject: [PATCH 03/11] fix: constrain shared KV prefix pages --- docs/design/ws2-cp-attention-contract.md | 8 ++- rl_engine/kernels/attention_contract.py | 61 ++++++++++++++++++ tests/test_attention_contract.py | 82 ++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 2 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index 0bc1be02..becd90a7 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -128,12 +128,16 @@ Decode additionally requires `KVCacheSpec` with: - a block/page table; - the physical page size; - global token positions for every logical cached token; -- a prefix-cache key when prefix caching is enabled. +- a prefix-cache key and explicit shared-prefix page count when prefix caching is enabled. Within each logical sequence, global token positions must be strictly increasing. Block-table padding must be trailing, the active page count must match `ceil(kv_seq_len / page_size)`, and a sequence cannot repeat one physical page id. Different sequences may share physical pages for an -equivalent prefix. Missing or inconsistent decode cache identity is an error at contract +equivalent prefix only when those pages are declared by `shared_prefix_page_count`, use the same +leading page ids and logical positions, and are fully populated. Declared shared prefix pages are +read-only; all suffix pages are exclusive to one sequence, providing the contract boundary needed +for copy-on-write before divergent decode. When prefix caching is disabled, no active page may be +shared across sequences. Missing or inconsistent decode cache identity is an error at contract construction time. ## Contract-Aware Dispatch diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 2607b737..cfb4e8f1 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -279,6 +279,7 @@ class KVCacheSpec: 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") @@ -305,6 +306,7 @@ def __post_init__(self) -> None: 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 @@ -317,6 +319,7 @@ def __post_init__(self) -> None: "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 try: @@ -329,6 +332,7 @@ def __post_init__(self) -> None: 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) ): @@ -360,9 +364,13 @@ def __post_init__(self) -> None: raise AttentionContractError( f"block_table row {row_index} contains duplicate active page ids" ) + active_block_rows.append(tuple(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" @@ -371,6 +379,58 @@ def __post_init__(self) -> 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(active_blocks) < shared_prefix_page_count for active_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, (active_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if active_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, active_blocks in enumerate(active_block_rows): + for page_id in active_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) @@ -495,6 +555,7 @@ def to_dict(self) -> dict[str, Any]: "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, } return { "semantic_operator": "standard_softmax_attention", diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index a07274be..9d0959c8 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -280,9 +280,91 @@ def test_prefix_pages_may_be_shared_across_sequences(): page_size=2, prefix_cache_enabled=True, prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, ) assert cache.block_table == ((3,), (3,)) + assert cache.shared_prefix_page_count == 1 + + +def test_non_prefix_cache_rejects_cross_sequence_page_sharing(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=False, + ) + + +def test_prefix_cache_requires_explicit_shared_page_count(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=0, + ) + + +def test_prefix_cache_rejects_shared_writable_suffix_pages(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 4)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_identity_must_match_pages_and_positions(): + with pytest.raises(AttentionContractError, match="page ids must match"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (5, 6)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + with pytest.raises(AttentionContractError, match="token positions must match"): + KVCacheSpec( + cache_positions=(3, 13), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 5)), + global_token_positions=(0, 1, 2, 3, 10, 11, 12, 13), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_pages_must_be_fully_populated(): + with pytest.raises(AttentionContractError, match="fully populated and read-only"): + KVCacheSpec( + cache_positions=(0, 0), + kv_seq_lens=(1, 1), + block_table=((3,), (3,)), + global_token_positions=(0, 0), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="partial-prefix-page", + shared_prefix_page_count=1, + ) def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): From 0f1fd760402126a4aaebe2620f95fb6c8012b751 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 15:59:08 +0000 Subject: [PATCH 04/11] fix: satisfy attention contract type checks --- .github/workflows/ci.yml | 3 +++ rl_engine/kernels/attention_contract.py | 24 +++++++++++------------- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..26f0575c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,9 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run WS2 Attention Contract Tests (CPU-safe) + run: python -m pytest tests/test_attention_contract.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index cfb4e8f1..2c68e7dd 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -336,7 +336,7 @@ def __post_init__(self) -> None: for row_index, (row, sequence_length) in enumerate( zip(block_table, kv_seq_lens, strict=True) ): - active_blocks: list[int] = [] + 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: @@ -352,19 +352,19 @@ def __post_init__(self) -> None: "block_table -1 padding must be trailing; " f"row {row_index} contains an active block after padding" ) - active_blocks.append(block) + row_active_blocks.append(block) expected_blocks = (sequence_length + page_size - 1) // page_size - if len(active_blocks) != expected_blocks: + 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(active_blocks)}" + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" ) - if len(set(active_blocks)) != len(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(active_blocks)) + 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") @@ -386,9 +386,7 @@ def __post_init__(self) -> None: shared_prefix_pages: tuple[int, ...] = () if shared_prefix_page_count > 0: - if any( - len(active_blocks) < shared_prefix_page_count for active_blocks in active_block_rows - ): + 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" ) @@ -400,10 +398,10 @@ def __post_init__(self) -> None: 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, (active_blocks, positions) in enumerate( + for sequence_index, (row_blocks, positions) in enumerate( zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 ): - if active_blocks[:shared_prefix_page_count] != shared_prefix_pages: + 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" @@ -416,8 +414,8 @@ def __post_init__(self) -> None: exclusive_page_owners: dict[int, int] = {} shared_prefix_page_ids = set(shared_prefix_pages) - for sequence_index, active_blocks in enumerate(active_block_rows): - for page_id in active_blocks[shared_prefix_page_count:]: + 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" From 6d826df393ba6d4972d1706cd5ffd704a44e4456 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Mon, 20 Jul 2026 09:31:13 +0000 Subject: [PATCH 05/11] fix: validate attention position metadata --- docs/design/ws2-cp-attention-contract.md | 8 +++++ rl_engine/kernels/attention_contract.py | 21 ++++++++++++- tests/test_attention_contract.py | 38 +++++++++++++++++++++++- 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index becd90a7..f9bb6ec1 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -39,6 +39,10 @@ materialize it. equal `len(packed_sequence_offsets) - 1`; it is not the physical leading dimension of a flattened token tensor. +For full `prefill`, `query_sequence_length` equals the local sequence length described by +`ShardingSpec`. Chunked prefill and decode may use shorter query lengths than their available KV +context. + ## Qwen3-8B TP=4 CP=4 Example ```python @@ -140,6 +144,10 @@ for copy-on-write before divergent decode. When prefix caching is disabled, no a shared across sequences. Missing or inconsistent decode cache identity is an error at contract construction time. +Each `cache_positions` entry is the terminal logical position already present in that sequence's +KV cache, so it must equal the final corresponding `global_token_positions` entry. It is not the +next position to be written. + ## Contract-Aware Dispatch Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 2c68e7dd..0f14fdea 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -322,6 +322,17 @@ def __post_init__(self) -> None: 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: @@ -458,12 +469,20 @@ def __post_init__(self) -> None: 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") - _positive_int(self.query_sequence_length, "query_sequence_length") + 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 ( + 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: diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 9d0959c8..2a14d856 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -67,6 +67,7 @@ def _contract( kv_cache: KVCacheSpec | None = None, causal_offsets: tuple[int, ...] = (0,), batch_size: int = 1, + query_sequence_length: int | None = None, ) -> AttentionContract: resolved_sharding = sharding or _sharding() return AttentionContract( @@ -74,7 +75,11 @@ def _contract( mode=mode, dtype="bf16", batch_size=batch_size, - query_sequence_length=(1 if mode == "decode" else resolved_sharding.local_sequence_length), + query_sequence_length=( + query_sequence_length + if query_sequence_length is not None + else (1 if mode == "decode" else resolved_sharding.local_sequence_length) + ), head_dim=128, causal=True, causal_offsets=causal_offsets, @@ -200,6 +205,26 @@ def test_causal_attention_requires_explicit_offset(): replace(contract, causal_offsets=None) +def test_full_prefill_query_length_must_match_local_sequence_length(): + with pytest.raises(AttentionContractError, match="prefill query_sequence_length must equal"): + _contract(mode="prefill", query_sequence_length=2048) + + chunked = _contract(mode="chunked_prefill", query_sequence_length=512) + decode = _contract( + mode="decode", + query_sequence_length=1, + kv_cache=KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ), + ) + assert chunked.query_sequence_length == 512 + assert decode.query_sequence_length == 1 + + def test_decode_requires_complete_kv_cache_identity(): with pytest.raises(AttentionContractError, match="kv_cache metadata is required"): _contract(mode="decode") @@ -240,6 +265,17 @@ def test_cache_positions_must_match_kv_sequence_count(): ) +def test_cache_position_must_match_terminal_global_token_position(): + with pytest.raises(AttentionContractError, match="terminal global token position"): + KVCacheSpec( + cache_positions=(999,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + @pytest.mark.parametrize("positions", [(7, 6), (7, 7)]) def test_kv_cache_positions_must_be_strictly_increasing_per_sequence(positions): with pytest.raises(AttentionContractError, match="strictly increasing"): From ed05a09f66f12fa85e2cb7f5dbf68410178ecb41 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 20 Jul 2026 22:59:48 +0800 Subject: [PATCH 06/11] feat(attention): add deterministic CP reference Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 22 + rl_engine/kernels/gtest/operator_inputs.py | 22 + .../ops/pytorch/attention/cp_attention.py | 548 ++++++++++++++++++ rl_engine/kernels/registry.py | 9 + tests/test_cp_attention.py | 385 ++++++++++++ tests/test_operator_inputs.py | 1 + 6 files changed, 987 insertions(+) create mode 100644 rl_engine/kernels/ops/pytorch/attention/cp_attention.py create mode 100644 tests/test_cp_attention.py diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3cb5f9b..97dfed29 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -83,6 +83,15 @@ path. When fused attention kernels land, they are prepended to the priority list op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. +`kernel_registry.get_op("cp_attention")` resolves to +`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel +reference. It emulates CP prefill and chunked-prefill by splitting logical query +and KV sequence blocks, computing per-block `(out, lse)` partial states, and +merging them in fp32 by global KV block index. This path is not a production +fused backend; it defines the CP/LSE merge behavior that downstream fused paths +must match. Optional per-batch `query_position_offsets` / `key_position_offsets` +cover varlen causal-mask metadata while keeping the dense tensor layout. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -135,6 +144,7 @@ memory. ```bash python -m pytest tests/test_attention.py -v +python -m pytest tests/test_cp_attention.py -v ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -144,15 +154,27 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below), gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. +`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard +attention, CP=2 prefill vs CP=1, chunked-prefill replay, global-position causal +masking across CP boundaries, order-independent LSE merge by global block index, +padding/all-masked stability, BF16 final-write behavior, input purity, argument +validation, and registry dispatch. +`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill +synthetic case for local harnesses. + ## Implementation Files - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` +- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` - `rl_engine/kernels/registry.py` - `tests/test_attention.py` +- `tests/test_cp_attention.py` ## Known Limitations - PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). +- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, + not a distributed runtime or fused kernel. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - The naive path materializes the full `[B, Hq, Sq, Skv]` scores tensor — no query-chunking, so the LARGE load point is memory-heavy and GPU-only. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f124cafb..8d26ac04 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -28,6 +28,7 @@ def make_operator_inputs( "rms_norm": _make_rms_norm_inputs, "matmul": _make_matmul_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "rope": _make_rope_inputs, @@ -50,6 +51,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", @@ -107,6 +109,26 @@ def _make_attention_inputs( } +def _make_cp_attention_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + return { + "q": _floating_tensor( + (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0 + ), + "k": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1 + ), + "v": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2 + ), + "causal": True, + "cp_world_size": 2, + "kv_chunk_size": max(1, seq // 2), + } + + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..c658b134 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,548 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(q.dtype if output_dtype is None else output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + """ + + _validate_qkv(q, k, v) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + scale_value = scale if scale is not None else (1.0 / math.sqrt(dim)) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, None, :, :], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching out/lse shapes") + if state.block_start < previous_end: + raise ValueError("partial state block ranges must not overlap") + previous_end = state.block_end + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have the same shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionPartialState", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "merge_attention_partial_states", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..d463bb5b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -72,6 +72,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + # WS2 correctness-first context-parallel attention reference. It emulates + # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. + PYTORCH_CP_ATTENTION = ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ) # WS1 pure-PyTorch ground-truth linear ops PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops @@ -165,6 +171,7 @@ def __init__(self): ], "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], @@ -188,6 +195,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], @@ -206,6 +214,7 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..dd0f53f7 --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + merge_attention_partial_states, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 + + +@contextlib.contextmanager +def _single_thread(): + prev = torch.get_num_threads() + torch.set_num_threads(1) + try: + yield + finally: + torch.set_num_threads(prev) + + +def _qkv( + batch, + sq, + skv, + *, + seed, + dtype=torch.float32, + heads=_N_HEADS, + kv_heads=_N_KV, + dim=_HEAD_DIM, +): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(batch, heads, sq, dim, generator=gen, dtype=dtype) + k = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + v = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + return q, k, v + + +def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): + qf, kf = q.float(), k.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hq % hkv != 0: + raise ValueError("invalid GQA shape") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * ( + scale if scale is not None else 1.0 / math.sqrt(dim) + ) + if causal: + query_pos = torch.arange(skv - sq, skv) + key_pos = torch.arange(skv) + scores = scores.masked_fill( + (key_pos[None, :] > query_pos[:, None])[None, None, :, :], + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + return torch.logsumexp(scores, dim=-1) + + +def test_cp1_matches_native_attention_and_exports_lse(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 8, 8, seed=1) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + want = native.forward_fp32(q, k, v, causal=True) + want_lse = _full_lse(q, k, causal=True) + + torch.testing.assert_close(out, want, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse, want_lse, atol=_ATOL, rtol=0.0) + assert lse.dtype == torch.float32 + assert lse.shape == q.shape[:3] + + +def test_cp2_prefill_matches_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 9, 9, seed=2) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + out2, lse2 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=2) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_chunked_prefill_replay_matches_unchunked_cp2(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 10, 10, seed=3) + + with _single_thread(): + unchunked_out, unchunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + ) + chunked_out, chunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=3, + ) + + torch.testing.assert_close(chunked_out, unchunked_out, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(chunked_lse, unchunked_lse, atol=_ATOL, rtol=0.0) + + +def test_causal_mask_uses_global_positions_across_cp_boundary(): + op = DeterministicCPAttentionReferenceOp() + batch, heads, kv_heads, seq, dim = 1, 2, 1, 5, 3 + q = torch.zeros(batch, heads, seq, dim) + k = torch.zeros(batch, kv_heads, seq, dim) + v = torch.arange(seq * dim, dtype=torch.float32).reshape(1, 1, seq, dim) + out = op.forward_fp32(q, k, v, causal=True, cp_world_size=2) + + expected = torch.stack([v[0, 0, : index + 1].mean(dim=0) for index in range(seq)]) + expected = expected.reshape(1, 1, seq, dim).repeat(1, heads, 1, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + + +def test_position_offsets_apply_varlen_causal_metadata_per_batch_row(): + op = DeterministicCPAttentionReferenceOp() + q = torch.zeros(2, 2, 2, 1) + k = torch.zeros(2, 1, 4, 1) + v = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4, 1) + + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=torch.tensor([0, 11]), + key_position_offsets=torch.tensor([0, 10]), + cp_world_size=2, + kv_chunk_size=1, + ) + + expected = torch.tensor([0.0, 0.5, 4.5, 5.0]).reshape(2, 1, 2, 1).repeat(1, 2, 1, 1) + expected_lse = torch.log(torch.tensor([1.0, 2.0, 2.0, 3.0])).reshape(2, 1, 2) + expected_lse = expected_lse.repeat(1, 2, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(lse, expected_lse, atol=1.0e-6, rtol=0.0) + + +def test_merge_order_uses_global_block_index_not_arrival_order(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 6, 6, seed=4) + first = op.local_partial_state( + q, + k[:, :, :3], + v[:, :, :3], + q_start=0, + k_start=0, + total_kv_len=6, + causal=True, + ) + second = op.local_partial_state( + q, + k[:, :, 3:], + v[:, :, 3:], + q_start=0, + k_start=3, + total_kv_len=6, + causal=True, + ) + + forward = merge_attention_partial_states([first, second]) + reversed_arrival = merge_attention_partial_states([second, first]) + assert torch.equal(forward.out, reversed_arrival.out) + assert torch.equal(forward.lse, reversed_arrival.lse) + + +def test_key_padding_mask_and_all_masked_rows_are_stable(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 6, 6, seed=5) + mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, False, False, False, False, False], + ], + dtype=torch.bool, + ) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=False, + key_padding_mask=mask, + cp_world_size=2, + kv_chunk_size=2, + ) + want = native.forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + torch.testing.assert_close(out[:1], want[:1], atol=_ATOL, rtol=0.0) + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + assert torch.isfinite(out).all() + + +def test_empty_query_and_empty_kv_edges_are_stable(): + op = DeterministicCPAttentionReferenceOp() + q_empty = torch.randn(1, 2, 0, 4, requires_grad=True) + k_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + v_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + out, lse = op.forward_fp32_with_lse(q_empty, k_empty, v_empty, cp_world_size=2) + assert out.shape == (1, 2, 0, 4) + assert lse.shape == (1, 2, 0) + assert out.requires_grad + out.sum().backward() + assert torch.equal(q_empty.grad, torch.zeros_like(q_empty)) + assert torch.equal(k_empty.grad, torch.zeros_like(k_empty)) + assert torch.equal(v_empty.grad, torch.zeros_like(v_empty)) + + q = torch.randn(1, 2, 3, 4) + out, lse = op.forward_fp32_with_lse(q, k_empty, v_empty, causal=False, cp_world_size=4) + assert torch.equal(out, torch.zeros_like(out)) + assert torch.isneginf(lse).all() + + +def test_empty_kv_backward_returns_zero_grads(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 2, 3, 4, requires_grad=True) + k = torch.randn(1, 1, 0, 4, requires_grad=True) + v = torch.randn(1, 1, 0, 4, requires_grad=True) + + out = op.forward_fp32(q, k, v, causal=False, cp_world_size=4) + assert out.requires_grad + out.sum().backward() + + assert torch.equal(q.grad, torch.zeros_like(q)) + assert torch.equal(k.grad, torch.zeros_like(k)) + assert torch.equal(v.grad, torch.zeros_like(v)) + + +def test_bf16_forward_uses_fp32_merge_then_final_write(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 8, 8, seed=6, dtype=torch.bfloat16) + + out, lse = op.forward_with_lse(q, k, v, causal=True, cp_world_size=2, kv_chunk_size=2) + fp32_out, fp32_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert torch.equal(out, fp32_out.to(torch.bfloat16)) + assert torch.equal(lse, fp32_lse) + + +def test_cp2_chunked_gradients_match_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 5, 5, seed=12, heads=4, kv_heads=2, dim=8) + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + q_cp = q.detach().clone().requires_grad_(True) + k_cp = k.detach().clone().requires_grad_(True) + v_cp = v.detach().clone().requires_grad_(True) + gen = torch.Generator().manual_seed(13) + dy = torch.randn(1, 4, 5, 8, generator=gen) + + with _single_thread(): + out_ref = op.forward_fp32(q_ref, k_ref, v_ref, causal=True, cp_world_size=1) + out_cp = op.forward_fp32( + q_cp, + k_cp, + v_cp, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + out_ref.backward(dy) + out_cp.backward(dy) + + torch.testing.assert_close(out_cp, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(q_cp.grad, q_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(k_cp.grad, k_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) + + +def test_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index bb1a2220..9e3eac34 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -36,6 +36,7 @@ def _args(**overrides): "rms_norm", "matmul", "attention", + "cp_attention", "logp", "linear_logp", "rope", From 8a4f9eb190ed19053d996583a034e5aaeb1ab502 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 27 Jul 2026 22:00:02 +0800 Subject: [PATCH 07/11] fix: align attention contract target with issue scope --- docs/design/ws2-cp-attention-contract.md | 26 ++++----- tests/test_attention_contract.py | 67 +++++++++++++----------- 2 files changed, 48 insertions(+), 45 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index f9bb6ec1..3e259fa8 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -43,7 +43,7 @@ For full `prefill`, `query_sequence_length` equals the local sequence length des `ShardingSpec`. Chunked prefill and decode may use shorter query lengths than their available KV context. -## Qwen3-8B TP=4 CP=4 Example +## Qwen3-8B TP=2 CP=2 Example ```python from rl_engine.kernels.attention_contract import ( @@ -54,20 +54,20 @@ from rl_engine.kernels.attention_contract import ( sharding = ShardingSpec( tp_rank=0, - tp_world_size=4, + tp_world_size=2, cp_rank=0, - cp_world_size=4, + cp_world_size=2, global_q_heads=32, global_kv_heads=8, local_q_head_start=0, - local_q_heads=8, + local_q_heads=16, local_kv_head_start=0, - local_kv_heads=2, + local_kv_heads=4, global_sequence_length=4096, - local_sequence_length=1024, + local_sequence_length=2048, global_block_indices=(0,), global_block_token_starts=(0,), - local_block_offsets=(0, 1024), + local_block_offsets=(0, 2048), ) contract = AttentionContract( @@ -75,7 +75,7 @@ contract = AttentionContract( mode="prefill", dtype="bf16", batch_size=1, - query_sequence_length=1024, + query_sequence_length=2048, head_dim=128, causal=True, causal_offsets=(0,), @@ -84,14 +84,14 @@ contract = AttentionContract( ) ``` -The TP fields preserve the global Qwen3 GQA mapping: each rank owns 8 of 32 query heads and 2 of +The TP fields preserve the global Qwen3 GQA mapping: each rank owns 16 of 32 query heads and 4 of 8 KV heads. The CP fields map local tensor slices to stable logical global block ids. A rank that owns non-contiguous blocks uses one global token start per block and one extra local boundary: ```python -global_block_indices=(0, 7) -global_block_token_starts=(0, 3584) -local_block_offsets=(0, 512, 1024) +global_block_indices=(0, 3) +global_block_token_starts=(0, 3072) +local_block_offsets=(0, 1024, 2048) ``` This metadata is sufficient for a later implementation to restore logical global order without @@ -183,6 +183,6 @@ Contract and dispatch behavior are covered by: python -m pytest tests/test_attention_contract.py -q ``` -The tests include Qwen3 TP=4/CP=4 construction, GQA ownership errors, non-contiguous CP blocks, +The tests include Qwen3 TP=2/CP=2 construction, GQA ownership errors, non-contiguous CP blocks, packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible fallback, and JSON-compatible provenance. diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 2a14d856..feab4c3a 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -27,14 +27,14 @@ def _sharding( *, tp_rank: int = 0, - tp_world_size: int = 4, + tp_world_size: int = 2, cp_rank: int = 0, - cp_world_size: int = 4, + cp_world_size: int = 2, global_sequence_length: int = 4096, - local_sequence_length: int = 1024, + local_sequence_length: int = 2048, global_block_indices: tuple[int, ...] = (0,), global_block_token_starts: tuple[int, ...] = (0,), - local_block_offsets: tuple[int, ...] = (0, 1024), + local_block_offsets: tuple[int, ...] = (0, 2048), packed_sequence_offsets: tuple[int, ...] | None = None, ) -> ShardingSpec: local_q_heads = 32 // tp_world_size @@ -97,8 +97,8 @@ def _declared_cp_backend() -> AttentionBackendCapability: {AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL, AttentionMode.DECODE} ), dtypes=frozenset({AttentionDType.BF16}), - tp_world_sizes=(4,), - cp_world_sizes=(1, 2, 4), + tp_world_sizes=(2,), + cp_world_sizes=(1, 2), exports_attention_lse=True, deterministic_cp_merge=True, supports_packed_varlen=True, @@ -107,11 +107,13 @@ def _declared_cp_backend() -> AttentionBackendCapability: ) -def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): +def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): contract = _contract() - assert contract.sharding.local_q_heads == 8 - assert contract.sharding.local_kv_heads == 2 + assert contract.sharding.local_q_heads == 16 + assert contract.sharding.local_kv_heads == 4 + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 assert contract.reduction.acc_dtype is AttentionDType.FP32 assert contract.to_dict()["reduction"] == { "merge": "online_softmax_lse", @@ -126,8 +128,8 @@ def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): @pytest.mark.parametrize( ("field", "value", "message"), [ - ("tp_rank", 4, "tp_rank=4"), - ("cp_rank", 4, "cp_rank=4"), + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), ("global_block_indices", (), "must not be empty"), ("global_block_indices", (1, 0), "strictly increasing"), ], @@ -135,20 +137,20 @@ def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): def test_invalid_rank_and_cp_order_metadata_fail_loudly(field, value, message): values = { "tp_rank": 0, - "tp_world_size": 4, + "tp_world_size": 2, "cp_rank": 0, - "cp_world_size": 4, + "cp_world_size": 2, "global_q_heads": 32, "global_kv_heads": 8, "local_q_head_start": 0, - "local_q_heads": 8, + "local_q_heads": 16, "local_kv_head_start": 0, - "local_kv_heads": 2, + "local_kv_heads": 4, "global_sequence_length": 4096, - "local_sequence_length": 1024, + "local_sequence_length": 2048, "global_block_indices": (0,), "global_block_token_starts": (0,), - "local_block_offsets": (0, 1024), + "local_block_offsets": (0, 2048), } values[field] = value @@ -171,26 +173,26 @@ def test_sequence_range_and_packed_offsets_are_validated(): with pytest.raises(AttentionContractError, match="final packed_sequence_offsets"): _sharding(packed_sequence_offsets=(0, 512)) - sharding = _sharding(packed_sequence_offsets=(0, 256, 1024)) - assert sharding.packed_sequence_offsets == (0, 256, 1024) + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + assert sharding.packed_sequence_offsets == (0, 512, 2048) def test_non_contiguous_cp_blocks_have_explicit_global_and_local_offsets(): sharding = _sharding( - global_block_indices=(0, 7), - global_block_token_starts=(0, 3584), - local_block_offsets=(0, 512, 1024), + global_block_indices=(0, 3), + global_block_token_starts=(0, 3072), + local_block_offsets=(0, 1024, 2048), ) - assert sharding.global_block_indices == (0, 7) - assert sharding.global_block_token_starts == (0, 3584) - assert sharding.local_block_offsets == (0, 512, 1024) + assert sharding.global_block_indices == (0, 3) + assert sharding.global_block_token_starts == (0, 3072) + assert sharding.local_block_offsets == (0, 1024, 2048) with pytest.raises(AttentionContractError, match="non-overlapping and ordered"): _sharding( global_block_indices=(0, 1), - global_block_token_starts=(0, 256), - local_block_offsets=(0, 512, 1024), + global_block_token_starts=(0, 512), + local_block_offsets=(0, 1024, 2048), ) @@ -207,7 +209,7 @@ def test_causal_attention_requires_explicit_offset(): def test_full_prefill_query_length_must_match_local_sequence_length(): with pytest.raises(AttentionContractError, match="prefill query_sequence_length must equal"): - _contract(mode="prefill", query_sequence_length=2048) + _contract(mode="prefill", query_sequence_length=1024) chunked = _contract(mode="chunked_prefill", query_sequence_length=512) decode = _contract( @@ -410,7 +412,7 @@ def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): registry.get_attention_op(_contract()) message = str(exc_info.value) - assert "CP=4 is unsupported" in message + assert "CP=2 is unsupported" in message assert "attention-domain LSE export is unsupported" in message assert "deterministic CP (out, lse) merge is unsupported" in message @@ -435,7 +437,8 @@ def test_declared_compatible_backend_resolves_and_records_provenance(): assert result.provenance["requested_backend"] == "deterministic" assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" assert result.provenance["fallback"] is False - assert result.provenance["contract"]["sharding"]["cp_world_size"] == 4 + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["cp_world_size"] == 2 json.dumps(result.provenance) @@ -455,7 +458,7 @@ def test_requested_stable_backend_id_is_enforced(): def test_packed_layout_requires_declared_backend_support(): capability = replace(_declared_cp_backend(), supports_packed_varlen=False) contract = _contract( - sharding=_sharding(packed_sequence_offsets=(0, 512, 1024)), + sharding=_sharding(packed_sequence_offsets=(0, 512, 2048)), causal_offsets=(0, 0), batch_size=2, ) @@ -464,7 +467,7 @@ def test_packed_layout_requires_declared_backend_support(): def test_packed_sequence_count_must_match_logical_batch_size(): - sharding = _sharding(packed_sequence_offsets=(0, 512, 1024)) + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) with pytest.raises(AttentionContractError, match="must equal logical batch_size"): _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=1) From d39f0a50a19559d97811684f01e671720f64a448 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 2 Aug 2026 23:39:25 +0800 Subject: [PATCH 08/11] feat(attention): add rope contract metadata Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/design/ws2-cp-attention-contract.md | 41 ++++++- rl_engine/kernels/attention_contract.py | 142 +++++++++++++++++++++++ tests/test_attention_contract.py | 75 ++++++++++++ 3 files changed, 257 insertions(+), 1 deletion(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index 3e259fa8..95320101 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -29,6 +29,7 @@ belong to later work in #235. - `ShardingSpec`: TP-local head ownership and CP block-to-token ownership; - `ReductionSpec`: fixed `(out, lse)` merge semantics; - `KVCacheSpec`: decode replay cache identity; +- `RoPESpec`: Qwen3 RoPE state, position identity, and fused/unfused boundary metadata; - `AttentionBackendCapability`: the layouts and semantics a backend explicitly supports. Construction performs validation immediately. A structurally valid contract means that the @@ -49,6 +50,7 @@ context. from rl_engine.kernels.attention_contract import ( AttentionContract, ReductionSpec, + RoPESpec, ShardingSpec, ) @@ -81,6 +83,18 @@ contract = AttentionContract( causal_offsets=(0,), sharding=sharding, reduction=ReductionSpec(), + rope=RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ), ) ``` @@ -97,6 +111,29 @@ local_block_offsets=(0, 1024, 2048) This metadata is sufficient for a later implementation to restore logical global order without using ring arrival order. +## RoPE / Position Semantics + +RoPE is part of the attention contract because rollout can materialize +`RoPE+Attention` as a fused or cache-aware path while training may materialize +`RoPE -> Attention` as separate operators. PR1 does not execute the RoPE kernel, +but it records the metadata required to prove both materializations use the same +model semantics. + +`RoPESpec` records: + +- whether Q, K, and cached K are `pre_rope` or `post_rope`; +- `theta`, optional `rope_scaling`, and `rotary_dim`; +- dense `position_ids` or per-sequence `query_position_offsets` / + `key_position_offsets`; +- the RoPE cast point and output dtype; +- `fusion_boundary`, either `unfused_rope_attention` or `fused_rope_attention`. + +When RoPE metadata is present, construction validates that rotary dimensions fit +the attention head dimension and that offset metadata matches the logical batch +shape. Backends must declare RoPE support through `AttentionBackendCapability`; +a backend that cannot consume RoPE/position metadata or cannot support a fused +RoPE+Attention boundary is rejected before dispatch. + ## Reduction Semantics The only PR1 reduction contract is: @@ -160,6 +197,8 @@ provenance = result.provenance Dispatch considers only backends with an `AttentionBackendCapability`. It checks role, attention mode, dtype, TP/CP degree, LSE export, deterministic CP merge, packed varlen, and KV-cache support. +When RoPE metadata is present, dispatch also checks whether the backend explicitly supports +RoPE/position metadata and fused RoPE+Attention boundaries. An undeclared or incompatible backend is skipped with an explicit rejection reason. The current WS1 PyTorch Attention implementations support local reference math but do not export @@ -185,4 +224,4 @@ python -m pytest tests/test_attention_contract.py -q The tests include Qwen3 TP=2/CP=2 construction, GQA ownership errors, non-contiguous CP blocks, packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible -fallback, and JSON-compatible provenance. +fallback, RoPE metadata validation, and JSON-compatible provenance. diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 0f14fdea..7deeddf0 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -55,6 +55,22 @@ class ReductionEngine(str, Enum): IN_OP_REFERENCE = "in_op_reference" +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) @@ -447,6 +463,65 @@ def __post_init__(self) -> None: 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.""" @@ -462,6 +537,7 @@ class AttentionContract: sharding: ShardingSpec reduction: ReductionSpec kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None export_lse: bool = True def __post_init__(self) -> None: @@ -520,6 +596,27 @@ def __post_init__(self) -> 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( @@ -574,6 +671,32 @@ def to_dict(self) -> dict[str, Any]: "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, @@ -591,6 +714,7 @@ def to_dict(self) -> dict[str, Any]: "sharding": sharding, "reduction": reduction, "kv_cache": kv_cache, + "rope": rope, } @@ -608,6 +732,8 @@ class AttentionBackendCapability: 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 implementation_kind: str = "production" def __post_init__(self) -> None: @@ -635,6 +761,8 @@ def __post_init__(self) -> None: "deterministic_cp_merge", "supports_packed_varlen", "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", ): if not isinstance(getattr(self, field), bool): raise AttentionContractError(f"{field} must be a bool") @@ -675,6 +803,14 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: 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") return tuple(reasons) def supports(self, contract: AttentionContract) -> bool: @@ -692,6 +828,8 @@ def to_dict(self) -> dict[str, Any]: "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, "implementation_kind": self.implementation_kind, } @@ -719,5 +857,9 @@ class AttentionDispatchResult: "ReductionEngine", "ReductionOrder", "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", "ShardingSpec", ] diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index feab4c3a..b986fc15 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -19,6 +19,8 @@ AttentionRole, KVCacheSpec, ReductionSpec, + RoPEFusionBoundary, + RoPESpec, ShardingSpec, ) from rl_engine.kernels.registry import KernelRegistry, OpBackend @@ -68,6 +70,7 @@ def _contract( causal_offsets: tuple[int, ...] = (0,), batch_size: int = 1, query_sequence_length: int | None = None, + rope: RoPESpec | None = None, ) -> AttentionContract: resolved_sharding = sharding or _sharding() return AttentionContract( @@ -86,6 +89,7 @@ def _contract( sharding=resolved_sharding, reduction=ReductionSpec(), kv_cache=kv_cache, + rope=rope, ) @@ -125,6 +129,55 @@ def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): json.dumps(contract.to_dict()) +def test_rope_metadata_is_part_of_attention_contract_provenance(): + rope = RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ) + + contract = _contract(rope=rope) + payload = contract.to_dict() + + assert payload["rope"] == { + "q_state": "post_rope", + "k_state": "post_rope", + "k_cache_state": "post_rope", + "theta": 1.0e6, + "rotary_dim": 128, + "rope_scaling": None, + "position_ids": None, + "query_position_offsets": [0], + "key_position_offsets": [0], + "cast_at": "after_rope", + "output_dtype": "bf16", + "fusion_boundary": "unfused_rope_attention", + } + json.dumps(payload) + + +def test_rope_position_metadata_is_validated_against_contract_shape(): + with pytest.raises(AttentionContractError, match="rotary_dim=256"): + _contract(rope=RoPESpec(rotary_dim=256)) + + with pytest.raises(AttentionContractError, match="query_position_offsets"): + _contract(batch_size=2, causal_offsets=(0, 0), rope=RoPESpec(query_position_offsets=(0,))) + + with pytest.raises(AttentionContractError, match="position_ids"): + _contract(rope=RoPESpec(position_ids=(0, 1, 2))) + + valid = _contract(rope=RoPESpec(position_ids=tuple(range(2048)))) + assert valid.rope is not None + assert valid.rope.position_ids == tuple(range(2048)) + + @pytest.mark.parametrize( ("field", "value", "message"), [ @@ -466,6 +519,28 @@ def test_packed_layout_requires_declared_backend_support(): assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) +def test_rope_contract_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec()) + capability = _declared_cp_backend() + + assert capability.incompatibilities(contract) == ("RoPE/position metadata is unsupported",) + + supported = replace(capability, supports_rope_metadata=True) + assert supported.incompatibilities(contract) == () + + +def test_fused_rope_attention_boundary_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec(fusion_boundary=RoPEFusionBoundary.FUSED_ROPE_ATTENTION)) + capability = replace(_declared_cp_backend(), supports_rope_metadata=True) + + assert capability.incompatibilities(contract) == ( + "fused RoPE+Attention boundary is unsupported", + ) + + supported = replace(capability, supports_fused_rope_attention=True) + assert supported.incompatibilities(contract) == () + + def test_packed_sequence_count_must_match_logical_batch_size(): sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) From 0480ce81fe2cf8a543f6aa6613a6ec2cf8550ede Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 2 Aug 2026 23:40:00 +0800 Subject: [PATCH 09/11] docs(attention): clarify rope boundary for cp reference Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 13 ++++++- .../ops/pytorch/attention/cp_attention.py | 8 ++++ tests/test_cp_attention.py | 38 +++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 9aac608e..988fad8c 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -94,6 +94,12 @@ fused backend; it defines the CP/LSE merge behavior that downstream fused paths must match. Optional per-batch `query_position_offsets` / `key_position_offsets` cover varlen causal-mask metadata while keeping the dense tensor layout. +For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not +call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The +position offsets passed to CP attention must describe the same absolute token +positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary +while PR7 can later validate production fused `RoPE+Attention` kernels. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -158,8 +164,9 @@ gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. `tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard -attention, CP=2 prefill vs CP=1, chunked-prefill replay, global-position causal -masking across CP boundaries, order-independent LSE merge by global block index, +attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared +global position metadata, chunked-prefill replay, global-position causal masking +across CP boundaries, order-independent LSE merge by global block index, padding/all-masked stability, BF16 final-write behavior, input purity, argument validation, and registry dispatch. `make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill @@ -246,6 +253,8 @@ for measured peak memory at representative shapes. - Full materialization of scores/P limits practical sequence length. - `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, not a distributed runtime or fused kernel. +- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused + `RoPE+Attention` backend alignment are outside PR3. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 734d97b9..302ad73f 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -84,6 +84,12 @@ def merge_attention_partial_states( class DeterministicCPAttentionReferenceOp: """Correctness-first CP attention reference for prefill and chunked prefill. + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + The op emulates CP by splitting query and KV sequence dimensions into logical CP shards. Each query shard computes one partial attention state per KV block, then merges those states in fixed global-block order using fp32 @@ -267,6 +273,8 @@ def local_partial_state( ``query_position_offsets`` and ``key_position_offsets`` are optional per-batch-row base positions. They let the reference express varlen or packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. """ _validate_qkv(q, k, v) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index dd0f53f7..03c861dd 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -19,6 +19,7 @@ merge_attention_partial_states, ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp from rl_engine.kernels.registry import kernel_registry _N_HEADS = 32 @@ -106,6 +107,43 @@ def test_cp2_prefill_matches_cp1_reference(): torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) +def test_cp2_consumes_post_rope_qk_with_shared_global_position_metadata(): + op = DeterministicCPAttentionReferenceOp() + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 7, 7, seed=14, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([17, 103], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + + assert not torch.equal(q, pre_rope_q.float()) + assert not torch.equal(k, pre_rope_k.float()) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=1, + ) + out2, lse2 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + def test_chunked_prefill_replay_matches_unchunked_cp2(): op = DeterministicCPAttentionReferenceOp() q, k, v = _qkv(2, 10, 10, seed=3) From f2c6acb5b59f3051eea646bedc831a6b99fd5628 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 4 Aug 2026 22:20:11 +0800 Subject: [PATCH 10/11] feat(alignment): bind rollout and training attention contracts (#235 PR4) Wire the CP attention path into the cross-configuration planner/runtime for the Qwen3-8B TP=2 CP=2 BF16 target. The PR4 criterion "rollout and training descriptors bind to the same semantic attention contract" cannot hold literally: training runs full-sequence prefill over a CP-sharded sequence while rollout runs vLLM paged-KV chunked prefill, so the two AttentionContract instances always differ. Binding is therefore split into three tiers -- identity must match bit for bit, reduction semantics must match each other and the WS2 mandate, and materialization differences are recorded and measured rather than rejected. reduction.engine stays in the recorded tier so a Transformer Engine merge oracle on one side does not fail the binding; reduction.order and acc_dtype stay in the semantic tier because that is the WS2 claim. Also adds the first two framework-shaped RuntimeMaterializer implementations. Before this the only one was CpuSmokeMaterializer over a synthetic CPU model, and every named scenario was planning-only. Neither adapter imports megatron or vllm, so the binding rules run on CPU in CI. Determinism is probed on both sides and compared, because the two frameworks mean different things by it: Megatron asserts NCCL_ALGO and leaves TF32 and BF16 reduced-precision reduction unmanaged, while vLLM hard-sets ten NCCL variables and disables both. Mismatches in NCCL_ALGO, NCCL_PROTO and CUBLAS_WORKSPACE_CONFIG are blocking; the rest are recorded. Fixes a latent break on the way: the planner normalizes dtype knobs to torch spellings (bfloat16) while AttentionDType uses short ones (bf16), so passing a normalized knob into the enum raised. Stacked on #236 (attention contract) and #238 (deterministic CP reference), on top of #230 (cross-configuration framework). Part of #235 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw --- .../ws2-attention-cross-config-integration.md | 135 ++++ ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 87 +++ .../cross_config/adapters/__init__.py | 33 + .../cross_config/adapters/_common.py | 263 ++++++++ .../alignment/cross_config/adapters/knobs.py | 160 +++++ .../cross_config/adapters/megatron.py | 381 +++++++++++ .../alignment/cross_config/adapters/vllm.py | 442 +++++++++++++ .../cross_config/attention_binding.py | 528 +++++++++++++++ .../alignment/cross_config/determinism.py | 305 +++++++++ tests/test_attention_cross_config_binding.py | 612 ++++++++++++++++++ 10 files changed, 2946 insertions(+) create mode 100644 docs/design/ws2-attention-cross-config-integration.md create mode 100644 examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json create mode 100644 rl_engine/alignment/cross_config/adapters/__init__.py create mode 100644 rl_engine/alignment/cross_config/adapters/_common.py create mode 100644 rl_engine/alignment/cross_config/adapters/knobs.py create mode 100644 rl_engine/alignment/cross_config/adapters/megatron.py create mode 100644 rl_engine/alignment/cross_config/adapters/vllm.py create mode 100644 rl_engine/alignment/cross_config/attention_binding.py create mode 100644 rl_engine/alignment/cross_config/determinism.py create mode 100644 tests/test_attention_cross_config_binding.py diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md new file mode 100644 index 00000000..06966ebe --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,135 @@ +# WS2 Attention Cross-Configuration Integration + +Implements PR4 of [#235](https://github.com/RL-Align/RL-Kernel/issues/235): wiring the +CP attention path into the cross-configuration planner/runtime for the Qwen3-8B +TP=2 CP=2 BF16 target. + +Builds on [#236](https://github.com/RL-Align/RL-Kernel/pull/236) (attention contract +and dispatch metadata), [#238](https://github.com/RL-Align/RL-Kernel/pull/238) +(deterministic CP reference) and [#230](https://github.com/RL-Align/RL-Kernel/pull/230) +(cross-configuration framework). + +## What "bind to the same contract" means here + +The PR4 acceptance criteria say rollout and training descriptors must "bind to the +same semantic attention contract". Under the frozen deployment the two sides can +never produce identical `AttentionContract` instances: + +| | training (Megatron) | rollout (vLLM) | +| --- | --- | --- | +| mode | full-sequence prefill | chunked prefill, later decode | +| CP | `context_parallel_size`, whole forward | `prefill_context_parallel_size`, prefill only | +| KV | no paging | paged KV with a block table | +| backend vocabulary | `AttnBackend{flash,fused,unfused,local,auto}` | `AttentionBackendEnum` | + +Read literally, the criterion is unsatisfiable. It is therefore implemented as three +tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: + +| tier | fields | rule | failure | +| --- | --- | --- | --- | +| `IDENTICAL` | checkpoint, model version, weight version, tokenizer, token ids, active mask, position ids, padding side, pre-update state, Q/KV heads, head dim, RoPE theta/scaling/rotary dim, QK-Norm, cached global token positions, KV sequence lengths | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | `reduction.merge`, `reduction.acc_dtype`, `reduction.order`, `reduction.downcast_at`, `export_lse`, cross-side determinism mode | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging, CP/TP world sizes, local sequence length | free to differ | none; recorded into provenance and measured | + +Two placements are load-bearing: + +* **`reduction.engine` is `RECORDED`, not `SEMANTIC`.** Training may run the in-op + deterministic reference while rollout runs a Transformer Engine merge oracle. + Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 + depend on. +* **`reduction.order` and `reduction.acc_dtype` are `SEMANTIC`.** This is the entire + WS2 claim: merge order and accumulation precision are decided by the contract, not + by whichever backend happens to be selected. + +`comparable` and `passed` are separate flags. A pair with mismatched identity is not +comparable. A pair that is comparable but violates the reduction mandate is still +rejected -- the drift would be real but attributable to the wrong thing. + +## Determinism is not one thing + +`rl_engine/alignment/cross_config/determinism.py` probes both sides and compares +them, because the two frameworks mean different things by "deterministic": + +| | Megatron `deterministic_mode` | vLLM `VLLM_BATCH_INVARIANT` | +| --- | --- | --- | +| `NCCL_ALGO` | asserts membership in a five-value set | hard-sets `allreduce:tree` | +| `NCCL_PROTO`, channels, threads | not managed | hard-set (`Simple`, `1`, `1`) | +| TF32 | **not managed at all** | disabled (`fp32_precision="ieee"`) | +| BF16 reduced-precision reduction | not managed | disabled | +| cuBLAS workspace / BLAS library | not managed | `:4096:8`, cuBLASLt | +| GEMM | cuBLAS / TE | Triton `matmul_persistent` | +| FlashAttention | forbidden | permitted | + +`NCCL_ALGO`, `NCCL_PROTO` and `CUBLAS_WORKSPACE_CONFIG` change arithmetic, so a +mismatch there is blocking. The remaining differences -- including the TF32 and +BF16-reduction asymmetry, which under a pure BF16 GEMM path does not fire -- are +recorded so the asymmetry appears in every artifact rather than being invisible. + +## Runtime adapters + +Before this PR the only `RuntimeMaterializer` in the repository was +`CpuSmokeMaterializer` over a synthetic CPU model, and every named scenario +(`S1`/`S2`/`S3`) was planning-only. This PR adds the first two framework-shaped +adapters: + +* `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and + distributed-context fingerprints, determinism probe, frozen-scope assertions) and + `MegatronAttentionMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (adds `kv_page_size` from + `CacheConfig.block_size` and `split_kv_policy` from + `AttentionConfig.flash_attn_max_num_splits_for_cuda_graph`) and + `VllmRolloutMaterializer`. + +Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding +rules are exercised on CPU in CI rather than only on a 2-node cluster. + +## Fail closed, never substitute + +`unsupported_reduction_reason` rejects requests that #236 cannot express, instead of +collapsing them onto the supported value: + +| request | status | why | +| --- | --- | --- | +| `attention.reduction_order=arrival` | `UNSUPPORTED` | the control group must stay distinguishable from the treatment | +| `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | +| `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | +| `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | +| `rollout.context_parallel_size>1` with `mode=decode` | `FALLBACK` | vLLM CP covers prefill only; recorded with the reason | + +## Knobs + +`adapters/knobs.py` extends `V1_KNOBS` additively. Added: training-side +`tensor_parallel_size` / `context_parallel_size` / `deterministic_mode` / +`cp_comm_type`, `rollout.batch_invariant` / `rollout.kv_block_size`, and the +reduction axis (`acc_dtype`, `order`, `downcast_at`, `engine`) plus +`attention.fusion_boundary` and `attention.split_kv_policy`. + +`training.attention_backend` keeps its path but its value domain is replaced with +Megatron's `AttnBackend`; the HuggingFace names have no Megatron counterpart, so this +is a replacement rather than a mapping. + +Not done here, because both change `V1_KNOBS` itself and would break existing +cross-config tests: removing `training.sharding` (Megatron has no such concept, and +DP=1 makes it moot) and renaming `rollout.context_parallel_size` to reflect that it +binds to `prefill_context_parallel_size`. + +## Scenario + +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json` supersedes +`cross_config_s1_distributed_smoke.json` and +`cross_config_s3_qwen3_8b_tp4_cp4_bf16.json`, whose training sides used `sdpa` / +`flash_attention_2` and `sharding: fsdp` -- none of which exist under Megatron -- and +whose TP=4/CP=4 topology does not match the target. +`cross_config_s2_vllm_tp_vs_fsdp.json` has no Megatron-only counterpart and should be +retired rather than rewritten. + +## Out of scope + +Deliberately not in this PR: + +* launching `torchrun`, initializing process groups, or executing attention; +* decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 + and is refused with that reference rather than stubbed; +* Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json new file mode 100644 index 00000000..e74d3e6c --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,87 @@ +{ + "experiment_id": "ws2-qwen3-8b-attention-tp2-cp2", + "scenario_id": "qwen3_8b_megatron_tp2_cp2_vllm", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "scenario": { + "issue": "https://github.com/RL-Align/RL-Kernel/issues/235", + "pull_request": "PR4 -- cross-config integration", + "model": "Qwen3-8B dense", + "training_framework": "megatron", + "rollout_framework": "vllm", + "topology": "2 nodes x 2 GPUs, TP=2 CP=2 PP=1 DP=1, BF16, SM90", + "notes": [ + "Supersedes cross_config_s1_distributed_smoke.json and", + "cross_config_s3_qwen3_8b_tp4_cp4_bf16.json, whose training side used", + "HuggingFace attention backends and FSDP sharding. Neither exists in", + "Megatron, and DP=1 makes the sharding knob meaningless.", + "cross_config_s2_vllm_tp_vs_fsdp.json has no Megatron-only counterpart at", + "all and should be retired rather than rewritten.", + "rollout.context_parallel_size binds to vLLM", + "ParallelConfig.prefill_context_parallel_size and therefore applies to", + "prefill only; a decode-mode contract runs at CP=1." + ] + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": false, + "enforce_eager": true, + "batch_invariant": true, + "kv_block_size": 16 + }, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "attention_backend": "unfused", + "compute_dtype": "bfloat16", + "deterministic_mode": true, + "cp_comm_type": "p2p", + "sharding": "unsharded" + }, + "attention": { + "reduction_acc_dtype": "fp32", + "reduction_order": "global_block_index", + "reduction_downcast_at": "final_write", + "reduction_engine": "in_op_reference", + "fusion_boundary": "unfused_rope_attention", + "split_kv_policy": 32 + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "training.context_parallel_size", + "values": [1, 2] + }, + { + "path": "training.tensor_parallel_size", + "values": [1, 2] + }, + { + "path": "attention.fusion_boundary", + "values": ["unfused_rope_attention", "fused_rope_attention"] + }, + { + "path": "training.cp_comm_type", + "values": ["p2p", "all_gather"] + }, + { + "path": "attention.reduction_order", + "values": ["global_block_index", "arrival"] + }, + { + "path": "attention.reduction_acc_dtype", + "values": ["fp32", "bf16"] + } + ] +} diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..c2db2134 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" + +from rl_engine.alignment.cross_config.adapters._common import QWEN3_8B, Qwen3ModelSpec +from rl_engine.alignment.cross_config.adapters.knobs import ( + MEGATRON_ATTENTION_BACKENDS, + WS2_ATTENTION_KNOB_DESCRIPTORS, + WS2_ATTENTION_KNOBS, + WS2_ATTENTION_NORMALIZERS, +) +from rl_engine.alignment.cross_config.adapters.megatron import ( + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, +) +from rl_engine.alignment.cross_config.adapters.vllm import ( + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", + "QWEN3_8B", + "Qwen3ModelSpec", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py new file mode 100644 index 00000000..33b4d91c --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared pieces for the Megatron and vLLM WS2 attention adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from rl_engine.alignment.cross_config.runtime import KnobApplication +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) +from rl_engine.kernels.attention_contract import ( + AttentionDType, + AttentionMerge, + DowncastPoint, + ReductionEngine, + ReductionOrder, + ReductionSpec, + ShardingSpec, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "unsupported_reduction_reason", +] + + +@dataclass(frozen=True) +class Qwen3ModelSpec: + """Architecture constants for the frozen dense target. + + These are *not* knobs. #235/#239/#241 all fix Qwen3-8B dense, so they belong to + the scenario, and both sides must agree on them or the comparison is void. + """ + + name: str = "qwen3-8b" + hidden_size: int = 4096 + ffn_hidden_size: int = 12288 + num_layers: int = 36 + q_heads: int = 32 + kv_heads: int = 8 + head_dim: int = 128 + real_vocab_size: int = 151936 + rope_theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + qk_layernorm: bool = True + + def identity_fields(self) -> dict[str, Any]: + """The subset of :data:`IDENTITY_FIELDS` this spec is responsible for.""" + + return { + "q_heads": self.q_heads, + "kv_heads": self.kv_heads, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + "rope_scaling": self.rope_scaling, + "rotary_dim": self.rotary_dim, + "qk_layernorm": self.qk_layernorm, + } + + +QWEN3_8B = Qwen3ModelSpec() + + +#: The planner normalizes dtype knobs to torch spellings (``bfloat16``), while +#: :class:`AttentionDType` uses short spellings (``bf16``). Passing a normalized knob +#: straight into the enum raises, so every adapter must translate here rather than +#: each inventing its own mapping. +_DTYPE_ALIASES: Mapping[str, AttentionDType] = { + "bf16": AttentionDType.BF16, + "bfloat16": AttentionDType.BF16, + "fp16": AttentionDType.FP16, + "float16": AttentionDType.FP16, + "half": AttentionDType.FP16, + "fp32": AttentionDType.FP32, + "float32": AttentionDType.FP32, + "float": AttentionDType.FP32, +} + + +def attention_dtype(value: Any, *, field: str) -> AttentionDType: + """Translate a normalized knob dtype into an :class:`AttentionDType`.""" + + if isinstance(value, AttentionDType): + return value + key = str(value).strip().lower().replace("torch.", "") + try: + return _DTYPE_ALIASES[key] + except KeyError as exc: + raise ValueError( + f"{field}={value!r} is not a supported attention dtype; " + f"expected one of {sorted(set(_DTYPE_ALIASES))}" + ) from exc + + +def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + """Flatten nested knob mappings into dotted paths.""" + + flat: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}{key}" + if isinstance(child, Mapping): + flat.update(flatten(child, f"{path}.")) + else: + flat[path] = child + return flat + + +def application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, + **evidence: Any, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason, **evidence}, + critical=descriptor.critical, + ) + + +def unsupported_reduction_reason(flat: Mapping[str, Any]) -> str | None: + """Return why the requested reduction cannot be materialized, if it cannot. + + #236 declares single-member enums for merge order, downcast point and reduction + engine, so the alternative knob values exist only as control groups. Requesting + one must fail loudly rather than quietly collapse onto the supported value -- + silently substituting ``global_block_index`` for a requested ``arrival`` would + make the control group indistinguishable from the treatment. + """ + + order = flat.get("attention.reduction_order") + if order is not None and order != ReductionOrder.GLOBAL_BLOCK_INDEX.value: + return ( + f"attention.reduction_order={order!r} has no backend; #236 ReductionOrder " + f"declares only {ReductionOrder.GLOBAL_BLOCK_INDEX.value!r}" + ) + downcast = flat.get("attention.reduction_downcast_at") + if downcast is not None and downcast != DowncastPoint.FINAL_WRITE.value: + return ( + f"attention.reduction_downcast_at={downcast!r} has no backend; #236 " + f"DowncastPoint declares only {DowncastPoint.FINAL_WRITE.value!r}" + ) + engine = flat.get("attention.reduction_engine") + if engine is not None and engine != ReductionEngine.IN_OP_REFERENCE.value: + return ( + f"attention.reduction_engine={engine!r} has no backend; the Transformer " + "Engine merge oracle lands in #235 PR2/PR3, not here" + ) + acc_dtype = flat.get("attention.reduction_acc_dtype") + if ( + acc_dtype is not None + and attention_dtype(acc_dtype, field="attention.reduction_acc_dtype") + is not AttentionDType.FP32 + ): + return ( + f"attention.reduction_acc_dtype={acc_dtype!r} violates the WS2 mandate; " + "the CP (out, lse) merge accumulates in fp32" + ) + return None + + +def build_reduction_spec(flat: Mapping[str, Any]) -> ReductionSpec: + """Build the reduction spec, having already rejected unsupported requests.""" + + return ReductionSpec( + merge=AttentionMerge.ONLINE_SOFTMAX_LSE, + acc_dtype=AttentionDType.FP32, + order=ReductionOrder.GLOBAL_BLOCK_INDEX, + downcast_at=DowncastPoint.FINAL_WRITE, + engine=ReductionEngine.IN_OP_REFERENCE, + ) + + +def build_sharding_spec( + *, + model: Qwen3ModelSpec, + tp_rank: int, + tp_world_size: int, + cp_rank: int, + cp_world_size: int, + global_sequence_length: int, +) -> ShardingSpec: + """Build a CP/TP sharding spec for one rank of the frozen layout. + + TP splits heads, CP splits the sequence. The #239 rank layout fixes + ``rank = cp_rank * tp_world_size + tp_rank`` for a 2-node x 2-GPU deployment, + but nothing here depends on that mapping: ownership is derived from the ranks + themselves so the same builder serves CP=1 baselines. + """ + + if model.q_heads % tp_world_size or model.kv_heads % tp_world_size: + raise ValueError( + f"Qwen3 GQA heads ({model.q_heads}/{model.kv_heads}) must divide evenly " + f"across tp_world_size={tp_world_size}" + ) + if global_sequence_length % cp_world_size: + raise ValueError( + f"global_sequence_length={global_sequence_length} must divide evenly " + f"across cp_world_size={cp_world_size}" + ) + + local_q_heads = model.q_heads // tp_world_size + local_kv_heads = model.kv_heads // tp_world_size + local_sequence_length = global_sequence_length // cp_world_size + + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=model.q_heads, + global_kv_heads=model.kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + # One contiguous CP block per rank. The merge order key is the global block + # index, never the arrival order of the CP exchange. + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * local_sequence_length,), + local_block_offsets=(0, local_sequence_length), + ) + + +def causal_offsets_for(sharding: ShardingSpec, batch_size: int) -> tuple[int, ...]: + """Causal offsets for one CP shard, one entry per batch entry. + + Under CP the local query block does not start at global position zero, so the + causal mask has to be shifted by the number of preceding global tokens. Taking + that from ``global_block_token_starts`` rather than recomputing + ``cp_rank * local_sequence_length`` keeps uneven CP splits correct. + """ + + offset = sharding.global_block_token_starts[0] + return (offset,) * batch_size + + +_PROCESS_SCOPES = (IsolationScope.PROCESS, IsolationScope.DISTRIBUTED_CONTEXT) diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py new file mode 100644 index 00000000..c1cf0b84 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 attention knobs for the Qwen3-8B TP=2 CP=2 Megatron + vLLM target. + +``V1_KNOBS`` was written against a HuggingFace/FSDP rollout-vs-training pair. Three +of its entries do not survive contact with the frozen Megatron + vLLM target: + +* ``training.sharding`` takes ``unsharded``/``fsdp``, neither of which exists in + Megatron, and is meaningless at DP=1 anyway; +* ``training.attention_backend`` takes HuggingFace names + (``flash_attention_2``/``sdpa``/``eager``/``model_default``) while Megatron's + ``AttnBackend`` is ``flash``/``fused``/``unfused``/``local``/``auto``; +* there is no training-side ``tensor_parallel_size`` or ``context_parallel_size`` + at all, so the target configuration cannot even be expressed. + +This module is deliberately **additive**: it extends ``V1_KNOBS`` rather than +editing it, and overrides only the normalizer for ``training.attention_backend``. +Deleting the two dead knobs changes ``V1_KNOBS`` itself and would break existing +cross-config tests, so it is left to a follow-up on the framework PR. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + Normalizer, + _normalize_choice, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.schema import IsolationScope, KnobDescriptor + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] + + +#: ``megatron.core.transformer.enums.AttnBackend``. +MEGATRON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "flash", + "fused", + "unfused", + "local", + "auto", +) + + +WS2_ATTENTION_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + # -- training-side parallelism: the target configuration itself ------------ + KnobDescriptor( + "training.tensor_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "training.context_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + # -- determinism switches, one per framework ------------------------------ + KnobDescriptor( + "training.deterministic_mode", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "rollout.batch_invariant", + IsolationScope.PROCESS, + ("rollout",), + ), + # -- reduction knobs: the "turn the noise sources on and off" axis --------- + # These are what make drift attributable. ``reduction.order=arrival`` in + # particular is a control group, not a supported production value. + KnobDescriptor( + "attention.reduction_acc_dtype", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("fp32", "bf16"), + ), + KnobDescriptor( + "attention.reduction_order", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("global_block_index", "arrival"), + ), + KnobDescriptor( + "attention.reduction_downcast_at", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("final_write", "per_block"), + ), + KnobDescriptor( + "attention.reduction_engine", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("in_op_reference", "te_oracle"), + ), + # -- materialization knobs: differences the experiment measures ------------ + KnobDescriptor( + "attention.fusion_boundary", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("unfused_rope_attention", "fused_rope_attention"), + ), + KnobDescriptor( + # vLLM: AttentionConfig.flash_attn_max_num_splits_for_cuda_graph + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + KnobDescriptor( + # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size + "rollout.kv_block_size", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + # The CP communication group cannot be reconfigured once built; it is bound to + # the distributed context, not merely to engine construction. + KnobDescriptor( + "training.cp_comm_type", + IsolationScope.DISTRIBUTED_CONTEXT, + ("training",), + allowed_values=("p2p", "all_gather", "a2a", "a2a+p2p"), + ), +) + + +WS2_ATTENTION_KNOBS: Mapping[str, KnobDescriptor] = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in WS2_ATTENTION_KNOB_DESCRIPTORS}, +} + + +WS2_ATTENTION_NORMALIZERS: Mapping[str, Normalizer] = { + **_NORMALIZERS, + # Replace, not map: the HuggingFace names have no Megatron counterpart. + "training.attention_backend": _normalize_choice(*MEGATRON_ATTENTION_BACKENDS), + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.deterministic_mode": _strict_bool, + "rollout.batch_invariant": _strict_bool, + # AttentionDType values, not torch dtype names -- these feed ReductionSpec directly. + "attention.reduction_acc_dtype": _normalize_choice("fp32", "bf16"), + "attention.reduction_order": _normalize_choice("global_block_index", "arrival"), + "attention.reduction_downcast_at": _normalize_choice("final_write", "per_block"), + "attention.reduction_engine": _normalize_choice("in_op_reference", "te_oracle"), + "attention.fusion_boundary": _normalize_choice( + "unfused_rope_attention", "fused_rope_attention" + ), + "attention.split_kv_policy": _positive_int, + "rollout.kv_block_size": _positive_int, + "training.cp_comm_type": _normalize_choice("p2p", "all_gather", "a2a", "a2a+p2p"), +} diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py new file mode 100644 index 00000000..8b96b159 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-side (Megatron) runtime adapter for WS2 attention cross-config. + +Two things live here: + +``MegatronProvenanceAdapter`` + Read-only. Turns a Megatron config object into the construction and + distributed-context fingerprints the cross-config framework already expects, + plus the determinism probe. It never imports ``megatron`` -- every accessor is + duck-typed -- so this module is importable and testable on a laptop. + +``MegatronAttentionMaterializer`` + Implements the ``RuntimeMaterializer`` protocol. Before this PR the only + implementation in the repository was ``CpuSmokeMaterializer`` over a synthetic + CPU model, so nothing had ever materialized a real distributed runtime. + +Scope boundary: materialization builds and validates the training-side +:class:`AttentionContract` and reports what would be constructed. It does not +launch ``torchrun``, initialize process groups, or execute attention. Binding a +constructed Megatron model to this contract is the next step and needs the 2-node +x 2-GPU environment that #239 fixes. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import ( + DeterminismProbe, + megatron_probe_from_config, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "MEGATRON_CONSTRUCTION_KEYS", + "MEGATRON_DISTRIBUTED_KEYS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", +] + + +#: ``TransformerConfig`` fields that change attention arithmetic. Hashed into the +#: construction fingerprint. Deliberately excludes MoE, Mamba, MLA and sparse +#: attention fields: the frozen target is Qwen3-8B dense, and those are asserted +#: off rather than recorded. +MEGATRON_CONSTRUCTION_KEYS: tuple[str, ...] = ( + "attention_backend", + "attention_softmax_in_fp32", + "apply_query_key_layer_scaling", + "apply_rope_fusion", + "masked_softmax_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "gradient_accumulation_fusion", + "cross_entropy_loss_fusion", + "cross_entropy_fusion_impl", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "rotary_base", + "rotary_percent", + "rotary_interleaved", + "rotary_scaling_factor", + "qk_layernorm", + "hidden_dropout", + "attention_dropout", + "params_dtype", + "bf16", + "fp16", + "fp8", + "deterministic_mode", +) + + +#: ``ModelParallelConfig`` fields that define the distributed context. +MEGATRON_DISTRIBUTED_KEYS: tuple[str, ...] = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "cp_comm_type", + "tp_comm_overlap", + "use_te_rng_tracker", +) + + +#: Fields that must hold these values for the frozen dense target. A mismatch is a +#: hard stop, not a recorded difference -- see the exclusion list in the WS2 scope. +MEGATRON_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "sequence_parallel": False, + "fp8": None, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class MegatronProvenanceAdapter: + """Extract fingerprints and determinism evidence from a Megatron config. + + ``config`` may be a real ``TransformerConfig``/``ModelParallelConfig``, a merged + namespace, or a test double. Missing attributes read as ``None`` and are + recorded as such rather than raising: an absent field is itself provenance. + """ + + framework = "megatron" + + def __init__(self, config: Any, *, env: Optional[Mapping[str, str]] = None): + self.config = config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_CONSTRUCTION_KEYS} + + def distributed_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_DISTRIBUTED_KEYS} + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return megatron_probe_from_config(self.config, env=self.env) + + def frozen_scope_violations(self) -> tuple[str, ...]: + """Return the frozen-scope assertions this config violates.""" + + violations: list[str] = [] + for name, expected in MEGATRON_FROZEN_ASSERTIONS.items(): + actual = _value(self.config, name) + if actual is None: + # Not declared. Treated as unknown rather than as satisfied, because + # a silently-absent MoE or FP8 setting is exactly the case that would + # otherwise slip past a dense-only claim. + violations.append(f"{name} is not declared (expected {expected!r})") + elif actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "determinism": self.determinism_probe().to_dict(), + } + + +class MegatronAttentionMaterializer: + """Materialize the training-side attention runtime for the WS2 target.""" + + runtime_kind = "megatron_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + backend_id: str = "rlkernel.cp_attention_reference", + provenance: Optional[MegatronProvenanceAdapter] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.backend_id = backend_id + self.provenance = provenance + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + """Build the training-side contract. Raises on an unusable request.""" + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=attention_dtype( + flat.get("training.compute_dtype", "bf16"), field="training.compute_dtype" + ), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "training" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the training-side attention contract", + frozen_scope_violations=list(scope_violations), + ) + ) + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "megatron", + "attention_backend": flat.get("training.attention_backend"), + "compute_dtype": flat.get("training.compute_dtype"), + "deterministic_mode": flat.get("training.deterministic_mode"), + "cp_comm_type": flat.get("training.cp_comm_type"), + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side_config, "rollout": {}}, + topology={ + "training": { + "world_size": tp_world_size * cp_world_size, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": cp_world_size, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "rollout": {"world_size": 1}, + }, + scorer={ + "mode": "teacher_forcing", + "framework": "megatron", + "export_lse": True, + }, + operator_backends={ + "training": self.backend_id, + "rollout": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py new file mode 100644 index 00000000..a8de61c9 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rollout-side (vLLM) runtime adapter for WS2 attention cross-config. + +Mirrors :mod:`.megatron`, with three differences that come straight from what vLLM +actually is: + +* vLLM's context parallelism is ``prefill_context_parallel_size`` -- it applies to + prefill only, so a decode-mode contract must declare ``cp_world_size == 1`` + regardless of what the prefill knob says. +* ``CacheConfig.block_size`` is the paged-KV page size, and it feeds + ``KVCacheSpec.page_size`` directly rather than being invented here. +* Determinism comes from the ``VLLM_BATCH_INVARIANT`` environment variable rather + than from a config field, because vLLM applies it inside + ``init_batch_invariance()`` at worker startup. + +Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so +the module is importable anywhere. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "VLLM_ATTENTION_KEYS", + "VLLM_CACHE_KEYS", + "VLLM_FROZEN_ASSERTIONS", + "VLLM_MODEL_KEYS", + "VLLM_PARALLEL_KEYS", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", +] + + +VLLM_MODEL_KEYS: tuple[str, ...] = ( + "dtype", + "seed", + "quantization", + "enforce_eager", + "max_logprobs", + "disable_cascade_attn", + "max_model_len", +) + +VLLM_CACHE_KEYS: tuple[str, ...] = ( + "block_size", + "cache_dtype", + "enable_prefix_caching", + "prefix_caching_hash_algo", + "calculate_kv_scales", + "sliding_window", +) + +VLLM_ATTENTION_KEYS: tuple[str, ...] = ( + "backend", + "flash_attn_version", + "use_prefill_decode_attention", + "flash_attn_max_num_splits_for_cuda_graph", + "use_cudnn_prefill", + "disable_flashinfer_prefill", + "use_non_causal", +) + +VLLM_PARALLEL_KEYS: tuple[str, ...] = ( + "tensor_parallel_size", + "pipeline_parallel_size", + "prefill_context_parallel_size", + "data_parallel_size", +) + + +#: Frozen dense-target assertions on the rollout side. ``cache_dtype`` must stay +#: ``auto`` because an FP8 KV cache is a representation-drift problem tracked +#: separately, and ``disable_cascade_attn`` must stay ``True`` because cascade +#: attention changes the block-merge structure the contract pins down. +VLLM_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "quantization": None, + "cache_dtype": "auto", + "calculate_kv_scales": False, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "sliding_window": None, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class VllmProvenanceAdapter: + """Extract fingerprints and determinism evidence from vLLM configs.""" + + framework = "vllm" + + def __init__( + self, + *, + model_config: Any = None, + cache_config: Any = None, + attention_config: Any = None, + parallel_config: Any = None, + env: Optional[Mapping[str, str]] = None, + ): + self.model_config = model_config + self.cache_config = cache_config + self.attention_config = attention_config + self.parallel_config = parallel_config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + view: dict[str, Any] = {} + for prefix, config, keys in ( + ("model", self.model_config, VLLM_MODEL_KEYS), + ("cache", self.cache_config, VLLM_CACHE_KEYS), + ("attention", self.attention_config, VLLM_ATTENTION_KEYS), + ): + for name in keys: + view[f"{prefix}.{name}"] = _value(config, name) + return view + + def distributed_view(self) -> dict[str, Any]: + return { + f"parallel.{name}": _value(self.parallel_config, name) for name in VLLM_PARALLEL_KEYS + } + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return vllm_probe_from_env(self.env, model_config=self.model_config) + + def frozen_scope_violations(self) -> tuple[str, ...]: + sources = { + "quantization": self.model_config, + "disable_cascade_attn": self.model_config, + "cache_dtype": self.cache_config, + "calculate_kv_scales": self.cache_config, + "sliding_window": self.cache_config, + "pipeline_parallel_size": self.parallel_config, + "data_parallel_size": self.parallel_config, + } + violations: list[str] = [] + for name, expected in VLLM_FROZEN_ASSERTIONS.items(): + config = sources.get(name) + if config is None: + continue + actual = _value(config, name) + if actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + @property + def kv_page_size(self) -> Optional[int]: + """vLLM's paged-KV block size, which is the contract's ``page_size``.""" + + block_size = _value(self.cache_config, "block_size") + return int(block_size) if block_size is not None else None + + @property + def split_kv_policy(self) -> Optional[int]: + """The split-KV knob #235 PR5/PR7 needs; #236 has no field for it yet.""" + + splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") + return int(splits) if splits is not None else None + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "kv_page_size": self.kv_page_size, + "split_kv_policy": self.split_kv_policy, + "determinism": self.determinism_probe().to_dict(), + } + + +class VllmRolloutMaterializer: + """Materialize the rollout-side attention runtime for the WS2 target.""" + + runtime_kind = "vllm_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, + backend_id: str = "vllm.flash_attn", + provenance: Optional[VllmProvenanceAdapter] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.mode = mode + self.backend_id = backend_id + self.provenance = provenance + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def effective_cp_world_size(self, flat: Mapping[str, Any]) -> int: + """CP applies to prefill only; decode always runs at CP=1.""" + + requested = int(flat.get("rollout.context_parallel_size", 1)) + if self.mode is AttentionMode.DECODE: + return 1 + return requested + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + if self.mode is AttentionMode.DECODE: + # Decode replay needs a validated KVCacheSpec (cache positions, page + # ownership, prefix-cache identity). That is #235 PR6's contract surface, + # and inventing a placeholder here would let an unvalidated decode case + # look bound. Fail instead. + raise AttentionContractError( + "decode-mode materialization requires KV-cache identity from #235 PR6; " + "this adapter covers prefill and chunked prefill" + ) + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + cp_world_size = self.effective_cp_world_size(flat) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank if cp_world_size > 1 else 0, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.FUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + # vLLM stores post-RoPE K in the cache; recorded, not asserted equal to + # the training side, because it is a materialization fact. + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.INFER, + mode=self.mode, + dtype=attention_dtype(flat.get("rollout.dtype", "bf16"), field="rollout.dtype"), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + requested_cp = int(flat.get("rollout.context_parallel_size", 1)) + effective_cp = self.effective_cp_world_size(flat) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "rollout" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if path == "rollout.context_parallel_size" and effective_cp != requested_cp: + applications.append( + application( + descriptor, + requested, + effective_cp, + effective_cp, + MaterializationStatus.FALLBACK, + ( + "vLLM context parallelism covers prefill only; a decode-mode " + f"contract runs at cp_world_size=1, not {requested_cp}" + ), + vllm_field="ParallelConfig.prefill_context_parallel_size", + ) + ) + continue + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the rollout-side attention contract", + frozen_scope_violations=list(scope_violations), + ) + ) + + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "vllm", + "dtype": flat.get("rollout.dtype"), + "enforce_eager": flat.get("rollout.enforce_eager"), + "enable_prefix_caching": flat.get("rollout.enable_prefix_caching"), + "batch_invariant": flat.get("rollout.batch_invariant"), + "kv_block_size": flat.get("rollout.kv_block_size"), + "split_kv_policy": flat.get("attention.split_kv_policy"), + "attention_mode": self.mode.value, + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"rollout": side_config, "training": {}}, + topology={ + "rollout": { + "world_size": tp_world_size * effective_cp, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": effective_cp, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "training": {"world_size": 1}, + }, + scorer={ + "mode": "rollout_logprob", + "framework": "vllm", + "export_lse": True, + }, + operator_backends={ + "rollout": self.backend_id, + "training": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py new file mode 100644 index 00000000..fcd4c203 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Three-tier binding between rollout-side and training-side attention contracts. + +Issue #235 PR4 requires that "rollout and training descriptors bind to the same +semantic attention contract". Under the frozen Megatron + vLLM deployment the two +sides can never produce *identical* :class:`AttentionContract` instances: training +runs full-sequence prefill over a CP-sharded sequence, while rollout runs vLLM +paged-KV chunked prefill and decode. Taking "same contract" literally would make +the target configuration permanently unbindable. + +This module therefore splits binding into three tiers: + +``IDENTICAL`` + Logical identity. Both sides must agree bit for bit, otherwise the pair is not + comparable at all and no drift number from it means anything. + +``SEMANTIC`` + The WS2 numerical claim: merge semantics, accumulation dtype, reduction order + and downcast point are decided by the contract, not by the implementation. + Both sides must carry the same values *and* those values must match the WS2 + mandate, otherwise the comparison fails closed. + +``RECORDED`` + Materialization facts that the two sides are expected to differ on -- attention + mode, RoPE fusion boundary, KV-cache paging, backend id, reduction engine. These + differences are exactly what the experiment measures, so they are recorded into + provenance rather than rejected. + +Deliberately *not* in ``SEMANTIC``: ``engine``. Training may run the in-op +deterministic reference while rollout runs a Transformer Engine merge oracle; forcing +those equal would defeat the purpose of the oracle comparison in #235 PR2/3/5/6. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, +) + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "first_blocking_issue", + "identity_fingerprint", + "summarize_binding", +] + + +class AttentionBindingError(ValueError): + """Raised when a caller supplies structurally unusable binding inputs.""" + + +class BindingTier(str, Enum): + """Which rule a field is governed by.""" + + IDENTICAL = "identical" + SEMANTIC = "semantic" + RECORDED = "recorded" + + +class BindingErrorCode(str, Enum): + """Stable, machine-readable reasons a binding is rejected. + + Callers branch on these; they are part of the artifact schema and must not be + renamed without a schema version bump. + """ + + IDENTITY_MISSING = "IDENTITY_MISSING" + IDENTITY_MISMATCH = "IDENTITY_MISMATCH" + REDUCTION_SEMANTIC_MISMATCH = "REDUCTION_SEMANTIC_MISMATCH" + REDUCTION_MANDATE_VIOLATION = "REDUCTION_MANDATE_VIOLATION" + LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" + ROLE_COLLISION = "ROLE_COLLISION" + DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + + +#: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). +#: Recorded explicitly so a future ``LogprobContract`` binding cannot be confused +#: with this one purely because both set ``export_lse=True``. +ATTENTION_LSE_DOMAIN = "attention" + + +#: Fields both sides must agree on bit for bit before any comparison is meaningful. +#: Sourced from #235 "Numerical Contract" preconditions plus the vime-owned rollout +#: provenance (weight version, sampling, padding) that the issue assumes but does +#: not enumerate. +IDENTITY_FIELDS: tuple[str, ...] = ( + "checkpoint_id", + "model_version", + "weight_version", + "tokenizer_fingerprint", + "token_ids_fingerprint", + "active_mask_fingerprint", + "position_ids_fingerprint", + "padding_side", + "pre_update_state", + # model semantics that decide what attention *means* + "q_heads", + "kv_heads", + "head_dim", + "rope_theta", + "rope_scaling", + "rotary_dim", + "qk_layernorm", + # decode replay identity (#235 PR6) + "global_token_positions_fingerprint", + "kv_seq_lens_fingerprint", +) + + +#: Reduction fields that decide the numerical result. Both sides must carry the +#: same value, and that value must satisfy :data:`WS2_ATTENTION_REDUCTION_MANDATE`. +SEMANTIC_REDUCTION_FIELDS: tuple[str, ...] = ( + "merge", + "acc_dtype", + "order", + "downcast_at", +) + + +#: The WS2 mandate itself. ``#236`` currently declares single-member enums for +#: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; +#: they are written out anyway so that widening any of those enums later fails here +#: instead of silently admitting a non-conforming backend. +WS2_ATTENTION_REDUCTION_MANDATE: Mapping[str, str] = { + "merge": AttentionMerge.ONLINE_SOFTMAX_LSE.value, + "acc_dtype": AttentionDType.FP32.value, + "order": ReductionOrder.GLOBAL_BLOCK_INDEX.value, + "downcast_at": DowncastPoint.FINAL_WRITE.value, +} + + +#: Materialization facts the two sides are expected to differ on. Recorded into +#: provenance; never a rejection reason. +RECORDED_FIELDS: tuple[str, ...] = ( + "mode", + "backend_id", + "reduction.engine", + "rope.fusion_boundary", + "rope.q_state", + "rope.k_state", + "rope.k_cache_state", + "rope.cast_at", + "rope.output_dtype", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", + "sharding.cp_world_size", + "sharding.tp_world_size", + "sharding.local_sequence_length", +) + + +@dataclass(frozen=True) +class BindingIssue: + """One reason a binding is not comparable or not admissible.""" + + code: BindingErrorCode + tier: BindingTier + field: str + rollout: Any = None + training: Any = None + message: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code.value, + "tier": self.tier.value, + "field": self.field, + "rollout": self.rollout, + "training": self.training, + "message": self.message, + } + + +@dataclass(frozen=True) +class AttentionBindingResult: + """Outcome of binding one rollout contract to one training contract. + + ``comparable`` and ``passed`` are deliberately separate. A pair whose identity + does not match is *not comparable* -- reporting a drift number for it would be + meaningless. A pair that is comparable but violates the reduction mandate *is* + comparable yet must still fail closed, because the whole WS2 claim is that + reduction order and accumulation precision come from the contract. + """ + + comparable: bool + passed: bool + issues: tuple[BindingIssue, ...] = () + identity_fingerprint: str = "" + reduction_fingerprint: str = "" + binding_fingerprint: str = "" + recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + provenance: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.attention_binding.v1" + + def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: + return tuple(issue for issue in self.issues if issue.code is code) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "comparable": self.comparable, + "passed": self.passed, + "issues": [issue.to_dict() for issue in self.issues], + "identity_fingerprint": self.identity_fingerprint, + "reduction_fingerprint": self.reduction_fingerprint, + "binding_fingerprint": self.binding_fingerprint, + "recorded_differences": { + key: dict(value) for key, value in self.recorded_differences.items() + }, + "provenance": dict(self.provenance), + } + + +def _canonical_fingerprint(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def identity_fingerprint(identity: Mapping[str, Any]) -> str: + """Fingerprint only the declared :data:`IDENTITY_FIELDS`, in a fixed order. + + Extra keys in ``identity`` are ignored on purpose: callers pass whole + provenance bundles, and the fingerprint must not drift when an unrelated + diagnostic field is added. + """ + + return _canonical_fingerprint({name: identity.get(name) for name in IDENTITY_FIELDS}) + + +def _reduction_view(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _recorded_view(contract: AttentionContract) -> dict[str, Any]: + rope = contract.rope + kv_cache = contract.kv_cache + view: dict[str, Any] = { + "mode": contract.mode.value, + "backend_id": None, + "reduction.engine": contract.reduction.engine.value, + "sharding.cp_world_size": contract.sharding.cp_world_size, + "sharding.tp_world_size": contract.sharding.tp_world_size, + "sharding.local_sequence_length": contract.sharding.local_sequence_length, + } + if rope is not None: + view.update( + { + "rope.fusion_boundary": rope.fusion_boundary.value, + "rope.q_state": rope.q_state.value, + "rope.k_state": rope.k_state.value, + "rope.k_cache_state": rope.k_cache_state.value, + "rope.cast_at": rope.cast_at.value, + "rope.output_dtype": rope.output_dtype.value, + } + ) + if kv_cache is not None: + view.update( + { + "kv_cache.page_size": kv_cache.page_size, + "kv_cache.prefix_cache_enabled": kv_cache.prefix_cache_enabled, + "kv_cache.block_table_shape": [ + len(kv_cache.block_table), + max((len(row) for row in kv_cache.block_table), default=0), + ], + } + ) + return view + + +#: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B +#: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- +#: both sides still have to agree on it, which the equality pass below handles. +NULLABLE_IDENTITY_FIELDS: frozenset[str] = frozenset({"rope_scaling"}) + + +def _missing_identity_fields(identity: Mapping[str, Any]) -> tuple[str, ...]: + return tuple( + name + for name in IDENTITY_FIELDS + if name not in NULLABLE_IDENTITY_FIELDS and identity.get(name) is None + ) + + +def bind_attention_contracts( + *, + rollout_contract: AttentionContract, + training_contract: AttentionContract, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), + require_full_identity: bool = True, +) -> AttentionBindingResult: + """Bind a rollout attention contract to a training attention contract. + + ``determinism_issues`` is threaded in from + :mod:`rl_engine.alignment.cross_config.determinism` rather than computed here, + so that this module stays free of framework probing and remains testable + without Megatron or vLLM present. + + ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which + legitimately has no KV-cache or decode identity to declare. Distributed callers + must leave it at ``True``. + """ + + if rollout_contract.role is not AttentionRole.INFER: + raise AttentionBindingError( + f"rollout_contract.role must be {AttentionRole.INFER.value!r}, " + f"got {rollout_contract.role.value!r}" + ) + if training_contract.role is not AttentionRole.TRAIN: + raise AttentionBindingError( + f"training_contract.role must be {AttentionRole.TRAIN.value!r}, " + f"got {training_contract.role.value!r}" + ) + + issues: list[BindingIssue] = [] + + # ---- tier 1: identity, bit for bit ------------------------------------- + if require_full_identity: + for side, identity in (("rollout", rollout_identity), ("training", training_identity)): + for name in _missing_identity_fields(identity): + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISSING, + tier=BindingTier.IDENTICAL, + field=f"{side}.{name}", + message=f"{side} identity does not declare {name!r}", + ) + ) + + for name in IDENTITY_FIELDS: + rollout_value = rollout_identity.get(name) + training_value = training_identity.get(name) + if rollout_value != training_value: + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=name, + rollout=rollout_value, + training=training_value, + message=( + f"{name!r} differs between sides; the pair is not comparable " + "and any drift computed from it is meaningless" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- + rollout_reduction = _reduction_view(rollout_contract) + training_reduction = _reduction_view(training_contract) + + for name in SEMANTIC_REDUCTION_FIELDS: + if rollout_reduction[name] != training_reduction[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"reduction.{name!r} must be decided by the contract, not by the " + "backend; the two sides disagree" + ), + ) + ) + mandated = WS2_ATTENTION_REDUCTION_MANDATE[name] + for side, view in (("rollout", rollout_reduction), ("training", training_reduction)): + if view[name] != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_MANDATE_VIOLATION, + tier=BindingTier.SEMANTIC, + field=f"{side}.reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"WS2 requires reduction.{name} == {mandated!r}; " + f"{side} declares {view[name]!r}" + ), + ) + ) + + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): + if not contract.export_lse: + issues.append( + BindingIssue( + code=BindingErrorCode.LSE_NOT_EXPORTED, + tier=BindingTier.SEMANTIC, + field=f"{side}.export_lse", + message=( + "attention-domain LSE must be exported; without it the deterministic " + "CP merge cannot be validated" + ), + ) + ) + + if rollout_backend_id == training_backend_id and rollout_backend_id: + # Not an error, but worth surfacing: an identical backend on both sides means + # the experiment is not actually measuring a cross-implementation difference. + pass + + issues.extend(determinism_issues) + + # ---- tier 3: recorded differences -------------------------------------- + rollout_recorded = _recorded_view(rollout_contract) + rollout_recorded["backend_id"] = rollout_backend_id + training_recorded = _recorded_view(training_contract) + training_recorded["backend_id"] = training_backend_id + + recorded_differences: dict[str, dict[str, Any]] = {} + for name in RECORDED_FIELDS: + rollout_value = rollout_recorded.get(name) + training_value = training_recorded.get(name) + if rollout_value != training_value: + recorded_differences[name] = { + "rollout": rollout_value, + "training": training_value, + } + + identity_fp = identity_fingerprint(training_identity if comparable else rollout_identity) + reduction_fp = _canonical_fingerprint( + {name: training_reduction[name] for name in SEMANTIC_REDUCTION_FIELDS} + ) + passed = comparable and not any(issue.tier is BindingTier.SEMANTIC for issue in issues) + + provenance = { + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout": { + "contract": rollout_contract.to_dict(), + "backend_id": rollout_backend_id, + "recorded": rollout_recorded, + }, + "training": { + "contract": training_contract.to_dict(), + "backend_id": training_backend_id, + "recorded": training_recorded, + }, + } + + return AttentionBindingResult( + comparable=comparable, + passed=passed, + issues=tuple(issues), + identity_fingerprint=identity_fp, + reduction_fingerprint=reduction_fp, + binding_fingerprint=_canonical_fingerprint( + { + "identity": identity_fp, + "reduction": reduction_fp, + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +def summarize_binding(result: AttentionBindingResult) -> str: + """One-line human summary for CLI output and failure messages.""" + + if result.passed: + return ( + f"attention binding OK " + f"(identity={result.identity_fingerprint[:12]}, " + f"{len(result.recorded_differences)} recorded difference(s))" + ) + if not result.comparable: + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.IDENTICAL}) + ) + return f"attention binding NOT COMPARABLE; identity problems: {fields}" + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.SEMANTIC}) + ) + return f"attention binding FAILED CLOSED; semantic problems: {fields}" + + +def first_blocking_issue( + result: AttentionBindingResult, +) -> Optional[BindingIssue]: + """Return the issue a caller should report, preferring identity over semantics.""" + + for tier in (BindingTier.IDENTICAL, BindingTier.SEMANTIC): + for issue in result.issues: + if issue.tier is tier: + return issue + return None diff --git a/rl_engine/alignment/cross_config/determinism.py b/rl_engine/alignment/cross_config/determinism.py new file mode 100644 index 00000000..be81654e --- /dev/null +++ b/rl_engine/alignment/cross_config/determinism.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-side determinism probing for the Megatron + vLLM cross-config target. + +Both frameworks ship a "make this deterministic" switch, but they mean different +things by it, and neither knows the other exists: + +``Megatron`` ``ModelParallelConfig.deterministic_mode`` + Asserts ``NCCL_ALGO`` is one of five values, forbids FlashAttention and fused + cross-entropy, calls ``torch.use_deterministic_algorithms(True)``, and requires + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO == 0``. It does **not** touch TF32, BF16 + reduced-precision reduction, cuBLAS workspace, NCCL protocol, or NCCL channel + counts. + +``vLLM`` ``VLLM_BATCH_INVARIANT`` + Replaces ``aten::mm/addmm/matmul/linear/bmm``, ``log_softmax``/``softmax``, + ``mean.dim`` and ``rms_norm`` with Triton kernels, disables TF32 and BF16/FP16 + reduced-precision reduction, pins cuBLAS workspace and the BLAS library, and + hard-sets ten NCCL environment variables. + +So a run can have both switches on and still be comparing two different notions of +determinism. This module makes that difference explicit and, where it changes the +numerics, blocking. It never imports Megatron or vLLM: probes are built from plain +mappings so the logic is testable on any machine. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from rl_engine.alignment.cross_config.attention_binding import ( + BindingErrorCode, + BindingIssue, + BindingTier, +) + +__all__ = [ + "COMPARED_NCCL_KEYS", + "DeterminismProbe", + "DeterminismReport", + "compare_determinism", + "megatron_probe_from_config", + "vllm_probe_from_env", +] + + +#: Environment keys whose value can change a reduction result. Compared across +#: sides; a difference is reported, and a difference in the *arithmetic* subset is +#: blocking. Ordering is fixed so the fingerprint is stable. +COMPARED_NCCL_KEYS: tuple[str, ...] = ( + "NCCL_ALGO", + "NCCL_PROTO", + "NCCL_MIN_NCHANNELS", + "NCCL_MAX_NCHANNELS", + "NCCL_NTHREADS", + "NCCL_SOCKET_NTHREADS", + "NCCL_COLLNET_ENABLE", + "NCCL_NVLS_ENABLE", + "NCCL_P2P_NET_DISABLE", + "NCCL_LAUNCH_MODE", + "CUBLAS_WORKSPACE_CONFIG", +) + + +#: The subset above that changes arithmetic rather than only scheduling. A mismatch +#: here fails the binding closed; a mismatch in the remainder is recorded only. +_ARITHMETIC_NCCL_KEYS: frozenset[str] = frozenset( + {"NCCL_ALGO", "NCCL_PROTO", "CUBLAS_WORKSPACE_CONFIG"} +) + + +@dataclass(frozen=True) +class DeterminismProbe: + """What one side actually has switched on. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are tri-state on + purpose: ``None`` means "the framework does not manage this", which is exactly + Megatron's situation and is itself the finding. + """ + + side: str + framework: str + mode_flag: str + enabled: bool + env: Mapping[str, Any] = field(default_factory=dict) + tf32_disabled: Optional[bool] = None + bf16_reduced_precision_reduction: Optional[bool] = None + forbids_flash_attention: Optional[bool] = None + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_probe.v1" + + def __post_init__(self) -> None: + if self.side not in ("rollout", "training"): + raise ValueError("side must be 'rollout' or 'training'") + if not self.framework: + raise ValueError("framework must not be empty") + object.__setattr__(self, "env", dict(self.env)) + object.__setattr__(self, "evidence", dict(self.evidence)) + + @property + def env_fingerprint(self) -> str: + payload = {key: self.env.get(key) for key in COMPARED_NCCL_KEYS} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "side": self.side, + "framework": self.framework, + "mode_flag": self.mode_flag, + "enabled": self.enabled, + "env": {key: self.env.get(key) for key in COMPARED_NCCL_KEYS}, + "env_fingerprint": self.env_fingerprint, + "tf32_disabled": self.tf32_disabled, + "bf16_reduced_precision_reduction": self.bf16_reduced_precision_reduction, + "forbids_flash_attention": self.forbids_flash_attention, + "evidence": dict(self.evidence), + } + + +@dataclass(frozen=True) +class DeterminismReport: + """Cross-side comparison result.""" + + rollout: DeterminismProbe + training: DeterminismProbe + issues: tuple[BindingIssue, ...] = () + differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_report.v1" + + @property + def compatible(self) -> bool: + return not self.issues + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "compatible": self.compatible, + "rollout": self.rollout.to_dict(), + "training": self.training.to_dict(), + "issues": [issue.to_dict() for issue in self.issues], + "differences": {key: dict(value) for key, value in self.differences.items()}, + } + + +def megatron_probe_from_config( + config: Any, + env: Optional[Mapping[str, str]] = None, +) -> DeterminismProbe: + """Build a training-side probe from a Megatron config object. + + ``config`` is duck-typed (anything exposing ``deterministic_mode`` and + optionally ``attention_backend`` / ``cross_entropy_loss_fusion``) so this works + against a real ``ModelParallelConfig``, a test double, or a plain namespace, + and so importing this module never requires Megatron. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are reported as + ``None`` because Megatron does not manage them -- a ``grep`` for ``allow_tf32`` + and ``fp32_precision`` across ``megatron/`` returns nothing. That asymmetry + against vLLM is the point of :func:`compare_determinism`. + """ + + environ = dict(env or {}) + enabled = bool(getattr(config, "deterministic_mode", False)) + return DeterminismProbe( + side="training", + framework="megatron", + mode_flag="deterministic_mode", + enabled=enabled, + env={key: environ.get(key) for key in COMPARED_NCCL_KEYS}, + tf32_disabled=None, + bf16_reduced_precision_reduction=None, + forbids_flash_attention=enabled, + evidence={ + "nvte_allow_nondeterministic_algo": environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO"), + "cross_entropy_loss_fusion": getattr(config, "cross_entropy_loss_fusion", None), + "attention_backend": _enum_value(getattr(config, "attention_backend", None)), + "tensor_model_parallel_size": getattr(config, "tensor_model_parallel_size", None), + "context_parallel_size": getattr(config, "context_parallel_size", None), + "sequence_parallel": getattr(config, "sequence_parallel", None), + "manages_tf32": False, + "manages_bf16_reduced_precision_reduction": False, + }, + ) + + +def vllm_probe_from_env( + env: Mapping[str, str], + *, + model_config: Any = None, +) -> DeterminismProbe: + """Build a rollout-side probe from the vLLM process environment. + + ``VLLM_BATCH_INVARIANT`` is read from ``env`` rather than ``vllm.envs`` so the + probe can be constructed from a remote worker's reported environment, which is + how vime's Ray actors expose it. + """ + + enabled = str(env.get("VLLM_BATCH_INVARIANT", "0")).strip() in ("1", "true", "True") + return DeterminismProbe( + side="rollout", + framework="vllm", + mode_flag="VLLM_BATCH_INVARIANT", + enabled=enabled, + env={key: env.get(key) for key in COMPARED_NCCL_KEYS}, + # vLLM sets both to "ieee"/disabled inside init_batch_invariance(). + tf32_disabled=enabled or None, + bf16_reduced_precision_reduction=(False if enabled else None), + forbids_flash_attention=False, + evidence={ + "vllm_allreduce_use_symm_mem": env.get("VLLM_ALLREDUCE_USE_SYMM_MEM"), + "vllm_use_aot_compile": env.get("VLLM_USE_AOT_COMPILE"), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "disable_cascade_attn": getattr(model_config, "disable_cascade_attn", None), + "quantization": getattr(model_config, "quantization", None), + "manages_tf32": True, + "manages_bf16_reduced_precision_reduction": True, + }, + ) + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def compare_determinism( + *, + rollout: DeterminismProbe, + training: DeterminismProbe, +) -> DeterminismReport: + """Compare two probes and produce blocking issues plus recorded differences.""" + + if rollout.side != "rollout" or training.side != "training": + raise ValueError("compare_determinism expects one rollout probe and one training probe") + + issues: list[BindingIssue] = [] + differences: dict[str, dict[str, Any]] = {} + + for probe in (rollout, training): + if not probe.enabled: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{probe.side}.{probe.mode_flag}", + rollout=rollout.enabled, + training=training.enabled, + message=( + f"{probe.framework} {probe.mode_flag} is not enabled; the " + f"{probe.side} side is not batch-invariant and cannot anchor a " + "cross-config comparison" + ), + ) + ) + + for key in COMPARED_NCCL_KEYS: + rollout_value = rollout.env.get(key) + training_value = training.env.get(key) + if rollout_value == training_value: + continue + differences[key] = {"rollout": rollout_value, "training": training_value} + if key in _ARITHMETIC_NCCL_KEYS: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"env.{key}", + rollout=rollout_value, + training=training_value, + message=( + f"{key} differs between sides; the two sides would reduce with " + "different arithmetic and the resulting drift is not attributable" + ), + ) + ) + + # Megatron reports None for these because it does not manage them at all. That is + # recorded rather than blocking: under a pure BF16 GEMM path TF32 does not fire, and + # forcing Megatron to manage it is out of scope for this PR. It is surfaced so the + # asymmetry appears in every artifact instead of being invisible. + for name in ("tf32_disabled", "bf16_reduced_precision_reduction"): + rollout_value = getattr(rollout, name) + training_value = getattr(training, name) + if rollout_value != training_value: + differences[name] = { + "rollout": rollout_value, + "training": training_value, + "note": ( + "megatron does not manage this setting; vllm sets it inside " + "init_batch_invariance()" + ), + } + + return DeterminismReport( + rollout=rollout, + training=training, + issues=tuple(issues), + differences=differences, + ) diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py new file mode 100644 index 00000000..058dfa44 --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,612 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for #235 PR4: rollout/training attention contract binding. + +Every test here runs on CPU without Megatron or vLLM installed. That is the point: +the binding rules are contract logic, and contract logic that can only be exercised +on a 2-node x 2-GPU cluster would never be exercised. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + first_blocking_issue, + identity_fingerprint, + summarize_binding, +) +from rl_engine.alignment.cross_config.determinism import ( + compare_determinism, + megatron_probe_from_config, + vllm_probe_from_env, +) +from rl_engine.alignment.cross_config.schema import MaterializationStatus +from rl_engine.kernels.attention_contract import AttentionContractError, AttentionMode + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 1, + "rollout.dtype": "bf16", +} + + +def _identity(**overrides): + identity = { + "checkpoint_id": "qwen3-8b", + "model_version": "v1", + "weight_version": 7, + "tokenizer_fingerprint": "tokenizer-abc", + "token_ids_fingerprint": "tokens-abc", + "active_mask_fingerprint": "mask-abc", + "position_ids_fingerprint": "pos-abc", + "padding_side": "right", + "pre_update_state": "pre_update", + "global_token_positions_fingerprint": "gtp-abc", + "kv_seq_lens_fingerprint": "kvlen-abc", + } + identity.update(QWEN3_8B.identity_fields()) + identity.update(overrides) + return identity + + +def _contracts(): + training = MegatronAttentionMaterializer().build_contract(TRAINING_KNOBS) + rollout = VllmRolloutMaterializer().build_contract(ROLLOUT_KNOBS) + return rollout, training + + +def _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + return bind_attention_contracts( + rollout_contract=kwargs.pop("rollout_contract", rollout), + training_contract=kwargs.pop("training_contract", training), + rollout_identity=rollout_identity if rollout_identity is not None else _identity(), + training_identity=training_identity if training_identity is not None else _identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_binds_despite_different_materialization(): + """The core claim of PR4: same identity + same reduction, different runtimes.""" + + result = _bind() + + assert result.comparable + assert result.passed + assert result.issues == () + # Training runs CP=2 full prefill, rollout runs CP=1 chunked prefill. Those + # differences are recorded, not rejected. + assert "mode" in result.recorded_differences + assert "sharding.cp_world_size" in result.recorded_differences + assert result.recorded_differences["mode"] == { + "rollout": "chunked_prefill", + "training": "prefill", + } + + +def test_weight_version_mismatch_is_not_comparable(): + result = _bind(rollout_identity=_identity(weight_version=6)) + + assert not result.comparable + assert not result.passed + codes = {issue.code for issue in result.issues} + assert BindingErrorCode.IDENTITY_MISMATCH in codes + blocking = first_blocking_issue(result) + assert blocking is not None and blocking.tier is BindingTier.IDENTICAL + assert "NOT COMPARABLE" in summarize_binding(result) + + +def test_rope_theta_mismatch_is_not_comparable(): + """RoPE math constants are identity, not materialization.""" + + result = _bind(training_identity=_identity(rope_theta=10000.0)) + + assert not result.comparable + assert any(issue.field == "rope_theta" for issue in result.issues) + + +def test_null_rope_scaling_is_a_value_not_an_omission(): + """Qwen3-8B applies no RoPE scaling; ``None`` must not read as undeclared.""" + + result = _bind() + + assert not result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + + +def test_missing_identity_field_is_reported_per_side(): + identity = _identity() + del identity["padding_side"] + result = _bind(rollout_identity=identity, training_identity=identity) + + missing = result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + assert {issue.field for issue in missing} == { + "rollout.padding_side", + "training.padding_side", + } + assert not result.comparable + + +def test_single_gpu_harness_may_waive_full_identity(): + """#235 PR2 has no KV-cache identity to declare; it opts out explicitly.""" + + identity = _identity() + del identity["global_token_positions_fingerprint"] + del identity["kv_seq_lens_fingerprint"] + + strict = _bind(rollout_identity=identity, training_identity=identity) + waived = _bind( + rollout_identity=identity, + training_identity=identity, + require_full_identity=False, + ) + + assert not strict.comparable + assert waived.comparable and waived.passed + + +def test_identity_fingerprint_ignores_undeclared_extra_keys(): + base = _identity() + decorated = dict(base, diagnostic_note="added later") + + assert identity_fingerprint(base) == identity_fingerprint(decorated) + + +# -------------------------------------------------------------------------- +# tier 2: reduction semantics +# -------------------------------------------------------------------------- + + +def test_reduction_semantics_are_bound_and_fingerprinted(): + result = _bind() + + reduction = result.provenance["training"]["contract"]["reduction"] + assert reduction["merge"] == "online_softmax_lse" + assert reduction["acc_dtype"] == "fp32" + assert reduction["order"] == "global_block_index" + assert reduction["downcast_at"] == "final_write" + assert result.reduction_fingerprint + + +def test_reduction_engine_difference_is_recorded_not_rejected(): + """A TE merge oracle on one side must not fail the binding.""" + + from rl_engine.alignment.cross_config.attention_binding import ( + RECORDED_FIELDS, + SEMANTIC_REDUCTION_FIELDS, + ) + + assert "reduction.engine" in RECORDED_FIELDS + assert "engine" not in SEMANTIC_REDUCTION_FIELDS + + +def test_lse_domain_is_recorded_as_attention_domain(): + """#235: attention exports attention-domain LSE, not vocab-logprob LSE.""" + + result = _bind() + + assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" + + +# -------------------------------------------------------------------------- +# role and input validation +# -------------------------------------------------------------------------- + + +def test_swapped_roles_are_rejected_outright(): + rollout, training = _contracts() + + with pytest.raises(AttentionBindingError): + bind_attention_contracts( + rollout_contract=training, + training_contract=rollout, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="a", + training_backend_id="b", + ) + + +# -------------------------------------------------------------------------- +# determinism cross-check +# -------------------------------------------------------------------------- + + +def _megatron_env(): + return {"NCCL_ALGO": "Tree", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0"} + + +def _vllm_env(**overrides): + env = { + "VLLM_BATCH_INVARIANT": "1", + "NCCL_ALGO": "allreduce:tree", + "NCCL_PROTO": "Simple", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + } + env.update(overrides) + return env + + +def test_nccl_algo_mismatch_blocks_the_binding(): + """Megatron asserts NCCL_ALGO; vLLM hard-sets a different value.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert not report.compatible + fields = {issue.field for issue in report.issues} + assert "env.NCCL_ALGO" in fields + assert "env.NCCL_PROTO" in fields + + +def test_matching_nccl_settings_are_compatible(): + shared = {"NCCL_ALGO": "allreduce:tree", "NCCL_PROTO": "Simple"} + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=dict(shared) + ) + rollout = vllm_probe_from_env({**_vllm_env(**shared), "CUBLAS_WORKSPACE_CONFIG": None}) + + report = compare_determinism(rollout=rollout, training=training) + + assert report.compatible, [issue.to_dict() for issue in report.issues] + + +def test_determinism_switch_off_on_either_side_blocks(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=False), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env(VLLM_BATCH_INVARIANT="0")) + + report = compare_determinism(rollout=rollout, training=training) + + fields = {issue.field for issue in report.issues} + assert "training.deterministic_mode" in fields + assert "rollout.VLLM_BATCH_INVARIANT" in fields + + +def test_tf32_asymmetry_is_recorded_not_blocking(): + """Megatron does not manage TF32 at all; vLLM disables it. Record the gap.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert training.tf32_disabled is None + assert rollout.tf32_disabled is True + assert "tf32_disabled" in report.differences + assert not any(issue.field == "tf32_disabled" for issue in report.issues) + + +def test_determinism_issues_flow_into_the_binding(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + report = compare_determinism(rollout=rollout, training=training) + + result = _bind(determinism_issues=report.issues) + + assert result.comparable # identity is fine + assert not result.passed # but the reduction environment is not + assert result.issues_by_code(BindingErrorCode.DETERMINISM_INCOMPATIBLE) + assert "FAILED CLOSED" in summarize_binding(result) + + +# -------------------------------------------------------------------------- +# sharding derived from the frozen #239 layout +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_cp_shards_cover_the_global_sequence_without_overlap(cp_rank): + contract = MegatronAttentionMaterializer( + cp_rank=cp_rank, global_sequence_length=4096 + ).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert sharding.local_sequence_length == 2048 + assert sharding.global_block_indices == (cp_rank,) + assert sharding.global_block_token_starts == (cp_rank * 2048,) + # The causal offset must be the number of preceding *global* tokens, otherwise + # rank 1 would mask as if its shard started at position zero. + assert contract.causal_offsets == (cp_rank * 2048, cp_rank * 2048) + + +def test_tp_head_shards_split_qwen3_gqa_evenly(): + contract = MegatronAttentionMaterializer(tp_rank=1).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert (sharding.global_q_heads, sharding.global_kv_heads) == (32, 8) + assert (sharding.local_q_heads, sharding.local_kv_heads) == (16, 4) + assert (sharding.local_q_head_start, sharding.local_kv_head_start) == (16, 4) + + +@pytest.mark.parametrize("tp_world_size", [2, 4, 8]) +def test_supported_tp_degrees_shard_qwen3_gqa(tp_world_size): + """Qwen3-8B has 32 Q heads and 8 KV heads, so TP in {2, 4, 8} all divide.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": tp_world_size} + ) + + sharding = contract.sharding + assert sharding.local_q_heads == 32 // tp_world_size + assert sharding.local_kv_heads == 8 // tp_world_size + + +@pytest.mark.parametrize( + ("knob_value", "expected"), + [("bfloat16", "bf16"), ("float16", "fp16"), ("float32", "fp32"), ("fp16", "fp16")], +) +def test_planner_normalized_dtypes_reach_the_contract(knob_value, expected): + """The planner emits torch spellings; AttentionDType uses short ones.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": knob_value} + ) + + assert contract.dtype.value == expected + + +def test_unknown_dtype_is_rejected_with_the_offending_field(): + with pytest.raises(ValueError, match="training.compute_dtype"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "int8"} + ) + + +def test_indivisible_tp_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": 3} + ) + + +def test_indivisible_cp_sequence_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer(global_sequence_length=4097).build_contract(TRAINING_KNOBS) + + +# -------------------------------------------------------------------------- +# materialization: fail closed rather than silently substitute +# -------------------------------------------------------------------------- + + +def _statuses(materialization, path): + return [app.status for app in materialization.applications if app.path == path] + + +def test_arrival_merge_order_is_unsupported_not_silently_corrected(): + """The control group must stay distinguishable from the treatment.""" + + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +def test_bf16_reduction_accumulation_is_unsupported(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_acc_dtype": "bf16"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_acc_dtype") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_te_oracle_engine_is_unsupported_until_pr2_pr3(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_engine": "te_oracle"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_engine") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): + materializer = VllmRolloutMaterializer(mode=AttentionMode.DECODE) + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + + assert materializer.effective_cp_world_size({"rollout.context_parallel_size": 2}) == 1 + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + contract_error = materialization.binding.side_configs["rollout"]["contract_error"] + assert "#235 PR6" in contract_error + + +def test_decode_contract_is_refused_without_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="PR6"): + VllmRolloutMaterializer(mode=AttentionMode.DECODE).build_contract(ROLLOUT_KNOBS) + + +def test_materializers_expose_distinct_implementation_fingerprints(): + megatron = MegatronAttentionMaterializer().implementation_fingerprint + vllm = VllmRolloutMaterializer().implementation_fingerprint + + assert megatron and vllm and megatron != vllm + + +def test_runtime_binding_reports_the_frozen_topology(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + binding = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS).binding + + topology = binding.topology["training"] + assert topology["tensor_parallel_size"] == 2 + assert topology["context_parallel_size"] == 2 + assert topology["world_size"] == 4 + assert topology["pipeline_parallel_size"] == 1 + assert topology["data_parallel_size"] == 1 + + +# -------------------------------------------------------------------------- +# provenance adapters +# -------------------------------------------------------------------------- + + +def test_megatron_provenance_flags_undeclared_frozen_scope_fields(): + adapter = MegatronProvenanceAdapter(SimpleNamespace(deterministic_mode=True)) + + violations = adapter.frozen_scope_violations() + + # Nothing is declared, so every assertion reads as unknown rather than as met. + assert any("expert_model_parallel_size" in text for text in violations) + assert any("fp8" in text for text in violations) + + +def test_megatron_provenance_accepts_a_conforming_dense_config(): + adapter = MegatronProvenanceAdapter( + SimpleNamespace( + deterministic_mode=True, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + sequence_parallel=False, + fp8=None, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + ) + + assert adapter.frozen_scope_violations() == ("fp8 is not declared (expected None)",) + + +def test_megatron_construction_fingerprint_tracks_fusion_changes(): + base = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=False) + fused = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=True) + + assert ( + MegatronProvenanceAdapter(base).construction_fingerprint + != MegatronProvenanceAdapter(fused).construction_fingerprint + ) + + +def test_vllm_provenance_reads_page_size_and_split_kv_policy(): + adapter = VllmProvenanceAdapter( + cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), + attention_config=SimpleNamespace(flash_attn_max_num_splits_for_cuda_graph=32), + ) + + assert adapter.kv_page_size == 16 + assert adapter.split_kv_policy == 32 + + +def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): + adapter = VllmProvenanceAdapter( + model_config=SimpleNamespace(quantization=None, disable_cascade_attn=False), + cache_config=SimpleNamespace( + cache_dtype="fp8", calculate_kv_scales=False, sliding_window=None + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, data_parallel_size=1), + ) + + violations = adapter.frozen_scope_violations() + + assert any("cache_dtype" in text for text in violations) + assert any("disable_cascade_attn" in text for text in violations) + + +# -------------------------------------------------------------------------- +# scenario config +# -------------------------------------------------------------------------- + + +SCENARIO = ( + Path(__file__).resolve().parents[1] + / "examples" + / "cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json" +) + + +def test_scenario_uses_megatron_vocabulary_only(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + training = config["baseline"]["training"] + + assert training["attention_backend"] in {"flash", "fused", "unfused", "local", "auto"} + assert training["tensor_parallel_size"] == 2 + assert training["context_parallel_size"] == 2 + assert config["baseline"]["rollout"]["batch_invariant"] is True + + +def test_scenario_knob_paths_all_exist(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + + def paths(mapping, prefix=""): + for key, value in mapping.items(): + path = f"{prefix}{key}" + if isinstance(value, dict): + yield from paths(value, f"{path}.") + else: + yield path + + declared = set(paths(config["baseline"])) + unknown = declared - set(WS2_ATTENTION_KNOBS) + assert not unknown, f"scenario declares unknown knobs: {sorted(unknown)}" + + for intervention in config["interventions"]: + assert intervention["path"] in WS2_ATTENTION_KNOBS From 4ad305b0829b21da0bbb610b325148c337ab7e3b Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 4 Aug 2026 22:23:29 +0800 Subject: [PATCH 11/11] fix(alignment): bind dtype, batch size and split-KV policy (#235 PR4) Three fields could differ between the two sides without the binding noticing. dtype was in no tier at all, so a BF16 rollout could bind to an FP16 training pass and produce a drift number attributable to nothing. It joins the semantic tier, with allow_dtype_difference for the #235 PR5 sweep that deliberately scores BF16 against an FP32 reference. batch_size was likewise unchecked. Batch invariance is a claim about results not changing with batch makeup, so two sides scoring different batches are not comparable and it belongs to identity. split_kv_policy has no field in the #236 contract, so it only reached side_configs and never took part in binding. Callers now pass it through rollout_recorded_extra / training_recorded_extra so the difference is at least visible in provenance; it can move into the contract once #236 grows the field. Part of #235 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw --- .../cross_config/attention_binding.py | 55 ++++++++++++++++- tests/test_attention_cross_config_binding.py | 59 +++++++++++++++++++ 2 files changed, 111 insertions(+), 3 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index fcd4c203..48e382b8 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -61,6 +61,7 @@ "IDENTITY_FIELDS", "NULLABLE_IDENTITY_FIELDS", "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", "SEMANTIC_REDUCTION_FIELDS", "WS2_ATTENTION_REDUCTION_MANDATE", "bind_attention_contracts", @@ -126,6 +127,9 @@ class BindingErrorCode(str, Enum): "rope_scaling", "rotary_dim", "qk_layernorm", + # batch composition: batch-invariance is a claim about results not changing with + # batch makeup, so two sides scoring different batches are not comparable at all + "batch_size", # decode replay identity (#235 PR6) "global_token_positions_fingerprint", "kv_seq_lens_fingerprint", @@ -142,6 +146,14 @@ class BindingErrorCode(str, Enum): ) +#: Contract fields outside ``ReductionSpec`` that still decide the numerical result. +#: ``dtype`` is here rather than in :data:`RECORDED_FIELDS` because comparing a BF16 +#: rollout against an FP16 training pass produces a real drift number attributable to +#: nothing. #235 PR5 does sweep BF16 against an FP32 reference; that sweep opts in via +#: ``allow_dtype_difference`` instead of loosening the default. +SEMANTIC_CONTRACT_FIELDS: tuple[str, ...] = ("dtype",) + + #: The WS2 mandate itself. ``#236`` currently declares single-member enums for #: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; #: they are written out anyway so that widening any of those enums later fails here @@ -172,6 +184,11 @@ class BindingErrorCode(str, Enum): "sharding.cp_world_size", "sharding.tp_world_size", "sharding.local_sequence_length", + # Supplied by the caller, not by the contract: #236 has no split-KV field yet, so + # the value comes from vLLM's flash_attn_max_num_splits_for_cuda_graph via the + # adapter. Recorded so split-KV differences are at least visible in provenance + # until #236 grows the field and it can move into the contract proper. + "split_kv_policy", ) @@ -264,7 +281,10 @@ def _reduction_view(contract: AttentionContract) -> dict[str, Any]: } -def _recorded_view(contract: AttentionContract) -> dict[str, Any]: +def _recorded_view( + contract: AttentionContract, + extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: rope = contract.rope kv_cache = contract.kv_cache view: dict[str, Any] = { @@ -297,6 +317,8 @@ def _recorded_view(contract: AttentionContract) -> dict[str, Any]: ], } ) + if extra: + view.update(extra) return view @@ -324,6 +346,9 @@ def bind_attention_contracts( training_backend_id: str, determinism_issues: Sequence[BindingIssue] = (), require_full_identity: bool = True, + allow_dtype_difference: bool = False, + rollout_recorded_extra: Optional[Mapping[str, Any]] = None, + training_recorded_extra: Optional[Mapping[str, Any]] = None, ) -> AttentionBindingResult: """Bind a rollout attention contract to a training attention contract. @@ -335,6 +360,13 @@ def bind_attention_contracts( ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which legitimately has no KV-cache or decode identity to declare. Distributed callers must leave it at ``True``. + + ``allow_dtype_difference`` exists for the #235 PR5 sweep that deliberately scores + a BF16 path against an FP32 reference. It must stay ``False`` everywhere else. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` carry materialization + facts that #236 does not yet model -- today that is ``split_kv_policy``. They are + merged into the recorded tier, never into identity or semantics. """ if rollout_contract.role is not AttentionRole.INFER: @@ -419,6 +451,22 @@ def bind_attention_contracts( ) ) + if not allow_dtype_difference and rollout_contract.dtype is not training_contract.dtype: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field="dtype", + rollout=rollout_contract.dtype.value, + training=training_contract.dtype.value, + message=( + "the two sides compute in different dtypes; the resulting drift is " + "not attributable. Pass allow_dtype_difference=True only for a " + "deliberate precision sweep" + ), + ) + ) + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): if not contract.export_lse: issues.append( @@ -441,9 +489,9 @@ def bind_attention_contracts( issues.extend(determinism_issues) # ---- tier 3: recorded differences -------------------------------------- - rollout_recorded = _recorded_view(rollout_contract) + rollout_recorded = _recorded_view(rollout_contract, rollout_recorded_extra) rollout_recorded["backend_id"] = rollout_backend_id - training_recorded = _recorded_view(training_contract) + training_recorded = _recorded_view(training_contract, training_recorded_extra) training_recorded["backend_id"] = training_backend_id recorded_differences: dict[str, dict[str, Any]] = {} @@ -464,6 +512,7 @@ def bind_attention_contracts( provenance = { "lse_domain": ATTENTION_LSE_DOMAIN, + "dtype": training_contract.dtype.value, "rollout": { "contract": rollout_contract.to_dict(), "backend_id": rollout_backend_id, diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 058dfa44..815daf12 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -71,6 +71,7 @@ def _identity(**overrides): "position_ids_fingerprint": "pos-abc", "padding_side": "right", "pre_update_state": "pre_update", + "batch_size": 2, "global_token_positions_fingerprint": "gtp-abc", "kv_seq_lens_fingerprint": "kvlen-abc", } @@ -224,6 +225,64 @@ def test_lse_domain_is_recorded_as_attention_domain(): assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" +def test_mixed_dtypes_fail_closed(): + """BF16 rollout against FP16 training produces an unattributable number.""" + + rollout = VllmRolloutMaterializer().build_contract( + {**ROLLOUT_KNOBS, "rollout.dtype": "float16"} + ) + result = _bind(rollout_contract=rollout) + + assert result.comparable # identity is fine + assert not result.passed + assert any(issue.field == "dtype" for issue in result.issues) + + +def test_precision_sweep_may_opt_into_mixed_dtypes(): + """#235 PR5 sweeps BF16 against an FP32 reference; it says so explicitly.""" + + training = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "float32"} + ) + result = _bind(training_contract=training, allow_dtype_difference=True) + + assert result.passed + assert result.provenance["dtype"] == "fp32" + + +def test_batch_size_mismatch_is_not_comparable(): + """Batch invariance is a claim about batch makeup, so it belongs to identity.""" + + result = _bind(rollout_identity=_identity(batch_size=4)) + + assert not result.comparable + assert any(issue.field == "batch_size" for issue in result.issues) + + +def test_split_kv_policy_difference_is_recorded(): + """#236 has no split-KV field, so the adapter supplies it to the recorded tier.""" + + result = _bind( + rollout_recorded_extra={"split_kv_policy": 8}, + training_recorded_extra={"split_kv_policy": None}, + ) + + assert result.passed + assert result.recorded_differences["split_kv_policy"] == { + "rollout": 8, + "training": None, + } + + +def test_matching_split_kv_policy_is_not_reported_as_a_difference(): + result = _bind( + rollout_recorded_extra={"split_kv_policy": 32}, + training_recorded_extra={"split_kv_policy": 32}, + ) + + assert "split_kv_policy" not in result.recorded_differences + + # -------------------------------------------------------------------------- # role and input validation # --------------------------------------------------------------------------