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 optional Transformer Engine paths. It reports output, LSE, and optional active-token log-probability drift, exports the API, documents usage, and adds pytest coverage. ChangesAttention comparison harness
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Test
participant compare_single_gpu_attention
participant AttentionPaths
participant TransformerEngine
participant AttentionComparisonReport
Test->>compare_single_gpu_attention: submit Q/K/V and comparison options
compare_single_gpu_attention->>AttentionPaths: run full, chunked, and paged-KV paths
AttentionPaths->>TransformerEngine: optionally merge paged-KV states
TransformerEngine-->>AttentionPaths: return merged output and LSE
AttentionPaths-->>compare_single_gpu_attention: return path results
compare_single_gpu_attention-->>AttentionComparisonReport: calculate drift statistics
AttentionComparisonReport-->>Test: return report and unavailable paths
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 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
| def _merge_partial_states_transformer_engine( | ||
| states: list[_PartialAttentionState], | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: | ||
| te_cp = _load_te_context_parallel() | ||
| merged_out = states[0].out.float() | ||
| merged_lse = states[0].lse.float() | ||
| for state in states[1:]: | ||
| previous_lse = merged_lse | ||
| merged_lse = previous_lse.clone() | ||
| te_cp.flash_attn_fwd_softmax_lse_correction(merged_lse, state.lse.float()) | ||
| merged_out = te_cp.flash_attn_fwd_out_correction_init( | ||
| merged_out, | ||
| merged_lse, | ||
| previous_lse, | ||
| seq_dim=2, | ||
| ) | ||
| te_cp.flash_attn_fwd_out_correction( | ||
| merged_out, | ||
| state.out.float(), | ||
| merged_lse, | ||
| state.lse.float(), | ||
| seq_dim=2, | ||
| ) | ||
| return merged_out, merged_lse | ||
|
|
||
|
|
||
| def _load_te_context_parallel() -> Any: | ||
| try: | ||
| return importlib.import_module(_TE_CONTEXT_PARALLEL_MODULE) | ||
| except (ImportError, OSError, RuntimeError) as exc: | ||
| raise TransformerEngineUnavailable(str(exc)) from exc |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file outline / size =="
wc -l rl_engine/testing/attention_comparison.py
ast-grep outline rl_engine/testing/attention_comparison.py --view compact || true
echo "== relevant snippets =="
sed -n '130,175p' rl_engine/testing/attention_comparison.py
printf '\n---\n'
sed -n '380,450p' rl_engine/testing/attention_comparison.py
echo "== search usages and tests =="
rg -n "TransformerEngineUnavailable|compare_single_gpu_attention|_merge_partial_states_transformer_engine|flash_attn_fwd_out_correction|flash_attn_fwd_softmax_lse_correction|_TE_CONTEXT_PARALLEL_MODULE" rl_engine/testing -S
echo "== git diff stat/name-only if available =="
git diff --name-only --stat || trueRepository: RL-Align/RL-Kernel
Length of output: 6205
🌐 Web query:
NVIDIA Transformer Engine Fused attention context parallel flash_attn_fwd_out_correction_init flash_attn_fwd_out_correction
💡 Result:
In the context of the NVIDIA Transformer Engine (TE), "fused attention" refers to optimized, high-performance kernels that execute dot-product attention—including key components like softmax, dropout, and masking—as a single fused operation on the GPU to minimize memory access and improve throughput [1][2][3]. The functions flash_attn_fwd_out_correction and flash_attn_fwd_out_correction_init are specialized utility functions used within the Transformer Engine's Context Parallelism (CP) implementation [4][5][6]. Because context parallelism distributes sequences across multiple GPU ranks (e.g., using a ring-based communication pattern), partial attention results are computed locally on different chunks of the sequence [5][6]. These functions are responsible for merging these partial results using the log-sum-exp (LSE) values from the softmax operation, which is necessary to maintain numerical correctness when combining attention outputs across different steps or ranks [4][5][6]. - flash_attn_fwd_out_correction_init: Specifically handles the merging of partial attention outputs for the first step of the attention process [4][7][5]. - flash_attn_fwd_out_correction: Merges partial attention outputs from subsequent steps by applying the correct scaling based on the difference between the partial LSE and the running, accumulated LSE [4][7][5]. Mathematically, these functions apply the standard technique for merging online softmax outputs, where the partial output (O) and LSE are combined using the formula [4][5]: O_new = O_old * exp(LSE_old - LSE_new) + O_step * exp(LSE_step - LSE_new) [5] These helpers are typically marked with @jit_fuser or @torch.compile in the Transformer Engine source code to ensure they are compiled into efficient fused kernels during execution [4][7].
Citations:
- 1: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/c/fused_attn.html
- 2: https://nvidia.github.io/TransformerEngine/api/c/fused_attn.html
- 3: https://docs.nvidia.com/deeplearning/transformer-engine-releases/release-2.10/user-guide/api/c/fused_attn.html
- 4: https://github.com/NVIDIA/TransformerEngine/blob/1df4a69f761672f633d40ea3605327087d1ea737/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 5: https://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 6: https://leeroopedia.com/index.php/Implementation:NVIDIA_TransformerEngine_Context_Parallel
- 7: https://github.com/meituan-longcat/LongCat-Video/blob/8334da93/longcat_video/block_sparse_attention/bsa_interface.py
🌐 Web query:
site:github.com "flash_attn_fwd_out_correction_init" "flash_attn_fwd_out_correction" transformer_engine
💡 Result:
In the NVIDIA Transformer Engine library, flash_attn_fwd_out_correction_init and flash_attn_fwd_out_correction are utility functions used to manage the merging of partial attention outputs when performing context-parallel training [1][2][3]. When training transformers with context parallelism (CP), the attention computation is split across multiple ranks or segments. These functions implement the necessary numerical correction to ensure the final attention output remains mathematically correct despite the partitioned computation [1][2]. Specifically: 1. flash_attn_fwd_out_correction_init: Initializes the correction process by merging the partial outputs of the first step in the attention computation [1][2]. 2. flash_attn_fwd_out_correction: Merges the partial outputs of subsequent steps in the attention computation [1][2]. Both functions perform this by adjusting the partial outputs using the Log-Sum-Exp (LSE) values of the respective steps, essentially re-normalizing the attention probabilities before adding them to the global output [1][2]. Technical Context and Recent Fixes: - Numerical Accuracy: These functions (and the broader implementation they belong to) were introduced to enable fp32 accumulation for attention operations during context-parallel forward and backward passes, which significantly improves numerical stability compared to lower-precision accumulation [4]. - Known Issues: A bug related to the movedim method in these functions was reported (issue #1235) when handling certain tensor dimensions [3]. This was resolved in Transformer Engine v1.12 by updating the handling of packed softmax_lse for compatibility with newer FlashAttention versions (v2.6+ and FA3) [3].
Citations:
- 1: https://github.com/NVIDIA/TransformerEngine/blob/1df4a69f761672f633d40ea3605327087d1ea737/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
- 2: https://github.com/meituan-longcat/LongCat-Video/blob/8334da93/longcat_video/block_sparse_attention/bsa_interface.py
- 3: [PyTorch] Failed running call_method movedim NVIDIA/TransformerEngine#1235
- 4: https://github.com/NVIDIA/TransformerEngine/pull/821/files
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== attention_comparison imports and TE module path =="
sed -n '1,60p' rl_engine/testing/attention_comparison.py
echo "== run_paged_kv_attention and exception handling =="
sed -n '270,330p' rl_engine/testing/attention_comparison.py
echo "== static probe of load/use relationship =="
python3 - <<'PY'
from pathlib import Path
p = Path("rl_engine/testing/attention_comparison.py")
src = p.read_text()
print("_load_te_context_parallel returns module assigned to te_cp:")
print(" module assignment" in src and "te_cp = _load_te_context_parallel()" in src)
print("API calls use attributes of returned module:")
for name in [
"flash_attn_fwd_softmax_lse_correction",
"flash_attn_fwd_out_correction_init",
"flash_attn_fwd_out_correction",
]:
print(f" call {name} in source:", name in src)
# Show whether the TE call sites are inside the same try/except block as _load_te_context_parallel.
try:
import ast
except Exception as exc:
print("ast unavailable:", repr(exc))
else:
tree = ast.parse(src)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == "_merge_partial_states_transformer_engine":
call_sites = []
raises = []
handles = []
for child in ast.walk(node):
if isinstance(child, ast.Call):
func = child.func
call_sites.append(
getattr(func, 'attr', None)
or (isinstance(func, ast.Attribute) and func.attr)
)
if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Store):
pass
# Find surrounding except handlers and try blocks that contain function body
# but use source-line membership since AST doesn't give textual try/except grouping directly.
lines = src.splitlines()
start, end = node.lineno, node.end_lineno
print(f" function range: {start}-{end}")
for other in tree.body:
if isinstance(other, ast.Try):
tstart, tend = other.lineno, other.end_lineno
if tstart <= start < tstart < (tstart+20) or (tstart <= start and start < tend):
print(f" TRY same scope/function? try spans {tstart}-{tend}, handlers={len(other.handlers)}")
if isinstance(node, ast.FunctionDef) and node.name == "compare_single_gpu_attention":
print("compare_single_gpu_attention range:", node.lineno, node.end_lineno)
PYRepository: RL-Align/RL-Kernel
Length of output: 4246
🌐 Web query:
NVIDIA TransformerEngine context_parallel.py flash_attn_fwd_out_correction_init source merge while out_correction_init called
💡 Result:
The function flash_attn_fwd_out_correction_init is a component of NVIDIA TransformerEngine's context parallelism (CP) implementation, located in transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py [1][2]. It is used to merge partial attention outputs generated during the first step of the forward pass when processing sequences across multiple ranks [1][3]. The function operates by applying a correction factor derived from the log-sum-exp (LSE) values to the partial output, according to the formula: O_corrected = O_init_step * exp(LSE_init_step - LSE) [1][3]. Regarding your inquiry about the "source merge" while out_correction_init is called: this function specifically handles the initial merging of the first step's partial results [1][3]. Subsequent steps in the sequence parallelism ring are handled by flash_attn_fwd_out_correction, which performs an in-place addition (out.add_) of the corrected partial outputs to the running output tensor [1][3]. Recent development in the TransformerEngine repository has addressed issues related to the conditional logic and output-indexing branches for different versions of FlashAttention (e.g., ensuring correct routing for FlashAttention-3 vs. earlier versions) [4]. These changes ensure that the output tensors are correctly identified and processed by the correction functions during the context-parallel forward pass [4].
Citations:
- 1: https://github.com/NVIDIA/TransformerEngine/blob/68493d2d/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://deepwiki.com/NVIDIA/TransformerEngine/7.1-context-parallelism
- 4: 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
Validate TE helper attributes at load time
_load_te_context_parallel only converts import failures into TransformerEngineUnavailable; the TE merge path later dereferences flash_attn_fwd_softmax_lse_correction, flash_attn_fwd_out_correction_init, and flash_attn_fwd_out_correction without a catch around these calls. If an installed TE version exposes the module but is missing/renames any helper, one of those calls propagates as AttributeError/TypeError instead of the optional TE path being marked unavailable. Normalize these errors to TransformerEngineUnavailable at the TE load/validation edge.
🤖 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 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.
| 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.
| def test_single_gpu_attention_harness_preserves_key_padding_mask(): | ||
| q, k, v = _qkv(seed=3) | ||
| key_padding_mask = torch.tensor( | ||
| [ | ||
| [True, True, True, False, False, False], | ||
| [True, False, True, True, False, False], | ||
| ], | ||
| dtype=torch.bool, | ||
| ) | ||
|
|
||
| report = compare_single_gpu_attention( | ||
| AttentionComparisonInputs( | ||
| q=q, | ||
| k=k, | ||
| v=v, | ||
| causal=True, | ||
| key_padding_mask=key_padding_mask, | ||
| ), | ||
| query_chunk_size=4, | ||
| kv_page_size=2, | ||
| ) | ||
|
|
||
| assert report.unavailable == () | ||
| for drift in report.drifts: | ||
| assert drift.out.max_abs <= 1.0e-6 | ||
| assert drift.lse.max_abs <= 1.0e-6 | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert key-padding semantics, not only candidate agreement.
Because the reference and all candidates receive the same mask, every path could ignore it and still produce low relative drift. Add an independent masked-reference check, such as comparing against an equivalent input with padded KV entries removed or neutralized.
🤖 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 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.
1 participant
Part of #235.
Summary
out, attention-domainlse, and optional active-tokendlogpdrift with structured provenance.Testing
Local validation:
pre-commit run --all-files- passed.mypy --ignore-missing-imports rl_engine/- passed.pytest rl_engine/tests/test_dispatch.py tests/test_attention_comparison.py tests/test_attention.py tests/test_kv_cache_attention.py tests/test_operator_inputs.py -q -k "not large and not gpu"- passed,67 passed, 4 deselected.PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest tests/test_attention_correctness.py -q -rs- collected GPU/ROCm correctness cases; all skipped locally because CUDA/ROCm is not available.python scripts/check_operator.py --op attention --candidate pytorch --dtype fp32 --device cpu --batch 1 --seq 3- passed.mkdocs build --strict -f mkdocs.yaml- passed.GitHub Actions:
CI-Pipeline / linting- passed.CI-Pipeline / docs- passed.CI-Pipeline / unit-tests- passed.GPU CI / gpu-tests- skipped as expected because the PR does not have theneeds-gpu-cilabel.CodeRabbit- review in progress at the time this body was updated.Notes
Signed-off-by: inaniloquentee <3051000145@qq.com>.Summary by CodeRabbit
New Features
Documentation
Tests