From cb359fe6c7f9d9f5900a9515d0fcf2fd811404d2 Mon Sep 17 00:00:00 2001
From: "yangbosong.ljx" <yangbosong.ljx@bytedance.com>
Date: Sun, 19 Jul 2026 14:49:37 +0000
Subject: [PATCH 01/13] 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" <yangbosong.ljx@bytedance.com>
Date: Sun, 19 Jul 2026 15:12:26 +0000
Subject: [PATCH 02/13] 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" <yangbosong.ljx@bytedance.com>
Date: Sun, 19 Jul 2026 15:28:36 +0000
Subject: [PATCH 03/13] 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" <yangbosong.ljx@bytedance.com>
Date: Sun, 19 Jul 2026 15:59:08 +0000
Subject: [PATCH 04/13] 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" <yangbosong.ljx@bytedance.com>
Date: Mon, 20 Jul 2026 09:31:13 +0000
Subject: [PATCH 05/13] 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 8a4f9eb190ed19053d996583a034e5aaeb1ab502 Mon Sep 17 00:00:00 2001
From: inaniloquentee <3051000145@qq.com>
Date: Mon, 27 Jul 2026 22:00:02 +0800
Subject: [PATCH 06/13] 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 07/13] 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 deab4ed47e44ec827ab5cd2e7856579e9cddaf8c Mon Sep 17 00:00:00 2001
From: lamentropetion <3051000145@qq.com>
Date: Wed, 12 Aug 2026 22:24:21 +0800
Subject: [PATCH 08/13] fix(attention): complete Split-KV contract validation

---
 docs/operators/attention.md                   | 108 ++-
 rl_engine/kernels/attention_contract.py       | 683 +++++++++++++++++-
 rl_engine/kernels/gtest/operator_inputs.py    |  79 +-
 rl_engine/kernels/gtest/operator_specs.py     | 133 ++++
 .../kernels/ops/cuda/attention/__init__.py    |  40 +
 rl_engine/kernels/registry.py                 | 179 ++++-
 tests/test_attention_contract.py              | 136 +++-
 7 files changed, 1310 insertions(+), 48 deletions(-)

diff --git a/docs/operators/attention.md b/docs/operators/attention.md
index 8bec7d22..bc7f4a38 100644
--- a/docs/operators/attention.md
+++ b/docs/operators/attention.md
@@ -48,7 +48,7 @@ The op exposes the WS1 dual-path contract:
 | Backend | Wrapper | Native symbol | Status |
 | --- | --- | --- | --- |
 | PyTorch fallback | `NativeAttentionOp` | None | fp32 ground-truth reference; CPU and any GPU. |
-| CUDA / ROCm / Triton | — | — | Planned: downstream fused attention kernels validate against this reference. |
+| CUDA deterministic | `DeterministicAttentionOp` | `_C.deterministic_attention_forward/backward` | Batch-invariant CUDA implementation (issue #147). |
 
 ## Tensor Contract
 
@@ -76,12 +76,14 @@ the inputs' device.
 ## Dispatch Behavior
 
 `kernel_registry.get_op("attention")` resolves through the `OpBackend` priority map. On
-`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op
-(`PYTORCH_NATIVE_ATTENTION`), so every device dispatches to this op. Calling it (`__call__` ->
-`forward(...)`) computes in the input dtype; `forward_fp32(...)` is the explicit fp32 golden
-path. When fused attention kernels land, they are prepended to the priority list and the native
-op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is
-a separate dispatch chain and is unaffected.
+`cuda` the priority is:
+
+1. `CUDA_DETERMINISTIC_ATTENTION` — `DeterministicAttentionOp` (batch-invariant, fixed-order).
+2. `PYTORCH_NATIVE_ATTENTION` — `NativeAttentionOp` (fallback).
+
+Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_fp32(...)` is
+the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type
+(SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected.
 
 ### WS2 CP-aware dispatch
 
@@ -94,6 +96,21 @@ Existing WS1 implementations do not yet export attention-domain LSE or implement
 CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as
 a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md).
 
+Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow
+`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block
+boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason.
+Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and
+validate the same actual plan.
+
+The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report:
+
+```bash
+python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json
+python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \
+  --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \
+  --output artifacts/ws2-cp-attention-drift.json
+```
+
 ## Accuracy
 
 Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled):
@@ -146,6 +163,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
@@ -157,14 +175,80 @@ GPU-only LARGE Qwen3-8B real-shape smoke test.
 
 ## Implementation Files
 
-- `rl_engine/kernels/ops/pytorch/attention/standard_attn.py`
+- `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference
+- `rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` — CUDA deterministic op
+- `csrc/cuda/attention/deterministic_attention.cu` — CUDA kernels
 - `rl_engine/kernels/registry.py`
 - `tests/test_attention.py`
+- `tests/test_deterministic_attention_cuda.py`
+
+## Fixed Reduction Order (CUDA Deterministic Backend)
+
+The CUDA `DeterministicAttentionOp` pins reduction order for batch-invariance:
+
+**QK kernel**: D-dimension FP32 accumulation, `d = 0 .. D-1`. Each `scores[b,hq,q,k]` has
+exactly one writer thread.
+
+**Softmax + LSE kernel**: Each `(b, hq, q)` row processed by one CTA with fixed 256 threads.
+Max and sum-exp use a fixed shared-memory tree reduction (power-of-two stride). No split by
+batch size, sequence length, or SM count.
+
+**PV kernel**: K-dimension FP32 accumulation, `k = 0 .. Skv-1`. Each `out[b,hq,q,d]` has
+exactly one writer thread.
+
+**Backward dK/dV**: Per `(b, hkv, k, d)` output element, a single thread accumulates over
+query heads in group order then query positions:
+```text
+for local = 0 .. g-1:       # g = Hq / Hkv
+  hq = hkv * g + local
+  for q = 0 .. Sq-1:
+    acc += ...
+```
+No cross-CTA atomics. No launch-order dependent accumulation.
+
+## Prefill / Decode / KV-cache Shared Contract
+
+All inference modes use **the same standard attention kernels** (only Sq/Skv differ):
+
+- **Prefill**: `Sq == Skv`. Causal mask `key_index <= Skv - Sq + query_index`.
+- **Chunked-prefill**: each chunk uses `Sq = chunk_size`, `Skv = past + chunk_size`.
+  Same causal offset formula produces identical results to full prefill at matching positions.
+- **Decode**: `Sq = 1` (or few), `Skv = full_context`. Same kernel, same offset.
+- **KV-cache**: caller does `k_full = cat([k_cache, k_new], dim=2)` then calls this op.
+  No separate KV-cache softmax implementation allowed.
+
+Hooks:
+- `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable.
+- `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging,
+  and future KV-cache / training integration.
+- `backward_reference(q, k, v, dout, ...)` — runs the deterministic training backward
+  validation path and returns `dq`, `dk`, `dv`, `out`, `lse`, and provenance.
+- `compare_cp_attention_backward(q, k, v, dout, ...)` — compares CP=1 backward against
+  CP/chunked-prefill backward and emits whole-tensor plus per-logical-rank drift stats.
+
+## Tolerance
+
+| Scenario | Comparison | Tolerance |
+| --- | --- | --- |
+| Same physical shape, varying batch position/size/chunk | bitwise | `batch_invariance` (atol=0, rtol=0) |
+| Chunked-prefill on/off at same position | bitwise | `batch_invariance` |
+| Prefill tail vs decode slice | bitwise | `batch_invariance` |
+| CUDA vs `forward_fp32` output/grad | tolerance | `accuracy.default.attention` |
+| Valid-only vs padded (reduction width differs) | near-equal | accuracy tolerance (NOT bitwise) |
+
+## Memory Tradeoff (First Version)
+
+The first version materializes full FP32 `scores [B, Hq, Sq, Skv]` and `P [B, Hq, Sq, Skv]`.
+Memory cost: `4 * B * Hq * Sq * Skv` bytes per tensor. For Qwen3-8B at B=8, Sq=Skv=4096,
+Hq=32: each tensor is ~17 GB. This is acceptable for correctness verification and moderate
+sequence lengths but OOM-prone for long sequences. See `benchmarks/benchmark_deterministic_attention.py`
+for measured peak memory at representative shapes.
 
 ## Known Limitations
 
-- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work).
+- First version: `D=128` only (Qwen3-8B alignment).
+- Supported dtypes: BF16, FP16.
+- Full materialization of scores/P limits practical sequence length.
 - `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.
-- Covers softmax attention only; QK-Norm and RoPE are applied before the call.
+- CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op).
+- No FP8, no multi-GPU / sequence-parallel.
diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py
index 7deeddf0..bb0104ad 100644
--- a/rl_engine/kernels/attention_contract.py
+++ b/rl_engine/kernels/attention_contract.py
@@ -11,7 +11,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Iterable, TypeVar
 
@@ -55,6 +55,12 @@ class ReductionEngine(str, Enum):
     IN_OP_REFERENCE = "in_op_reference"
 
 
+class SplitKVMode(str, Enum):
+    DISABLED = "disabled"
+    FIXED = "fixed"
+    AUTO = "auto"
+
+
 class RoPEState(str, Enum):
     PRE_ROPE = "pre_rope"
     POST_ROPE = "post_rope"
@@ -284,6 +290,645 @@ def __post_init__(self) -> None:
             )
 
 
+@dataclass(frozen=True)
+class SplitKVExecutionPlan:
+    """Actual backend-local Split-KV schedule emitted by a runtime.
+
+    CP ownership is intentionally separate from this plan.  ``boundaries``
+    describe the canonical logical KV ranges reduced by one backend invocation;
+    CP may transport those partial states between ranks but must not change the
+    FP32 merge contract recorded here.
+    """
+
+    requested_mode: SplitKVMode
+    requested_split_size: int | None
+    actual_mode: SplitKVMode | None
+    actual_split_size: int | None
+    boundaries: tuple[tuple[int, int], ...]
+    merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX
+    acc_dtype: AttentionDType = AttentionDType.FP32
+    downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE
+    backend: str = "reference"
+    source: str = "contract_exact"
+    fallback: bool = False
+    fallback_reason: str | None = None
+
+    def __post_init__(self) -> None:
+        object.__setattr__(
+            self,
+            "requested_mode",
+            _enum_value(SplitKVMode, self.requested_mode, "requested_mode"),
+        )
+        if self.actual_mode is not None:
+            object.__setattr__(
+                self,
+                "actual_mode",
+                _enum_value(SplitKVMode, self.actual_mode, "actual_mode"),
+            )
+        for field_name in ("requested_split_size", "actual_split_size"):
+            value = getattr(self, field_name)
+            if value is not None:
+                _positive_int(value, field_name)
+        if self.requested_mode is SplitKVMode.FIXED:
+            if self.requested_split_size is None:
+                raise AttentionContractError(
+                    "requested fixed Split-KV mode requires requested_split_size"
+                )
+        elif self.requested_split_size is not None:
+            raise AttentionContractError(
+                "requested_split_size is only valid for requested fixed Split-KV mode"
+            )
+        try:
+            boundaries = tuple(tuple(boundary) for boundary in self.boundaries)
+        except TypeError as exc:
+            raise AttentionContractError(
+                "Split-KV boundaries must be an iterable of (start, end) pairs"
+            ) from exc
+        previous_end = 0
+        for index, boundary in enumerate(boundaries):
+            if len(boundary) != 2:
+                raise AttentionContractError(
+                    f"Split-KV boundary {index} must contain exactly start and end"
+                )
+            start, end = boundary
+            if (
+                isinstance(start, bool)
+                or isinstance(end, bool)
+                or not isinstance(start, int)
+                or not isinstance(end, int)
+                or start < 0
+                or end <= start
+            ):
+                raise AttentionContractError(
+                    "Split-KV boundaries must satisfy 0 <= start < end"
+                )
+            if index > 0 and start != previous_end:
+                raise AttentionContractError(
+                    "Split-KV boundaries must be contiguous and in logical KV order"
+                )
+            previous_end = end
+        if self.actual_mode is None and boundaries:
+            raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries")
+        if self.actual_mode is not None and not boundaries:
+            raise AttentionContractError("actual Split-KV plan must declare logical boundaries")
+        if self.actual_mode is SplitKVMode.FIXED:
+            if self.actual_split_size is None:
+                raise AttentionContractError(
+                    "actual fixed Split-KV mode requires actual_split_size"
+                )
+            widths = tuple(end - start for start, end in boundaries)
+            if any(width != self.actual_split_size for width in widths[:-1]) or (
+                widths and widths[-1] > self.actual_split_size
+            ):
+                raise AttentionContractError(
+                    "fixed Split-KV boundaries must use actual_split_size except "
+                    "for a shorter final split"
+                )
+        elif self.actual_split_size is not None:
+            raise AttentionContractError(
+                "actual_split_size is only valid for actual fixed Split-KV mode"
+            )
+        if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1:
+            raise AttentionContractError(
+                "disabled Split-KV execution must contain exactly one boundary"
+            )
+        object.__setattr__(self, "boundaries", boundaries)
+        object.__setattr__(
+            self,
+            "merge_order",
+            _enum_value(ReductionOrder, self.merge_order, "merge_order"),
+        )
+        object.__setattr__(
+            self,
+            "acc_dtype",
+            _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"),
+        )
+        object.__setattr__(
+            self,
+            "downcast_at",
+            _enum_value(DowncastPoint, self.downcast_at, "downcast_at"),
+        )
+        if self.acc_dtype is not AttentionDType.FP32:
+            raise AttentionContractError("Split-KV partial states must be merged in fp32")
+        if not isinstance(self.backend, str) or not self.backend.strip():
+            raise AttentionContractError("Split-KV backend must be a non-empty string")
+        if not isinstance(self.source, str) or not self.source.strip():
+            raise AttentionContractError("Split-KV plan source must be a non-empty string")
+        if not isinstance(self.fallback, bool):
+            raise AttentionContractError("Split-KV fallback must be a bool")
+        if self.fallback and not self.fallback_reason:
+            raise AttentionContractError("Split-KV fallback_reason is required for a fallback")
+        if not self.fallback and self.fallback_reason is not None:
+            raise AttentionContractError(
+                "Split-KV fallback_reason must be None when fallback=False"
+            )
+        if not self.fallback and self.actual_mode is not None:
+            if self.actual_mode is not self.requested_mode:
+                raise AttentionContractError(
+                    "actual Split-KV mode may differ from requested mode only for a fallback"
+                )
+            if self.actual_split_size != self.requested_split_size:
+                raise AttentionContractError(
+                    "actual Split-KV size may differ from requested size only for a fallback"
+                )
+
+    @property
+    def actual_split_count(self) -> int | None:
+        return None if self.actual_mode is None else len(self.boundaries)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "requested_split_kv_policy": self.requested_mode.value,
+            "requested_split_kv_size": self.requested_split_size,
+            "actual_split_kv_policy": (
+                None if self.actual_mode is None else self.actual_mode.value
+            ),
+            "actual_split_kv_size": self.actual_split_size,
+            "actual_split_kv_count": self.actual_split_count,
+            "actual_split_boundaries": [list(boundary) for boundary in self.boundaries],
+            "split_kv_merge_order": self.merge_order.value,
+            "split_kv_accum_dtype": self.acc_dtype.value,
+            "split_kv_downcast_at": self.downcast_at.value,
+            "split_kv_backend": self.backend,
+            "split_kv_plan_source": self.source,
+            "split_kv_fallback": self.fallback,
+            "split_kv_fallback_reason": self.fallback_reason,
+        }
+
+
+@dataclass(frozen=True)
+class SplitKVSpec:
+    """Requested Split-KV policy shared by training and rollout paths."""
+
+    mode: SplitKVMode = SplitKVMode.DISABLED
+    fixed_split_size: int | None = None
+    strict_consistency: bool = True
+
+    @classmethod
+    def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec":
+        return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency)
+
+    @classmethod
+    def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec":
+        return cls(
+            mode=SplitKVMode.FIXED,
+            fixed_split_size=split_size,
+            strict_consistency=strict_consistency,
+        )
+
+    @classmethod
+    def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec":
+        return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency)
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode"))
+        if not isinstance(self.strict_consistency, bool):
+            raise AttentionContractError("split_kv.strict_consistency must be a bool")
+        if self.mode is SplitKVMode.FIXED:
+            if self.fixed_split_size is None:
+                raise AttentionContractError(
+                    "fixed Split-KV policy requires fixed_split_size"
+                )
+            _positive_int(self.fixed_split_size, "split_kv.fixed_split_size")
+        elif self.fixed_split_size is not None:
+            raise AttentionContractError(
+                "fixed_split_size is only valid for fixed Split-KV policy"
+            )
+        if self.strict_consistency and self.mode is SplitKVMode.AUTO:
+            raise AttentionContractError(
+                "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency"
+            )
+
+    def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan:
+        """Resolve policies whose logical schedule is fully known by contract."""
+
+        total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens")
+        if self.mode is SplitKVMode.AUTO:
+            return SplitKVExecutionPlan(
+                requested_mode=self.mode,
+                requested_split_size=None,
+                actual_mode=None,
+                actual_split_size=None,
+                boundaries=(),
+                backend=backend,
+                source="runtime_required",
+            )
+        split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size
+        assert split_size is not None
+        boundaries = tuple(
+            (start, min(start + split_size, total_kv_tokens))
+            for start in range(0, total_kv_tokens, split_size)
+        )
+        return SplitKVExecutionPlan(
+            requested_mode=self.mode,
+            requested_split_size=self.fixed_split_size,
+            actual_mode=self.mode,
+            actual_split_size=self.fixed_split_size,
+            boundaries=boundaries,
+            backend=backend,
+            source="contract_exact",
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "mode": self.mode.value,
+            "fixed_split_size": self.fixed_split_size,
+            "strict_consistency": self.strict_consistency,
+        }
+
+
+def validate_split_kv_alignment(
+    training: SplitKVExecutionPlan,
+    rollout: SplitKVExecutionPlan,
+) -> None:
+    """Fail closed unless train and rollout executed the same logical plan."""
+
+    if training.actual_mode is None or rollout.actual_mode is None:
+        raise AttentionContractError(
+            "strict Split-KV alignment requires actual runtime plans from both sides"
+        )
+    fields = (
+        "requested_mode",
+        "requested_split_size",
+        "actual_mode",
+        "actual_split_size",
+        "boundaries",
+        "merge_order",
+        "acc_dtype",
+        "downcast_at",
+        "fallback",
+    )
+    mismatches = [
+        field_name
+        for field_name in fields
+        if getattr(training, field_name) != getattr(rollout, field_name)
+    ]
+    if mismatches:
+        raise AttentionContractError(
+            "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches)
+        )
+
+
+@dataclass(frozen=True, order=True)
+class SplitKVRuntimeCoordinate:
+    """Identity of one batch/rank/owner Split-KV runtime plan."""
+
+    batch_index: int
+    tp_rank: int
+    cp_rank: int
+    owner_cp_rank: int
+
+    def validate(
+        self,
+        *,
+        batch_size: int,
+        tp_world_size: int,
+        cp_world_size: int,
+    ) -> None:
+        batch_index = _non_negative_int(self.batch_index, "batch_index")
+        tp_rank = _non_negative_int(self.tp_rank, "tp_rank")
+        cp_rank = _non_negative_int(self.cp_rank, "cp_rank")
+        owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank")
+        if batch_index >= batch_size:
+            raise AttentionContractError("Split-KV batch_index is outside batch_size")
+        if tp_rank >= tp_world_size:
+            raise AttentionContractError("Split-KV tp_rank is outside tp_world_size")
+        if cp_rank >= cp_world_size:
+            raise AttentionContractError("Split-KV cp_rank is outside cp_world_size")
+        if owner_cp_rank >= cp_world_size:
+            raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size")
+
+    def to_dict(self) -> dict[str, int]:
+        return {
+            "batch_index": self.batch_index,
+            "tp_rank": self.tp_rank,
+            "cp_rank": self.cp_rank,
+            "owner_cp_rank": self.owner_cp_rank,
+        }
+
+
+@dataclass(frozen=True)
+class SplitKVRuntimePlanEntry:
+    """Actual plan for one batch/TP/CP consumer and logical KV owner."""
+
+    coordinate: SplitKVRuntimeCoordinate
+    expected_kv_range: tuple[int, int]
+    execution: SplitKVExecutionPlan
+
+    def validate(
+        self,
+        *,
+        batch_size: int,
+        tp_world_size: int,
+        cp_world_size: int,
+        total_kv_tokens: int,
+    ) -> None:
+        self.coordinate.validate(
+            batch_size=batch_size,
+            tp_world_size=tp_world_size,
+            cp_world_size=cp_world_size,
+        )
+        try:
+            start, end = self.expected_kv_range
+        except (TypeError, ValueError) as exc:
+            raise AttentionContractError(
+                "expected_kv_range must contain exactly (start, end)"
+            ) from exc
+        if (
+            isinstance(start, bool)
+            or isinstance(end, bool)
+            or not isinstance(start, int)
+            or not isinstance(end, int)
+            or start < 0
+            or end <= start
+            or end > total_kv_tokens
+        ):
+            raise AttentionContractError(
+                "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens"
+            )
+        if self.execution.actual_mode is None:
+            raise AttentionContractError(
+                "complete Split-KV plan sets require actual runtime plans"
+            )
+        if (
+            self.execution.boundaries[0][0] != start
+            or self.execution.boundaries[-1][1] != end
+        ):
+            raise AttentionContractError(
+                "Split-KV execution boundaries must exactly cover expected_kv_range"
+            )
+        if any(
+            boundary_start < start or boundary_end > end
+            for boundary_start, boundary_end in self.execution.boundaries
+        ):
+            raise AttentionContractError(
+                "Split-KV execution boundary escapes expected_kv_range"
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            **self.coordinate.to_dict(),
+            "expected_kv_range": list(self.expected_kv_range),
+            **self.execution.to_dict(),
+        }
+
+
+@dataclass(frozen=True)
+class SplitKVRuntimePlanSet:
+    """Complete actual Split-KV plans across batch, TP, CP, and KV owners.
+
+    Every CP consumer must report the plan used for every CP-owned KV range.
+    This duplicates owner plans across consumers intentionally: it detects one
+    rank silently choosing a different Split-K schedule or merge policy.
+    """
+
+    batch_size: int
+    tp_world_size: int
+    cp_world_size: int
+    total_kv_tokens: tuple[int, ...]
+    entries: tuple[SplitKVRuntimePlanEntry, ...]
+
+    def __post_init__(self) -> None:
+        batch_size = _positive_int(self.batch_size, "batch_size")
+        tp_world_size = _positive_int(self.tp_world_size, "tp_world_size")
+        cp_world_size = _positive_int(self.cp_world_size, "cp_world_size")
+        totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens")
+        if len(totals) != batch_size or any(total <= 0 for total in totals):
+            raise AttentionContractError(
+                "total_kv_tokens must contain one positive length per batch item"
+            )
+        object.__setattr__(self, "total_kv_tokens", totals)
+        entries = tuple(self.entries)
+        object.__setattr__(self, "entries", entries)
+        expected_coordinates = {
+            SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank)
+            for batch_index in range(batch_size)
+            for tp_rank in range(tp_world_size)
+            for cp_rank in range(cp_world_size)
+            for owner_cp_rank in range(cp_world_size)
+        }
+        actual_coordinates = [entry.coordinate for entry in entries]
+        if len(set(actual_coordinates)) != len(actual_coordinates):
+            raise AttentionContractError(
+                "Split-KV runtime plan set contains duplicate coordinates"
+            )
+        missing = expected_coordinates.difference(actual_coordinates)
+        extra = set(actual_coordinates).difference(expected_coordinates)
+        if missing or extra:
+            raise AttentionContractError(
+                "Split-KV runtime plan set coordinate coverage is incomplete; "
+                f"missing={_format_split_kv_coordinates(missing)}, "
+                f"extra={_format_split_kv_coordinates(extra)}"
+            )
+        for entry in entries:
+            entry.validate(
+                batch_size=batch_size,
+                tp_world_size=tp_world_size,
+                cp_world_size=cp_world_size,
+                total_kv_tokens=totals[entry.coordinate.batch_index],
+            )
+        self._validate_owner_coverage()
+        self._validate_rank_invariance()
+
+    def _validate_owner_coverage(self) -> None:
+        for batch_index, total in enumerate(self.total_kv_tokens):
+            for tp_rank in range(self.tp_world_size):
+                ranges = []
+                for owner_cp_rank in range(self.cp_world_size):
+                    matches = [
+                        entry
+                        for entry in self.entries
+                        if entry.coordinate.batch_index == batch_index
+                        and entry.coordinate.tp_rank == tp_rank
+                        and entry.coordinate.cp_rank == 0
+                        and entry.coordinate.owner_cp_rank == owner_cp_rank
+                    ]
+                    ranges.append(matches[0].expected_kv_range)
+                previous_end = 0
+                for start, end in ranges:
+                    if start != previous_end:
+                        raise AttentionContractError(
+                            "Split-KV owner ranges must be gap-free in CP owner order"
+                        )
+                    previous_end = end
+                if previous_end != total:
+                    raise AttentionContractError(
+                        "Split-KV owner ranges do not cover total_kv_tokens"
+                    )
+
+    def _validate_rank_invariance(self) -> None:
+        for batch_index in range(self.batch_size):
+            for owner_cp_rank in range(self.cp_world_size):
+                entries = sorted(
+                    (
+                        entry
+                        for entry in self.entries
+                        if entry.coordinate.batch_index == batch_index
+                        and entry.coordinate.owner_cp_rank == owner_cp_rank
+                    ),
+                    key=lambda entry: (
+                        entry.coordinate.tp_rank,
+                        entry.coordinate.cp_rank,
+                    ),
+                )
+                reference = entries[0]
+                for entry in entries[1:]:
+                    if entry.expected_kv_range != reference.expected_kv_range:
+                        raise AttentionContractError(
+                            "Split-KV owner range differs across TP/CP consumers"
+                        )
+                    try:
+                        validate_split_kv_alignment(
+                            reference.execution,
+                            entry.execution,
+                        )
+                    except AttentionContractError as exc:
+                        raise AttentionContractError(
+                            "Split-KV plan differs across TP/CP consumers for "
+                            f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}"
+                        ) from exc
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "batch_size": self.batch_size,
+            "tp_world_size": self.tp_world_size,
+            "cp_world_size": self.cp_world_size,
+            "total_kv_tokens": list(self.total_kv_tokens),
+            "entries": [
+                entry.to_dict()
+                for entry in sorted(self.entries, key=lambda entry: entry.coordinate)
+            ],
+            "coverage": "complete_batch_tp_cp_owner_cartesian_product",
+        }
+
+
+def validate_split_kv_plan_set_alignment(
+    training: SplitKVRuntimePlanSet,
+    rollout: SplitKVRuntimePlanSet,
+) -> None:
+    """Fail closed unless complete train/rollout runtime plan sets align."""
+
+    topology_fields = (
+        "batch_size",
+        "tp_world_size",
+        "cp_world_size",
+        "total_kv_tokens",
+    )
+    topology_mismatches = [
+        field_name
+        for field_name in topology_fields
+        if getattr(training, field_name) != getattr(rollout, field_name)
+    ]
+    if topology_mismatches:
+        raise AttentionContractError(
+            "training/rollout Split-KV plan-set topology differs: "
+            + ", ".join(topology_mismatches)
+        )
+    training_by_coordinate = {
+        entry.coordinate: entry for entry in training.entries
+    }
+    rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries}
+    if training_by_coordinate.keys() != rollout_by_coordinate.keys():
+        raise AttentionContractError(
+            "training/rollout Split-KV plan-set coordinates differ"
+        )
+    for coordinate in sorted(training_by_coordinate):
+        train_entry = training_by_coordinate[coordinate]
+        rollout_entry = rollout_by_coordinate[coordinate]
+        if train_entry.expected_kv_range != rollout_entry.expected_kv_range:
+            raise AttentionContractError(
+                f"training/rollout expected KV range differs at {coordinate}"
+            )
+        try:
+            validate_split_kv_alignment(
+                train_entry.execution,
+                rollout_entry.execution,
+            )
+        except AttentionContractError as exc:
+            raise AttentionContractError(
+                f"training/rollout Split-KV plan differs at {coordinate}: {exc}"
+                ) from exc
+
+
+def build_split_kv_runtime_plan_set(
+    total_kv_tokens: Iterable[int],
+    *,
+    tp_world_size: int,
+    cp_world_size: int,
+    split_kv: SplitKVSpec,
+    backend: str = "contract_reference",
+) -> SplitKVRuntimePlanSet:
+    """Build a complete owner-local plan set for contract tests and adapters."""
+
+    totals = _integer_tuple(total_kv_tokens, "total_kv_tokens")
+    if not totals or any(total < cp_world_size for total in totals):
+        raise AttentionContractError(
+            "contract plan sets require at least one KV token per CP owner"
+        )
+    tp_world_size = _positive_int(tp_world_size, "tp_world_size")
+    cp_world_size = _positive_int(cp_world_size, "cp_world_size")
+    if not isinstance(split_kv, SplitKVSpec):
+        raise AttentionContractError("split_kv must be a SplitKVSpec")
+
+    entries: list[SplitKVRuntimePlanEntry] = []
+    for batch_index, total in enumerate(totals):
+        base = total // cp_world_size
+        remainder = total % cp_world_size
+        owner_ranges: list[tuple[int, int]] = []
+        start = 0
+        for owner_cp_rank in range(cp_world_size):
+            end = start + base + (1 if owner_cp_rank < remainder else 0)
+            owner_ranges.append((start, end))
+            start = end
+        for tp_rank in range(tp_world_size):
+            for cp_rank in range(cp_world_size):
+                for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges):
+                    local_total = owner_end - owner_start
+                    local = split_kv.resolve(local_total, backend=backend)
+                    execution = SplitKVExecutionPlan(
+                        requested_mode=local.requested_mode,
+                        requested_split_size=local.requested_split_size,
+                        actual_mode=local.actual_mode,
+                        actual_split_size=local.actual_split_size,
+                        boundaries=tuple(
+                            (owner_start + start, owner_start + end)
+                            for start, end in local.boundaries
+                        ),
+                        merge_order=local.merge_order,
+                        acc_dtype=local.acc_dtype,
+                        downcast_at=local.downcast_at,
+                        backend=local.backend,
+                        source=local.source,
+                        fallback=local.fallback,
+                        fallback_reason=local.fallback_reason,
+                    )
+                    entries.append(
+                        SplitKVRuntimePlanEntry(
+                            coordinate=SplitKVRuntimeCoordinate(
+                                batch_index=batch_index,
+                                tp_rank=tp_rank,
+                                cp_rank=cp_rank,
+                                owner_cp_rank=owner_cp_rank,
+                            ),
+                            expected_kv_range=(owner_start, owner_end),
+                            execution=execution,
+                        )
+                    )
+    return SplitKVRuntimePlanSet(
+        batch_size=len(totals),
+        tp_world_size=tp_world_size,
+        cp_world_size=cp_world_size,
+        total_kv_tokens=totals,
+        entries=tuple(entries),
+    )
+
+
+def _format_split_kv_coordinates(
+    coordinates: Iterable[SplitKVRuntimeCoordinate],
+) -> list[dict[str, int]]:
+    return [coordinate.to_dict() for coordinate in sorted(coordinates)]
+
+
 @dataclass(frozen=True)
 class KVCacheSpec:
     """Logical identity of the paged/block KV cache used for replay."""
@@ -536,6 +1181,7 @@ class AttentionContract:
     causal_offsets: tuple[int, ...] | None
     sharding: ShardingSpec
     reduction: ReductionSpec
+    split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled)
     kv_cache: KVCacheSpec | None = None
     rope: RoPESpec | None = None
     export_lse: bool = True
@@ -551,6 +1197,8 @@ 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 not isinstance(self.split_kv, SplitKVSpec):
+            raise AttentionContractError("split_kv must be a SplitKVSpec")
         if (
             self.mode is AttentionMode.PREFILL
             and query_sequence_length != self.sharding.local_sequence_length
@@ -713,6 +1361,7 @@ def to_dict(self) -> dict[str, Any]:
             "lse_domain": "attention",
             "sharding": sharding,
             "reduction": reduction,
+            "split_kv": self.split_kv.to_dict(),
             "kv_cache": kv_cache,
             "rope": rope,
         }
@@ -734,6 +1383,10 @@ class AttentionBackendCapability:
     supports_kv_cache: bool = False
     supports_rope_metadata: bool = False
     supports_fused_rope_attention: bool = False
+    supports_split_kv_disabled: bool = True
+    supports_split_kv_fixed: bool = False
+    supports_split_kv_auto: bool = False
+    reports_actual_split_kv_plan: bool = False
     implementation_kind: str = "production"
 
     def __post_init__(self) -> None:
@@ -763,6 +1416,10 @@ def __post_init__(self) -> None:
             "supports_kv_cache",
             "supports_rope_metadata",
             "supports_fused_rope_attention",
+            "supports_split_kv_disabled",
+            "supports_split_kv_fixed",
+            "supports_split_kv_auto",
+            "reports_actual_split_kv_plan",
         ):
             if not isinstance(getattr(self, field), bool):
                 raise AttentionContractError(f"{field} must be a bool")
@@ -811,6 +1468,17 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]:
             and not self.supports_fused_rope_attention
         ):
             reasons.append("fused RoPE+Attention boundary is unsupported")
+        split_support = {
+            SplitKVMode.DISABLED: self.supports_split_kv_disabled,
+            SplitKVMode.FIXED: self.supports_split_kv_fixed,
+            SplitKVMode.AUTO: self.supports_split_kv_auto,
+        }
+        if not split_support[contract.split_kv.mode]:
+            reasons.append(
+                f"Split-KV policy={contract.split_kv.mode.value} is unsupported"
+            )
+        if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan:
+            reasons.append("actual Split-KV execution-plan provenance is unsupported")
         return tuple(reasons)
 
     def supports(self, contract: AttentionContract) -> bool:
@@ -830,6 +1498,10 @@ def to_dict(self) -> dict[str, Any]:
             "supports_kv_cache": self.supports_kv_cache,
             "supports_rope_metadata": self.supports_rope_metadata,
             "supports_fused_rope_attention": self.supports_fused_rope_attention,
+            "supports_split_kv_disabled": self.supports_split_kv_disabled,
+            "supports_split_kv_fixed": self.supports_split_kv_fixed,
+            "supports_split_kv_auto": self.supports_split_kv_auto,
+            "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan,
             "implementation_kind": self.implementation_kind,
         }
 
@@ -862,4 +1534,13 @@ class AttentionDispatchResult:
     "RoPESpec",
     "RoPEState",
     "ShardingSpec",
+    "build_split_kv_runtime_plan_set",
+    "SplitKVExecutionPlan",
+    "SplitKVMode",
+    "SplitKVRuntimeCoordinate",
+    "SplitKVRuntimePlanEntry",
+    "SplitKVRuntimePlanSet",
+    "SplitKVSpec",
+    "validate_split_kv_alignment",
+    "validate_split_kv_plan_set_alignment",
 ]
diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py
index f124cafb..78a9cfed 100644
--- a/rl_engine/kernels/gtest/operator_inputs.py
+++ b/rl_engine/kernels/gtest/operator_inputs.py
@@ -27,9 +27,12 @@ def make_operator_inputs(
     builders = {
         "rms_norm": _make_rms_norm_inputs,
         "matmul": _make_matmul_inputs,
+        "det_gemm": _make_det_gemm_inputs,
         "attention": _make_attention_inputs,
+        "cp_attention": _make_attention_inputs,
         "logp": _make_logp_inputs,
         "linear_logp": _make_linear_logp_inputs,
+        "batch_invariant_logp": _make_batch_invariant_logp_inputs,
         "rope": _make_rope_inputs,
         "silu": _make_silu_inputs,
         "swiglu": _make_swiglu_inputs,
@@ -49,14 +52,17 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str:
     names = {
         "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}",
         "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}",
+        "det_gemm": 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}",
+        "batch_invariant_logp": f"{batch}x{seq}x{vocab}",
         "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}",
         "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}",
         "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}",
-        "embedding": f"{batch}x{seq}x{vocab}x{DEFAULT_HIDDEN}",
-        "lm_head": f"{batch}x{seq}x{vocab}",
+        "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}",
+        "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}",
         "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}",
     }
     try:
@@ -89,23 +95,51 @@ def _make_matmul_inputs(
     }
 
 
-def _make_attention_inputs(
+def _make_det_gemm_inputs(
     args: argparse.Namespace, dtype: torch.dtype, device: torch.device
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
+    k_dim = _matmul_k(args)
+    n_dim = _matmul_n(args)
+    m_dim = batch * seq
     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,
+        "a": _floating_tensor((m_dim, k_dim), args, dtype, device, offset=0),
+        "b": _floating_tensor((k_dim, n_dim), args, dtype, device, offset=1),
+    }
+
+
+def _make_attention_inputs(
+    args: argparse.Namespace, dtype: torch.dtype, device: torch.device
+) -> dict[str, Any]:
+    batch, seq = _batch_seq(args)
+    skv = _arg_int(args, "skv", seq)
+    n_heads = _arg_int(args, "n_heads", DEFAULT_N_HEADS)
+    n_kv_heads = _arg_int(args, "n_kv_heads", DEFAULT_N_KV_HEADS)
+    causal = bool(_arg_int(args, "causal", 1))
+    use_padding = bool(_arg_int(args, "use_padding", 0))
+    scale_mode = _arg_str(args, "scale_mode", "default")
+
+    inputs: dict[str, Any] = {
+        "q": _floating_tensor((batch, n_heads, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0),
+        "k": _floating_tensor((batch, n_kv_heads, skv, DEFAULT_HEAD_DIM), args, dtype, device, 1),
+        "v": _floating_tensor((batch, n_kv_heads, skv, DEFAULT_HEAD_DIM), args, dtype, device, 2),
+        "causal": causal,
     }
 
+    if scale_mode == "zero":
+        inputs["scale"] = 0.0
+    elif scale_mode == "custom":
+        inputs["scale"] = 0.05
+    # else: scale_mode == "default" -> no scale kwarg (uses 1/sqrt(D))
+
+    if use_padding:
+        generator = _generator(args, device, offset=42)
+        key_padding_mask = torch.rand((batch, skv), generator=generator, device=device) > 0.3
+        key_padding_mask[:, 0] = True
+        inputs["key_padding_mask"] = key_padding_mask
+
+    return inputs
+
 
 def _make_logp_inputs(
     args: argparse.Namespace, dtype: torch.dtype, device: torch.device
@@ -132,6 +166,17 @@ def _make_linear_logp_inputs(
     }
 
 
+def _make_batch_invariant_logp_inputs(
+    args: argparse.Namespace, dtype: torch.dtype, device: torch.device
+) -> dict[str, Any]:
+    batch, seq = _batch_seq(args)
+    vocab = _arg_int(args, "vocab", DEFAULT_VOCAB)
+    return {
+        "logits": _floating_tensor((batch, seq, vocab), args, dtype, device, offset=0),
+        "target_ids": _token_ids((batch, seq), vocab, args, device),
+    }
+
+
 def _make_rope_inputs(
     args: argparse.Namespace, dtype: torch.dtype, device: torch.device
 ) -> dict[str, Any]:
@@ -169,9 +214,10 @@ def _make_embedding_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     vocab = _arg_int(args, "vocab", DEFAULT_VOCAB)
+    hidden_dim = _normalized_dim(args)
     return {
         "token_ids": _token_ids((batch, seq), vocab, args, device),
-        "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 0),
+        "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 0),
     }
 
 
@@ -180,9 +226,10 @@ def _make_lm_head_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     vocab = _arg_int(args, "vocab", DEFAULT_VOCAB)
+    hidden_dim = _normalized_dim(args)
     return {
-        "hidden": _floating_tensor((batch, seq, DEFAULT_HIDDEN), args, dtype, device, 0),
-        "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 1),
+        "hidden": _floating_tensor((batch, seq, hidden_dim), args, dtype, device, 0),
+        "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 1),
         "bias": None,
     }
 
diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py
index 55a4a203..454f00ba 100644
--- a/rl_engine/kernels/gtest/operator_specs.py
+++ b/rl_engine/kernels/gtest/operator_specs.py
@@ -32,6 +32,53 @@ def _load_object(path: str) -> Any:
 
 
 OP_SPECS = {
+    "rms_norm": OperatorSpec(
+        name="rms_norm",
+        op_class="reduction",
+        gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp",
+            "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp",
+            "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp",
+        },
+        grad_input_names=("x", "weight"),
+    ),
+    "attention": OperatorSpec(
+        name="attention",
+        op_class="attention",
+        gold_path="rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp",
+            "triton": (
+                "rl_engine.kernels.ops.triton.attention.standard_attn."
+                "TritonBatchInvariantAttentionOp"
+            ),
+            "cuda": (
+                "rl_engine.kernels.ops.cuda.attention.deterministic_attn."
+                "DeterministicAttentionOp"
+            ),
+        },
+        grad_input_names=("q", "k", "v"),
+    ),
+    "cp_attention": OperatorSpec(
+        name="cp_attention",
+        op_class="attention",
+        gold_path=(
+            "rl_engine.kernels.ops.pytorch.attention.cp_attention."
+            "DeterministicCPAttentionReferenceOp"
+        ),
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": (
+                "rl_engine.kernels.ops.pytorch.attention.cp_attention."
+                "DeterministicCPAttentionReferenceOp"
+            ),
+        },
+        grad_input_names=("q", "k", "v"),
+    ),
     "logp": OperatorSpec(
         name="logp",
         op_class="logprob",
@@ -57,6 +104,92 @@ def _load_object(path: str) -> Any:
         },
         grad_input_names=("hidden", "lm_head_weight"),
     ),
+    "embedding": OperatorSpec(
+        name="embedding",
+        op_class="elementwise",
+        gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp",
+        },
+        grad_input_names=("weight",),
+    ),
+    "lm_head": OperatorSpec(
+        name="lm_head",
+        op_class="reduction",
+        gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp",
+        },
+        grad_input_names=("hidden", "weight"),
+    ),
+    "det_gemm": OperatorSpec(
+        name="det_gemm",
+        op_class="reduction",
+        gold_path="rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp",
+        gold_method="__call__",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp",
+            "cuda": "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp",
+            "triton": "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp",
+        },
+        grad_input_names=("a", "b"),
+    ),
+    "rope": OperatorSpec(
+        name="rope",
+        op_class="elementwise",
+        gold_path="rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp",
+            "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op",
+        },
+        grad_input_names=("x",),
+    ),
+    "silu": OperatorSpec(
+        name="silu",
+        op_class="elementwise",
+        gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp",
+            "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp",
+            "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp",
+        },
+        grad_input_names=("x",),
+    ),
+    "swiglu": OperatorSpec(
+        name="swiglu",
+        op_class="elementwise",
+        gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp",
+            "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp",
+            "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp",
+        },
+        grad_input_names=("gate", "up"),
+    ),
+    "batch_invariant_logp": OperatorSpec(
+        name="batch_invariant_logp",
+        op_class="logprob",
+        gold_path="rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp."
+        "NativeBatchInvariantLogpOp",
+        gold_method="apply",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp."
+            "NativeBatchInvariantLogpOp",
+            "triton": "rl_engine.kernels.ops.triton.loss.batch_invariant_logp."
+            "TritonBatchInvariantLogpOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp."
+            "BatchInvariantLogpSM90Op",
+        },
+        grad_input_names=("logits",),
+    ),
 }
 
 
diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py
index 8d6addd9..2c3a1f1b 100644
--- a/rl_engine/kernels/ops/cuda/attention/__init__.py
+++ b/rl_engine/kernels/ops/cuda/attention/__init__.py
@@ -1,9 +1,49 @@
 # File: rl_engine/kernels/ops/cuda/attention/__init__.py
 
+from .cp_comm import (
+    AttentionCPBlockMetadata,
+    AttentionCPCommunication,
+    AttentionCPCommunicationPlan,
+    AttentionCPCommunicationUnavailable,
+    AttentionCPMergedState,
+    AttentionCPPartialState,
+    AttentionParallelSpec,
+    CPCommunicationBackend,
+    CPCommunicationStatus,
+    CUDAAGRSAttentionCPCommunication,
+    P2PNCCLAttentionCPCommunication,
+    sort_attention_cp_partial_states,
+)
+from .deterministic_attn import DeterministicAttentionOp
 from .flash_attn import FlashAttentionOp
+from .flashinfer_paged_attention import (
+    FlashInferPagedAttentionConfig,
+    FlashInferQwen3PagedAttentionOp,
+    FlashInferRoPEFusionConfig,
+    FlashInferSplitKVPolicy,
+    FlashInferUnavailable,
+)
 from .prefix_shared_attn import PrefixSharedAttentionOp
 
 __all__ = [
+    "AttentionCPBlockMetadata",
+    "AttentionCPCommunication",
+    "AttentionCPCommunicationPlan",
+    "AttentionCPCommunicationUnavailable",
+    "AttentionCPMergedState",
+    "AttentionCPPartialState",
+    "AttentionParallelSpec",
+    "CPCommunicationBackend",
+    "CPCommunicationStatus",
+    "CUDAAGRSAttentionCPCommunication",
+    "P2PNCCLAttentionCPCommunication",
+    "DeterministicAttentionOp",
     "FlashAttentionOp",
+    "FlashInferPagedAttentionConfig",
+    "FlashInferQwen3PagedAttentionOp",
+    "FlashInferRoPEFusionConfig",
+    "FlashInferSplitKVPolicy",
+    "FlashInferUnavailable",
     "PrefixSharedAttentionOp",
+    "sort_attention_cp_partial_states",
 ]
diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py
index cab339fd..edd5efd5 100644
--- a/rl_engine/kernels/registry.py
+++ b/rl_engine/kernels/registry.py
@@ -6,6 +6,8 @@
 from enum import Enum, EnumMeta
 from typing import Any, Dict, Optional, Set, Type
 
+import torch
+
 from rl_engine.kernels.attention_contract import (
     AttentionBackendCapability,
     AttentionContract,
@@ -39,6 +41,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     CUDA_FUSED_LOGP_SM90 = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op"
     CUDA_FUSED_LOGP_GENERIC = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp"
     CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp"
+    # Deterministic standard-softmax attention (issue #147); not FlashAttention.
+    CUDA_DETERMINISTIC_ATTENTION = (
+        "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp"
+    )
 
     # AMD ROCm optimized stack
     ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp"
@@ -59,6 +65,25 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     TRITON_RATIO_KL = "rl_engine.kernels.ops.triton.loss.ratio_kl.TritonRatioKLOp"
     PYTORCH_RATIO_KL = "rl_engine.kernels.ops.pytorch.loss.ratio_kl.NativeRatioKLOp"
 
+    # Variable-length packing (pack-and-pad), [B,S,...] -> [Total_Active,...]
+    PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp"
+    # Batch-invariant deterministic GEMM (WS1 #146)
+    CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp"
+    TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp"
+    # NON-deterministic reference (torch.matmul); reference/benchmark ONLY,
+    # intentionally excluded from det_gemm dispatch (cuBLAS breaks invariance).
+    PYTORCH_GEMM = "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp"
+    # Batch-invariant selected-logprob (WS1 #148: locked reduction order)
+    TRITON_BATCH_INVARIANT_LOGP = (
+        "rl_engine.kernels.ops.triton.loss.batch_invariant_logp.TritonBatchInvariantLogpOp"
+    )
+    PYTORCH_BATCH_INVARIANT_LOGP = (
+        "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvariantLogpOp"
+    )
+    CUDA_BATCH_INVARIANT_LOGP_SM90 = (
+        "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op"
+    )
+
     # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth)
     PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp"
 
@@ -68,8 +93,14 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp"
     PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp"
     PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp"
+    TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp"
+    CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op"
     PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp"
     PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp"
+    CUDA_SILU = "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp"
+    CUDA_SWIGLU = "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp"
+    TRITON_SILU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp"
+    TRITON_SWIGLU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp"
 
     # WS1 pure-PyTorch ground-truth attention reference (hand-written fp32 softmax).
     # Distinct from PYTORCH_ATTN above, which is the production SDPA fallback.
@@ -81,10 +112,18 @@ 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
     PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp"
+    CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp"
+    CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp"
 
 
 def resolve_logp_op_type(
@@ -175,6 +214,27 @@ def __init__(self):
                 supports_kv_cache=False,
                 implementation_kind="reference",
             ),
+            OpBackend.PYTORCH_CP_ATTENTION: AttentionBackendCapability(
+                backend_id="pytorch-deterministic-cp-attention-reference",
+                roles=common_roles,
+                modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}),
+                dtypes=common_dtypes,
+                tp_world_sizes=(1, 2),
+                cp_world_sizes=(1, 2),
+                exports_attention_lse=True,
+                deterministic_cp_merge=True,
+                supports_packed_varlen=False,
+                supports_kv_cache=False,
+                # PR3 consumes attention-ready post-RoPE Q/K; RoPE execution
+                # and fused boundary validation remain in PR2/PR7 harnesses.
+                supports_rope_metadata=False,
+                supports_fused_rope_attention=False,
+                supports_split_kv_disabled=True,
+                supports_split_kv_fixed=True,
+                supports_split_kv_auto=False,
+                reports_actual_split_kv_plan=True,
+                implementation_kind="deterministic",
+            ),
         }
 
         self._priority_map = {
@@ -206,22 +266,55 @@ def __init__(self):
                     OpBackend.PYTORCH_NATIVE,
                 ],
                 "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN],
-                "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION],
+                "attention": [
+                    OpBackend.CUDA_DETERMINISTIC_ATTENTION,
+                    OpBackend.PYTORCH_NATIVE_ATTENTION,
+                ],
+                "ws2_attention": [
+                    OpBackend.PYTORCH_CP_ATTENTION,
+                    OpBackend.CUDA_DETERMINISTIC_ATTENTION,
+                    OpBackend.PYTORCH_NATIVE_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],
+                "linear_logp": [
+                    OpBackend.TRITON_LINEAR_LOGP,
+                    OpBackend.PYTORCH_LINEAR_LOGP,
+                ],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
+                "pack": [OpBackend.PYTORCH_PACK],
+                "det_gemm": [OpBackend.CUDA_DET_GEMM, OpBackend.TRITON_DET_GEMM],
+                "batch_invariant_logp": [
+                    OpBackend.TRITON_BATCH_INVARIANT_LOGP,
+                    OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
+                ],
                 "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM],
                 "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
                 "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING],
-                "silu": [OpBackend.PYTORCH_NATIVE_SILU],
-                "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU],
+                "silu": [
+                    OpBackend.CUDA_SILU,
+                    OpBackend.TRITON_SILU,
+                    OpBackend.PYTORCH_NATIVE_SILU,
+                ],
+                "swiglu": [
+                    OpBackend.CUDA_SWIGLU,
+                    OpBackend.TRITON_SWIGLU,
+                    OpBackend.PYTORCH_NATIVE_SWIGLU,
+                ],
                 # Default dispatch logic for new operators
                 "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL],
-                "rope": [OpBackend.PYTORCH_NATIVE_ROPE],
+                "rope": [
+                    OpBackend.CUDA_ROPE_SM90,
+                    OpBackend.TRITON_ROPE,
+                    OpBackend.PYTORCH_NATIVE_ROPE,
+                ],
             },
             "rocm": {
-                "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE],
+                "logp": [
+                    OpBackend.ROCM_AITER,
+                    OpBackend.TRITON_GENERIC,
+                    OpBackend.PYTORCH_NATIVE,
+                ],
                 "logp_deterministic": [OpBackend.PYTORCH_NATIVE],
                 "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE],
                 "attn": [
@@ -230,17 +323,27 @@ def __init__(self):
                     OpBackend.TRITON_GENERIC,
                 ],
                 "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION],
+                "ws2_attention": [
+                    OpBackend.PYTORCH_CP_ATTENTION,
+                    OpBackend.PYTORCH_NATIVE_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],
+                "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE],
                 "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
+                "pack": [OpBackend.PYTORCH_PACK],
+                "det_gemm": [OpBackend.TRITON_DET_GEMM],
+                "batch_invariant_logp": [
+                    OpBackend.TRITON_BATCH_INVARIANT_LOGP,
+                    OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
+                ],
                 "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL],
                 "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM],
                 "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
                 "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING],
-                "silu": [OpBackend.PYTORCH_NATIVE_SILU],
-                "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU],
+                "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU],
+                "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU],
             },
             "cpu": {
                 "logp": [OpBackend.PYTORCH_NATIVE],
@@ -248,11 +351,18 @@ def __init__(self):
                 "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE],
                 "attn": [OpBackend.PYTORCH_ATTN],
                 "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION],
+                "ws2_attention": [
+                    OpBackend.PYTORCH_CP_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],
                 "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP],
                 "ratio_kl": [OpBackend.PYTORCH_RATIO_KL],
+                "pack": [OpBackend.PYTORCH_PACK],
+                "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP],
                 "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL],
                 "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM],
                 "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
@@ -279,7 +389,11 @@ def _adjust_priority_from_env(self):
                 OpBackend.ROCM_FLASH_ATTN,
                 OpBackend.TRITON_GENERIC,
             ]
-        elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}:
+        elif rocm_attn_backend and rocm_attn_backend not in {
+            "native",
+            "pytorch",
+            "sdpa",
+        }:
             logger.warning(
                 "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.",
                 rocm_attn_backend,
@@ -315,24 +429,38 @@ def _adjust_priority_for_hardware(self):
                 ll_list = self._priority_map["cuda"]["linear_logp"]
                 if OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90 not in ll_list:
                     ll_list.insert(0, OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90)
+
+            # Batch-invariant logp SM90 kernel: same sm_90a TMA gating (Hopper only).
+            batch_inv_compiled = _EXT_AVAILABLE and hasattr(_C, "batch_invariant_logp_sm90")
+            if batch_inv_compiled and cc_major == 9:
+                bi_list = self._priority_map["cuda"]["batch_invariant_logp"]
+                if OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90 not in bi_list:
+                    bi_list.insert(0, OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90)
             elif cc >= 90:
                 logger.debug(
                     f"SM{cc}: fused linear-logp SM90 kernel not compiled into _C; "
                     "using generic linear-logp backend."
                 )
+
+            sm90_embedding_compiled = _EXT_AVAILABLE and hasattr(_C, "embedding_sm90_forward")
+            if sm90_embedding_compiled and cc_major == 9:
+                embedding_list = self._priority_map["cuda"]["embedding"]
+                if OpBackend.CUDA_SM90_EMBEDDING not in embedding_list:
+                    embedding_list.insert(0, OpBackend.CUDA_SM90_EMBEDDING)
+
+            sm90_lm_head_compiled = _EXT_AVAILABLE and hasattr(_C, "lm_head_sm90_forward")
+            if sm90_lm_head_compiled and cc_major == 9:
+                lm_head_list = self._priority_map["cuda"]["lm_head"]
+                if OpBackend.CUDA_SM90_LM_HEAD not in lm_head_list:
+                    lm_head_list.insert(0, OpBackend.CUDA_SM90_LM_HEAD)
         except Exception as e:
             logger.warning(f"Failed to probe device capability: {e}")
 
-    def get_op(self, op_type: str) -> Any:
+    def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any:
         """Core distribution logic: Automatically select the best operator
         based on hardware and priority.
         """
-        if device_ctx.is_rocm:
-            platform = "rocm"
-        elif device_ctx.device_type == "cuda":
-            platform = "cuda"
-        else:
-            platform = "cpu"
+        platform = self._platform_for_device(device)
         candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE])
 
         for backend in candidates:
@@ -356,6 +484,21 @@ def get_op(self, op_type: str) -> Any:
 
         raise RuntimeError(f"No functional backend found for {op_type} on {platform}")
 
+    def _platform_for_device(self, device: torch.device | str | None) -> str:
+        if device is None:
+            if device_ctx.is_rocm:
+                return "rocm"
+            if device_ctx.device_type == "cuda":
+                return "cuda"
+            return "cpu"
+
+        resolved = torch.device(device)
+        if resolved.type == "cuda":
+            return "rocm" if torch.version.hip is not None else "cuda"
+        if resolved.type in self._priority_map:
+            return resolved.type
+        return "cpu"
+
     def get_attention_op(
         self,
         contract: AttentionContract,
@@ -377,7 +520,7 @@ def get_attention_op(
 
         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, [])
+        candidates = self._priority_map.get(platform, {}).get("ws2_attention", [])
         rejected: list[str] = []
 
         for backend in candidates:
diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py
index b986fc15..47c94c40 100644
--- a/tests/test_attention_contract.py
+++ b/tests/test_attention_contract.py
@@ -22,6 +22,12 @@
     RoPEFusionBoundary,
     RoPESpec,
     ShardingSpec,
+    SplitKVExecutionPlan,
+    SplitKVRuntimePlanSet,
+    SplitKVSpec,
+    build_split_kv_runtime_plan_set,
+    validate_split_kv_alignment,
+    validate_split_kv_plan_set_alignment,
 )
 from rl_engine.kernels.registry import KernelRegistry, OpBackend
 
@@ -107,6 +113,8 @@ def _declared_cp_backend() -> AttentionBackendCapability:
         deterministic_cp_merge=True,
         supports_packed_varlen=True,
         supports_kv_cache=True,
+        supports_split_kv_fixed=True,
+        reports_actual_split_kv_plan=True,
         implementation_kind="deterministic",
     )
 
@@ -254,6 +262,126 @@ def test_reduction_requires_fp32_accumulation():
         ReductionSpec(acc_dtype="bf16")
 
 
+def test_split_kv_policy_is_a_first_class_strict_contract():
+    contract = _contract()
+    assert contract.to_dict()["split_kv"] == {
+        "mode": "disabled",
+        "fixed_split_size": None,
+        "strict_consistency": True,
+    }
+
+    fixed = replace(contract, split_kv=SplitKVSpec.fixed(128))
+    assert fixed.to_dict()["split_kv"]["mode"] == "fixed"
+    assert fixed.to_dict()["split_kv"]["fixed_split_size"] == 128
+
+    with pytest.raises(AttentionContractError, match="auto Split-KV"):
+        SplitKVSpec.auto(strict_consistency=True)
+
+
+def test_split_kv_execution_plan_records_actual_logical_schedule():
+    plan = SplitKVSpec.fixed(4).resolve(10, backend="training-reference")
+
+    assert plan.actual_split_count == 3
+    assert plan.to_dict()["actual_split_boundaries"] == [[0, 4], [4, 8], [8, 10]]
+    assert plan.to_dict()["split_kv_merge_order"] == "global_block_index"
+    assert plan.to_dict()["split_kv_accum_dtype"] == "fp32"
+    assert plan.to_dict()["split_kv_downcast_at"] == "final_write"
+
+
+def test_strict_split_kv_alignment_rejects_unknown_or_mismatched_actual_plan():
+    training = SplitKVSpec.fixed(4).resolve(10, backend="training")
+    rollout = SplitKVSpec.fixed(4).resolve(10, backend="rollout")
+    validate_split_kv_alignment(training, rollout)
+
+    unknown = SplitKVSpec.auto().resolve(10, backend="rollout")
+    with pytest.raises(AttentionContractError, match="actual runtime plans"):
+        validate_split_kv_alignment(training, unknown)
+
+    mismatched = SplitKVSpec.fixed(5).resolve(10, backend="rollout")
+    with pytest.raises(AttentionContractError, match="differ"):
+        validate_split_kv_alignment(training, mismatched)
+
+    with pytest.raises(AttentionContractError, match="contiguous"):
+        SplitKVExecutionPlan(
+            requested_mode="fixed",
+            requested_split_size=4,
+            actual_mode="fixed",
+            actual_split_size=4,
+            boundaries=((0, 4), (5, 10)),
+        )
+
+
+def test_complete_split_kv_plan_set_covers_batch_tp_cp_and_owner_coordinates():
+    plan_set = build_split_kv_runtime_plan_set(
+        (8, 10),
+        tp_world_size=2,
+        cp_world_size=2,
+        split_kv=SplitKVSpec.fixed(2),
+        backend="training-reference",
+    )
+
+    assert len(plan_set.entries) == 16
+    assert plan_set.to_dict()["coverage"] == (
+        "complete_batch_tp_cp_owner_cartesian_product"
+    )
+    assert {
+        tuple(entry["expected_kv_range"])
+        for entry in plan_set.to_dict()["entries"]
+        if entry["batch_index"] == 0
+    } == {(0, 4), (4, 8)}
+
+
+def test_split_kv_plan_set_alignment_rejects_missing_and_mismatched_rank_plans():
+    training = build_split_kv_runtime_plan_set(
+        (8,),
+        tp_world_size=2,
+        cp_world_size=2,
+        split_kv=SplitKVSpec.fixed(2),
+        backend="training",
+    )
+    rollout = build_split_kv_runtime_plan_set(
+        (8,),
+        tp_world_size=2,
+        cp_world_size=2,
+        split_kv=SplitKVSpec.fixed(2),
+        backend="rollout",
+    )
+    validate_split_kv_plan_set_alignment(training, rollout)
+
+    with pytest.raises(AttentionContractError, match="coordinate coverage is incomplete"):
+        SplitKVRuntimePlanSet(
+            batch_size=training.batch_size,
+            tp_world_size=training.tp_world_size,
+            cp_world_size=training.cp_world_size,
+            total_kv_tokens=training.total_kv_tokens,
+            entries=training.entries[:-1],
+        )
+
+    mismatched = build_split_kv_runtime_plan_set(
+        (8,),
+        tp_world_size=2,
+        cp_world_size=2,
+        split_kv=SplitKVSpec.fixed(1),
+        backend="rollout",
+    )
+    with pytest.raises(AttentionContractError, match="plan differs"):
+        validate_split_kv_plan_set_alignment(training, mismatched)
+
+
+def test_backend_must_support_policy_and_actual_plan_provenance():
+    fixed = replace(_contract(), split_kv=SplitKVSpec.fixed(128))
+    capability = replace(
+        _declared_cp_backend(),
+        supports_split_kv_fixed=False,
+        reports_actual_split_kv_plan=False,
+    )
+
+    assert capability.incompatibilities(fixed)[-2:] == (
+        "Split-KV policy=fixed is unsupported",
+        "actual Split-KV execution-plan provenance is unsupported",
+    )
+
+
 def test_causal_attention_requires_explicit_offset():
     contract = _contract()
     with pytest.raises(AttentionContractError, match="causal_offsets are required"):
@@ -460,6 +588,8 @@ def test_shared_prefix_pages_must_be_fully_populated():
 
 def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback():
     registry = KernelRegistry()
+    platform = registry._platform()
+    registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION]
 
     with pytest.raises(RuntimeError) as exc_info:
         registry.get_attention_op(_contract())
@@ -473,7 +603,7 @@ def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback():
 def test_undeclared_backend_capability_is_never_selected():
     registry = KernelRegistry()
     platform = registry._platform()
-    registry._priority_map[platform]["attention"] = [OpBackend.PYTORCH_ATTN]
+    registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_ATTN]
 
     with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"):
         registry.get_attention_op(_contract())
@@ -481,6 +611,8 @@ def test_undeclared_backend_capability_is_never_selected():
 
 def test_declared_compatible_backend_resolves_and_records_provenance():
     registry = KernelRegistry()
+    platform = registry._platform()
+    registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION]
     registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend()
 
     result = registry.get_attention_op(_contract(), requested_backend="deterministic")
@@ -497,6 +629,8 @@ def test_declared_compatible_backend_resolves_and_records_provenance():
 
 def test_requested_stable_backend_id_is_enforced():
     registry = KernelRegistry()
+    platform = registry._platform()
+    registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION]
     registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend()
 
     with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"):

From 7f0e2914ef3d59ea2d3bbe33aa8ae48c9e0ec70f Mon Sep 17 00:00:00 2001
From: lamentropetion <3051000145@qq.com>
Date: Thu, 13 Aug 2026 21:19:13 +0800
Subject: [PATCH 09/13] style(attention): satisfy full PR lint

Signed-off-by: lamentropetion <3051000145@qq.com>
---
 rl_engine/kernels/attention_contract.py | 70 +++++++++----------------
 rl_engine/kernels/registry.py           |  1 -
 tests/test_attention_contract.py        |  4 +-
 3 files changed, 26 insertions(+), 49 deletions(-)

diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py
index bb0104ad..206663e7 100644
--- a/rl_engine/kernels/attention_contract.py
+++ b/rl_engine/kernels/attention_contract.py
@@ -359,9 +359,7 @@ def __post_init__(self) -> None:
                 or start < 0
                 or end <= start
             ):
-                raise AttentionContractError(
-                    "Split-KV boundaries must satisfy 0 <= start < end"
-                )
+                raise AttentionContractError("Split-KV boundaries must satisfy 0 <= start < end")
             if index > 0 and start != previous_end:
                 raise AttentionContractError(
                     "Split-KV boundaries must be contiguous and in logical KV order"
@@ -486,14 +484,10 @@ def __post_init__(self) -> None:
             raise AttentionContractError("split_kv.strict_consistency must be a bool")
         if self.mode is SplitKVMode.FIXED:
             if self.fixed_split_size is None:
-                raise AttentionContractError(
-                    "fixed Split-KV policy requires fixed_split_size"
-                )
+                raise AttentionContractError("fixed Split-KV policy requires fixed_split_size")
             _positive_int(self.fixed_split_size, "split_kv.fixed_split_size")
         elif self.fixed_split_size is not None:
-            raise AttentionContractError(
-                "fixed_split_size is only valid for fixed Split-KV policy"
-            )
+            raise AttentionContractError("fixed_split_size is only valid for fixed Split-KV policy")
         if self.strict_consistency and self.mode is SplitKVMode.AUTO:
             raise AttentionContractError(
                 "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency"
@@ -647,13 +641,8 @@ def validate(
                 "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens"
             )
         if self.execution.actual_mode is None:
-            raise AttentionContractError(
-                "complete Split-KV plan sets require actual runtime plans"
-            )
-        if (
-            self.execution.boundaries[0][0] != start
-            or self.execution.boundaries[-1][1] != end
-        ):
+            raise AttentionContractError("complete Split-KV plan sets require actual runtime plans")
+        if self.execution.boundaries[0][0] != start or self.execution.boundaries[-1][1] != end:
             raise AttentionContractError(
                 "Split-KV execution boundaries must exactly cover expected_kv_range"
             )
@@ -661,9 +650,7 @@ def validate(
             boundary_start < start or boundary_end > end
             for boundary_start, boundary_end in self.execution.boundaries
         ):
-            raise AttentionContractError(
-                "Split-KV execution boundary escapes expected_kv_range"
-            )
+            raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range")
 
     def to_dict(self) -> dict[str, Any]:
         return {
@@ -709,9 +696,7 @@ def __post_init__(self) -> None:
         }
         actual_coordinates = [entry.coordinate for entry in entries]
         if len(set(actual_coordinates)) != len(actual_coordinates):
-            raise AttentionContractError(
-                "Split-KV runtime plan set contains duplicate coordinates"
-            )
+            raise AttentionContractError("Split-KV runtime plan set contains duplicate coordinates")
         missing = expected_coordinates.difference(actual_coordinates)
         extra = set(actual_coordinates).difference(expected_coordinates)
         if missing or extra:
@@ -821,17 +806,12 @@ def validate_split_kv_plan_set_alignment(
     ]
     if topology_mismatches:
         raise AttentionContractError(
-            "training/rollout Split-KV plan-set topology differs: "
-            + ", ".join(topology_mismatches)
+            "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches)
         )
-    training_by_coordinate = {
-        entry.coordinate: entry for entry in training.entries
-    }
+    training_by_coordinate = {entry.coordinate: entry for entry in training.entries}
     rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries}
     if training_by_coordinate.keys() != rollout_by_coordinate.keys():
-        raise AttentionContractError(
-            "training/rollout Split-KV plan-set coordinates differ"
-        )
+        raise AttentionContractError("training/rollout Split-KV plan-set coordinates differ")
     for coordinate in sorted(training_by_coordinate):
         train_entry = training_by_coordinate[coordinate]
         rollout_entry = rollout_by_coordinate[coordinate]
@@ -847,7 +827,7 @@ def validate_split_kv_plan_set_alignment(
         except AttentionContractError as exc:
             raise AttentionContractError(
                 f"training/rollout Split-KV plan differs at {coordinate}: {exc}"
-                ) from exc
+            ) from exc
 
 
 def build_split_kv_runtime_plan_set(
@@ -1148,14 +1128,16 @@ def __post_init__(self) -> None:
             not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip()
         ):
             raise AttentionContractError("rope_scaling must be a non-empty string when provided")
-        for field in ("position_ids", "query_position_offsets", "key_position_offsets"):
-            values = getattr(self, field)
+        for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"):
+            values = getattr(self, position_field)
             if values is None:
                 continue
-            normalized = _integer_tuple(values, field)
+            normalized = _integer_tuple(values, position_field)
             if not normalized or any(value < 0 for value in normalized):
-                raise AttentionContractError(f"{field} must contain non-negative positions")
-            object.__setattr__(self, field, normalized)
+                raise AttentionContractError(
+                    f"{position_field} must contain non-negative positions"
+                )
+            object.__setattr__(self, position_field, normalized)
         object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at"))
         object.__setattr__(
             self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype")
@@ -1259,11 +1241,11 @@ def __post_init__(self) -> None:
                     "position_ids must describe the local query sequence or full local "
                     "sequence length"
                 )
-            for field in ("query_position_offsets", "key_position_offsets"):
-                offsets = getattr(self.rope, field)
+            for position_field in ("query_position_offsets", "key_position_offsets"):
+                offsets = getattr(self.rope, position_field)
                 if offsets is not None and len(offsets) != batch_size:
                     raise AttentionContractError(
-                        f"{field} must contain one entry per logical batch entry"
+                        f"{position_field} must contain one entry per logical batch entry"
                     )
         if self.mode is AttentionMode.DECODE and self.kv_cache is not None:
             if len(self.kv_cache.kv_seq_lens) != batch_size:
@@ -1409,7 +1391,7 @@ def __post_init__(self) -> None:
                 raise AttentionContractError("tp_world_sizes must contain positive values")
             if len(set(tp_world_sizes)) != len(tp_world_sizes):
                 raise AttentionContractError("tp_world_sizes must not contain duplicates")
-        for field in (
+        for capability_field in (
             "exports_attention_lse",
             "deterministic_cp_merge",
             "supports_packed_varlen",
@@ -1421,8 +1403,8 @@ def __post_init__(self) -> None:
             "supports_split_kv_auto",
             "reports_actual_split_kv_plan",
         ):
-            if not isinstance(getattr(self, field), bool):
-                raise AttentionContractError(f"{field} must be a bool")
+            if not isinstance(getattr(self, capability_field), bool):
+                raise AttentionContractError(f"{capability_field} must be a bool")
         if self.implementation_kind not in {"production", "reference", "deterministic"}:
             raise AttentionContractError(
                 "implementation_kind must be production, reference, or deterministic"
@@ -1474,9 +1456,7 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]:
             SplitKVMode.AUTO: self.supports_split_kv_auto,
         }
         if not split_support[contract.split_kv.mode]:
-            reasons.append(
-                f"Split-KV policy={contract.split_kv.mode.value} is unsupported"
-            )
+            reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported")
         if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan:
             reasons.append("actual Split-KV execution-plan provenance is unsupported")
         return tuple(reasons)
diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py
index edd5efd5..6472bf08 100644
--- a/rl_engine/kernels/registry.py
+++ b/rl_engine/kernels/registry.py
@@ -519,7 +519,6 @@ def get_attention_op(
         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("ws2_attention", [])
         rejected: list[str] = []
 
diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py
index 47c94c40..350545a1 100644
--- a/tests/test_attention_contract.py
+++ b/tests/test_attention_contract.py
@@ -321,9 +321,7 @@ def test_complete_split_kv_plan_set_covers_batch_tp_cp_and_owner_coordinates():
     )
 
     assert len(plan_set.entries) == 16
-    assert plan_set.to_dict()["coverage"] == (
-        "complete_batch_tp_cp_owner_cartesian_product"
-    )
+    assert plan_set.to_dict()["coverage"] == ("complete_batch_tp_cp_owner_cartesian_product")
     assert {
         tuple(entry["expected_kv_range"])
         for entry in plan_set.to_dict()["entries"]

From c8971732daae5c3de064754cb94af26f94482ae8 Mon Sep 17 00:00:00 2001
From: lamentropetion <3051000145@qq.com>
Date: Sun, 16 Aug 2026 15:02:08 +0800
Subject: [PATCH 10/13] feat(attention): contract QKV and output projection
 semantics

Signed-off-by: lamentropetion <3051000145@qq.com>
---
 rl_engine/kernels/attention_contract.py | 144 ++++++++++++++++++++++++
 tests/test_attention_contract.py        |  51 +++++++++
 2 files changed, 195 insertions(+)

diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py
index 206663e7..cbaa4890 100644
--- a/rl_engine/kernels/attention_contract.py
+++ b/rl_engine/kernels/attention_contract.py
@@ -61,6 +61,13 @@ class SplitKVMode(str, Enum):
     AUTO = "auto"
 
 
+class ProjectionCollective(str, Enum):
+    NONE = "none"
+    ALL_REDUCE = "all_reduce"
+    ALL_GATHER = "all_gather"
+    REDUCE_SCATTER = "reduce_scatter"
+
+
 class RoPEState(str, Enum):
     PRE_ROPE = "pre_rope"
     POST_ROPE = "post_rope"
@@ -134,12 +141,16 @@ class ShardingSpec:
     global_block_token_starts: tuple[int, ...]
     local_block_offsets: tuple[int, ...]
     packed_sequence_offsets: tuple[int, ...] | None = None
+    sp_rank: int = 0
+    sp_world_size: int = 1
 
     def __post_init__(self) -> None:
         tp_world_size = _positive_int(self.tp_world_size, "tp_world_size")
         cp_world_size = _positive_int(self.cp_world_size, "cp_world_size")
+        sp_world_size = _positive_int(self.sp_world_size, "sp_world_size")
         tp_rank = _non_negative_int(self.tp_rank, "tp_rank")
         cp_rank = _non_negative_int(self.cp_rank, "cp_rank")
+        sp_rank = _non_negative_int(self.sp_rank, "sp_rank")
         if tp_rank >= tp_world_size:
             raise AttentionContractError(
                 f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}"
@@ -148,6 +159,10 @@ def __post_init__(self) -> None:
             raise AttentionContractError(
                 f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}"
             )
+        if sp_rank >= sp_world_size:
+            raise AttentionContractError(
+                f"sp_rank={sp_rank} must be smaller than sp_world_size={sp_world_size}"
+            )
 
         global_q_heads = _positive_int(self.global_q_heads, "global_q_heads")
         global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads")
@@ -1149,6 +1164,98 @@ def __post_init__(self) -> None:
         )
 
 
+@dataclass(frozen=True)
+class AttentionProjectionSpec:
+    """Deterministic QKV or output-projection execution contract."""
+
+    name: str
+    input_dtype: AttentionDType = AttentionDType.BF16
+    output_dtype: AttentionDType = AttentionDType.BF16
+    acc_dtype: AttentionDType = AttentionDType.FP32
+    split_kv: SplitKVMode = SplitKVMode.DISABLED
+    k_order: str = "ascending"
+    backend_policy: str = "native_verified_then_common_deterministic"
+    deterministic_backend: str = "rlkernel.cuda.det_gemm"
+    tp_forward_collective: ProjectionCollective = ProjectionCollective.NONE
+    tp_backward_dgrad_collective: ProjectionCollective = ProjectionCollective.NONE
+    sp_forward_collective: ProjectionCollective = ProjectionCollective.NONE
+    sp_backward_collective: ProjectionCollective = ProjectionCollective.NONE
+    qkv_split_order: tuple[str, ...] = ()
+    require_runtime_readback: bool = True
+
+    @classmethod
+    def qkv(cls) -> "AttentionProjectionSpec":
+        return cls(
+            name="qkv",
+            tp_backward_dgrad_collective=ProjectionCollective.ALL_REDUCE,
+            sp_forward_collective=ProjectionCollective.ALL_GATHER,
+            sp_backward_collective=ProjectionCollective.REDUCE_SCATTER,
+            qkv_split_order=("q", "k", "v"),
+        )
+
+    @classmethod
+    def output(cls) -> "AttentionProjectionSpec":
+        return cls(
+            name="o_proj",
+            tp_forward_collective=ProjectionCollective.ALL_REDUCE,
+            sp_forward_collective=ProjectionCollective.REDUCE_SCATTER,
+            sp_backward_collective=ProjectionCollective.ALL_GATHER,
+        )
+
+    def __post_init__(self) -> None:
+        if self.name not in {"qkv", "o_proj"}:
+            raise AttentionContractError("projection name must be qkv or o_proj")
+        for field_name in ("input_dtype", "output_dtype", "acc_dtype"):
+            object.__setattr__(
+                self,
+                field_name,
+                _enum_value(AttentionDType, getattr(self, field_name), field_name),
+            )
+        object.__setattr__(
+            self,
+            "split_kv",
+            _enum_value(SplitKVMode, self.split_kv, "projection.split_kv"),
+        )
+        for field_name in (
+            "tp_forward_collective",
+            "tp_backward_dgrad_collective",
+            "sp_forward_collective",
+            "sp_backward_collective",
+        ):
+            object.__setattr__(
+                self,
+                field_name,
+                _enum_value(
+                    ProjectionCollective,
+                    getattr(self, field_name),
+                    field_name,
+                ),
+            )
+        if self.input_dtype is not AttentionDType.BF16:
+            raise AttentionContractError("Attention projections require BF16 input")
+        if self.output_dtype is not AttentionDType.BF16:
+            raise AttentionContractError("Attention projections require BF16 output")
+        if self.acc_dtype is not AttentionDType.FP32:
+            raise AttentionContractError("Attention projections require FP32 accumulation")
+        if self.split_kv is not SplitKVMode.DISABLED:
+            raise AttentionContractError("Attention projection GEMMs must disable Split-K")
+        if self.k_order != "ascending":
+            raise AttentionContractError("Attention projection GEMMs require ascending K order")
+        for field_name in ("backend_policy", "deterministic_backend"):
+            value = getattr(self, field_name)
+            if not isinstance(value, str) or not value.strip():
+                raise AttentionContractError(f"{field_name} must be a non-empty string")
+        if not isinstance(self.require_runtime_readback, bool) or not self.require_runtime_readback:
+            raise AttentionContractError("projection runtime readback must be required")
+        split_order = tuple(self.qkv_split_order)
+        if self.name == "qkv":
+            if split_order != ("q", "k", "v"):
+                raise AttentionContractError("QKV projection must split in Q, K, V order")
+        elif split_order:
+            raise AttentionContractError("o_proj must not declare a QKV split order")
+        object.__setattr__(self, "qkv_split_order", split_order)
+
+
 @dataclass(frozen=True)
 class AttentionContract:
     """Complete semantic request consumed by contract-aware dispatch."""
@@ -1167,6 +1274,10 @@ class AttentionContract:
     kv_cache: KVCacheSpec | None = None
     rope: RoPESpec | None = None
     export_lse: bool = True
+    qkv_projection: AttentionProjectionSpec = field(default_factory=AttentionProjectionSpec.qkv)
+    output_projection: AttentionProjectionSpec = field(
+        default_factory=AttentionProjectionSpec.output
+    )
 
     def __post_init__(self) -> None:
         object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role"))
@@ -1181,6 +1292,16 @@ def __post_init__(self) -> None:
             raise AttentionContractError("reduction must be a ReductionSpec")
         if not isinstance(self.split_kv, SplitKVSpec):
             raise AttentionContractError("split_kv must be a SplitKVSpec")
+        if not isinstance(self.qkv_projection, AttentionProjectionSpec):
+            raise AttentionContractError("qkv_projection must be an AttentionProjectionSpec")
+        if self.qkv_projection.name != "qkv":
+            raise AttentionContractError("qkv_projection must use name='qkv'")
+        if not isinstance(self.output_projection, AttentionProjectionSpec):
+            raise AttentionContractError(
+                "output_projection must be an AttentionProjectionSpec"
+            )
+        if self.output_projection.name != "o_proj":
+            raise AttentionContractError("output_projection must use name='o_proj'")
         if (
             self.mode is AttentionMode.PREFILL
             and query_sequence_length != self.sharding.local_sequence_length
@@ -1265,6 +1386,8 @@ def to_dict(self) -> dict[str, Any]:
             "tp_world_size": self.sharding.tp_world_size,
             "cp_rank": self.sharding.cp_rank,
             "cp_world_size": self.sharding.cp_world_size,
+            "sp_rank": self.sharding.sp_rank,
+            "sp_world_size": self.sharding.sp_world_size,
             "global_q_heads": self.sharding.global_q_heads,
             "global_kv_heads": self.sharding.global_kv_heads,
             "local_q_head_start": self.sharding.local_q_head_start,
@@ -1327,6 +1450,24 @@ def to_dict(self) -> dict[str, Any]:
                 "output_dtype": self.rope.output_dtype.value,
                 "fusion_boundary": self.rope.fusion_boundary.value,
             }
+        projections = {
+            spec.name: {
+                "input_dtype": spec.input_dtype.value,
+                "output_dtype": spec.output_dtype.value,
+                "acc_dtype": spec.acc_dtype.value,
+                "split_kv": spec.split_kv.value,
+                "k_order": spec.k_order,
+                "backend_policy": spec.backend_policy,
+                "deterministic_backend": spec.deterministic_backend,
+                "tp_forward_collective": spec.tp_forward_collective.value,
+                "tp_backward_dgrad_collective": spec.tp_backward_dgrad_collective.value,
+                "sp_forward_collective": spec.sp_forward_collective.value,
+                "sp_backward_collective": spec.sp_backward_collective.value,
+                "qkv_split_order": list(spec.qkv_split_order),
+                "require_runtime_readback": spec.require_runtime_readback,
+            }
+            for spec in (self.qkv_projection, self.output_projection)
+        }
         return {
             "semantic_operator": "standard_softmax_attention",
             "role": self.role.value,
@@ -1346,6 +1487,7 @@ def to_dict(self) -> dict[str, Any]:
             "split_kv": self.split_kv.to_dict(),
             "kv_cache": kv_cache,
             "rope": rope,
+            "projections": projections,
         }
 
 
@@ -1501,6 +1643,7 @@ class AttentionDispatchResult:
     "AttentionBackendCapability",
     "AttentionDispatchResult",
     "AttentionDType",
+    "AttentionProjectionSpec",
     "AttentionMerge",
     "AttentionMode",
     "AttentionRole",
@@ -1509,6 +1652,7 @@ class AttentionDispatchResult:
     "ReductionEngine",
     "ReductionOrder",
     "ReductionSpec",
+    "ProjectionCollective",
     "RoPECastPoint",
     "RoPEFusionBoundary",
     "RoPESpec",
diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py
index 350545a1..0c3347be 100644
--- a/tests/test_attention_contract.py
+++ b/tests/test_attention_contract.py
@@ -16,8 +16,10 @@
     AttentionContractError,
     AttentionDType,
     AttentionMode,
+    AttentionProjectionSpec,
     AttentionRole,
     KVCacheSpec,
+    ProjectionCollective,
     ReductionSpec,
     RoPEFusionBoundary,
     RoPESpec,
@@ -44,6 +46,8 @@ def _sharding(
     global_block_token_starts: tuple[int, ...] = (0,),
     local_block_offsets: tuple[int, ...] = (0, 2048),
     packed_sequence_offsets: tuple[int, ...] | None = None,
+    sp_rank: int = 0,
+    sp_world_size: int = 1,
 ) -> ShardingSpec:
     local_q_heads = 32 // tp_world_size
     local_kv_heads = 8 // tp_world_size
@@ -64,6 +68,8 @@ def _sharding(
         global_block_token_starts=global_block_token_starts,
         local_block_offsets=local_block_offsets,
         packed_sequence_offsets=packed_sequence_offsets,
+        sp_rank=sp_rank,
+        sp_world_size=sp_world_size,
     )
 
 
@@ -681,3 +687,48 @@ def test_packed_sequence_count_must_match_logical_batch_size():
 
     contract = _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=2)
     assert contract.batch_size == 2
+
+
+def test_projection_contract_pins_gemm_and_tp_sp_collectives():
+    contract = _contract(sharding=_sharding(sp_rank=1, sp_world_size=2))
+    provenance = contract.to_dict()
+
+    assert provenance["sharding"]["sp_rank"] == 1
+    assert provenance["sharding"]["sp_world_size"] == 2
+    assert provenance["projections"]["qkv"] == {
+        "input_dtype": "bf16",
+        "output_dtype": "bf16",
+        "acc_dtype": "fp32",
+        "split_kv": "disabled",
+        "k_order": "ascending",
+        "backend_policy": "native_verified_then_common_deterministic",
+        "deterministic_backend": "rlkernel.cuda.det_gemm",
+        "tp_forward_collective": "none",
+        "tp_backward_dgrad_collective": "all_reduce",
+        "sp_forward_collective": "all_gather",
+        "sp_backward_collective": "reduce_scatter",
+        "qkv_split_order": ["q", "k", "v"],
+        "require_runtime_readback": True,
+    }
+    assert provenance["projections"]["o_proj"]["tp_forward_collective"] == "all_reduce"
+    assert provenance["projections"]["o_proj"]["sp_forward_collective"] == "reduce_scatter"
+    assert provenance["projections"]["o_proj"]["sp_backward_collective"] == "all_gather"
+
+
+def test_projection_contract_rejects_split_k_and_wrong_collective_identity():
+    with pytest.raises(AttentionContractError, match="disable Split-K"):
+        replace(AttentionProjectionSpec.qkv(), split_kv="fixed")
+
+    with pytest.raises(AttentionContractError, match="qkv_projection"):
+        replace(
+            _contract(),
+            qkv_projection=replace(
+                AttentionProjectionSpec.output(),
+                tp_forward_collective=ProjectionCollective.NONE,
+            ),
+        )
+
+
+def test_sharding_rejects_out_of_range_sp_rank():
+    with pytest.raises(AttentionContractError, match="sp_rank=2"):
+        _sharding(sp_rank=2, sp_world_size=2)

From a5a01e05c36db5bd2cb9b6cf30d11e7f022c0ed6 Mon Sep 17 00:00:00 2001
From: lamentropetion <3051000145@qq.com>
Date: Mon, 17 Aug 2026 18:07:28 +0800
Subject: [PATCH 11/13] feat(attention): identify shared deterministic core

---
 rl_engine/kernels/attention_contract.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py
index cbaa4890..e3298fa9 100644
--- a/rl_engine/kernels/attention_contract.py
+++ b/rl_engine/kernels/attention_contract.py
@@ -18,6 +18,11 @@
 _EnumT = TypeVar("_EnumT", bound=Enum)
 
 
+# Stable identity shared by the training and rollout deterministic Attention
+# core.  Backend adapters may differ, but strict mode must report this ID.
+STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1"
+
+
 class AttentionContractError(ValueError):
     """Raised when attention metadata does not describe a valid invocation."""
 
@@ -1297,9 +1302,7 @@ def __post_init__(self) -> None:
         if self.qkv_projection.name != "qkv":
             raise AttentionContractError("qkv_projection must use name='qkv'")
         if not isinstance(self.output_projection, AttentionProjectionSpec):
-            raise AttentionContractError(
-                "output_projection must be an AttentionProjectionSpec"
-            )
+            raise AttentionContractError("output_projection must be an AttentionProjectionSpec")
         if self.output_projection.name != "o_proj":
             raise AttentionContractError("output_projection must use name='o_proj'")
         if (
@@ -1665,6 +1668,7 @@ class AttentionDispatchResult:
     "SplitKVRuntimePlanEntry",
     "SplitKVRuntimePlanSet",
     "SplitKVSpec",
+    "STRICT_ATTENTION_CORE_ID",
     "validate_split_kv_alignment",
     "validate_split_kv_plan_set_alignment",
 ]

From 2e208acebea4cab8a1de81e8653baed0bffe7a2c Mon Sep 17 00:00:00 2001
From: lamentropetion <3051000145@qq.com>
Date: Tue, 18 Aug 2026 21:48:14 +0800
Subject: [PATCH 12/13] feat(attention): pin strict arithmetic schedule
 identity

---
 rl_engine/kernels/attention_contract.py | 2 ++
 tests/test_attention_contract.py        | 7 +++++++
 2 files changed, 9 insertions(+)

diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py
index e3298fa9..bf2bf329 100644
--- a/rl_engine/kernels/attention_contract.py
+++ b/rl_engine/kernels/attention_contract.py
@@ -21,6 +21,7 @@
 # Stable identity shared by the training and rollout deterministic Attention
 # core.  Backend adapters may differ, but strict mode must report this ID.
 STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1"
+STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks"
 
 
 class AttentionContractError(ValueError):
@@ -1669,6 +1670,7 @@ class AttentionDispatchResult:
     "SplitKVRuntimePlanSet",
     "SplitKVSpec",
     "STRICT_ATTENTION_CORE_ID",
+    "STRICT_ATTENTION_SCHEDULE_ID",
     "validate_split_kv_alignment",
     "validate_split_kv_plan_set_alignment",
 ]
diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py
index 0c3347be..6f1587b3 100644
--- a/tests/test_attention_contract.py
+++ b/tests/test_attention_contract.py
@@ -11,6 +11,8 @@
 import pytest
 
 from rl_engine.kernels.attention_contract import (
+    STRICT_ATTENTION_CORE_ID,
+    STRICT_ATTENTION_SCHEDULE_ID,
     AttentionBackendCapability,
     AttentionContract,
     AttentionContractError,
@@ -34,6 +36,11 @@
 from rl_engine.kernels.registry import KernelRegistry, OpBackend
 
 
+def test_strict_attention_identity_pins_core_and_arithmetic_schedule():
+    assert STRICT_ATTENTION_CORE_ID == "rlkernel.attention.deterministic_core.v1"
+    assert STRICT_ATTENTION_SCHEDULE_ID == "single_batch_single_query_global_kv_blocks"
+
+
 def _sharding(
     *,
     tp_rank: int = 0,

From 51b8f52286818755f4bab04c6bd4bf3fa2762c08 Mon Sep 17 00:00:00 2001
From: lamentropetion <3051000145@qq.com>
Date: Tue, 18 Aug 2026 23:25:26 +0800
Subject: [PATCH 13/13] feat(attention): expose strict core identity and
 optional CP imports

---
 .../kernels/ops/cuda/attention/__init__.py    | 103 ++++++++++--------
 .../ops/cuda/attention/deterministic_attn.py  |   8 ++
 tests/test_attention.py                       |  12 ++
 3 files changed, 78 insertions(+), 45 deletions(-)

diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py
index 2c3a1f1b..8ca8df7b 100644
--- a/rl_engine/kernels/ops/cuda/attention/__init__.py
+++ b/rl_engine/kernels/ops/cuda/attention/__init__.py
@@ -1,49 +1,62 @@
-# File: rl_engine/kernels/ops/cuda/attention/__init__.py
-
-from .cp_comm import (
-    AttentionCPBlockMetadata,
-    AttentionCPCommunication,
-    AttentionCPCommunicationPlan,
-    AttentionCPCommunicationUnavailable,
-    AttentionCPMergedState,
-    AttentionCPPartialState,
-    AttentionParallelSpec,
-    CPCommunicationBackend,
-    CPCommunicationStatus,
-    CUDAAGRSAttentionCPCommunication,
-    P2PNCCLAttentionCPCommunication,
-    sort_attention_cp_partial_states,
-)
 from .deterministic_attn import DeterministicAttentionOp
 from .flash_attn import FlashAttentionOp
-from .flashinfer_paged_attention import (
-    FlashInferPagedAttentionConfig,
-    FlashInferQwen3PagedAttentionOp,
-    FlashInferRoPEFusionConfig,
-    FlashInferSplitKVPolicy,
-    FlashInferUnavailable,
-)
 from .prefix_shared_attn import PrefixSharedAttentionOp
 
-__all__ = [
-    "AttentionCPBlockMetadata",
-    "AttentionCPCommunication",
-    "AttentionCPCommunicationPlan",
-    "AttentionCPCommunicationUnavailable",
-    "AttentionCPMergedState",
-    "AttentionCPPartialState",
-    "AttentionParallelSpec",
-    "CPCommunicationBackend",
-    "CPCommunicationStatus",
-    "CUDAAGRSAttentionCPCommunication",
-    "P2PNCCLAttentionCPCommunication",
-    "DeterministicAttentionOp",
-    "FlashAttentionOp",
-    "FlashInferPagedAttentionConfig",
-    "FlashInferQwen3PagedAttentionOp",
-    "FlashInferRoPEFusionConfig",
-    "FlashInferSplitKVPolicy",
-    "FlashInferUnavailable",
-    "PrefixSharedAttentionOp",
-    "sort_attention_cp_partial_states",
-]
+__all__ = ["DeterministicAttentionOp", "FlashAttentionOp", "PrefixSharedAttentionOp"]
+
+# CP communication and FlashInfer are optional layers owned by later WS2 PRs.
+# Keep the base Attention package importable while those PRs are developed or
+# tested independently, then expose their symbols automatically when present.
+try:
+    from .cp_comm import (
+        AttentionCPBlockMetadata,
+        AttentionCPCommunication,
+        AttentionCPCommunicationPlan,
+        AttentionCPCommunicationUnavailable,
+        AttentionCPMergedState,
+        AttentionCPPartialState,
+        AttentionParallelSpec,
+        CPCommunicationBackend,
+        CPCommunicationStatus,
+        CUDAAGRSAttentionCPCommunication,
+        P2PNCCLAttentionCPCommunication,
+        sort_attention_cp_partial_states,
+    )
+except ModuleNotFoundError as exc:
+    if exc.name != f"{__package__}.cp_comm":
+        raise
+else:
+    __all__ += [
+        "AttentionCPBlockMetadata",
+        "AttentionCPCommunication",
+        "AttentionCPCommunicationPlan",
+        "AttentionCPCommunicationUnavailable",
+        "AttentionCPMergedState",
+        "AttentionCPPartialState",
+        "AttentionParallelSpec",
+        "CPCommunicationBackend",
+        "CPCommunicationStatus",
+        "CUDAAGRSAttentionCPCommunication",
+        "P2PNCCLAttentionCPCommunication",
+        "sort_attention_cp_partial_states",
+    ]
+
+try:
+    from .flashinfer_paged_attention import (
+        FlashInferPagedAttentionConfig,
+        FlashInferQwen3PagedAttentionOp,
+        FlashInferRoPEFusionConfig,
+        FlashInferSplitKVPolicy,
+        FlashInferUnavailable,
+    )
+except ModuleNotFoundError as exc:
+    if exc.name != f"{__package__}.flashinfer_paged_attention":
+        raise
+else:
+    __all__ += [
+        "FlashInferPagedAttentionConfig",
+        "FlashInferQwen3PagedAttentionOp",
+        "FlashInferRoPEFusionConfig",
+        "FlashInferSplitKVPolicy",
+        "FlashInferUnavailable",
+    ]
diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py
index 01c51c99..1e3d01e2 100644
--- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py
+++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py
@@ -17,6 +17,10 @@
 from torch.autograd import Function
 from torch.autograd.function import once_differentiable
 
+from rl_engine.kernels.attention_contract import (
+    STRICT_ATTENTION_CORE_ID,
+    STRICT_ATTENTION_SCHEDULE_ID,
+)
 from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
 from rl_engine.utils.logger import logger
 
@@ -82,6 +86,10 @@ class DeterministicAttentionOp:
     so #108 harness can call forward(**inputs) with key_padding_mask.
     """
 
+    core_id = STRICT_ATTENTION_CORE_ID
+    strict_schedule = STRICT_ATTENTION_SCHEDULE_ID
+    backend_id = "rlkernel.cuda.deterministic_attention"
+
     def __init__(self) -> None:
         if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_attention_forward"):
             raise RuntimeError(
diff --git a/tests/test_attention.py b/tests/test_attention.py
index 469c6d30..40ee6fd5 100644
--- a/tests/test_attention.py
+++ b/tests/test_attention.py
@@ -443,6 +443,18 @@ def test_registry_dispatches_native_attention_op():
     assert isinstance(op, (NativeAttentionOp, DeterministicAttentionOp))
 
 
+def test_deterministic_attention_op_exposes_shared_strict_identity():
+    from rl_engine.kernels.attention_contract import (
+        STRICT_ATTENTION_CORE_ID,
+        STRICT_ATTENTION_SCHEDULE_ID,
+    )
+    from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp
+
+    assert DeterministicAttentionOp.core_id == STRICT_ATTENTION_CORE_ID
+    assert DeterministicAttentionOp.strict_schedule == STRICT_ATTENTION_SCHEDULE_ID
+    assert DeterministicAttentionOp.backend_id == "rlkernel.cuda.deterministic_attention"
+
+
 # --------------------------------------------------------------------------- #
 # Qwen3-8B LARGE real-scale GPU smoke test
 # --------------------------------------------------------------------------- #