feat(attention): add CP backward drift validation

[inaniloquentee]

Summary

Implements PR8 of #235 on top of PR3 / #238.

This PR extends the deterministic CP attention reference with training-side backward validation for causal prefill and chunked-prefill. It does not add a production fused backward kernel; it adds a correctness/reporting path that future fused training backends must match before they claim WS2 alignment.

Changes

• Add DeterministicCPAttentionReferenceOp.backward_reference(...) to materialize deterministic training backward over the same CP forward graph.
• Add compare_cp_attention_backward(...) to compare CP=1 backward against a CP/chunked-prefill candidate.
• Report whole-tensor and per-logical-CP-rank drift for dq, dk, and dv.
• Include out and attention-domain lse drift in the same report.
• Record backward provenance: saved forward state, CP world size, KV chunk size, fixed global_block_index merge order, fp32 accumulation, final-write downcast, dtype metadata, and unsupported decode backward.
• Register cp_attention in the generic operator checker so scripts/check_operator.py --op cp_attention --check-grad can run.
• Document the PR8 backward validation boundary in the attention operator docs.

RoPE / TE Boundary

• The backward reference consumes the same attention-ready Q/K boundary as PR3. For Qwen3 WS2 that means post-QK-Norm, post-RoPE Q/K.
• Tests cover post-RoPE Q/K with shared global position offsets.
• Transformer Engine backward is intentionally not claimed in this PR because compatible saved forward state is not exposed here. The report records te_backward_oracle=not_used.
• Decode replay remains forward-only; decode backward is recorded as not_supported.

Testing

Local validation:

• PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_cp_attention.py -q - passed, 23 passed.
• PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_cp_attention.py tests/test_cp_attention_transformer_engine.py tests/test_operator_inputs.py -q - passed, 41 passed, 1 skipped.
• PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_attention.py -q -k "not large and not gpu" - passed, 24 passed, 2 deselected.
• PYTHONPATH=. pytest rl_engine/tests/test_dispatch.py tests/test_attention.py tests/test_kv_cache_attention.py tests/test_cp_attention.py tests/test_cp_attention_transformer_engine.py tests/test_operator_inputs.py -q -k "not large and not gpu" - passed, 85 passed, 1 skipped, 2 deselected.
• PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_attention_correctness.py -q -rs - passed/collected, 45 passed, 82 skipped locally because CUDA FlashAttention and ROCm are unavailable.
• PYTHONPATH=. mypy --ignore-missing-imports rl_engine/ - passed.
• mkdocs build --strict -f mkdocs.yaml - passed.
• pre-commit run --all-files - passed.
• PYTHONPATH=. py -3.13 scripts/check_operator.py --op cp_attention --candidate pytorch --dtype fp32 --device cpu --batch 1 --seq 4 --check-grad - passed.

Qwen3-8B TP=2 CP=2 synthetic Attention smoke:

local TP shard: Hq=16, Hkv=4, D=128, BF16, CP=2, seq=4
forward_out_max_abs=0.00000000e+00
forward_lse_max_abs=2.38418579e-07
backward_dq_max_abs=1.56250000e-02
backward_dk_max_abs=1.56250000e-02
backward_dv_max_abs=3.12500000e-02

DCO

• Commit is signed off: Signed-off-by: inaniloquentee <3051000145@qq.com>.

Stack

Base: feat/ws2-cp-attention-reference-pr3 / #238.

Related: #235, #236, #238.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: a0c7d81c-c4c2-498c-b8b7-f27885c4d71c

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}