[WS2][Cross-config][Attention] Integrate the Qwen3-8B TP=2 CP=2 BF16 Attention setting to cross-config alignment framework

docs/operators/attention.md

@@ -13,6 +13,12 @@ friends) are validated against.
 This op covers **only** the softmax attention. Qwen3's QK-Norm and RoPE are applied *before*
 the call (see the chain), so the `q`, `k` passed in are already normalized and rotated.
 
+For the WS2 Attention experiment, the measured boundary also includes QKV and `o_proj`
+projections plus their TP/SP communication contracts. Those projections use native TE or
+vLLM callables only after an H100 bitwise probe; otherwise both sides use the deterministic
+`DetGemmOp` path with BF16 I/O, FP32 accumulation, ascending-K reduction, and Split-K disabled.
+The model input RMSNorm and residual add remain outside this boundary.
+
 ```text
 q --\
 k ----softmax(QKᵀ/√d + mask)·V--> out
@@ -95,6 +101,34 @@ 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).
 
+### WS2 deterministic CP reference
+
+### 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
+```
+
 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.

rl_engine/kernels/gtest/tolerance.py

@@ -202,6 +202,62 @@ def load_contract(
     return contract
 
 
+def resolve_logprob_threshold(dtype: Any) -> float:
+    """Return the fixed WS1 selected-logprob absolute-difference threshold.
+
+    The contract path is intentionally not configurable through this accessor.
+    Cross-configuration experiment definitions may select a dtype, but they cannot
+    inject or override a numerical threshold.
+    """
+
+    dtype_name = _normalize_dtype_name(dtype)
+    contract = load_contract()
+    try:
+        values = contract["accuracy"]["default"]["logprob"][dtype_name]
+        raw_threshold = values["atol"]
+    except (KeyError, TypeError) as exc:
+        raise ValueError(f"WS1 has no logprob threshold for dtype {dtype_name!r}") from exc
+    if isinstance(raw_threshold, bool) or not isinstance(raw_threshold, (int, float)):
+        raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}")
+    threshold = float(raw_threshold)
+    if not math.isfinite(threshold) or threshold < 0.0:
+        raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}")
+    return threshold
+
+
+def tolerance_contract_fingerprint() -> str:
+    """Return a deterministic fingerprint of the current WS1 contract contents."""
+
+    canonical = json.dumps(
+        load_contract(),
+        ensure_ascii=True,
+        separators=(",", ":"),
+        sort_keys=True,
+    ).encode("utf-8")
+    return hashlib.sha256(canonical).hexdigest()
+
+
+def _normalize_dtype_name(dtype: Any) -> str:
+    normalized = str(dtype).strip().lower().replace("torch.", "").replace("-", "")
+    aliases = {
+        "bf16": "bfloat16",
+        "bfloat16": "bfloat16",
+        "fp16": "float16",
+        "float16": "float16",
+        "half": "float16",
+        "fp32": "float32",
+        "float32": "float32",
+        "float": "float32",
+    }
+    try:
+        return aliases[normalized]
+    except KeyError as exc:
+        valid = ", ".join(sorted(set(aliases.values())))
+        raise ValueError(
+            f"unsupported WS1 logprob dtype {dtype!r}; expected one of: {valid}"
+        ) from exc
+
+
 def validate_contract_schema(contract: Mapping[str, Any]) -> None:
     """Validate four-judgment schema, dtype policy, roles, and aggregates."""
 
@@ -1056,6 +1112,8 @@ def tolerance_contract_fingerprint() -> str:
     "tolerance_contract_fingerprint",
     "resolve_tolerance",
     "resolve_tolerance_support",
+    "resolve_logprob_threshold",
+    "tolerance_contract_fingerprint",
     "validate_backend_provenance",
     "validate_contract_schema",
 ]

tests/test_attention_ablation.py

@@ -139,6 +139,44 @@ def forward_with_lse(self, q, k, v, *, causal, scale):
     assert torch.equal(result.out, q)
 
 
+def test_cp_production_configuration_fails_closed_without_ag_rs_backend():
+    q, k, v = _qkv()
+    cp_sharding = ShardingSpec(
+        tp_rank=0,
+        tp_world_size=1,
+        cp_rank=0,
+        cp_world_size=2,
+        global_q_heads=2,
+        global_kv_heads=1,
+        local_q_head_start=0,
+        local_q_heads=2,
+        local_kv_head_start=0,
+        local_kv_heads=1,
+        global_sequence_length=4,
+        local_sequence_length=2,
+        global_block_indices=(0,),
+        global_block_token_starts=(0,),
+        local_block_offsets=(0, 2),
+    )
+    contract = AttentionContract(
+        role=AttentionRole.TRAIN,
+        mode=AttentionMode.PREFILL,
+        dtype=AttentionDType.BF16,
+        batch_size=1,
+        query_sequence_length=2,
+        head_dim=4,
+        causal=True,
+        causal_offsets=(0,),
+        sharding=cp_sharding,
+        reduction=ReductionSpec(),
+        split_kv=SplitKVSpec.disabled(),
+    )
+    with pytest.raises(AttentionContractError, match="injected AG/RS backend"):
+        AttentionAblationOp(communication_backend="self_owned_cuda_ag_rs")(
+            q[:, :, :2], k[:, :, :2], v[:, :, :2], contract=contract
+        )
+
+
 def test_wrapper_owned_deterministic_core_does_not_require_external_provenance():
     q, k, v = _qkv()
 
@@ -279,6 +317,13 @@ def __call__(self, q, k, v, *, causal, scale):
     assert result.provenance["production_ready"] is True
 
 
+def test_deterministic_attention_rejects_runtime_split_kv_auto():
+    q, k, v = _qkv()
+    contract = _contract(split_kv=SplitKVSpec.auto(strict_consistency=False))
+    with pytest.raises(AttentionContractError, match="Split-KV"):
+        AttentionAblationOp()(q, k, v, contract=contract)
+
+
 @pytest.mark.parametrize(
     ("missing_field", "replacement"),
     [

tests/test_attention_preprocess.py

@@ -152,6 +152,7 @@ def _inputs():
 @requires_h100_preprocess
 def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope():
     q, k, q_weight, k_weight, positions = _inputs()
+    result = H100AttentionPreprocessor()(q, k, q_weight, k_weight, positions)
     result = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False)(
         q, k, q_weight, k_weight, positions
     )
@@ -179,6 +180,7 @@ def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope():
 @requires_h100_preprocess
 def test_h100_preprocessor_is_bitwise_batch_invariant_for_2d_positions():
     q, k, q_weight, k_weight, positions = _inputs()
+    op = H100AttentionPreprocessor()
     op = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False)
     full = op(q, k, q_weight, k_weight, positions)
 

tests/test_cp_attention.py

@@ -745,18 +745,15 @@ def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype):
     op = DeterministicCPAttentionReferenceOp(strict_bitwise=True)
 
     cp1_out, cp1_lse = op.forward_with_lse(
-        q,
-        k,
-        v,
-        cp_world_size=1,
-        kv_chunk_size=None,
+        q, k, v, cp_world_size=1, kv_chunk_size=3
     )
     cp2_out, cp2_lse = op.forward_with_lse(
-        q,
-        k,
-        v,
-        cp_world_size=2,
-        kv_chunk_size=3,
+        q, k, v, cp_world_size=2, kv_chunk_size=3
+    )
+    cp1_out, cp1_lse = op.forward_with_lse(q, k, v, cp_world_size=1, kv_chunk_size=3)
+    cp2_out, cp2_lse = op.forward_with_lse(q, k, v, cp_world_size=2, kv_chunk_size=3)
+    single_out, single_lse = op.forward_with_lse(
+        q[:1], k[:1], v[:1], cp_world_size=1, kv_chunk_size=3
     )
     assert torch.equal(cp1_out, cp2_out)
     assert torch.equal(cp1_lse, cp2_lse)

tests/test_rms_norm.py

@@ -5,14 +5,17 @@
 import torch
 import torch.nn.functional as F
 
-from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda
+from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp, rmsnorm_cuda
 from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp
 from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton
 
 try:
     from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
 
-    _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward")
+    _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and all(
+        hasattr(_C, name)
+        for name in ("rmsnorm_forward", "rmsnorm_backward_dx", "rmsnorm_backward_dw")
+    )
 except ImportError:  # pragma: no cover - import can fail when the extension is not built.
     _HAS_CUDA_RMSNORM = False
 
@@ -236,8 +239,12 @@ def test_registry_dispatches_rms_norm():
     from rl_engine.kernels.registry import kernel_registry
 
     op = kernel_registry.get_op("rms_norm")
-    assert isinstance(op, NativeRMSNormOp)
-    assert hasattr(op, "forward") and hasattr(op, "forward_fp32")
+    if torch.cuda.is_available() and _HAS_CUDA_RMSNORM:
+        assert isinstance(op, RMSNormCudaOp)
+        assert hasattr(op, "forward")
+    else:
+        assert isinstance(op, NativeRMSNormOp)
+        assert hasattr(op, "forward") and hasattr(op, "forward_fp32")
 
 
 @requires_cuda

tests/test_tolerance_contract.py

@@ -6,11 +6,15 @@
 from __future__ import annotations
 
 import copy
+import hashlib
+import inspect
+import json
 import math
 
 import pytest
 import torch
 
+from rl_engine.kernels.gtest import tolerance as tolerance_module
 from rl_engine.kernels.gtest.tolerance import (
     CHAIN_AGGREGATE_METRICS,
     JUDGMENTS,
@@ -26,8 +30,10 @@
     resolve_chain_aggregate_thresholds,
     resolve_comparison_roles,
     resolve_dtype_policy,
+    resolve_logprob_threshold,
     resolve_tolerance,
     resolve_tolerance_support,
+    tolerance_contract_fingerprint,
     validate_backend_provenance,
     validate_contract_schema,
 )