[Performance][Attention] Fuse local SFA LSE formation into DCP pack

[Ruiqiu-Zheng]

What this PR does / why we need it?

This PR removes a local SparseFlashAttention DCP post-processing boundary by forming LSE inside the existing pack kernel instead of materializing a standalone FP32 LSE tensor before packing.

Before:

• native FP32 softmax_max / softmax_sum
• materialize LSE = softmax_max + log(softmax_sum)
• permute / reshape
• DCP pack
• HCCL All2All
• existing stable combine

After:

• native FP32 softmax_max / softmax_sum
• max/sum-aware DCP pack, which forms the same LSE while writing the packed layout
• same HCCL All2All
• unchanged stable combine

The change preserves the existing row/head/token mapping, packed payload semantics, HCCL communication pattern, and downstream stable-combine ABI.

The legacy materialized-LSE custom-op path is kept for compatibility and regression coverage.

Invalid local statistics are handled explicitly. Non-finite softmax_max, non-finite softmax_sum, or softmax_sum <= 0 produce an invalid (-inf) LSE contribution, which is then handled by the existing combine semantics. Focused tests cover these cases.

Performance evidence

The performance experiment used a preregistered paired/interleaved A/B protocol with fixed inputs and ordering:

• protocol seed: 20260916
• BF16 producer seed: 20260913
• FP16 correctness seed: 20260914
• 10 warmup iterations
• 5 timed samples per arm/block
• 8 paired A/B blocks
• fixed base/candidate order
• pre/post baseline-only null controls
• no post-hoc performance samples

The positive decision rule was fixed before measurement:

at least 75% of paired blocks favor candidate AND paired median improvement > max(1%, 2 * relative_noise_floor)

Matched measurements on the recorded Ascend test environment were:

Interval	Baseline median	Candidate median	Paired median effect	Noise floor	Result
local composition: precomputed output/max/sum -> pack	310.342 us	242.750 us	21.655% faster	4.237%	positive, 8/8 pairs favor candidate
world-size-one SparseFlashAttention producer -> pack	435.846 us	375.199 us	12.945% faster	5.847%	positive, 8/8 pairs favor candidate
same-host two-rank producer -> HCCL -> stable combine	5.043 ms	5.068 ms	0.385% paired effect in favor of candidate	4.410%	neutral, 4/8 pairs favor candidate

For the two-rank lane, the standalone arm medians and paired median point in slightly different directions, and both effects are far below the 8.819% decision threshold. It is therefore classified as neutral, not as either a speedup or a regression.

These numbers are scoped to the measured local/laboratory intervals. They are not end-to-end model, serving-throughput, or production-performance claims.

A sanitized reproducibility packet is available on an auxiliary branch of the fork. It is not part of this production PR and is not a CI dependency:

• protocol, seeds, fixture, raw paired/null-control measurements, and correctness sentinels:
  https://github.com/Ruiqiu-Zheng/vllm-ascend/tree/repro/sfa-local-lse-r016-20260902/repro/sfa-local-lse-r016

• README and reproduction instructions:
  https://github.com/Ruiqiu-Zheng/vllm-ascend/blob/repro/sfa-local-lse-r016-20260902/repro/sfa-local-lse-r016/README.md

• independent reducer that recomputes the table above from raw JSON:
  https://github.com/Ruiqiu-Zheng/vllm-ascend/blob/repro/sfa-local-lse-r016-20260902/repro/sfa-local-lse-r016/reduce_results.py

The exact binary fixture is:

SHA256 63785fa8f4592b7f8f9e06d345a20cf29a4ac8d19adc1e89dc8f90e60eb7db19

The packet also contains the exact candidate reconstruction patch and source/runtime lineage.

Absolute latency is not expected to reproduce bit-for-bit on a different hardware/runtime state. The intended reproduction target is the paired effect relative to a freshly observed null-control noise floor under matched lineage and the same protocol.

Does this PR introduce any user-facing change?

No.

This is an internal DCP post-processing optimization. It does not add or change a public API, user configuration, environment variable, or user-visible model behavior.

How was this patch tested?

Focused post-rebase regression on the PR branch:

• host / fake-op / routing tests: 13 passed
• single-device NPU tests: 32 passed
• same-host two-rank HCCL eager + NPUGraph tests: 2 passed
• Ruff check: passed
• Ruff format check: passed
• py_compile: passed
• git diff --check: passed

NPU correctness coverage includes:

• BF16 and FP16
• scatter_dim = 0 and scatter_dim = 1
• head dimension 512 on the measured specialization
• non-contiguous output / max / sum strides
• zero, negative, NaN, and Inf softmax sums
• non-finite softmax maxima
• all-invalid rows
• legacy materialized-LSE regression

The two-rank graph regression covers eager execution, graph capture, first replay, and changed-input replay while keeping input/output storage addresses fixed. This checks that replay consumes updated tensor contents rather than returning stale captured results.

The PR was rebased as upstream main advanced. Before each rebase, the decision-relevant Local-SFA callsite/pack seam was checked for upstream drift. No relevant source drift was observed, and the focused post-rebase regressions remained passing.

• vLLM main: vllm-project/vllm@ba07e4a

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-02T09:01:09.534463Z	398f837	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 optimizes the SparseFlashAttention (SFA) DCP post-processing pipeline by integrating LSE formation directly into the DCP pack kernel. By avoiding the materialization of a standalone FP32 LSE tensor, the change reduces memory overhead and improves latency in local composition and producer-to-pack intervals. The implementation maintains existing communication patterns and ABI compatibility while introducing robust handling for non-finite or invalid statistics.

Highlights

• Kernel Optimization: Fused LSE formation directly into the DCP pack kernel, eliminating the need to materialize a standalone FP32 LSE tensor.
• Performance Improvement: Achieved significant latency reductions in local composition and producer-to-pack intervals, with measured improvements exceeding 12% in key segments.
• Robustness: Added explicit handling for invalid local statistics, including non-finite values and non-positive softmax sums, ensuring stable behavior.
• Testing: Introduced comprehensive regression tests, including NPU graph capture and replay scenarios, to ensure correctness across various input conditions.

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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

Suggested PR Title:

[Attention][Feature] Support fused All2All with max/sum statistics in SFA CP

Suggested PR Summary:

### What this PR does / why we need it?
This PR introduces a new custom operator `sfa_dcp_a2a_fused_max_sum` to support context-parallel All2All communication directly using `softmax_max` and `softmax_sum` statistics. Previously, LSE was materialized in Python before the All2All operation. By computing LSE on-the-fly within the Triton kernel (`_pack_sfa_dcp_output_lse_kernel`), we avoid extra memory allocation and overhead.

Key changes:
- Updated Triton kernel `_pack_sfa_dcp_output_lse_kernel` to support `INPUT_MAX_SUM` mode.
- Added `pack_sfa_dcp_output_max_sum` and registered the `sfa_dcp_a2a_fused_max_sum` custom op.
- Integrated the new op into `vllm_ascend/attention/context_parallel/sfa_cp.py`.
- Added comprehensive unit and end-to-end tests.

### Does this PR introduce _any_ user-facing change?
No. This is an internal optimization for context-parallel attention.

### How was this patch tested?
- Added unit tests in `tests/ut/ops/test_sfa_cp_max_sum.py` and `tests/ut/attention/test_sfa_cp_precision.py`.
- Added Triton-level tests in `tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_sfa_cp_a2a.py`.
- Added multi-rank integration tests in `tests/e2e/nightly/single_node/ops/multicard_ops_a3/test_sfa_cp_a2a.py`.

I have no feedback to provide as there are no review comments.

[Ruiqiu-Zheng]

CI triage: pre-commit, coverage recommendation, and select-tests all passed. The selected test jobs were skipped only because this PR does not yet have a ready label, so ci-gate failed without any test case actually failing.

Could a maintainer please add ready-precise to run the selected relevant test set? ready-all is not needed at this stage.