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 input generation, documentation, and correctness 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 issues
Suggested labels: 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. |
3 participants
Summary
(out, lse)partial-state merge in deterministic global KV-block ordercp_attentionand add CP=2 synthetic inputs for local harnessesPart of #235. This PR intentionally does not add a CUDA/fused backend, real distributed runtime, or decode KV-cache replay.
Checks
git commit -saddedSigned-off-bypre-commit run --all-filespython -m mypy --ignore-missing-imports rl_engine/python -m pytest rl_engine/tests/test_dispatch.py -vPYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rspython -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 -qmkdocs build --strict -f mkdocs.yamlpython -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.pySummary by CodeRabbit
cp_attention) attention reference for prefill and chunked-prefill, including causal masking, key padding masks, and position offsets.cp_attentionoperator registration/dispatch and synthetic test input generation for CP setups.cp_attentionis a PyTorch reference (not a fused/distributed runtime).