feat(attention): add single-gpu comparison harness - #253
Conversation
Signed-off-by: inaniloquentee <3051000145@qq.com>
📝 WalkthroughWalkthroughAdds a single-GPU attention comparison harness with full, chunked, paged-KV, and RoPE paths. It reports output, LSE, selected-logprob, and post-RoPE drift. It adds optional Transformer Engine handling, public exports, design documents, and pytest coverage. ChangesAttention comparison harness
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Test
participant AttentionComparator
participant AttentionPaths
participant TransformerEngine
participant AttentionComparisonReport
Test->>AttentionComparator: submit Q/K/V and comparison options
AttentionComparator->>AttentionPaths: run full, chunked, paged-KV, and RoPE paths
AttentionPaths->>TransformerEngine: optionally merge partial states
TransformerEngine-->>AttentionPaths: return merged output and LSE
AttentionPaths-->>AttentionComparator: return path results
AttentionComparator-->>AttentionComparisonReport: calculate drift statistics
AttentionComparisonReport-->>Test: return serialized report
Possibly related issues
Possibly related PRs
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 |
![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.
1 participant
Part of #235.
Summary
out, attention-domainlse, optional active-tokendlogp, RoPE Q/K drift, and runtime provenance.RoPE -> Attentionwith fused-likeRoPE+Attentionmaterialization.docs/design/ws2-attention-transformer-engine-reuse-plan.mdas the local design mirror for the issue-level TE reuse plan.Scope
Implemented in this PR:
full_prefill: training-style full-sequence softmax attention.chunked_prefill: rollout-style query chunk replay over full KV.rl_kernel_paged_kv: rollout-style KV page replay with RL-Kernel FP32(out, lse)merge.transformer_engine_paged_kv: optional TE merge-oracle path when TE is installed and compatible.unfused_rope_attention: canonicalRoPE -> Attentionpath.fused_like_rope_attention: semantic fused-boundary path that applies the same canonical RoPE rules, then recordsfusion_boundaryprovenance.Boundary:
Transformer Engine Reuse
TE is lazy-loaded from:
The optional oracle uses only these helpers:
The merge contract remains RL-Kernel-owned:
TE hardening added in the latest update:
lse = -infandout = 0do not become NaN;-infvs-infis zero drift;te_available,te_version,te_module,te_symbols,te_capability_probe,te_signature_checked,te_numeric_selftest,actual_backend,actual_backend_source,accum_dtype, anddowncast_at.DCO
Signed-off-by.cbb5e2c fix(attention): harden transformer engine merge oracle.Local Validation
Environment note: local pytest/pre-commit tools ran under Windows with Python 3.13.7. The GitHub CI workflow will rerun on Ubuntu/Python 3.10.
pre-commit run --files rl_engine/testing/attention_comparison.py tests/test_attention_comparison.py docs/design/ws2-attention-transformer-engine-reuse-plan.md- passed.ruff check rl_engine/testing/attention_comparison.py tests/test_attention_comparison.py- passed.mypy --ignore-missing-imports rl_engine/- passed,Success: no issues found in 90 source files.pytest tests/test_attention_comparison.py -q- passed,10 passed.pytest rl_engine/tests/test_dispatch.py -v- passed,6 passed.PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_attention_correctness.py -q -rs- passed,45 passed, 82 skipped.pytest tests/test_attention.py -v -k "not large and not gpu"- passed,24 passed, 2 deselected.pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu"- passed,15 passed.mkdocs build --strict -f mkdocs.yaml- passed. It emitted existing MkDocs/Material and non-nav page informational warnings only.git diff --check- passed with Windows CRLF conversion warnings only.GitHub Actions
Latest push:
cbb5e2c.CI-Pipeline / linting- passed,52s.CI-Pipeline / docs- passed,54s.CI-Pipeline / unit-tests- passed,1m22s.GPU CI / gpu-tests- skipped as expected because the PR does not have theneeds-gpu-cilabel.CodeRabbit- review completed.