feat(attention): add deterministic CP reference - #238
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a deterministic PyTorch context-parallel attention reference with chunked-prefill support, FP32 partial-state merging, WS2 contracts, 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 and KV-block partials
DeterministicCPAttentionReferenceOp->>merge_attention_partial_states: Merge partial states by global KV block order
merge_attention_partial_states-->>DeterministicCPAttentionReferenceOp: Return output and LSE
DeterministicCPAttentionReferenceOp-->>Caller: Return attention result and provenance
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
🛠️ Fix failing CI checks 💡
🧪 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.
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (3)
rl_engine/kernels/ops/pytorch/attention/cp_attention.py (1)
767-767: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
strict=Trueto thezipcall.Ruff B905 flags this call.
_split_boundsreturnscp_world_sizeentries for both bounds, sostrict=Truedoes not change behavior and documents the invariant.♻️ Proposed change
- for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate( + zip(q_bounds, kv_bounds, strict=True) + ):🤖 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` at line 767, Update the zip call in the bounds iteration within the attention processing logic to pass strict=True, preserving the existing loop behavior while explicitly validating that q_bounds and kv_bounds have matching lengths.Source: Linters/SAST tools
rl_engine/kernels/attention_contract.py (2)
1153-1162: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the identical
causal_offsetscount branches.Both branches assign
expected_causal_offsets = batch_size. Onlyoffset_ownerdiffers, and it is used for the message text. Simplify to one assignment plus a conditional label.♻️ Proposed simplification
- if self.sharding.packed_sequence_offsets is not None: - expected_causal_offsets = batch_size - offset_owner = "packed sequence" - else: - expected_causal_offsets = batch_size - offset_owner = "batch entry" - if len(causal_offsets) != expected_causal_offsets: + offset_owner = ( + "packed sequence" + if self.sharding.packed_sequence_offsets is not None + else "batch entry" + ) + if len(causal_offsets) != batch_size: raise AttentionContractError( f"causal_offsets must contain one entry per {offset_owner}" )🤖 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/attention_contract.py` around lines 1153 - 1162, In the causal-offset validation block, collapse the duplicated branches assigning expected_causal_offsets: assign it once to batch_size, then compute offset_owner conditionally from packed_sequence_offsets while preserving the existing labels and error behavior.
733-789: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winIndex entries by coordinate before the coverage and invariance scans.
Both
_validate_owner_coverageand_validate_rank_invariancere-scanself.entriesinside nested loops. The entry count isbatch_size * tp_world_size * cp_world_size**2, so the total work is quadratic in the entry count. For a moderately large topology, such asbatch_size=64,tp_world_size=8, andcp_world_size=8, this performs on the order of 10^8 comparisons per validation. Build one dictionary keyed bySplitKVRuntimeCoordinatein__post_init__and look entries up directly.♻️ Sketch of the indexed lookup
def __post_init__(self) -> None: ... actual_coordinates = [entry.coordinate for entry in entries] + by_coordinate = {entry.coordinate: entry for entry in entries} ... - self._validate_owner_coverage() - self._validate_rank_invariance() + self._validate_owner_coverage(by_coordinate) + self._validate_rank_invariance(by_coordinate)Then replace each list comprehension with
by_coordinate[SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank)].🤖 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/attention_contract.py` around lines 733 - 789, Build a coordinate-to-entry dictionary keyed by SplitKVRuntimeCoordinate in __post_init__. Update _validate_owner_coverage and _validate_rank_invariance to retrieve entries directly from this index instead of filtering self.entries inside nested loops, preserving the existing validation behavior and error handling.
🤖 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/attention_contract.py`:
- Around line 1-2: Run Black on attention_contract.py using the repository’s
standard formatter command and commit the resulting formatting changes, without
modifying behavior.
- Around line 1078-1085: Rename the loop variable field to a non-conflicting
name in the shown normalization loop and the corresponding loops in
AttentionContract.__post_init__ and AttentionBackendCapability.__post_init__.
Update all references within each loop, preserving the existing behavior while
avoiding shadowing the imported dataclasses.field.
In `@rl_engine/kernels/ops/pytorch/attention/cp_attention.py`:
- Around line 1019-1035: Remove the CPAttentionReferenceOp alias assignment and
its corresponding __all__ export; use DeterministicCPAttentionReferenceOp
directly and verify no remaining importers reference the removed alias.
- Around line 457-466: Update split_kv_execution_plan_provenance to return an
empty plan list when the KV sequence length is zero, before validation rejects
the input. Preserve existing planning behavior for positive lengths, and add a
regression test covering the backward_reference path with empty KV tensors.
- Around line 794-804: Update the percentile calculations in the
GradientDriftStats construction to avoid torch.quantile for tensors exceeding
2**24 elements, while preserving default linear interpolation semantics for
p95_abs and p99_abs. Use a large-tensor-safe reduction based on sorting and
interpolation, and keep the existing empty-tensor handling and other statistics
unchanged.
---
Nitpick comments:
In `@rl_engine/kernels/attention_contract.py`:
- Around line 1153-1162: In the causal-offset validation block, collapse the
duplicated branches assigning expected_causal_offsets: assign it once to
batch_size, then compute offset_owner conditionally from packed_sequence_offsets
while preserving the existing labels and error behavior.
- Around line 733-789: Build a coordinate-to-entry dictionary keyed by
SplitKVRuntimeCoordinate in __post_init__. Update _validate_owner_coverage and
_validate_rank_invariance to retrieve entries directly from this index instead
of filtering self.entries inside nested loops, preserving the existing
validation behavior and error handling.
In `@rl_engine/kernels/ops/pytorch/attention/cp_attention.py`:
- Line 767: Update the zip call in the bounds iteration within the attention
processing logic to pass strict=True, preserving the existing loop behavior
while explicitly validating that q_bounds and kv_bounds have matching lengths.
🪄 Autofix
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: cefd3690-2125-47b9-8634-25af36822583
📒 Files selected for processing (4)
docs/operators/attention.mdrl_engine/kernels/attention_contract.pyrl_engine/kernels/ops/pytorch/attention/cp_attention.pytests/test_cp_attention.py
| # SPDX-License-Identifier: Apache-2.0 | ||
| # Copyright (c) 2026 RL-Kernel Contributors |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Run black on this file.
The linting job reports that Black reformats this file. Run black rl_engine/kernels/attention_contract.py and commit the result.
🧰 Tools
🪛 GitHub Actions: CI-Pipeline / 2_linting.txt
[error] 1-1: Black formatting check failed and reformatted this file. Run 'pre-commit run --all-files' or 'black' locally, then commit the changes.
🪛 GitHub Actions: CI-Pipeline / linting
[error] 1-1: Black formatting check failed and reformatted this file. Run 'black rl_engine/kernels/attention_contract.py' and commit the 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/attention_contract.py` around lines 1 - 2, Run Black on
attention_contract.py using the repository’s standard formatter command and
commit the resulting formatting changes, without modifying behavior.
Source: Pipeline failures
| for field in ("position_ids", "query_position_offsets", "key_position_offsets"): | ||
| values = getattr(self, field) | ||
| if values is None: | ||
| continue | ||
| normalized = _integer_tuple(values, field) | ||
| if not normalized or any(value < 0 for value in normalized): | ||
| raise AttentionContractError(f"{field} must contain non-negative positions") | ||
| object.__setattr__(self, field, normalized) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Rename the loop variable field to unblock CI.
Flake8 F402 and Ruff F402 fail here and at Line 1189 and Line 1339. The loop variable field shadows the dataclasses.field import at Line 14. The linting job fails the build. Rename the loop variable at all three sites.
🔧 Proposed fix for Line 1078
- for field in ("position_ids", "query_position_offsets", "key_position_offsets"):
- values = getattr(self, field)
+ for field_name in ("position_ids", "query_position_offsets", "key_position_offsets"):
+ values = getattr(self, field_name)
if values is None:
continue
- normalized = _integer_tuple(values, field)
+ normalized = _integer_tuple(values, field_name)
if not normalized or any(value < 0 for value in normalized):
- raise AttentionContractError(f"{field} must contain non-negative positions")
- object.__setattr__(self, field, normalized)
+ raise AttentionContractError(f"{field_name} must contain non-negative positions")
+ object.__setattr__(self, field_name, normalized)Also apply the same rename in AttentionContract.__post_init__ at Line 1189 and in AttentionBackendCapability.__post_init__ at Line 1339.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for field in ("position_ids", "query_position_offsets", "key_position_offsets"): | |
| values = getattr(self, field) | |
| if values is None: | |
| continue | |
| normalized = _integer_tuple(values, field) | |
| if not normalized or any(value < 0 for value in normalized): | |
| raise AttentionContractError(f"{field} must contain non-negative positions") | |
| object.__setattr__(self, field, normalized) | |
| for field_name in ("position_ids", "query_position_offsets", "key_position_offsets"): | |
| values = getattr(self, field_name) | |
| if values is None: | |
| continue | |
| normalized = _integer_tuple(values, field_name) | |
| if not normalized or any(value < 0 for value in normalized): | |
| raise AttentionContractError(f"{field_name} must contain non-negative positions") | |
| object.__setattr__(self, field_name, normalized) |
🧰 Tools
🪛 Ruff (0.16.1)
[error] 1078-1078: Import field from line 14 shadowed by loop variable
(F402)
🤖 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/attention_contract.py` around lines 1078 - 1085, Rename the
loop variable field to a non-conflicting name in the shown normalization loop
and the corresponding loops in AttentionContract.__post_init__ and
AttentionBackendCapability.__post_init__. Update all references within each
loop, preserving the existing behavior while avoiding shadowing the imported
dataclasses.field.
Sources: Linters/SAST tools, Pipeline failures
| "requested_split_kv_policy": ( | ||
| "disabled" if kv_chunk_size is None else "fixed" | ||
| ), | ||
| "requested_split_kv_size": kv_chunk_size, | ||
| "actual_split_kv_plans": split_kv_execution_plan_provenance( | ||
| k.size(2), | ||
| cp_world_size=cp_world_size, | ||
| kv_chunk_size=kv_chunk_size, | ||
| backend="deterministic_cp_backward_reference", | ||
| ), |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C6 'empty_kv|size\(2\) == 0|skv == 0' rl_engine/kernels/ops/pytorch/attention/cp_attention.py tests/test_cp_attention.py
# Inspect the shared QKV validator to see whether an empty KV sequence is rejected earlier.
ast-grep run --pattern 'def _validate_qkv($$$) {
$$$
}' --lang python rl_engine || rg -n -A30 'def _validate_qkv' rl_engineRepository: RL-Align/RL-Kernel
Length of output: 6865
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- empty-KV backward test ---'
sed -n '271,315p' tests/test_cp_attention.py
printf '%s\n' '--- backward_reference and provenance ---'
rg -n -A45 -B12 'def backward_reference|split_kv_execution_plan_provenance|actual_split_kv_plans' \
rl_engine/kernels/ops/pytorch/attention/cp_attention.py
printf '%s\n' '--- all provenance callers ---'
rg -n -C5 'split_kv_execution_plan_provenance\(' rl_engine testsRepository: RL-Align/RL-Kernel
Length of output: 15777
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '409,480p' rl_engine/kernels/ops/pytorch/attention/cp_attention.py
sed -n '530,565p' rl_engine/kernels/ops/pytorch/attention/cp_attention.py
rg -n -C8 'backward_reference\(' tests rl_engineRepository: RL-Align/RL-Kernel
Length of output: 10355
Handle empty-KV provenance in backward
backward_reference supports empty KV tensors through the forward and autograd steps, but split_kv_execution_plan_provenance rejects length 0. Return an empty plan list for empty KV input and add a regression test.
🤖 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 457 -
466, Update split_kv_execution_plan_provenance to return an empty plan list when
the KV sequence length is zero, before validation rejects the input. Preserve
existing planning behavior for positive lengths, and add a regression test
covering the backward_reference path with empty KV tensors.
| diff = (candidate.float() - reference.float()).abs().reshape(-1) | ||
| active_count = int(diff.numel()) | ||
| if active_count == 0: | ||
| return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) | ||
| return GradientDriftStats( | ||
| max_abs=float(diff.max().item()), | ||
| mean_abs=float(diff.mean().item()), | ||
| p95_abs=float(torch.quantile(diff, 0.95).item()), | ||
| p99_abs=float(torch.quantile(diff, 0.99).item()), | ||
| active_count=active_count, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
torch.quantile input tensor too large limit 16777216 PyTorch 2.4
💡 Result:
The RuntimeError: quantile input tensor is too large error occurs in PyTorch because torch.quantile historically imposed a limit of 2^24 (16,777,216) elements when processing float32 tensors [1][2]. This limitation existed because the calculation of quantile ranks—specifically the formula rank = q * (size - 1)—was performed using the input tensor's own dtype (float32) [1]. Since float32 can only exactly represent integers up to 2^24, values exceeding this limit led to loss of precision, where distinct ranks would alias to the same index and interpolation weights would quantize incorrectly [1]. To prevent producing silent, inaccurate results, PyTorch implemented a safety check that raised this error [1]. In PyTorch 2.4 and later versions, this behavior has been updated to support larger float32 tensors by performing rank calculations in float64 (double precision), which maintains precision well beyond the 2^24 limit [1]. Key points: - Historical Limit: The 16,777,216 element cap was a safety mechanism to avoid floating-point precision errors during rank calculations [1]. - Updated Behavior: Recent updates (targeting versions including and following the 2.4 release cycle) allow float32 tensors larger than 2^24 to be processed by using float64 for intermediate rank arithmetic [1]. - Device Specifics: The update is device-aware; while most devices benefit from this, hardware or backends lacking float64 support (such as certain MPS configurations) may still maintain the original cap [1]. If you are encountering this error on an older version of PyTorch, the standard workarounds included using float64 inputs (if memory allows) or manually splitting the tensor into smaller chunks [2]. However, upgrading to the latest stable version of PyTorch is the primary resolution for this issue [1].
Citations:
- 1: Support large float32 tensors in quantile by computing ranks in float64 pytorch/pytorch#187574
- 2: Quantile is limited to 16 million elements and have poor performance. pytorch/pytorch#64947
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate files ---'
git ls-files | rg '(^|/)(pyproject\.toml|setup\.cfg|setup\.py|requirements[^/]*|environment[^/]*|Dockerfile|[^/]*\.ya?ml)$|rl_engine/kernels/ops/pytorch/attention/cp_attention\.py'
printf '%s\n' '--- PyTorch version references ---'
rg -n -i 'torch(==|>=|<=|~=|>|<)|pytorch(==|>=|<=|~=|>|<)|torch:[[:space:]]|pytorch:' . --glob '!**/.git/**' --glob '!**/node_modules/**' | head -200
printf '%s\n' '--- relevant source context ---'
sed -n '760,820p' rl_engine/kernels/ops/pytorch/attention/cp_attention.pyRepository: RL-Align/RL-Kernel
Length of output: 3637
🌐 Web query:
PyTorch 2.4.1 torch.quantile "input tensor is too large" 2^24 source
💡 Result:
The RuntimeError: quantile input tensor is too large in PyTorch is caused by a historical limitation where the internal rank calculation for torch.quantile was performed using the input tensor's own dtype [1]. When using float32, ranks above 2^24 (approximately 16.7 million elements) could not be represented exactly, leading to potential inaccuracies in interpolation and index gathering [1]. To prevent silently incorrect results, PyTorch implemented an explicit cap that raised this error [1][2]. Recent developments, including pull request #187574, have addressed this by switching the internal rank calculations to float64, which allows for exact representation of indices up to 2^53 [1]. While this update improves support for large tensors, the following context remains important: 1. Device-specific limitations: On devices like MPS (Metal Performance Shaders) that lack float64 support, the 2^24 element cap remains in place for float32 tensors [1]. 2. Workarounds: If you encounter this error in older versions of PyTorch or on unsupported hardware/dtypes, common workarounds include: * Using torch.kthvalue, which does not share this limitation and is often more performant for single quantile calculations [2][3]. * Implementing custom quantile functions using torch.kthvalue or sorting (though sorting can be memory-intensive) [2][3]. * Reducing the size of the input tensor (e.g., through sampling or processing in chunks) if an exact full-tensor quantile is not required for the application [4][5]. As of August 2026, ensure your environment is updated to a recent PyTorch version to benefit from the improved large-tensor support [1].
Citations:
- 1: Support large float32 tensors in quantile by computing ranks in float64 pytorch/pytorch#187574
- 2: Quantile is limited to 16 million elements and have poor performance. pytorch/pytorch#64947
- 3:
torch.quantileedge case pytorch/pytorch#157431 - 4: [algo] fix: remove torch.quantile-based percentile metrics to resolve tensor size limit error verl-project/verl#3810
- 5: Bug when calculating the rollout_is metrics verl-project/verl#3787
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
from pathlib import Path
for path in [Path("pyproject.toml"), Path("setup.py"), Path("ci/run_gpu_ci.sh")]:
print(f"--- {path} ---")
lines = path.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "torch" in line.lower() or "TORCH_SPEC" in line:
lo, hi = max(1, i - 3), min(len(lines), i + 3)
for n in range(lo, hi + 1):
print(f"{n}: {lines[n-1]}")
PY
printf '%s\n' '--- percentile-related tests and callers ---'
rg -n '_drift_stats|p95_abs|p99_abs|GradientDriftStats|attention_backward_drift' . --glob '!**/.git/**'Repository: RL-Align/RL-Kernel
Length of output: 15944
🏁 Script executed:
#!/bin/bash
set -eu
tmp=$(mktemp -d)
trap 'rm -rf "$tmp"' EXIT
for ref in v2.4.1 v2.5.0 v2.6.0; do
url="https://raw.githubusercontent.com/pytorch/pytorch/$ref/aten/src/ATen/native/Sorting.cpp"
echo "--- $ref Sorting.cpp quantile guards ---"
curl -fsSL "$url" | rg -n -C 8 'quantile input tensor is too large|2\^24|16777216|quantile|kthvalue' | head -180 || true
done
echo '--- upstream PR metadata ---'
curl -fsSL https://api.github.com/repos/pytorch/pytorch/pulls/187574 |
jq '{number,title,state,merged_at,merge_commit_sha,base:.base.ref,head:.head.ref}'
echo '--- v2.4.1 quantile references in source tree listing ---'
curl -fsSL 'https://api.github.com/repos/pytorch/pytorch/git/trees/v2.4.1?recursive=1' |
jq -r '.tree[].path' | rg 'quantile|Sorting\.(cpp|cu|h)$' | head -100Repository: RL-Align/RL-Kernel
Length of output: 14645
Make percentile calculation safe for tensors larger than 2²⁴.
PyTorch 2.4.1 raises RuntimeError: quantile() input tensor is too large above 2**24 elements. CI installs torch==2.4.1, and realistic gradients can exceed this limit. Replace torch.quantile with a large-tensor-safe reduction. Preserve its default linear interpolation; rounded indexing after torch.sort changes p95 and p99 values.
🤖 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 794 -
804, Update the percentile calculations in the GradientDriftStats construction
to avoid torch.quantile for tensors exceeding 2**24 elements, while preserving
default linear interpolation semantics for p95_abs and p99_abs. Use a
large-tensor-safe reduction based on sorting and interpolation, and keep the
existing empty-tensor handling and other statistics unchanged.
| CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp | ||
|
|
||
| __all__ = [ | ||
| "AttentionBackwardComparisonReport", | ||
| "AttentionBackwardGradients", | ||
| "AttentionBackwardPathDrift", | ||
| "AttentionBackwardPathResult", | ||
| "AttentionBackwardRankDrift", | ||
| "AttentionPartialState", | ||
| "build_reference_split_kv_runtime_plan_set", | ||
| "CPAttentionReferenceOp", | ||
| "DeterministicCPAttentionReferenceOp", | ||
| "GradientDriftStats", | ||
| "compare_cp_attention_backward", | ||
| "merge_attention_partial_states", | ||
| "split_kv_execution_plan_provenance", | ||
| ] |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
The alias CPAttentionReferenceOp is still present.
A previous review asked to delete the alias, and the PR description states that the alias was removed. Line 1019 still defines it, and Line 1029 still exports it. Remove both lines, or state why the alias must stay.
♻️ Proposed removal
-CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp
-
__all__ = [
"AttentionBackwardComparisonReport",
"AttentionBackwardGradients",
"AttentionBackwardPathDrift",
"AttentionBackwardPathResult",
"AttentionBackwardRankDrift",
"AttentionPartialState",
"build_reference_split_kv_runtime_plan_set",
- "CPAttentionReferenceOp",
"DeterministicCPAttentionReferenceOp",Run the following script to find any remaining importers before removal:
#!/bin/bash
set -euo pipefail
rg -n --type=py -C2 '\bCPAttentionReferenceOp\b'🧰 Tools
🪛 Ruff (0.16.1)
[warning] 1021-1035: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
🤖 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 1019 -
1035, Remove the CPAttentionReferenceOp alias assignment and its corresponding
__all__ export; use DeterministicCPAttentionReferenceOp directly and verify no
remaining importers reference the removed alias.
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com>
3 participants
Scope
PR3 provides the independent deterministic CP Attention reference for prefill and chunked-prefill.
Contract
Validation
Out/LSEmax abs0.0across 2/4/8 ranksBoundary
This PR validates reference arithmetic and fixed logical merge semantics. It does not claim real NCCL/AG/RS execution, fused CUDA arithmetic, decode replay, or performance.
Commit: 988d954