Conversation
Signed-off-by: Jiao Ziang <2720649216@qq.com>
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a bug in the Sparse Flash Attention (SFA) implementation for Ascend A2/A3 hardware where empty sparse rows incorrectly reported non-zero softmax sums. By forcing these rows to return zero attention output and zero softmax sum, the system now correctly produces a negative infinity LSE, maintaining consistency in DCP decode operations. The changes include kernel-level logic updates and a new end-to-end regression test suite to ensure robustness across various inference configurations. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
github-actions
Bot
added
documentation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [BugFix] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the Sparse Flash Attention documentation and kernel implementation to correctly handle empty local sparse rows on Atlas A2/A3, ensuring they return a zero softmax sum and negative infinity LSE. Specifically, it adds a check in the MLA kernel for empty rows under the TND layout with page attention when returning softmax LSE, and introduces comprehensive end-to-end regression tests. The feedback provides a structured PR title and summary in compliance with the repository's style guide.
| <td>returnSoftmaxLse</td> | ||
| <td>可选属性</td> | ||
| <td>用于表示是否返回softmax_max和softmax_sum。True表示返回，False表示不返回，默认值为False。该参数仅在训练且layout_kv不为PA_BSND场景支持。</td> | ||
| <td>用于表示是否返回softmax_max和softmax_sum。True表示返回，False表示不返回，默认值为False。支持推理的PA_BSND场景；DCP decode使用softmax_max + log(softmax_sum)合并各rank的局部注意力结果。</td> |
There was a problem hiding this comment.
According to the Repository Style Guide (Pull Request Summary Style Guide), a PR review should include a suggested PR Title and PR Summary formatted in markdown code blocks.
Suggested PR Title:
[Ops][BugFix] Return zero softmax sums for empty A2/A3 DCP rowsSuggested PR Summary:
### What this PR does / why we need it?
This PR fixes an issue where empty local sparse rows in the scoped A2/A3 unquantized SFA kernel return a nonzero softmax sum (e.g., 512) instead of zero, causing `softmax_max + log(softmax_sum)` to be a finite sentinel instead of `-inf` LSE.
It detects empty prefixes in the TND/PA_BSND, token-wise, `sparse_mode=0`, return-LSE path and reuses the existing zero-output path.
### Does this PR introduce _any_ user-facing change?
Yes, empty local sparse rows in the scoped A2/A3 return-LSE path now return zero attention output and zero softmax sum, giving `-inf` when reconstructing LSE. The public signature is unchanged.
### How was this patch tested?
Tested on Ascend 910B3 with CANN 9.1.0. Added regressions in `tests/e2e/nightly/single_node/ops/singlecard_ops/test_sparse_flash_attention_lse_a2_a3.py` covering FP16/BF16, RoPE/NoPE, permuted physical pages, variable local KV lengths, multiple KV tiles, and graph replay.References
- The style guide requires generating a suggested PR Title and PR Summary in markdown code blocks during PR review. (link)
LiPu-jpg
changed the title
1 participant
What this PR does / why we need it?
An A2/A3 unquantized SFA row with a nonempty local KV cache but only
-1sparse indices returns a zero attention output while reporting a nonzero softmax sum (512 in the regression). Consequently,softmax_max + log(softmax_sum)is a finite sentinel instead of the empty row's-infLSE.AscendSFADCPImplcompacts valid local indices before padding and uses these statistics to merge local attention results. Detect an empty prefix in its TND/PA_BSND, token-wise,sparse_mode=0, return-LSE path and reuse the existing zero-output path. Valid rows retain the existing computation; pipeline draining is preserved. Add independent FP64-reference regressions for FP16/BF16, RoPE/NoPE, permuted physical pages, variable local KV lengths, multiple KV tiles, and graph replay from nonempty to empty and back. Correct the operator documentation's inference-LSE and NoPE restrictions.This complements merged #16656 (A5). It does not trim nonempty prefixes as proposed by #17653; that draft explicitly preserves the old empty-prefix path. No competition source was copied.
Does this PR introduce any user-facing change?
Empty local sparse rows in the scoped A2/A3 return-LSE path now return zero attention output and zero softmax sum, giving
-infwhen reconstructing LSE. The public signature is unchanged.How was this patch tested?
Draft limitations: A2 hardware correctness and paired operator performance are verified. A3 compilation/hardware validation and model-level validation remain pending. No model throughput improvement is claimed.
bash format.sh ci: passed (including the new test file). GitHub pre-commit, DCO, and CPU UT on both the pinned vLLM main revision and v0.30.0 also passed. The CI gate is currently red because a maintainer must add a precision/full-test label; selected NPU tests have not run.bash build.sh --opkernel --soc=ascend910b --ops=sparse_flash_attention -j2fromcsrc.e86df70; baseline and candidate use identical Python, host adapter and runtime libraries in separate private package copies. Only the generated SFA.oand.jsonfiles are replaced for the candidate. Both installed candidate object hashes match the fresh build and differ from baseline.torch_npu; their full logs and partial measurements are retained separately and not used in the paired comparison. These detections do not prove concurrent NPU kernel execution. No cases were removed from any completed run.Reproduce in a supported CANN/torch_npu/vLLM environment with the custom operators built from this branch:
To reproduce the baseline, run the same added test file with a custom-op installation built from
e86df70. Keep the installations in separate environments so a cached native library or vendor OPP package cannot hide which kernel is loaded.Verified generated/installed candidate kernel SHA256:
Operator-only performance evidence
Ascend 910B3 (physical device 3), CANN 9.1.0, torch 2.10.0+cpu with torch_npu 2.10.0.post4. Baseline and candidate both use Release/O3 kernel compilation. Identical framework/host libraries and deterministic inputs; separate vendor package copies. Test file SHA256:
89bd64c40d1e53a3b50d09998122e77212b70b608a663da5f16bd45fa6bdb206. Runs completed consecutively on 2026-09-30, approximately 13:55–13:59 Asia/Shanghai.Each graph contains 8 SFA invocations; warm up 5 eager calls and 10 graph replays. Each sample times 100 replays with NPU events and divides by 800 invocations. Keep 5 samples per process, repeat each variant in 3 processes. Checks before/after every configuration reject other torch_npu/compilation processes. This measures amortized device graph execution for a repeatedly reused input/cache; it excludes setup/reference work and is not serving or model throughput. Query heads=64, D=512, page size=128, sparse capacity=2048, mode=0, return-LSE=True. Single-row cases use KV=512; varlen uses Q lengths [1,3,2], KV [0,129,512]; pipeline uses Q [7,28], KV [4096,4096], counts [0,1,513,0,1025,2048,0] repeated five times; graph-replay-input uses Q=2 and counts [1,129].
Positive latency change means slowdown.
All 36 configurations: min/median/max and latency changes
All 1,080 measured samples and paired process ratios
Reproducible benchmark script
Save this script as
benchmark_sfa.pyin a fresh benchmark directory. Copy the regression file intotests/test_sparse_flash_attention_lse_a2_a3.pybeneath that directory. In each isolated baseline/candidate environment, with its matching custom-op package on PYTHONPATH, runpython benchmark_sfa.py OUTPUT.json LABEL. Use the B/C/C/B/B/C process order above; ensure the same NPU is free of other work.