Repository navigation
[Bugfix][NPU] Support piecewise block-causal attention for full_attn_spans - #9244
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR appears to belong to: docs/design/module/diffusion/index.md. Module owners: @Isotr0py @david6666666 @princepride Routing: @Isotr0py via module of the changed files, CODEOWNERS; @david6666666 via module of the changed files; @princepride via module of the changed files @YawningPants, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Omni ReviewBot triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
@YawningPants, the source change follows the shared segment plan, and the new tests exercise the real NPU wrapper. The following validation evidence is still needed for 4be8385d4c8b020c2496af9607aa8bd8c58396a3:
- Attach the command and output for the three
test_npu_piecewise_*tests in test_flash_attn.py, using the same tests against the baseline implementation717f8a3499fb7dc6b3128b58a8baeacded1a7bbband this head. Acceptance: the original span-dropping behavior fails the relevant assertions, while the fixed implementation passes. - Make the reported A5 Before/Golden/After comparison reproducible: link the script/commands and tested source revisions, software/model versions, prompt/reference inputs and seed, warmup/measured counts, and timing scope. Keep the existing BF16/batch-1/two-step setup matched, and state the quality acceptance criterion for the residual I2I difference. Acceptance: the supplied record supports the stated kernel-time comparison and output-quality conclusion; a bounded smoke observation should be labeled as such.
Source inspection only on my side; I did not execute the CPU or A5 tests. These are evidence requests, not a demonstrated new code defect.
Dong1017
left a comment
Dong1017
left a comment
There was a problem hiding this comment.
Reviewed 4be8385: no findings. The NPU path now mirrors forward_cuda's piecewise dispatch (same shared segment plan; spans take precedence over masks and packed metadata on both platforms), the quantized path already rejects spans, and resolve_execution_path is descriptive-only, so the new npu_piecewise string has no runtime consumers.
Independent verification of head 4be8385:
- CPU contract suite on an H200 host (vllm 0.31.0 venv): 55 passed, 5 deselected. Running the same head tests against the pre-fix implementation (4a910f2) fails exactly the three new piecewise tests (52 passed), so they reproduce the silent full-attention defect instead of passing unconditionally.
- Kernel semantics on Ascend 910B1 (torch_npu 2.10.0.post4, CANN 9.1.0),
npu_fusion_attentionvs a dense fp32 reference:sparse_mode=3with the shared compressed mask matches bottom-right-aligned causal for Sq==Skv, Sq<Skv, and Skv=4100 beyond the 2048 mask width (max_abs ≤ 4.6e-3);sparse_mode=0withatten_mask=Nonematches full attention (≤ 3.8e-3). Misalignment would show as O(1) error. Note this is 910B1, not the A5 the PR validates on. Scripts:/data/dxw/tmp/pr9244-verify/(H200),/dataset/home/dxw/pr9244/npu_smoke.py(910B1).
Two gates still block merge:
- Make the DCO check pass on the head commit (currently
action_required). - Close superseded #8876 so Mergify can evaluate this PR.
The A5 before/golden/after tables remain author-provided; they are internally consistent, and the residual I2I After-vs-Golden difference is already disclosed in the description.
| # ``forward_cuda``: the caller emitted ``full_attn_spans`` instead of a | ||
| # mask, so the generic mask handling below cannot express it. | ||
| if attn_metadata is not None and attn_metadata.full_attn_spans is not None: | ||
| return self._forward_piecewise_npu(query, key, value, attn_metadata) |
There was a problem hiding this comment.
Mirrors forward_cuda (flash_attn.py:697-721): spans take precedence over attn_mask and packed metadata on both platforms, and query_ranges is forwarded the same way. Verified on the head tree: CPU contract suite 55 passed; the same head tests against the pre-fix impl (4a910f2) fail exactly the three new piecewise tests (52 passed), so the defect is reproduced on base.
| Inputs follow the model-side ``(B, S, H, D)`` convention == BSND. | ||
| """ | ||
| if causal: | ||
| atten_mask = _get_npu_compressed_causal_mask(query.device) |
There was a problem hiding this comment.
Device-checked this contract on Ascend 910B1 (torch_npu 2.10.0.post4): sparse_mode=3 with the shared (2048, 2048) strict-upper mask reproduces bottom-right-aligned causal against a dense reference for Sq==Skv (64/64), Sq<Skv (64/2048), and Skv=4100 beyond the mask width, max_abs ≤ 4.6e-3 (bf16 scale; misalignment would be O(1)). inner_precise=0 is right here since piecewise causal segments always have Sq ≤ Skv. Note 910B1, not the A5 this PR validates on.
| assert full_call.args[0].shape == (1, 3, 2, 4) | ||
| assert full_call.args[1].shape == (1, 5, 2, 4) | ||
| assert full_call.kwargs["sparse_mode"] == 0 | ||
| assert full_call.kwargs["atten_mask"] is None |
There was a problem hiding this comment.
Same contract device-checked on 910B1: sparse_mode=0 with atten_mask=None matches plain full attention (128×512 and 4096×4096, max_abs ≤ 3.8e-3 vs an fp32 reference), so asserting the absent mask pins the right producer contract.
| if attn_metadata is not None and attn_metadata.full_attn_spans is not None: | ||
| return replace( | ||
| result, | ||
| path="npu_piecewise", |
There was a problem hiding this comment.
Non-blocking question: forward_fa_npu now takes the piecewise branch before its self.causal check, but this resolver still returns npu_unverified for self.causal (the early-return above fires first) — so for causal ∧ spans, routing reports unverified while execution runs piecewise. No current producer combines them (diffusion DiT impls are non-causal). Should the spans check move above the causal early-return, or is leaving that combination unverified intentional?
There was a problem hiding this comment.
Thanks for catching this. Fixed in e19cc3d: the resolver now checks full_attn_spans before the self.causal early return, matching forward_fa_npu.
The three piecewise tests now cover both causal settings, with an additional causal-without-spans regression test. CPU regression: 121 passed, 5 deselected, using mocked NPU kernels. The new causal+spans routing test fails before the fix and passes afterward. Kernel execution is unchanged by this follow-up. DCO now passes.
4be8385 to
d23a875
Compare
``forward_fa_npu`` never read ``AttentionMetadata.full_attn_spans``, so a
model emitting block-causal span metadata (e.g. Qwen-Image-2.1 prefill)
silently ran full bidirectional attention on Ascend: condition images
attended to the future target image and text attended to image blocks.
Because the prefill K/V is written into the cross-step KV cache, the
wrong results propagated through the whole denoising trajectory.
Mirror ``forward_cuda`` and route spans through the shared ``piecewise_attn``
segment plan, so segment construction is identical on CUDA and NPU and only
the per-segment kernel differs. Both segment kinds run on
``torch_npu.npu_fusion_attention``:
- causal segments (text runs) reuse the shared compressed (2048, 2048) mask
with ``sparse_mode=3`` (rightDownCausal). ``piecewise_attn`` gives causal
segments ``kv_end == q_end``, so they always satisfy ``Sq <= Skv``, and
bottom-right alignment reproduces FlashAttention's
``attn(Q[:, s:e], K[:, :e], causal=True)`` contract.
- full segments (condition/target image blocks) pass no mask at all --
``sparse_mode`` is ignored when ``atten_mask`` is ``None`` -- so they are
plain bidirectional attention over the truncated K/V prefix.
The op call moves into ``_npu_fusion_attention`` so the pre-existing explicit
causal path and the new piecewise path share one call site.
``_resolve_npu_execution_path`` now reports ``npu_piecewise`` instead of
bailing out on span metadata, and the FP8/MX quantized path already rejects
spans ("piecewise attention requires the float attention path").
Tests (CPU, fake torch_npu, matching the file's existing NPU contract
section): the execution path resolves to ``npu_piecewise``; a two-segment
input issues one kernel call per segment with the expected mask/mode; and a
four-segment input matches a dense ``(q >= kv) or same_image_block``
reference.
Keep execution-path reporting aligned with forward_fa_npu when both
causal=True and full_attn_spans are present. Extend the CPU contracts to
both causal settings and retain the causal-without-spans boundary.
Resubmitted by YawningPants with A5 validation and CPU regression testing.
Original implementation: vllm-project#8876.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: YawningPants <54137998+YawningPants@users.noreply.github.com>
Signed-off-by: YawningPants <54137998+YawningPants@users.noreply.github.com>
d23a875 to
e19cc3d
Compare
|
@hsliuustc0106 Thanks. The A5 configurations, reference input, comparison variants, run counts and timing scope are now documented in the PR description. For the requested CPU comparison, the exact same test file from Equivalent pytest selection: python -m pytest tests/diffusion/attention/test_flash_attn.py \
-k 'test_npu_piecewise_' \
-m 'core_model and cpu' \
-vv --tb=short -p no:cacheproviderBoth runs were executed through a shared CPU harness. It supplied CPU kernel stand-ins, including the baseline MindIE dense fallback, and isolated an unrelated optional Mooncake shutdown issue. The command above shows the pytest selection, not the complete harness setup. Production routing and the test assertions were unchanged. Observed output summary: The baseline failures were:
Separately, CPU regression on the final revision |
|
@amy-why-3459 PTAL |
Bounty-hunter
added
diffusion
amy-why-3459
left a comment
amy-why-3459
left a comment
There was a problem hiding this comment.
Reviewed tip 32e9ea48 (feature commit e19cc3dc; tip itself is merge-from-main only). Covered forward_fa_npu / _forward_piecewise_npu / _npu_piecewise_attn_func / _npu_fusion_attention, _resolve_npu_execution_path, the quantized-path span rejection, and the new CPU contracts in test_flash_attn.py.
Dispatch
- Spans take precedence in
forward_fa_npubefore the causal / mask branches, matchingforward_cuda(flash_attn.py~735–760): same sharedpiecewise_attnsegment plan andquery_rangesforwarding; only the per-segment kernel differs. _resolve_npu_execution_pathnow checksfull_attn_spansbefore theself.causalearly-return, socausal ∧ spansreportsnpu_piecewiseinstead ofnpu_unverified. That closes the mismatch Dong flagged on4be8385;test_npu_piecewise_reports_its_own_execution_path[True/False]andtest_npu_causal_without_spans_stays_unverifiedpin both sides.- FP8/MX quantized NPU path still rejects spans (
"piecewise attention requires the float attention path"); no change needed there.
Kernel contract
- Causal segments: shared compressed
(2048, 2048)mask +sparse_mode=3.piecewise_attnsetskv_end == q_endfor those segments, soSq <= Skvalways holds and defaultinner_precise=0is consistent with the explicit-causal branch (which only bumpsinner_precisewhenSq > Skv). - Full segments:
atten_mask=None,sparse_mode=0— Ascend ignores sparse mode when the mask is absent, so this is plain bidirectional attention over the truncated K/V prefix. _npu_fusion_attentioncorrectly shares one call site between the pre-existing explicit-causal path and the new piecewise path.
Tests
- Per-segment call shapes/modes (
test_npu_piecewise_runs_one_fusion_attention_per_segment) and dense block-causal numerical agreement (test_npu_piecewise_matches_block_causal_reference) look right for the intended Qwen-Image-2.1 prefill pattern. I did not re-run the suite or the A5 measurements here; Dong’s 910B1 kernel smokes and the author’s A5 Before/Golden/After tables are consistent with the claimed contract.
Process (not code)
- Mergify still cannot evaluate while superseded #8876 shares the head SHA — that needs #8876 closed/merged before this PR can proceed through Mergify.
No new code defect found in this pass.
|
@amy-why-3459 @Bounty-hunter the ci shows failure, is this an obstacle? |
…spans (vllm-project#9244) Signed-off-by: YawningPants <54137998+YawningPants@users.noreply.github.com> Co-authored-by: changetheway <guotaoyuan1@h-partners.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: dengyunyang <584797741@qq.com> (cherry picked from commit 0be1493) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This PR supersedes #8876, preserving its original implementation and adding Ascend A5 validation results.
Purpose
FlashAttentionImpl.forward_fa_npunever readAttentionMetadata.full_attn_spans. A model that emits block-causal spanmetadata — e.g. Qwen-Image-2.1 prefill — therefore ran full bidirectional
attention on Ascend instead of the block-causal pattern, silently producing
wrong results: condition images attended to the future target image and text
attended to image blocks. Because the prefill K/V is written into the
cross-step KV cache, the corruption propagated through the entire denoising
trajectory.
forward_cudaalready dispatches onfull_attn_spans; only the NPU path wasmissing it.
Changes
Mirror
forward_cuda: route spans through the sharedpiecewise_attnsegmentplan, so segment construction is identical on CUDA and NPU and only the
per-segment kernel differs. Both segment kinds run on
torch_npu.npu_fusion_attention:(2048, 2048)mask withsparse_mode=3(rightDownCausal).piecewise_attngives causal segments
kv_end == q_endin global coordinates, so they alwayssatisfy
Sq <= Skv, and bottom-right alignment reproduces FlashAttention'sattn(Q[:, s:e], K[:, :e], causal=True)contract. No new mask is created —this is the same
_get_npu_compressed_causal_maskthe native causal path uses.sparse_modeis ignored whenatten_maskisNone— so they are plainbidirectional attention over the truncated K/V prefix.
The operator call is factored into
_npu_fusion_attentionso the pre-existingexplicit-causal branch and the new piecewise path share one call site.
_resolve_npu_execution_pathnow reportsnpu_piecewiseinstead of fallingthrough to
npu_unverifiedon span metadata. The FP8/MX quantized path alreadyrejects spans (
"piecewise attention requires the float attention path"), sono change was needed there.
Test Plan
CPU regression on the final revision (
e19cc3d): 121 passed, 5 deselected. Tests use mocked NPU kernels and cover execution-path resolution, per-segment kernel/mask arguments, numerical agreement with a dense block-causal reference, both causal settings for piecewise attention, and the causal-without-spans routing boundary.A5 experiments used the
quay.io/ascend/vllm-omni:v0.31.0rc1-a5base image.A5 comparison variants
full_attn_spansand performed full bidirectional attention instead of the intended block-causal pattern.4a910f239079attn_path="mask".4be8385d4c8b+ mask override4be8385d4c8bThese revisions identify the sources used for the recorded A5 measurements. The final revision (
e19cc3d) additionally fixes execution-path reporting and expands CPU coverage without changing attention kernel execution.Shared experiment settings
qwen_bear.png, RGBA; downloaded from the public reference imageGeneration follows the Qwen-Image 2.1 offline examples, using the inputs and settings listed above. Profiling and image generation use the same generation settings except for the denoising step count.
I2I reference image (saved locally as
qwen_bear.png):Test Result
Timings are accumulated device-kernel durations from a single two-step profiling run per variant, without additional warmup. They are bounded observations, not end-to-end generation latency.
Performance — I2I, accumulated kernel durations in ms
After reduces prefill attention time by 30.03% versus Golden, or 33.86% including adjacent conversions. Decode remains approximately unchanged. Total recorded kernel time decreases by 3.54%.
Image similarity — RGB pixel range 0–255, 50 denoising steps
After is substantially closer to Golden than Before in both tasks. These are descriptive similarity measurements; no predefined image-quality acceptance threshold was used.