feat(attention): add CP backward drift validation

docs/operators/attention.md

@@ -13,12 +13,6 @@ 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
@@ -90,18 +84,6 @@ 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).
-
-### WS2 deterministic CP reference
 
 ### WS2 CP-aware dispatch
 
@@ -129,76 +111,6 @@ python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2
   --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.
-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
-```
-
-### 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 the
-[WS2 CP-aware Attention contract in PR #236](https://github.com/RL-Align/RL-Kernel/blob/feat/issue-235-attention-cp-contract/docs/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
-```
-
-### 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 the [WS2 CP-aware Attention contract PR][attention-contract-pr].
-
-[attention-contract-pr]: https://github.com/RL-Align/RL-Kernel/pull/236
-
-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):
@@ -309,8 +221,12 @@ 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.
+- `save_forward_state(q, k, v, ...)` captures FP32 `out/lse`, masks, position metadata,
+  topology, Split-KV boundaries, and content fingerprints for the exact forward invocation.
+- `backward_reference(q, k, v, dout, ..., saved_forward_state=state)` consumes that saved
+  state in canonical global KV-block order. It fails closed if Q/K/V, saved tensors, masks,
+  offsets, topology, or Split-KV metadata changed, 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.
 

rl_engine/kernels/gtest/tolerance.py

@@ -202,62 +202,6 @@ 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."""
 
@@ -1112,8 +1056,6 @@ 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",
 ]