feat(attention): add deterministic CP reference - #238
Conversation
📝 WalkthroughWalkthroughAdds a deterministic PyTorch context-parallel attention reference with chunked-prefill support, FP32 partial-state merging, registry dispatch, synthetic inputs, documentation, and validation tests. ChangesContext-parallel attention
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant DeterministicCPAttentionReferenceOp
participant AttentionPartialState
participant merge_attention_partial_states
Caller->>DeterministicCPAttentionReferenceOp: Submit Q, K, V and CP parameters
DeterministicCPAttentionReferenceOp->>AttentionPartialState: Compute query-shard/KV-block partials
DeterministicCPAttentionReferenceOp->>merge_attention_partial_states: Merge partial states by global KV block order
merge_attention_partial_states-->>DeterministicCPAttentionReferenceOp: FP32 output and LSE
DeterministicCPAttentionReferenceOp-->>Caller: Return output and optional LSE
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
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. Comment |
Signed-off-by: inaniloquentee <3051000145@qq.com>
8e3c6aa to
ed05a09
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/ops/pytorch/attention/cp_attention.py`:
- Around line 313-329: Run Black on
rl_engine/kernels/ops/pytorch/attention/cp_attention.py and commit its
formatting changes, specifically the torch.zeros call in the skv == 0 early
return and the None if key_padding_mask is None conditional in
local_partial_state(...).
- Around line 427-444: Update the zero-output tensor created in the states-empty
branch of the attention operation to explicitly use torch.float32, matching the
lse tensor and the analogous no-chunks branch. Keep the existing device, shape,
and zero-dependency behavior unchanged so all entries in out_chunks have a
consistent dtype for torch.cat.
- Around line 340-353: In local_partial_state and _merge_two_states, replace
non-finite LSE values with a finite placeholder before any subtraction or
exponentiation. Use the guarded LSE for scores - lse and lse_a - merged_lse,
while preserving zero outputs for fully masked rows and the existing merge
behavior for finite rows.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1055b627-fdd1-486c-bf2d-d99ae146f9fa
📒 Files selected for processing (6)
docs/operators/attention.mdrl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/ops/pytorch/attention/cp_attention.pyrl_engine/kernels/registry.pytests/test_cp_attention.pytests/test_operator_inputs.py
| if skv == 0: | ||
| zero_dep = _zero_dependency(qf, kf, vf) | ||
| return AttentionPartialState( | ||
| out=torch.zeros( | ||
| q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32 | ||
| ) | ||
| + zero_dep, | ||
| lse=torch.full( | ||
| (q.size(0), hq, sq), | ||
| float("-inf"), | ||
| device=q.device, | ||
| dtype=torch.float32, | ||
| ) | ||
| + zero_dep, | ||
| block_start=k_start, | ||
| block_end=k_start, | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Black formatting not applied — CI linting fails.
Per pipeline failures, the black pre-commit hook reformats this file at two spots that were committed unformatted: the multi-line torch.zeros(...) call inside the skv == 0 early-return (around lines 316-319), and the None if key_padding_mask is None else ... conditional inside the local_partial_state(...) call (around lines 415-419). Run black locally and commit the reformatted file.
Also applies to: 405-426
🧰 Tools
🪛 GitHub Actions: CI-Pipeline / 1_linting.txt
[error] 313-315: Pre-commit hook 'black' failed: reformatting required. The hook modified the file (e.g., changed multi-line torch.zeros(...) call to a single line).
🪛 GitHub Actions: CI-Pipeline / linting
[error] 313-313: pre-commit hook "black" failed. File was reformatted by Black (CI requires committed formatting changes).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/ops/pytorch/attention/cp_attention.py` around lines 313 -
329, Run Black on rl_engine/kernels/ops/pytorch/attention/cp_attention.py and
commit its formatting changes, specifically the torch.zeros call in the skv == 0
early return and the None if key_padding_mask is None conditional in
local_partial_state(...).
Source: Pipeline failures
| if key_padding_mask is not None: | ||
| scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) | ||
|
|
||
| lse = torch.logsumexp(scores, dim=-1) | ||
| finite_lse = torch.isfinite(lse) | ||
| weights = torch.exp(scores - lse.unsqueeze(-1)) | ||
| weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) | ||
| out = torch.matmul(weights, vf) | ||
| return AttentionPartialState( | ||
| out=out, | ||
| lse=lse, | ||
| block_start=k_start, | ||
| block_end=k_start + skv, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and nearby symbols first.
ast-grep outline rl_engine/kernels/ops/pytorch/attention/cp_attention.py --view expanded || true
# Read the implementation around the cited lines.
sed -n '300,500p' rl_engine/kernels/ops/pytorch/attention/cp_attention.py
# Find the tests mentioned in the review comment and nearby masking/gradient coverage.
rg -n "test_key_padding_mask_and_all_masked_rows_are_stable|test_empty_kv_backward_returns_zero_grads|chunked_gradients_match_cp1_reference|all_masked|key_padding_mask" -S rl_engine tests . || trueRepository: RL-Align/RL-Kernel
Length of output: 26910
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the cp_attention tests around the cited cases.
sed -n '180,340p' tests/test_cp_attention.py
# Read the standard attention implementation's masked-row handling for comparison.
sed -n '150,190p' rl_engine/kernels/ops/pytorch/attention/standard_attn.py
# Run a minimal autograd repro for the exact pattern.
python3 - <<'PY'
import torch
torch.set_printoptions(sci_mode=False)
scores = torch.tensor([[[float('-inf'), float('-inf')]]], requires_grad=True)
lse = torch.logsumexp(scores, dim=-1)
finite = torch.isfinite(lse)
weights = torch.exp(scores - lse.unsqueeze(-1))
weights = torch.where(finite.unsqueeze(-1), weights, torch.zeros_like(weights))
loss = weights.sum()
loss.backward()
print("lse:", lse)
print("weights:", weights)
print("scores.grad:", scores.grad)
print("has_nan_grad:", torch.isnan(scores.grad).any().item())
scores2 = torch.tensor([[[0.0, float('-inf')]]], requires_grad=True)
lse2 = torch.logsumexp(scores2, dim=-1)
finite2 = torch.isfinite(lse2)
weights2 = torch.exp(scores2 - lse2.unsqueeze(-1))
weights2 = torch.where(finite2.unsqueeze(-1), weights2, torch.zeros_like(weights2))
loss2 = weights2.sum()
loss2.backward()
print("mixed lse:", lse2)
print("mixed weights:", weights2)
print("mixed grad:", scores2.grad)
print("mixed has_nan_grad:", torch.isnan(scores2.grad).any().item())
# Repro for the merge path.
lse_a = torch.tensor([float('-inf')], requires_grad=True)
lse_b = torch.tensor([float('-inf')], requires_grad=True)
merged = torch.logaddexp(lse_a, lse_b)
finite_m = torch.isfinite(merged)
wa = torch.where(finite_m, torch.exp(lse_a - merged), torch.zeros_like(merged))
wb = torch.where(finite_m, torch.exp(lse_b - merged), torch.zeros_like(merged))
loss3 = (wa + wb).sum()
loss3.backward()
print("merged:", merged)
print("weights:", wa, wb)
print("grad_a:", lse_a.grad, "grad_b:", lse_b.grad)
print("merge_has_nan_grad:", torch.isnan(lse_a.grad).any().item() or torch.isnan(lse_b.grad).any().item())
PYRepository: RL-Align/RL-Kernel
Length of output: 8383
Guard the -inf LSE before subtracting in both attention merge paths.
Fully masked rows/blocks produce nan gradients here: torch.where(...) only fixes the forward value, while exp(scores - lse) and exp(lse_a - merged_lse) still backpropagate through -inf - (-inf). Use a finite placeholder LSE before the subtraction in local_partial_state and _merge_two_states.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/ops/pytorch/attention/cp_attention.py` around lines 340 -
353, In local_partial_state and _merge_two_states, replace non-finite LSE values
with a finite placeholder before any subtraction or exponentiation. Use the
guarded LSE for scores - lse and lse_a - merged_lse, while preserving zero
outputs for fully masked rows and the existing merge behavior for finite rows.
zhangj1an
left a comment
zhangj1an
left a comment
There was a problem hiding this comment.
LGTM. Thanks for your work!
This FP32 Ring CP golden reference is numerically equivalent with the PyTorch implementation of flash_attn_fwd_softmax_lse_correction / flash_attn_fwd_out_correction in Transformer Engine, https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py. The softmax merging order is by KV block order.
Also the unit tests covered both forward and backward test cases (gradients via autograd).
| CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp | ||
|
|
||
| __all__ = [ | ||
| "AttentionPartialState", | ||
| "CPAttentionReferenceOp", | ||
| "DeterministicCPAttentionReferenceOp", | ||
| "merge_attention_partial_states", | ||
| ] |
There was a problem hiding this comment.
delete the alias DeterministicCPAttentionReferenceOp please
|
please resolve the code conflicts first, Thank you. |
Thanks for the review! Removed the alias and kept only DeterministicCPAttentionReferenceOp. Also agreed on the TE point: this PR keeps the PyTorch golden reference self-contained, while matching the same FP32 (out, lse) correction semantics as Transformer Engine. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/test_cp_attention_transformer_engine.py (1)
18-24: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winGuard the optional Transformer engine oracle against API drift.
This test only skips when the internal module cannot be imported. If an installed Transformer Engine release exposes the module but removes or renames the helper used at lines 43–45, the optional test raises
AttributeErrorinstead of skipping. Check that the requiredFlash attentionhelpers are present before calling them, or pin the supported Transformer Engine version.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_cp_attention_transformer_engine.py` around lines 18 - 24, Update _te_context_parallel_module to validate that the imported context-parallel module exposes the required Flash attention helpers used by the optional test before returning it. If any helper is missing or renamed, skip with a clear unavailable message instead of allowing AttributeError during the test; preserve the existing import-error skips.Source: MCP tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/test_cp_attention_transformer_engine.py`:
- Around line 18-24: Update _te_context_parallel_module to validate that the
imported context-parallel module exposes the required Flash attention helpers
used by the optional test before returning it. If any helper is missing or
renamed, skip with a clear unavailable message instead of allowing
AttributeError during the test; preserve the existing import-error skips.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 9b396caf-830f-45b7-8f6d-e4cd55c71ee5
📒 Files selected for processing (6)
docs/operators/attention.mdrl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/ops/pytorch/attention/cp_attention.pyrl_engine/kernels/registry.pytests/test_cp_attention_transformer_engine.pytests/test_operator_inputs.py
💤 Files with no reviewable changes (1)
- rl_engine/kernels/ops/pytorch/attention/cp_attention.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/test_operator_inputs.py
- rl_engine/kernels/registry.py
- rl_engine/kernels/gtest/operator_inputs.py
Fixed code conflicts. |
Signed-off-by: inaniloquentee <3051000145@qq.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
docs/operators/attention.md (1)
256-257: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClarify the compound modifier in the limitation.
Rewrite the sentence so that the relationship between production backends and fused
RoPE+Attentionkernels is explicit.Proposed wording
-- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused -- `RoPE+Attention` backend alignment are outside PR3. +- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2. RoPE execution and alignment + with production `RoPE+Attention`-fused backends are outside PR3.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/operators/attention.md` around lines 256 - 257, Rewrite the cp_attention limitation sentence in docs/operators/attention.md to explicitly state the relationship between production backends and fused RoPE+Attention kernels, while preserving that it consumes post-RoPE Q/K for Qwen3 WS2 and that RoPE execution is outside PR3.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_cp_attention.py`:
- Around line 114-144: Strengthen the test around rope.forward_fp32 and
op.forward_fp32_with_lse by using different global starts for query and key
positions, then pass those distinct offsets to the attention call. Add a
continuation-case assertion against an independent position-aware reference,
rather than relying only on the CP1-versus-CP2 comparison, so implementations
that ignore query_position_offsets or key_position_offsets cannot pass.
---
Nitpick comments:
In `@docs/operators/attention.md`:
- Around line 256-257: Rewrite the cp_attention limitation sentence in
docs/operators/attention.md to explicitly state the relationship between
production backends and fused RoPE+Attention kernels, while preserving that it
consumes post-RoPE Q/K for Qwen3 WS2 and that RoPE execution is outside PR3.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 365e8ce0-81ac-4717-8aa8-a9717322436f
📒 Files selected for processing (3)
docs/operators/attention.mdrl_engine/kernels/ops/pytorch/attention/cp_attention.pytests/test_cp_attention.py
🚧 Files skipped from review as they are similar to previous changes (1)
- rl_engine/kernels/ops/pytorch/attention/cp_attention.py
| position_offsets = torch.tensor([17, 103], dtype=torch.long) | ||
| positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) | ||
| q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) | ||
| k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) | ||
|
|
||
| assert not torch.equal(q, pre_rope_q.float()) | ||
| assert not torch.equal(k, pre_rope_k.float()) | ||
|
|
||
| with _single_thread(): | ||
| out1, lse1 = op.forward_fp32_with_lse( | ||
| q, | ||
| k, | ||
| v, | ||
| causal=True, | ||
| query_position_offsets=position_offsets, | ||
| key_position_offsets=position_offsets, | ||
| cp_world_size=1, | ||
| ) | ||
| out2, lse2 = op.forward_fp32_with_lse( | ||
| q, | ||
| k, | ||
| v, | ||
| causal=True, | ||
| query_position_offsets=position_offsets, | ||
| key_position_offsets=position_offsets, | ||
| cp_world_size=2, | ||
| kv_chunk_size=2, | ||
| ) | ||
|
|
||
| torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) | ||
| torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make position-offset handling observable in this test.
Lines 114-117 apply the same offset to Q and K. Lines 128-129 pass the same offset to both attention inputs. Adding the same value to both sides does not change the causal relation. An implementation that ignores both offset arguments can therefore pass this test.
The torch.equal checks only prove that RoPE changed the tensors. The CP1-versus-CP2 comparison only proves partition equivalence. Add a continuation case with different query and key global starts, and compare the result with an independent position-aware reference.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_cp_attention.py` around lines 114 - 144, Strengthen the test
around rope.forward_fp32 and op.forward_fp32_with_lse by using different global
starts for query and key positions, then pass those distinct offsets to the
attention call. Add a continuation-case assertion against an independent
position-aware reference, rather than relying only on the CP1-versus-CP2
comparison, so implementations that ignore query_position_offsets or
key_position_offsets cannot pass.
3 participants
Summary
(out, lse)partial-state merge in deterministic global KV-block order.cp_attentionand add CP=2 synthetic inputs for local harnesses.Part of #235. This PR intentionally does not add a CUDA/fused backend, real distributed runtime, or decode KV-cache replay.
RoPE Scope Update
Implemented in this PR:
DeterministicCPAttentionReferenceOpconsumes attention-ready Q/K; for Qwen3 WS2 this means post-QK-Norm, post-RoPE Q/K;RoPE+Attentioncan be validated at the input boundary;NativeRoPEOpfirst, then runs CP attention with the same global position offsets used to build RoPE positions;(out, lse)merge oracle.Boundary: PR3 validates deterministic CP partial-state merge over post-RoPE Q/K. Production fused
RoPE+Attentionalignment remains PR7 scope.Checks
RoPE amendment local validation:
PYTHONPATH=. PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_cp_attention.py tests/test_cp_attention_transformer_engine.py -q- passed,18 passed, 1 skipped.pre-commit run --files rl_engine/kernels/ops/pytorch/attention/cp_attention.py tests/test_cp_attention.py docs/operators/attention.md- passed.Previous PR3 checks before the RoPE amendment:
git commit -saddedSigned-off-by.pre-commit run --all-files.python -m mypy --ignore-missing-imports rl_engine/.python -m pytest rl_engine/tests/test_dispatch.py -v.PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs.python -m pytest tests/test_attention.py -v -k "not large and not gpu".python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu".python -m pytest tests/test_cp_attention.py tests/test_operator_inputs.py -q.mkdocs build --strict -f mkdocs.yaml.python -m ruff check rl_engine/kernels/ops/pytorch/attention/cp_attention.py tests/test_cp_attention.py rl_engine/kernels/gtest/operator_inputs.py tests/test_operator_inputs.py rl_engine/kernels/registry.py.GitHub Actions after the RoPE amendment:
CI-Pipeline / linting- passed.CI-Pipeline / docs- passed.CI-Pipeline / unit-tests- passed.GPU CI / gpu-tests- skipped as expected because the PR does not have theneeds-gpu-cilabel.CodeRabbit- review completed.