[WS2] Add CP-aware Attention contract and dispatch metadata

.github/workflows/ci.yml

@@ -96,6 +96,10 @@ jobs:
         run: |
           python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu"
 
+<<<<<<< HEAD
+      - name: Run WS2 Attention Contract Tests (CPU-safe)
+        run: python -m pytest tests/test_attention_contract.py -v
+=======
       - name: Run WS2 Logprob Contract Tests (CPU-safe)
         run: python -m pytest tests/test_logprob_contract.py -v
 
@@ -107,6 +111,7 @@ jobs:
           python -m pytest -q \
             tests/test_alignment_wrapper_interfaces.py \
             tests/test_qwen_ffn.py
+>>>>>>> origin/test
 
   docs:
     runs-on: ubuntu-latest

docs/design/runtime-dispatch.md

@@ -11,6 +11,11 @@ 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).
 WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In
 addition to platform priority, this path requires a backend capability descriptor and checks
 the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token

docs/operators/attention.md

@@ -84,6 +84,31 @@ the inputs' device.
 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
+
+WS2 distributed callers use a separate contract-aware entry point,
+`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed
+`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity
+before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged.
+
+Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic
+CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as
+a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md).
+
+Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow
+`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block
+boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason.
+Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and
+validate the same actual plan.
+
+The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report:
+
+```bash
+python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json
+python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \
+  --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \
+  --output artifacts/ws2-cp-attention-drift.json
+```
 
 ## Accuracy
 
@@ -137,6 +162,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
@@ -194,6 +220,10 @@ Hooks:
 - `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable.
 - `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging,
   and future KV-cache / training integration.
+- `backward_reference(q, k, v, dout, ...)` — runs the deterministic training backward
+  validation path and returns `dq`, `dk`, `dv`, `out`, `lse`, and provenance.
+- `compare_cp_attention_backward(q, k, v, dout, ...)` — compares CP=1 backward against
+  CP/chunked-prefill backward and emits whole-tensor plus per-logical-rank drift stats.
 
 ## Tolerance
 

rl_engine/kernels/attention_contract.py

@@ -18,6 +18,8 @@
 _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"
 STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks"
 
@@ -65,6 +67,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"
@@ -138,10 +147,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")
         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:
@@ -152,6 +167,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")
@@ -1153,6 +1172,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."""
@@ -1171,6 +1282,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"))
@@ -1185,6 +1300,14 @@ def __post_init__(self) -> None:
             raise AttentionContractError("reduction must be a ReductionSpec")
         if not isinstance(self.split_kv, SplitKVSpec):
             raise AttentionContractError("split_kv must be a SplitKVSpec")
+        if not isinstance(self.qkv_projection, AttentionProjectionSpec):
+            raise AttentionContractError("qkv_projection must be an AttentionProjectionSpec")
+        if self.qkv_projection.name != "qkv":
+            raise AttentionContractError("qkv_projection must use name='qkv'")
+        if not isinstance(self.output_projection, AttentionProjectionSpec):
+            raise AttentionContractError("output_projection must be an AttentionProjectionSpec")
+        if self.output_projection.name != "o_proj":
+            raise AttentionContractError("output_projection must use name='o_proj'")
         if (
             self.mode is AttentionMode.PREFILL
             and query_sequence_length != self.sharding.local_sequence_length
@@ -1269,6 +1392,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,
@@ -1331,6 +1456,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,
@@ -1350,6 +1493,7 @@ def to_dict(self) -> dict[str, Any]:
             "split_kv": self.split_kv.to_dict(),
             "kv_cache": kv_cache,
             "rope": rope,
+            "projections": projections,
         }
 
 
@@ -1505,6 +1649,7 @@ class AttentionDispatchResult:
     "AttentionBackendCapability",
     "AttentionDispatchResult",
     "AttentionDType",
+    "AttentionProjectionSpec",
     "AttentionMerge",
     "AttentionMode",
     "AttentionRole",
@@ -1513,6 +1658,7 @@ class AttentionDispatchResult:
     "ReductionEngine",
     "ReductionOrder",
     "ReductionSpec",
+    "ProjectionCollective",
     "RoPECastPoint",
     "RoPEFusionBoundary",
     "RoPESpec",

rl_engine/kernels/gtest/operator_inputs.py

@@ -31,6 +31,7 @@ def make_operator_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,
@@ -58,6 +59,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str:
         "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}",

rl_engine/kernels/gtest/operator_specs.py

@@ -76,6 +76,22 @@ def _load_object(path: str) -> Any:
         },
         grad_input_names=("q", "k", "v"),
     ),
+    "cp_attention": OperatorSpec(
+        name="cp_attention",
+        op_class="attention",
+        gold_path=(
+            "rl_engine.kernels.ops.pytorch.attention.cp_attention."
+            "DeterministicCPAttentionReferenceOp"
+        ),
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": (
+                "rl_engine.kernels.ops.pytorch.attention.cp_attention."
+                "DeterministicCPAttentionReferenceOp"
+            ),
+        },
+        grad_input_names=("q", "k", "v"),
+    ),
     "logp": OperatorSpec(
         name="logp",
         op_class="logprob",