feat(attention): add deterministic CP reference

docs/operators/attention.md

@@ -85,6 +85,21 @@ 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.
 
+`kernel_registry.get_op("cp_attention")` resolves to
+`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel
+reference. It emulates CP prefill and chunked-prefill by splitting logical query
+and KV sequence blocks, computing per-block `(out, lse)` partial states, and
+merging them in fp32 by global KV block index. This path is not a production
+fused backend; it defines the CP/LSE merge behavior that downstream fused paths
+must match. Optional per-batch `query_position_offsets` / `key_position_offsets`
+cover varlen causal-mask metadata while keeping the dense tensor layout.
+
+For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not
+call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The
+position offsets passed to CP attention must describe the same absolute token
+positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary
+while PR7 can later validate production fused `RoPE+Attention` kernels.
+
 ## Accuracy
 
 Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled):
@@ -137,6 +152,8 @@ memory.
 
 ```bash
 python -m pytest tests/test_attention.py -v
+python -m pytest tests/test_cp_attention.py -v
+python -m pytest tests/test_cp_attention_transformer_engine.py -v  # optional TE oracle
 ```
 
 Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile
@@ -146,14 +163,30 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below),
 gradient flow, registry dispatch, and a
 GPU-only LARGE Qwen3-8B real-shape smoke test.
 
+`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard
+attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared
+global position metadata, chunked-prefill replay, global-position causal masking
+across CP boundaries, order-independent LSE merge by global block index,
+padding/all-masked stability, BF16 final-write behavior, input purity, argument
+validation, and registry dispatch.
+`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill
+synthetic case for local harnesses.
+`tests/test_cp_attention_transformer_engine.py` optionally imports NVIDIA
+Transformer Engine's context-parallel PyTorch correction helpers and checks that
+RL-Kernel's fp32 `(out, lse)` merge matches those helpers; the test skips when
+Transformer Engine is not installed.
+
 ## Implementation Files
 
 - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference
+- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` — CP prefill/chunked 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`
+- `tests/test_cp_attention.py`
+- `tests/test_cp_attention_transformer_engine.py`
 
 ## Fixed Reduction Order (CUDA Deterministic Backend)
 
@@ -218,6 +251,10 @@ for measured peak memory at representative shapes.
 - First version: `D=128` only (Qwen3-8B alignment).
 - Supported dtypes: BF16, FP16.
 - Full materialization of scores/P limits practical sequence length.
+- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics,
+  not a distributed runtime or fused kernel.
+- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused
+  `RoPE+Attention` backend alignment are outside PR3.
 - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise).
 - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op).
 - No FP8, no multi-GPU / sequence-parallel.

rl_engine/kernels/gtest/operator_inputs.py

@@ -29,6 +29,7 @@ def make_operator_inputs(
         "matmul": _make_matmul_inputs,
         "det_gemm": _make_det_gemm_inputs,
         "attention": _make_attention_inputs,
+        "cp_attention": _make_cp_attention_inputs,
         "logp": _make_logp_inputs,
         "linear_logp": _make_linear_logp_inputs,
         "batch_invariant_logp": _make_batch_invariant_logp_inputs,
@@ -53,6 +54,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}",
@@ -139,6 +141,26 @@ def _make_attention_inputs(
     return inputs
 
 
+def _make_cp_attention_inputs(
+    args: argparse.Namespace, dtype: torch.dtype, device: torch.device
+) -> dict[str, Any]:
+    batch, seq = _batch_seq(args)
+    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,
+        "cp_world_size": 2,
+        "kv_chunk_size": max(1, seq // 2),
+    }
+
+
 def _make_logp_inputs(
     args: argparse.Namespace, dtype: torch.dtype, device: torch.device
 ) -> dict[str, Any]: