[iboiko-habana]

Contributor intent (required human answers)

1. What outcome do you want to achieve?

HunyuanImage-3.0 runs correctly on Intel XPU with DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN instead of failing every request with HTTP 500.

2. What is the highest-leverage variable for that outcome?

The XPU FlashAttention backend honoring full_attn_spans, the mixed causal/full attention contract the backend already advertises.

3. How will this work move that variable?

forward_xpu now dispatches spans through piecewise_attn with the varlen wrapper, exactly as forward_cuda does, so models that pass spans get correct mixed attention on XPU.

4. Whose customer or user needs does this address, and what are those needs?

Users serving HunyuanImage-3.0 (and other span-based models such as Qwen-Image 2.1 and HiDream-O1) on Intel GPUs. They need the default fast attention backend to work and produce correct images, without falling back to TORCH_SDPA.

Purpose

Before: FlashAttentionImpl.forward_xpu ignores AttentionMetadata.full_attn_spans, although FlashAttentionBackend advertises supports_piecewise_spans = True on every platform.

• HunyuanImage-3.0 with DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN fails on XPU. On the dense prompt-KV path (_forward_dense_legacy), the model passes a 4D attn_mask together with full_attn_spans. forward_xpu sends the 4D mask to _forward_varlen_masked, which asserts attention_mask must be 2D, so every request returns HTTP 500. XPU uses this path because it has no Scheduler-paged KV.
• Callers that pass only spans lose the mixed causal/full pattern. Qwen-Image 2.1 (piecewise path) and HiDream-O1 pass spans only after checking supports_piecewise_spans. On XPU they fall through to _forward_varlen_dense, which runs plain full attention, so the causal segments are not applied. This is from code inspection; I did not run these models on XPU.

After: when full_attn_spans is present, forward_xpu dispatches through piecewise_attn with the existing _flash_varlen_wrapper, like forward_cuda (XPU only provides flash_attn_varlen_func). Spans take precedence over attn_mask, as on CUDA. Other XPU paths are unchanged.

Contracts and docs: this brings XPU in line with the existing piecewise-span contract of the Flash backend (see Attention execution contract and the span inputs in Diffusion paged KV cache). The user guides (Diffusion Attention Backends, Dense backends) need no update: they don't say XPU lacks span support, and the user-facing configuration is unchanged.

Alternatives considered

• Rework the HunyuanImage-3.0 dense path to pass spans only: the model already passes spans; the defect is in the XPU backend, and fixing it there also covers the other span-based models.
• Default HunyuanImage-3.0 on XPU to Scheduler-paged KV: this is a larger feature, needing XPU native paged attention and memory-profile validation. The paged path also does not support imported AR KV (_forward_paged raises NotImplementedError). It would still depend on this span dispatch.

Known limitations (same as CUDA): spans must be identical across batch rows, otherwise piecewise_attn raises ValueError. attn_mask is ignored when spans are present.

AI assistance was used to prepare this change. The author reviewed every changed line and ran the tests below.

Test Plan

vLLM Version: v0.31.0 (XPU image)

vLLM-Omni Commit:

• Unit tests: da572cddb (v0.32.0rc1) + this patch
• End-to-end: previous base 9bbfb61d5 + the identical patch. The rebased tree requires vLLM v0.32 serving APIs (e.g. vllm.renderers.chat_utils), which the v0.31.0 XPU image lacks. The patched code path is identical on both bases; the only upstream change in it ([Bugfix][Diffusion] Pass the resolved FA version into dense varlen attention #9100, fa_version kwargs) applies only to vLLM's CUDA varlen dispatcher.

Hardware: 4x Intel Arc Pro B70 (Battlemage)

1. Regression test, without and with the fix:

   pytest -q -m cpu tests/diffusion/attention/test_flash_attn.py -k test_xpu_piecewise_spans_ignore_4d_mask
   pytest -q -m cpu tests/diffusion/attention/test_flash_attn.py tests/diffusion/attention/test_piecewise_attn.py

2. End-to-end HunyuanImage-3.0-Instruct-Distil, online serving: TP=4 + expert parallel, FP8, 512x512, 3 identical requests per server start, denoise graph off. DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN vs TORCH_SDPA.
3. Operator-level equivalence (temporary instrumentation, not part of this PR): for each distinct attention shape in a real HunyuanImage-3.0 request, compared
   • the mask implied by full_attn_spans (bottom-right causal + full spans) with the model's real 4D mask, and
   • the piecewise FlashAttention output with F.scaled_dot_product_attention on the same Q/K/V with the real 4D mask.

Test Result

Regression test: test_xpu_piecewise_spans_ignore_4d_mask

• Without the fix (v0.32.0rc1): 1 failed, AssertionError: attention_mask must be 2D, (batch_size, seq_len)
• With the fix (v0.32.0rc1): test_flash_attn.py + test_piecewise_attn.py: 80 passed, 7 deselected

End-to-end (HunyuanImage-3.0 on XPU, previous base + identical patch):

Backend	Before	After
FLASH_ATTN	HTTP 500 on every request (attention_mask must be 2D)	3/3 requests OK, outputs bit-identical across requests
TORCH_SDPA	OK	OK (unchanged)

• FLASH_ATTN vs TORCH_SDPA image PSNR: 19.4 dB. For reference, two TORCH_SDPA server starts differ by 23.8 dB. The remaining gap is numeric drift over the denoising steps; the operator-level check below shows the attention semantics are identical.
• FLASH_ATTN is about 4% faster end-to-end than TORCH_SDPA in this configuration.

Operator-level equivalence (real HunyuanImage-3.0 shapes):

Call	q × k	Spans	Mask vs spans mismatches	max |FA − SDPA| (output abs max)
Prefill	1121 × 1121	[(94, 1118)]	0	2.4e-4 (6.1e-2)
Denoising step (prompt-KV reuse, q < k)	1027 × 1118	[(94, 1118)]	0	1.2e-4 (2.3e-2)

Not run:

• End-to-end on v0.32.0rc1 (needs a vLLM v0.32 XPU image).
• Qwen-Image 2.1 / HiDream-O1 on XPU.
• CUDA/ROCm/NPU paths (unchanged by this PR).
• Multi-row batches with heterogeneous spans (unsupported, same as CUDA).

[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	⚠️ Failed 2026-10-19T18:07:09.454756Z	e92137f	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.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

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, module named in the PR description, semantic router, CODEOWNERS; @david6666666 via module of the changed files, module named in the PR description, semantic router; @princepride via module of the changed files, module named in the PR description, semantic router

@iboiko-habana, 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.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 8538c449858c produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@iboiko-habana this PR reads as a performance or value claim:

• claim: - FLASH_ATTN is about 4% faster end-to-end than TORCH_SDPA in this configuration.

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[xuechendi]

LGTM