[WS2][PR5][Attention] Add distributed CP drift benchmark artifacts

docs/design/ws2-attention-pr5-distributed-drift-benchmark.md

@@ -0,0 +1,129 @@
+# WS2 Attention PR5 Drift Benchmark
+
+PR5 adds the report artifact path for issue #235. It does not introduce a
+production communication kernel. The benchmark is a rank-aware, torchrun-style
+driver around the deterministic CP attention reference. Under a matching
+two-rank CUDA/NCCL launch it executes the P2P reference transport; CPU/Gloo
+remains a report-generation smoke path.
+
+## Scope
+
+The benchmark covers the Qwen3-8B Attention target:
+
+- global heads: `Hq=32`, `Hkv=8`, `D=128`
+- TP sweep: `TP=1/2`; TP only changes the local head shard shape
+- CP sweep: `CP=1/2`
+- modes: full prefill and chunked-prefill replay
+- dtype path: BF16 candidate path compared with FP32 reference
+- optional backward: `dq`, `dk`, `dv` drift from the PR8 reference
+- optional RoPE composition before Attention, while CP Attention still consumes
+  post-RoPE Q/K
+
+The report separates two drift classes:
+
+| Field | Meaning |
+| --- | --- |
+| `drift.cp_merge_fp32` | CP/chunked candidate with FP32 output vs CP=1 FP32 prefill. This isolates CP merge and split-KV order. |
+| `drift.dtype_path_vs_fp32` | BF16 candidate path vs FP32 reference. This exposes arithmetic/final-write drift. |
+| `merge_order_probe` | Reversed-arrival partial states vs canonical sorted merge. This verifies that arrival order is ignored. |
+| `te_merge_oracle` | Optional Transformer Engine merge-oracle drift when TE is installed and passes capability probes. |
+| `backward` | Optional PR8 `dq/dk/dv` drift report when `--include-backward` is used. |
+| `distributed_p2p_reference` | Real NCCL P2P partial-state gather, FP32 merge, and query scatter drift. |
+
+Selected-logprob `dlogp` remains `not_available` here because the full logprob
+chain integration is outside PR5. PR4/WS2 runtime integration should fill that
+field once Attention is wired into the chain.
+
+## Commands
+
+Local smoke:
+
+```bash
+python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json
+```
+
+Qwen3 TP=2 / CP=2 with backward drift and a JSON artifact:
+
+```bash
+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
+```
+
+Two-GPU NCCL transport check:
+
+```bash
+torchrun --standalone --nproc-per-node=2 \
+  scripts/ws2_p2p_nccl_attention_reference_check.py
+```
+
+Two-GPU benchmark report with real P2P transport:
+
+```bash
+torchrun --standalone --nproc-per-node=2 \
+  benchmarks/benchmark_ws2_cp_attention_drift.py \
+  --smoke \
+  --device cuda \
+  --init-process-group \
+  --tp-world-sizes 2 \
+  --cp-world-sizes 2 \
+  --json
+```
+
+Rank 0 prints or writes the shared report. Other ranks can run the same
+rank-aware benchmark without changing the numerical reducer.  The recommended
+container is the repository CUDA image built from `docker/Dockerfile.cuda`
+(`ghcr.io/rl-align/rl-kernel/rl-kernel-ci:cuda` when using the repository image
+workflow). It is based on PyTorch 2.4 / CUDA 12.4 and includes NCCL support.
+
+## Transformer Engine Reuse
+
+PR5 reuses Transformer Engine only as an optional merge oracle, not as the
+source of truth. The adapter imports:
+
+```text
+transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
+```
+
+and uses these APIs when available:
+
+```text
+flash_attn_fwd_softmax_lse_correction
+flash_attn_fwd_out_correction_init
+flash_attn_fwd_out_correction
+```
+
+The benchmark first builds RL-Kernel partial states:
+
+```text
+state_i = (out_i, lse_i, global_block_index_i)
+```
+
+then sorts them by `global_block_index`. TE is allowed to perform only the
+online-softmax correction arithmetic for those already-sorted states. If TE is
+missing, incompatible, or fails the numeric self-test, the report records a
+provenance fallback and continues with the deterministic RL-Kernel merge.
+
+## Report Contract
+
+The JSON root contains:
+
+```text
+schema_version
+issue / pr
+launch.command
+runtime.rank_env
+target
+te_context_parallel_merge
+dlogp
+cases[]
+```
+
+Each case records topology, RoPE/cache provenance, split-KV policy, block
+metadata hash, drift summaries, per-logical-CP-rank metrics, and optional
+backward drift. The merge order is always `global_block_index`, and
+`downcast_at` is always `final_write`.

docs/operators/attention.md

@@ -85,6 +85,32 @@ Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_
 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
 
 Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled):
@@ -137,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
@@ -194,6 +221,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