feat(attention): add single-gpu comparison harness - #253
Conversation
Signed-off-by: inaniloquentee <3051000145@qq.com>
📝 WalkthroughWalkthroughAdds typed attention contracts and a single-GPU comparison harness. The harness supports standard, decode, paged-KV, chunked-query, RoPE, and optional Transformer Engine paths. It reports output, LSE, selected-logprob, and RoPE drift with provenance. ChangesAttention comparison harness
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant Test
participant compare_single_gpu_attention
participant AttentionPaths
participant TransformerEngine
participant AttentionComparisonReport
Test->>compare_single_gpu_attention: submit attention comparison inputs
compare_single_gpu_attention->>AttentionPaths: execute reference and candidate paths
AttentionPaths->>TransformerEngine: probe and optionally merge partial states
TransformerEngine-->>AttentionPaths: return merged output and LSE
AttentionPaths-->>compare_single_gpu_attention: return path results and provenance
compare_single_gpu_attention->>AttentionComparisonReport: calculate output, LSE, dlogp, and RoPE drift
AttentionComparisonReport-->>Test: return structured comparison report
Possibly related PRs
Suggested labels: 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 |
![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/testing/attention_comparison.py`:
- Around line 411-441: The _load_te_context_parallel helper must validate that
the imported module provides callable flash_attn_fwd_softmax_lse_correction,
flash_attn_fwd_out_correction_init, and flash_attn_fwd_out_correction
attributes. If any helper is missing or unusable, raise
TransformerEngineUnavailable from the load/validation boundary, preserving the
existing normalization of import failures.
In `@tests/test_attention_comparison.py`:
- Around line 73-76: Update the p99 assertion in the report serialization test
to validate the actual expected tolerance or compare
payload["drifts"][0]["out"]["p99_abs"] with the corresponding drift.out.p99_abs
value, rather than asserting it is merely nonnegative.
- Around line 79-105: Strengthen
test_single_gpu_attention_harness_preserves_key_padding_mask with an independent
assertion that masked key/value positions do not affect the result, rather than
relying only on agreement among masked candidates. Compare the masked
computation against an equivalent reference where padded KV entries are removed
or neutralized, and assert the outputs and log-sum-exp values remain within the
existing tolerance.
🪄 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: 535a60a1-79a3-4d30-92b3-5d8d7977e9cb
📒 Files selected for processing (4)
docs/design/ws2-attention-single-gpu-harness.mdrl_engine/testing/__init__.pyrl_engine/testing/attention_comparison.pytests/test_attention_comparison.py
| payload = report.to_dict() | ||
| assert payload["reference_name"] == "full_prefill" | ||
| assert payload["drifts"][0]["out"]["p99_abs"] >= 0.0 | ||
| json.dumps(payload) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make the p99 assertion meaningful.
p99_abs >= 0.0 is tautological for an absolute-error metric, so it will not catch an incorrect serialized value. Assert the expected tolerance or compare the payload value with drift.out.p99_abs.
🧰 Tools
🪛 ast-grep (0.45.0)
[info] 75-75: use jsonify instead of json.dumps for JSON output
Context: json.dumps(payload)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 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_attention_comparison.py` around lines 73 - 76, Update the p99
assertion in the report serialization test to validate the actual expected
tolerance or compare payload["drifts"][0]["out"]["p99_abs"] with the
corresponding drift.out.p99_abs value, rather than asserting it is merely
nonnegative.
Signed-off-by: inaniloquentee <3051000145@qq.com>
There was a problem hiding this comment.
🧹 Nitpick comments (3)
docs/design/ws2-attention-single-gpu-harness.md (1)
27-29: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFix the compound-modifier hyphenation.
"production fused kernels" needs a hyphen when used as a compound modifier before "kernels".
✏️ Proposed fix
The RoPE path is still single-GPU attribution. It proves that both sides agree on post-RoPE Q/K, `out`, attention-domain `lse`, and optional active-token -`dlogp` before CP communication or production fused kernels are introduced. +`dlogp` before CP communication or production-fused kernels are introduced.🤖 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/design/ws2-attention-single-gpu-harness.md` around lines 27 - 29, Update the RoPE path description to hyphenate the compound modifier before “kernels,” changing “production fused kernels” to the grammatically correct form while preserving the surrounding technical content.Source: Linters/SAST tools
rl_engine/testing/attention_comparison.py (1)
223-284: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider consolidating the identical reference and candidate RoPE paths.
run_unfused_rope_attention(Lines 223-251) andrun_fused_like_rope_attention(Lines 254-282) run the identical body:_apply_rope_to_qkfollowed by_attention_with_lsewith the same arguments. Only thename,materialization, andfusion_boundarystrings differ. If this is an intentional placeholder for a future fused kernel implementation, add a short comment stating that intent. Otherwise, extract a shared private helper that takesname,materialization, andfusion_boundaryas parameters, so a future change to the shared RoPE/attention computation does not need to be applied twice.♻️ Proposed refactor
+def _run_rope_attention_path( + inputs: AttentionComparisonInputs, + *, + name: str, + materialization: str, + fusion_boundary: str, +) -> AttentionPathResult: + post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) + out, lse = _attention_with_lse( + post_rope_q, + post_rope_k, + inputs.v, + causal=inputs.causal, + scale=inputs.scale, + key_padding_mask=inputs.key_padding_mask, + q_start=0, + k_start=0, + total_query_len=post_rope_q.size(2), + total_kv_len=post_rope_k.size(2), + output_dtype=inputs.output_dtype, + ) + return AttentionPathResult( + name=name, + out=out, + lse=lse, + provenance=_rope_attention_provenance( + inputs, + materialization=materialization, + fusion_boundary=fusion_boundary, + ), + post_rope_q=post_rope_q, + post_rope_k=post_rope_k, + ) + + def run_unfused_rope_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: """Canonical ``RoPE -> Attention`` reference materialization.""" - post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) - out, lse = _attention_with_lse( - post_rope_q, - post_rope_k, - inputs.v, - causal=inputs.causal, - scale=inputs.scale, - key_padding_mask=inputs.key_padding_mask, - q_start=0, - k_start=0, - total_query_len=post_rope_q.size(2), - total_kv_len=post_rope_k.size(2), - output_dtype=inputs.output_dtype, - ) - return AttentionPathResult( - name="unfused_rope_attention", - out=out, - lse=lse, - provenance=_rope_attention_provenance( - inputs, - materialization="rope_then_attention", - fusion_boundary="unfused_rope_attention", - ), - post_rope_q=post_rope_q, - post_rope_k=post_rope_k, - ) + return _run_rope_attention_path( + inputs, + name="unfused_rope_attention", + materialization="rope_then_attention", + fusion_boundary="unfused_rope_attention", + ) def run_fused_like_rope_attention(inputs: AttentionComparisonInputs) -> AttentionPathResult: """Semantic fused ``RoPE+Attention`` path using the same canonical RoPE rules.""" - post_rope_q, post_rope_k = _apply_rope_to_qk(inputs) - out, lse = _attention_with_lse( - post_rope_q, - post_rope_k, - inputs.v, - causal=inputs.causal, - scale=inputs.scale, - key_padding_mask=inputs.key_padding_mask, - q_start=0, - k_start=0, - total_query_len=post_rope_q.size(2), - total_kv_len=post_rope_k.size(2), - output_dtype=inputs.output_dtype, - ) - return AttentionPathResult( - name="fused_like_rope_attention", - out=out, - lse=lse, - provenance=_rope_attention_provenance( - inputs, - materialization="fused_like_rope_attention", - fusion_boundary="fused_rope_attention", - ), - post_rope_q=post_rope_q, - post_rope_k=post_rope_k, - ) + return _run_rope_attention_path( + inputs, + name="fused_like_rope_attention", + materialization="fused_like_rope_attention", + fusion_boundary="fused_rope_attention", + )🤖 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/testing/attention_comparison.py` around lines 223 - 284, Consolidate the duplicated computation in run_unfused_rope_attention and run_fused_like_rope_attention by extracting a shared private helper that accepts name, materialization, and fusion_boundary, while preserving the existing _apply_rope_to_qk and _attention_with_lse arguments and returned provenance. Have both public paths delegate to the helper; if duplication is intentionally retained for a future fused kernel, add a brief comment documenting that intent instead.tests/test_attention_comparison.py (1)
150-155: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend RoPE validation coverage.
This test only exercises the "missing
rope_positions" branch of_validate_rope_inputs. That function also rejects mismatchedSq/Skv, an invalidrope_rotary_dim,rope_cast_at != "after_rope", a device mismatch, an invalid position dtype, and a malformed position shape. Add parametrized negative tests for these branches to guard the validation logic added in this PR.
Do you want me to generate the additional test cases?🤖 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_attention_comparison.py` around lines 150 - 155, Add parametrized negative tests alongside test_single_gpu_rope_attention_requires_position_metadata covering each remaining _validate_rope_inputs rejection: mismatched Sq/Skv, invalid rope_rotary_dim, rope_cast_at other than "after_rope", position/device mismatch, invalid rope_positions dtype, and malformed rope_positions shape. Build each case from _comparison_inputs, assert ValueError, and match the corresponding validation error while preserving the existing missing-metadata test.
🤖 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 `@docs/design/ws2-attention-single-gpu-harness.md`:
- Around line 27-29: Update the RoPE path description to hyphenate the compound
modifier before “kernels,” changing “production fused kernels” to the
grammatically correct form while preserving the surrounding technical content.
In `@rl_engine/testing/attention_comparison.py`:
- Around line 223-284: Consolidate the duplicated computation in
run_unfused_rope_attention and run_fused_like_rope_attention by extracting a
shared private helper that accepts name, materialization, and fusion_boundary,
while preserving the existing _apply_rope_to_qk and _attention_with_lse
arguments and returned provenance. Have both public paths delegate to the
helper; if duplication is intentionally retained for a future fused kernel, add
a brief comment documenting that intent instead.
In `@tests/test_attention_comparison.py`:
- Around line 150-155: Add parametrized negative tests alongside
test_single_gpu_rope_attention_requires_position_metadata covering each
remaining _validate_rope_inputs rejection: mismatched Sq/Skv, invalid
rope_rotary_dim, rope_cast_at other than "after_rope", position/device mismatch,
invalid rope_positions dtype, and malformed rope_positions shape. Build each
case from _comparison_inputs, assert ValueError, and match the corresponding
validation error while preserving the existing missing-metadata test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: d8fe70ab-d037-4e5d-966a-f46b00eb15d6
📒 Files selected for processing (4)
docs/design/ws2-attention-single-gpu-harness.mdrl_engine/testing/__init__.pyrl_engine/testing/attention_comparison.pytests/test_attention_comparison.py
Signed-off-by: inaniloquentee <3051000145@qq.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_attention_comparison.py (1)
307-337: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for the numeric self-test failure path.
_probe_te_context_parallelhas three hardening layers: acallablecheck, a signature check, and a numeric self-test. This file tests the first two (test_transformer_engine_path_reports_missing_helpers,test_transformer_engine_path_reports_incompatible_helper_signature), but no test exercises a helper that has a correct name and signature yet produces wrong numeric output. Add a sibling test that monkeypatches helpers with matching signatures but incorrect arithmetic, and assertreport.unavailablereports the numeric self-test failure.def test_transformer_engine_path_reports_numeric_selftest_failure(monkeypatch): def lse_correction(softmax_lse, softmax_lse_per_step): softmax_lse.copy_(softmax_lse + softmax_lse_per_step) # wrong: not logaddexp def out_correction_init(out_init_step, softmax_lse, softmax_lse_init_step, seq_dim): return out_init_step def out_correction(out, out_per_step, softmax_lse, softmax_lse_per_step, seq_dim): out.add_(out_per_step) monkeypatch.setitem( sys.modules, _TE_CONTEXT_PARALLEL_MODULE, types.SimpleNamespace( flash_attn_fwd_softmax_lse_correction=lse_correction, flash_attn_fwd_out_correction_init=out_correction_init, flash_attn_fwd_out_correction=out_correction, ), ) report = compare_single_gpu_attention( _comparison_inputs(), query_chunk_size=3, kv_page_size=2, include_transformer_engine=True, ) assert len(report.unavailable) == 1 assert "numeric self-test failed" in report.unavailable[0]🤖 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_attention_comparison.py` around lines 307 - 337, The existing test only covers incompatible helper signatures; add a sibling test for the numeric self-test failure in the Transformer Engine path. In the new test, monkeypatch the helpers exposed by _TE_CONTEXT_PARALLEL_MODULE with matching signatures but intentionally incorrect arithmetic, invoke compare_single_gpu_attention with include_transformer_engine=True, and assert exactly one unavailable result containing “numeric self-test failed”.
🤖 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/testing/attention_comparison.py`:
- Around line 636-698: Update the numeric checks in _probe_te_context_parallel
so compatible Transformer Engine FP32 implementations are not rejected for minor
accumulation differences. Keep the existing torch.allclose validations and
absolute tolerance, but add a small relative tolerance appropriate for TE FP32
accuracy to both the LSE and output comparisons.
---
Nitpick comments:
In `@tests/test_attention_comparison.py`:
- Around line 307-337: The existing test only covers incompatible helper
signatures; add a sibling test for the numeric self-test failure in the
Transformer Engine path. In the new test, monkeypatch the helpers exposed by
_TE_CONTEXT_PARALLEL_MODULE with matching signatures but intentionally incorrect
arithmetic, invoke compare_single_gpu_attention with
include_transformer_engine=True, and assert exactly one unavailable result
containing “numeric self-test failed”.
🪄 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: 9eb554da-43cd-408b-bc34-022091ca4095
📒 Files selected for processing (3)
docs/design/ws2-attention-transformer-engine-reuse-plan.mdrl_engine/testing/attention_comparison.pytests/test_attention_comparison.py
| def _probe_te_context_parallel(module: Any) -> None: | ||
| missing = [ | ||
| name for name in _TE_CONTEXT_PARALLEL_HELPERS if not callable(getattr(module, name, None)) | ||
| ] | ||
| if missing: | ||
| raise TransformerEngineUnavailable( | ||
| f"{_TE_CONTEXT_PARALLEL_MODULE} missing required helpers: {', '.join(missing)}" | ||
| ) | ||
|
|
||
| for name, expected in _TE_CONTEXT_PARALLEL_HELPERS.items(): | ||
| helper = getattr(module, name) | ||
| try: | ||
| parameters = tuple(inspect.signature(helper).parameters) | ||
| except (TypeError, ValueError) as exc: | ||
| raise TransformerEngineUnavailable( | ||
| f"{_TE_CONTEXT_PARALLEL_MODULE}.{name} signature is not inspectable" | ||
| ) from exc | ||
| if parameters[: len(expected)] != expected: | ||
| raise TransformerEngineUnavailable( | ||
| f"{_TE_CONTEXT_PARALLEL_MODULE}.{name} has incompatible signature " | ||
| f"{parameters}; expected prefix {expected}" | ||
| ) | ||
|
|
||
| try: | ||
| lse_a = torch.tensor([[[0.0, -1.0]]], dtype=torch.float32) | ||
| lse_b = torch.tensor([[[1.0, -3.0]]], dtype=torch.float32) | ||
| out_a = torch.tensor([[[[1.0, -2.0], [0.5, 2.0]]]], dtype=torch.float32) | ||
| out_b = torch.tensor([[[[-1.0, 4.0], [3.0, -0.5]]]], dtype=torch.float32) | ||
| expected_lse = torch.logaddexp(lse_a, lse_b) | ||
| expected_out = ( | ||
| torch.exp(lse_a - expected_lse).unsqueeze(-1) * out_a | ||
| + torch.exp(lse_b - expected_lse).unsqueeze(-1) * out_b | ||
| ) | ||
|
|
||
| probed_lse = lse_a.clone() | ||
| module.flash_attn_fwd_softmax_lse_correction(probed_lse, lse_b) | ||
| probed_out = module.flash_attn_fwd_out_correction_init( | ||
| out_a.clone(), | ||
| probed_lse, | ||
| lse_a, | ||
| seq_dim=2, | ||
| ) | ||
| module.flash_attn_fwd_out_correction( | ||
| probed_out, | ||
| out_b, | ||
| probed_lse, | ||
| lse_b, | ||
| seq_dim=2, | ||
| ) | ||
| except Exception as exc: | ||
| raise TransformerEngineUnavailable( | ||
| f"{_TE_CONTEXT_PARALLEL_MODULE} helper numeric self-test failed: {exc}" | ||
| ) from exc | ||
|
|
||
| if not torch.allclose(probed_lse, expected_lse, atol=1.0e-6, rtol=0.0): | ||
| raise TransformerEngineUnavailable( | ||
| f"{_TE_CONTEXT_PARALLEL_MODULE} LSE helper numeric self-test failed" | ||
| ) | ||
| if not torch.allclose(probed_out, expected_out, atol=1.0e-6, rtol=0.0): | ||
| raise TransformerEngineUnavailable( | ||
| f"{_TE_CONTEXT_PARALLEL_MODULE} out helper numeric self-test failed" | ||
| ) | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
NVIDIA Transformer Engine flash_attn_fwd_out_correction numerical precision fp32 accumulation
💡 Result:
In the NVIDIA Transformer Engine, the flash_attn_fwd_out_correction function is a JIT-fused helper used during context parallelism (CP) to merge partial attention outputs across distributed ranks [1][2]. Regarding numerical precision and FP32 accumulation: 1. Purpose of Accumulation: When employing context parallelism, the attention output is computed in pieces across different ranks. The flash_attn_fwd_out_correction function performs the online softmax correction required to merge these partial outputs accurately [1][2]. 2. FP32 Accumulation Improvements: To enhance numerical stability and accuracy—particularly for BF16 and FP16 precision—the Transformer Engine was updated (notably in PR #821) to implement FP32 accumulation for these attention operations [3][4]. 3. Impact: Adopting FP32 buffers for the accumulation of attention outputs (in the forward pass) and the gradient computations (in the backward pass) significantly reduced numerical mismatches [3][4]. While earlier implementations or lower-precision accumulations led to higher error rates in output comparisons, the transition to FP32 buffers provided stricter numerical adherence with a negligible increase in memory usage (typically less than 2%) [3][4]. The function itself performs the following operation: O_new = O_old * exp(LSE_old - LSE_new) + O_step * exp(LSE_step - LSE_new), which facilitates the correct combination of distributed partial attention results [1].
Citations:
- 1: https://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 2: https://leeroopedia.com/index.php/Implementation:NVIDIA_TransformerEngine_Context_Parallel
- 3: [Pytorch] Implement fp32 accumulation for attention with context parallel in both forward and backward pass. NVIDIA/TransformerEngine#821
- 4: https://github.com/NVIDIA/TransformerEngine/pull/821/files
🌐 Web query:
TransformerEngine fp32 accumulation flash_attn_fwd_softmax_lse_correction precision tolerance
💡 Result:
In NVIDIA Transformer Engine, specifically within the context parallelism (CP) implementation, the flash_attn_fwd_softmax_lse_correction function is used to merge partial softmax log-sum-exp (LSE) stats across different sequence chunks [1][2]. Regarding precision and accumulation: 1. FP32 Accumulation for Softmax LSE Correction: The flash_attn_fwd_softmax_lse_correction function has been updated to perform LSE correction in FP32 precision [3]. Previously, it utilized higher precision (FP64), but was lowered to FP32 to align with cuDNN's native softmax LSE representation, as this maintains sufficient accuracy for the combined CP ranks while staying within float limits [3]. The implementation uses torch.log1p for improved numerical stability when handling values close to zero [3][4]. 2. FP32 Accumulation in Attention Operations: Beyond the LSE correction, Transformer Engine incorporates FP32 accumulation buffers for attention operations (both forward and backward passes) when context parallelism is enabled [5][6]. This practice, similar to techniques used in Megatron-LM for bf16 precision, was extended to fp16 to improve numerical accuracy [5][6]. 3. Precision Tolerance: When validating these implementations, Transformer Engine developers have used strict tolerances, such as rtol = 1.3e-6 and atol = 1e-5, with torch.testing.assert_close to measure the effectiveness of using FP32 accumulation [5][6]. These tolerances are consistent with standard requirements for FP32 computations and have been used to demonstrate reductions in numerical mismatches during both forward and backward passes [5][6].
Citations:
- 1: https://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 2: https://leeroopedia.com/index.php/Implementation:NVIDIA_TransformerEngine_Context_Parallel
- 3: change softmax_lse correction of CP to FP32 NVIDIA/TransformerEngine#1546
- 4: https://github.com/NVIDIA/TransformerEngine/blob/1df4a69f761672f633d40ea3605327087d1ea737/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 5: [Pytorch] Implement fp32 accumulation for attention with context parallel in both forward and backward pass. NVIDIA/TransformerEngine#821
- 6: https://github.com/NVIDIA/TransformerEngine/pull/821/files
🌐 Web query:
"flash_attn_fwd_out_correction" "torch.allclose" "TransformerEngine"
💡 Result:
The function flash_attn_fwd_out_correction is a JIT-fused utility within NVIDIA's Transformer Engine used specifically to support context parallelism [1][2][3]. It merges partial attention outputs and log-sum-exp (LSE) values collected from different sequence chunks during the attention forward pass [2][3]. The mathematical operation performed by this function is: O_new = O_old * exp(LSE_old - LSE_new) + O_step * exp(LSE_step - LSE_new) [2] Key technical details: - Implementation: It is defined in transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py and uses the @jit_fuser decorator to optimize execution [1][4]. - Context: It was previously associated with issues related to torch.compile and movedim (e.g., Issue #1235), which were resolved in later versions of Transformer Engine (e.g., v1.12+) [5][4]. - Precision: To improve numerical stability in context-parallel attention, Transformer Engine implements FP32 accumulation for these types of operations [6]. Regarding torch.allclose, it is a standard PyTorch utility used to verify that two tensors are numerically close within a specified absolute (atol) and relative (rtol) tolerance [7]. In the context of Transformer Engine development, torch.testing.assert_close (often with FP32-specific tolerances) is frequently used in test suites to validate that these fused operations or FP32-accumulated attention passes match expected outputs, especially after performance-oriented modifications [6]. There is no direct functional coupling between the flash_attn_fwd_out_correction code and the torch.allclose function, other than the latter being used to verify the numerical correctness of the former during testing [6][7].
Citations:
- 1: https://github.com/NVIDIA/TransformerEngine/blob/1df4a69f761672f633d40ea3605327087d1ea737/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 2: https://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 3: https://leeroopedia.com/index.php/Implementation:NVIDIA_TransformerEngine_Context_Parallel
- 4: Fix bug in torch compile and seqdim is integer NVIDIA/TransformerEngine#1217
- 5: [PyTorch] Failed running call_method movedim NVIDIA/TransformerEngine#1235
- 6: https://github.com/NVIDIA/TransformerEngine/pull/821/files
- 7: https://docs.pytorch.org/docs/2.1/generated/torch.allclose.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant probe code and nearby TE availability usage without executing repository code.
if [ -f rl_engine/testing/attention_comparison.py ]; then
sed -n '600,710p' rl_engine/testing/attention_comparison.py
printf '\n--- TE unavailable usages ---\n'
rg -n "TransformerEngineUnavailable|load_and_probe_te_context_parallel|_probe_te_context_parallel|attention_comparison" -S .
else
echo "rl_engine/testing/attention_comparison.py not found"
git ls-files | rg 'attention_comparison|transformer|attention' || true
fiRepository: RL-Align/RL-Kernel
Length of output: 6382
🌐 Web query:
TransformerEngine context_parallel.py flash_attn_fwd_softmax_lse_correction implementation source
💡 Result:
The function flash_attn_fwd_softmax_lse_correction is implemented in the NVIDIA TransformerEngine library within the file transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py [1][2][3]. It is a JIT-fused helper function designed to merge the Log-Sum-Exp (LSE) statistics from different steps of attention computation when using context parallelism [4][5]. The implementation performs an in-place update on the provided softmax_lse tensor using the following logic [1][2]: @jit_fuser def flash_attn_fwd_softmax_lse_correction( softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor,): """Merge softmax stats of each step in Attention with context parallelism""" max_scale = torch.max(softmax_lse, softmax_lse_per_step) min_scale = torch.min(softmax_lse, softmax_lse_per_step) new_scale = max_scale + torch.log1p(torch.exp(min_scale - max_scale)) softmax_lse.copy_(new_scale) This approach computes the updated LSE value numerically stably by identifying the maximum of the two LSE values and using the log1p(exp(...)) transformation to calculate the contribution of the smaller value [1][2]. The in-place copy_(new_scale) operation was specifically implemented to resolve issues where non-in-place updates resulted in incorrect final LSE values or numerical instability when handling values outside the range of double-precision floating-point numbers [6].
Citations:
- 1: https://github.com/NVIDIA/TransformerEngine/blob/80ea3133efa4c3a3679845b8ee46dfc06e2792f0/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 2: https://github.com/NVIDIA/TransformerEngine/blob/1df4a69f761672f633d40ea3605327087d1ea737/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 3: https://github.com/NVIDIA/TransformerEngine/blob/6a68c733/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 4: https://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 5: https://leeroopedia.com/index.php/Implementation:NVIDIA_TransformerEngine_Context_Parallel
- 6: [Pytorch] Update context parallel softmax lse correction func NVIDIA/TransformerEngine#716
🌐 Web query:
TransformerEngine context_parallel.py flash_attn_fwd_out_correction implementation source
💡 Result:
The function flash_attn_fwd_out_correction is defined in the TransformerEngine library within the file transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py [1][2]. Its primary purpose is to merge partial attention outputs generated during different steps of context parallelism (CP) [1][3]. When using context parallelism, the attention operation is split across multiple ranks/steps, and this function computes the corrected output by accounting for the log-sum-exp (LSE) differences between the partial step and the global state [1][3]. The implementation follows this logic [1][2]: @jit_fuser def flash_attn_fwd_out_correction( out: torch.Tensor, out_per_step: torch.Tensor, softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor, seq_dim: int,): """Merge partial outputs of each step in Attention with context parallelism""" softmax_lse_corrected_exp = torch.exp(softmax_lse_per_step - softmax_lse).movedim(2, seq_dim) softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1) out_corrected = out_per_step * softmax_lse_corrected_exp out.add_(out_corrected) Key implementation details include: - Decorator: It is typically decorated with @jit_fuser to optimize the execution via TorchScript or torch.compile [4][2]. - Calculation: It uses the formula O_new = O_old + O_step * exp(LSE_step - LSE_global), where LSE_step and LSE_global are the partial and accumulated log-sum-exp values respectively [3]. - Data Handling: The function adjusts the dimensions of the LSE correction factor using movedim and unsqueeze to ensure it aligns with the output tensor's shape for element-wise multiplication and accumulation [1][4][5]. Historical context: The function has undergone refinements, including fixes to resolve compatibility issues with torch.compile, handling of integer arguments, and ensuring correct output-indexing logic when switching between different versions of FlashAttention (e.g., FA2 vs FA3) [6][4][5].
Citations:
- 1: https://github.com/NVIDIA/TransformerEngine/blob/1df4a69f761672f633d40ea3605327087d1ea737/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 2: https://github.com/NVIDIA/TransformerEngine/blob/80ea3133efa4c3a3679845b8ee46dfc06e2792f0/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 3: https://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 4: Fix bug in torch compile and seqdim is integer NVIDIA/TransformerEngine#1217
- 5: [PyTorch] Failed running call_method movedim NVIDIA/TransformerEngine#1235
- 6: fix(CP, FA): the conditional logic in the FA version contains a vulnerability when processing the output of Flash Attn forward pass NVIDIA/TransformerEngine#2825
Avoid using the probe as a strict numeric compatibility gate.
torch.allclose(..., atol=1.0e-6, rtol=0.0) can reject a functionally compatible Transformer Engine build when TE’s @jit_fuser implementation accumulates these small fp32 operands differently than the CPU reference. Use a tolerance that matches TE FP32 accuracy expectations, such as adding a small relative component; otherwise TransformerEngineUnavailable masks a working context-parallel adapter.
🤖 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/testing/attention_comparison.py` around lines 636 - 698, Update the
numeric checks in _probe_te_context_parallel so compatible Transformer Engine
FP32 implementations are not rejected for minor accumulation differences. Keep
the existing torch.allclose validations and absolute tolerance, but add a small
relative tolerance appropriate for TE FP32 accuracy to both the LSE and output
comparisons.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/test_attention_comparison.py (1)
519-521: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winPin the exact offset position and check
lse.
provenance["global_token_positions"][0][0] >= offsetpasses for every offset slot, so it does not detect a wrong provenance value. The first logical slot holds position0 + offset, so assert equality. Add thelsebound for parity with the other decode tests.♻️ Proposed test tightening
report = compare_decode_kv_replay(inputs) assert report.drifts[0].out.max_abs <= 1.0e-6 - assert report.drifts[0].provenance["global_token_positions"][0][0] >= offset + assert report.drifts[0].lse.max_abs <= 1.0e-6 + assert report.drifts[0].provenance["global_token_positions"][0][0] == offset🤖 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_attention_comparison.py` around lines 519 - 521, In the decode replay assertions around compare_decode_kv_replay, require provenance["global_token_positions"][0][0] to equal offset rather than merely be greater than or equal to it, and add the expected lse bound matching the other decode tests.rl_engine/kernels/attention_contract.py (2)
1445-1472: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to clear the Ruff RUF022 warning.Ruff reports
__all__is not sorted.AttentionBackendCapabilityandAttentionDispatchResultcome afterAttentionContractError. Apply isort-style ordering.🤖 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 1445 - 1472, Sort the __all__ entries in attention_contract.py using isort-style ordering, placing AttentionBackendCapability and AttentionDispatchResult before AttentionContract and AttentionContractError while preserving all exported symbols.Source: Linters/SAST tools
1149-1163: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the duplicated
causal_offsetscount branch.Both branches set
expected_causal_offsets = batch_size. Only the message noun differs. Simplify to one expression so the intent stays clear.♻️ 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:🤖 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 1149 - 1163, In the causal_offsets validation within the attention contract initializer, remove the duplicated if/else assignment of expected_causal_offsets and set it directly to batch_size. Preserve the existing offset_owner selection so the validation error continues to use “packed sequence” or “batch entry” appropriately.
🤖 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`:
- Line 1078: Rename the loop variable field to field_name in all three loops
within the attention contract module, including the loops near the
position-offset fields and the corresponding loops near lines 1189 and 1339,
while updating their references. Preserve the imported dataclasses.field usage,
then run Black on the module so its formatting matches the linting checks.
---
Nitpick comments:
In `@rl_engine/kernels/attention_contract.py`:
- Around line 1445-1472: Sort the __all__ entries in attention_contract.py using
isort-style ordering, placing AttentionBackendCapability and
AttentionDispatchResult before AttentionContract and AttentionContractError
while preserving all exported symbols.
- Around line 1149-1163: In the causal_offsets validation within the attention
contract initializer, remove the duplicated if/else assignment of
expected_causal_offsets and set it directly to batch_size. Preserve the existing
offset_owner selection so the validation error continues to use “packed
sequence” or “batch entry” appropriately.
In `@tests/test_attention_comparison.py`:
- Around line 519-521: In the decode replay assertions around
compare_decode_kv_replay, require provenance["global_token_positions"][0][0] to
equal offset rather than merely be greater than or equal to it, and add the
expected lse bound matching the other decode tests.
🪄 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: 06bdd82a-385a-4436-887c-8f845c4c166e
📒 Files selected for processing (4)
rl_engine/kernels/attention_contract.pyrl_engine/testing/__init__.pyrl_engine/testing/attention_comparison.pytests/test_attention_comparison.py
| not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() | ||
| ): | ||
| raise AttentionContractError("rope_scaling must be a non-empty string when provided") | ||
| for field in ("position_ids", "query_position_offsets", "key_position_offsets"): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔴 Critical | ⚡ Quick win
Rename the loop variable to stop shadowing the field import.
Line 14 imports field from dataclasses, and line 1111 uses it for default_factory. Lines 1078, 1189, and 1339 bind field as a loop variable, so flake8 reports F402 and the linting job fails. Rename the loop variable in all three loops, for example to field_name.
Also run black rl_engine/kernels/attention_contract.py and commit the result. The Black check currently reformats this file and fails the pipeline.
🔧 Proposed fix for the three shadowing loops
- 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)- for field in ("query_position_offsets", "key_position_offsets"):
- offsets = getattr(self.rope, field)
+ for field_name in ("query_position_offsets", "key_position_offsets"):
+ offsets = getattr(self.rope, field_name)
if offsets is not None and len(offsets) != batch_size:
raise AttentionContractError(
- f"{field} must contain one entry per logical batch entry"
+ f"{field_name} must contain one entry per logical batch entry"
)- for field in (
+ for field_name in (
"exports_attention_lse",
...
"reports_actual_split_kv_plan",
):
- if not isinstance(getattr(self, field), bool):
- raise AttentionContractError(f"{field} must be a bool")
+ if not isinstance(getattr(self, field_name), bool):
+ raise AttentionContractError(f"{field_name} must be a bool")🧰 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` at line 1078, Rename the loop
variable field to field_name in all three loops within the attention contract
module, including the loops near the position-offset fields and the
corresponding loops near lines 1189 and 1339, while updating their references.
Preserve the imported dataclasses.field usage, then run Black on the module so
its formatting matches the linting checks.
Sources: Linters/SAST tools, Pipeline failures
Signed-off-by: lamentropetion <3051000145@qq.com>
…live # Conflicts: # rl_engine/testing/__init__.py
1 participant
Scope
PR2 adds the single-GPU full, chunked, and paged-KV comparison harness used to attribute Attention drift.
Contract
Validation
Out/LSEmax abs0.0Boundary
This PR is single-GPU attribution coverage. It does not prove distributed communication or end-to-end Qwen3 checkpoint dlogp equality.
Commit: 1ce2349