[tthakkal]

FIX #6882

Purpose

HiDream-ai/HiDream-O1-Image generates pure noise on XPU — exit 0, no warning. supports_piecewise_spans was a plain class attribute: False on the AttentionBackend base, but hardcoded True on all three flash backends (flash_attn.py:175, flash_attn_hub.py:85/:253) regardless of platform — and FLASH_ATTN is what HiDream-O1 selects on XPU. Only forward_cuda dispatches piecewise_attn over AttentionMetadata.full_attn_spans. hidream_o1_image_transformer.py trusted the flag, skipped building the equivalent attn_mask, and XPU/NPU then ran fully dense unmasked attention — text rows that must stay causal attended to future image tokens. Full analysis in #6882.

A flag describes the forward_* that actually runs, not the backend class, so the mask-capability flags become platform-gated classmethods. CUDA / ROCm / MUSA flag values are identical before and after.

• abstract.py — supports_piecewise_spans() becomes a classmethod defaulting to False; new supports_dense_attention_mask() defaults to supports_attention_mask(), which only promises some mask is read (several kernels take just a 2D [B, Skv] padding mask).
• flash_attn.py, flash_attn_hub.py — in all three backends spans are True only for the CUDA family, exactly the set that reaches forward_cuda (ROCm/MUSA via base-class forward_hip / forward_musa delegation). The hub backends are converted along with the plain one because the base declaration is now a classmethod: leaving their supports_piecewise_spans: bool = True in place would shadow it with a non-callable and break every caller.
• cudnn_attn.py, vdnh3_attn.py — the same conversion for two flag declarations that landed on main after this PR was opened (CUDNN_ATTN's via [New model] Qwen-Image-2.1 support (rebase + review follow-ups) #8099, the new VDNH3_ATTN backend via [Model] Add VDN-H3 hybrid attention (VDNH3_ATTN) for MiniMax-H3 #8439). Both are CUDA-only and keep their old values (True / False); left as bool attributes, every supports_piecewise_spans() call on them raised TypeError: 'bool' object is not callable.
• flash_attn.py — only this backend overrides supports_dense_attention_mask(), returning False on XPU alone, where its one mask route is _forward_varlen_masked, which asserts 2D. The two hub backends inherit the abstract.py default and keep reporting True.
• layer.py — _run_local_attention reroutes a >2D mask to the existing sdpa_fallback with a warning_once, scoped to backends that claim mask support so it cannot preempt TRTLLM_ATTN's deliberate ValueError. _assert_piecewise_compatible now calls the flag.
• hidream_o1_image_transformer.py — HiDream-O1's probe now calls the flag. This () is what turns the noise into an image.
• qwen_image_21_transformer.py — the other model-side consumer, from the same [New model] Qwen-Image-2.1 support (rebase + review follow-ups) #8099, gets the same (). Without it the bound method is always truthy, so attn_path_hint="auto" picked "piecewise" for every backend — including TORCH_SDPA, which then hit the _assert_piecewise_compatible ValueError.

Platform	supports_piecewise_spans()	Mask reroute	Net effect
CUDA / ROCm / MUSA	True (was True)	no	No change
XPU	False (was True)	yes	Fixed — model builds the mask, layer runs it on SDPA; verified e2e
NPU	False (was True)	no	Unverified — a 4D mask now reaches aclnnFlashAttentionScore instead of being skipped, which should be a strict improvement over the dense unmasked run, but no NPU hardware was available. Needs confirmation from someone who has one; see the follow-up on mindiesd mask conventions below before assuming it is fixed.

The reroute fires for exactly one backend on one platform: FLASH_ATTN on XPU. It needs supports_attention_mask() and not supports_dense_attention_mask(), and flash_attn.py is the only backend that overrides the latter — every other backend inherits the abstract.py default, where the two are the same value, so the pair can never both hold. CUDNN_ATTN, FLASHINFER_ATTN, TORCH_SDPA and the two hub backends cannot reroute.

Scope

Per review, this PR is only the XPU correctness fix: abstract.py + the backend flags + layer.py + the two model-side () fixes. Unrelated fixes that were riding along have been reverted out; the one still open is listed under Follow-ups.

Reviewer attention

• The reroute is not HiDream-only. On XPU, every model that hands FLASH_ATTN a >2D mask now runs that layer on SDPA after one warning_once — HunyuanImage3, glm_image, flux, flux2 and ernie_image all unsqueeze to 4D before Attention. Previously those masks hit forward_xpu's 2D assert, so this turns a hard failure into a working run. CUDA keeps HunyuanImage3's 4D mask + spans on piecewise_attn.
• supports_dense_attention_mask() is deliberately True on CUDA, though only conditionally honored there: forward_cuda honors a >2D mask only via piecewise_attn, so a >2D mask without spans still hits the 2D assert — reachable by lingbot_video and glm_image. Returning False fixes both but drags HunyuanImage3 (4D mask with spans) off piecewise_attn onto SDPA, so it is a follow-up. Same for the hub backends.
• build_hidream_o1_attention_mask is newly load-bearing on XPU/NPU. It was dead code there before this PR, so its test file is new coverage for existing unexercised code rather than for anything this PR rewrites; the function itself is untouched.
• The 4D additive bf16 mask never reaches forward_xpu. layer.py diverts it to SDPA first, so the reverted torch.any(~attention_mask) on that path is not on the HiDream-O1 route. Were it ever reached with an additive mask it would raise on ~ over bf16 rather than fail silently.
• mypy does not catch the bool-shadows-classmethod regression, which is how CUDNN_ATTN / VDNH3_ATTN slipped through the merge. The registry-wide test below is the guard.

Test Plan

136 new cases: 134 CPU (core_model), 2 hardware-gated.

• test_mask_capability_flags.py (new, 116)
  • The platform table above asserted rather than described: both flags × three flash backends × five platforms, through a real OmniPlatform subclass.
  • The forward_hip / forward_musa → forward_cuda delegation the ROCm/MUSA True rests on, and that the flags are callables rather than truthy attributes.
  • The _mask_needs_sdpa truth table over 2D/3D/4D masks × mask-aware / dense-capable / mask-blind backends.
  • _run_local_attention dispatch with FLASH_ATTN on a patched platform, recording which impl actually runs: XPU sends 4D (with or without spans) and 3D masks to sdpa_fallback and keeps 2D / no mask on flash; CUDA keeps HunyuanImage3's 4D mask + spans on flash; NPU keeps 4D on flash.
  • Every backend in DiffusionAttentionBackendEnum × five platforms answers both flags as a callable returning bool, so a future supports_piecewise_spans: bool = ... fails CI. Backends that cannot be imported on the host are skipped (SAGE_ATTN_3 on XPU).
• test_hidream_o1_attention_mask.py (new, 18) — rows causal or full at seq_len 129 / 1000 / 4231 across bf16 / fp16 / fp32 against an independent torch.where reference, asserting the fully-unmasked row count and not just the values, since a partially-applied mask writes correct values into the rows it reaches and silently leaves the rest causal. Also HiDreamO1ImageTransformer.supports_piecewise_attention over stub backends (all on / all off / mixed), which is the property the pipeline reads to decide whether to build the mask. The 2 hardware_test cases build the mask on the accelerator itself and skip rather than false-pass when none is visible; both ran on XPU here.
• test_attn_path_probe.py (new, 2) — Qwen-Image-2.1's _resolve_attn_path("auto") returns "piecewise" / "mask" according to the backend's supports_piecewise_spans().
• test_attention_config.py — no new cases, only the stubs the flag change invalidates: three SimpleNamespace backends whose supports_piecewise_spans becomes a callable, a _mask_needs_sdpa stub for the one test that drives _run_local_attention directly, and a FlashAttentionBackend subclass with both flags pinned to the CUDA-family answers so those CPU dispatch assertions do not depend on the host platform.

Each new guard was checked against its regression: dropping the () in either model probe, reverting CUDNN_ATTN / VDNH3_ATTN to bool attributes, or disabling the reroute in layer.py each fails the corresponding tests. Across tests/diffusion/attention/, tests/diffusion/models/hidream_o1_image/ and tests/diffusion/models/qwen_image_21/ on XPU, the failing set is identical to a clean main worktree on the same container (54 pre-existing environment failures, none new).

vLLM Version: 0.32.0 (torch 2.14.0+xpu)

Test Result

HiDream-O1-Image, 2048x2048, --seed 42, on Intel Graphics [0xe223], 30.3 GiB each. TORCHDYNAMO_DISABLE=1 is required on this box for an unrelated Triton/Intel defect. Timings exclude a discarded warm-up generation (first-run kernel JIT).

TP	Steps	Generate (s)	Peak memory*	Before this PR	After
1x	8	14.9	15.2 GiB	noise ❌	coherent image, legible sign text ✅
2x	8	9.4	8.3 GiB/card	noise ❌	coherent image, legible sign text ✅
2x	50	57.4	8.4 GiB/card	noise ❌	coherent image, legible sign text ✅

* Peak memory from the earlier run on vLLM 0.29.0 / torch 2.13.0; not re-measured on 0.32.0.

"Before" is main at the merge base, same script and seed, re-run on 0.32.0 for all three rows: SSIM vs the baseline is 0.053 (TP1/8), 0.029 (TP2/8) and 0.032 (TP2/50), against 0.49 / 0.47 / 0.82 after.

The reroute is observable in the log — layer.py:776 emits Attention backend 'FLASH_ATTN' cannot honor a 4D attn_mask on this platform; running this layer on SDPA once per run.

At 50 steps against the CUDA-generated tests/assets/hidream_o1/hidream_o1_baseline.png, using the same LANCZOS downscale the accuracy test applies: SSIM 0.8174 (passes the 0.8 gate), PSNR 18.577 (just under the CUDA-tuned 19.0), visually indistinguishable. FLASH_ATTN and an explicit DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA are bit-identical (max abs diff 0 over the full 2048x2048 RGB image), which is what the reroute guarantees, so the PSNR gap is XPU-vs-CUDA numerics rather than a masking difference. These numbers reproduce exactly on both the 0.29.0 and 0.32.0 stacks.

• Not covered: no CUDA or NPU hardware, so the CUDA/ROCm/MUSA "no change" claim rests on the flag values being identical before and after rather than on an e2e run, and the NPU row of the table above is explicitly unverified. The same applies to CUDNN_ATTN, VDNH3_ATTN and Qwen-Image-2.1: their behavior is restored to what main had, verified on CPU only.

Follow-ups (not in this PR)

Split out of this PR at review request — an independent fix to a pre-existing bug, not needed for the XPU noise fix:

• torch.any(~attention_mask) mis-routes int keep-masks at four flash call sites (forward_cuda / forward_xpu in flash_attn.py, and both hub backends): ~1 == -2, so every int keep-mask is unconditionally treated as padded and takes the slower masked-varlen path. Qwen-Image is the affected in-tree model. The same helper should reject an additive float mask instead of inverting its meaning inside _unpad_input.

Pre-existing and untouched by this PR:

• lingbot_video and glm_image cannot use FLASH_ATTN on CUDA with a >2D mask and no spans. Make the reroute predicate metadata-aware: reroute unless spans are present and supports_piecewise_spans(). Needs CUDA hardware, since it changes which kernel HunyuanImage3 runs on.
• HiDream-O1 on NPU now builds an additive bf16 [1,1,S,S] mask that forward_fa_npu passes unchanged into mindiesd.attention_forward, where every other producer supplies a bool keep-mask — so if mindiesd reads it as one the semantics invert. This is why the NPU row above is marked unverified rather than fixed. That path ran dense and silently wrong before this PR too.
• SAGE_ATTN, SAGE_ATTN_3 and RAINFUSION_ATTN silently ignore an attn_mask instead of raising; they should either raise like TRTLLM_ATTN or opt into the reroute. And the CUDA-tuned PSNR threshold in tests/e2e/accuracy/test_hidream_o1_image.py needs a per-platform value before that test can gate XPU.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.

[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 @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi

@tthakkal, 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 575d2f4fad51 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[tthakkal]

Self-review @vllm-omni-review-bot

layer.py:597 drops attn_spec. _mask_needs_sdpa calls supports_attention_mask() spec-less while _assert_metadata_compatible (:610) passes it, and supports_dense_attention_mask has no spec parameter at all. Safe today — FLASHINFER_ATTN (flashinfer_attn.py:41) answers True spec-less / False spec-ful, so the assert raises first — but the asymmetry will bite a backend that answers the other way. Pass self.attn_spec and add the parameter.

Checked and cleared

Ring path has no reroute gap (layer.py:635 rejects masks outright). sdpa_fallback is never None where the reroute runs (:516 returns first for custom_attention). forward_hip/forward_musa do delegate to forward_cuda (abstract.py:368-386), so ROCm/MUSA True is correct. No docs/ page documents these capability flags, so none needs updating.

Evidence (4× Intel Arc Pro B70, torch 2.13.0+xpu, vLLM 0.29.0)

Verified natively with FLASH_ATTN resolved: flags correct on XPU; raw FlashAttentionImpl.forward with a [2,1,32,32] mask really does AssertionError; a real Attention layer reroutes it to output bit-identical to plain SDPA (maxdiff 0.0); a 2D mask is correctly not rerouted.

[tthakkal]

Could someone add the ready label so L1/L2 run? Fix is XPU-only (@yma11 @xuechendi FYI) — CUDA/ROCm/MUSA flag values are unchanged, verified by test_mask_capability_flags.py.

[xuechendi]

@Joshna-Medisetty , please help to initial scan this PR

[Joshna-Medisetty]

scanned it. the xpu part looks correct to me - supports_piecewise_spans really is hardcoded True at flash_attn.py:24 and forward_xpu never reads full_attn_spans, so the flag gating plus the sdpa reroute is the right fix.

main ask is whether this can be split. the actual xpu fix is abstract.py + the flash_attn.py flags + layer.py + the hidream one-liner. mask_excludes_tokens (cuda routing), the npu kv-quant raise, and the sequence_parallel gate are three separate fixes riding along.

also the table says npu is Fixed but the follow-up section says mindiesd may invert that mask. can we mark npu unverified until CI passes?

[tthakkal]

scanned it. the xpu part looks correct to me - supports_piecewise_spans really is hardcoded True at flash_attn.py:24 and forward_xpu never reads full_attn_spans, so the flag gating plus the sdpa reroute is the right fix.

main ask is whether this can be split. the actual xpu fix is abstract.py + the flash_attn.py flags + layer.py + the hidream one-liner. mask_excludes_tokens (cuda routing), the npu kv-quant raise, and the sequence_parallel gate are three separate fixes riding along.

also the table says npu is Fixed but the follow-up section says mindiesd may invert that mask. can we mark npu unverified until CI passes?

@Joshna-Medisetty Thank you for the review, both asks done. The PR is now only the XPU correctness fix. @xuechendi

Split out mask_excludes_tokens (+ its four call sites), the NPU KV-quant raise, and the sequence_parallel.py gate, along with their tests; all three are now Follow-ups, and utils/fa.py / sequence_parallel.py / test_flash_attn.py are byte-identical to 97912de8d. NPU is marked Unverified and points at the mindiesd mask-convention question.

One addition to your list: flash_attn_hub.py has to stay in, since once the base declaration is a classmethod the hub backends' supports_piecewise_spans: bool = True shadows it with a non-callable. Same platform gate, nothing else.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 7 days

@tthakkal this pull request has had no human commit, comment or review since 2026-09-17. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot review

2 actionable finding(s).

CI at 2ca811e00f16 (2026-10-10T01:42:52.024242+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Full review analysis

Scan:

Category	Result
Tests / verification	2 finding(s) below
Security	no finding reported
Docs / comments	no finding reported
Behavior / compatibility	no finding reported
Correctness	no finding reported

Validated:

• [resolved] NPU 4D mask → mindiesd: PR table marks unverified after Joshna; HiDream uses causal=False (hidream_o1_image_transformer.py:154) so forward_fa_npu takes the mindiesd path (flash_attn.py:962), not the torch_npu keep-mask invert (916).
• [resolved] split request: net diff matches the stated XPU-only set.
• [resolved] sdpa_fallback None: custom_attention returns at layer.py:584-585 before the reroute; sdpa_fallback is constructed at layer.py:226 when custom is None.
• [resolved] attn_spec dropped in _mask_needs_sdpa — author self-review already named FLASHINFER True/False-with-spec; residual: a future backend that answers False spec-less / True spec-ful would still reroute. Not republished.
• [resolved] Riding-along mask_excludes_tokens / NPU kv-quant raise / sequence_parallel truthiness gate — at this head sequence_parallel.py has no supports_piecewise_spans; grep confirms.
• [claim-verified] 'flag's one consumer' besides layer.py: repo-wide grep of supports_piecewise_spans in vllm_omni/ hits only hidream transformer + flash_attn*.py + abstract.py + layer.py.

2 actionable finding(s).

Verdict: COMMENT

Findings

• [P2] HiDream's piecewise probe is now a call (supports_piecewise_spans() at hidrea… — vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py

Evidence for HiDream's piecewise probe is now a call (`supports_piecewise_spans()` at hidrea…

HiDream's piecewise probe is now a call (supports_piecewise_spans() at hidream_o1_image_transformer.py:518), which is what pipeline_hidream_o1_image.py:646 (if not self.model.supports_piecewise_attention) uses to build the 4D mask. No test exercises that property: test_hidream_o1_attention_mask.py only imports build_hidream_o1_attention_mask, and test_piecewise_spans_is_a_call_not_a_truthy_attribute only checks AttentionBackend.supports_piecewise_spans. Dropping the () makes all(...) True on a bound method, so the pipeline skips the mask and every added test still passes. Add a CPU test that stubs a transformer whose backends return False and asserts supports_piecewise_attention is False.

Evidence: hidream_o1_image_transformer.py:518 layer.self_attn.attn.attn_backend.supports_piecewise_spans() for layer in self.model.language_model.layers — only production probe for HiDream piecewise. Unchanged by this diff, present in the PR-time tree: pipeline_hidream_o1_image.py:646 if not self.model.supports_piecewise_attention: then :647-650 attention_mask = build_hidream_o1_attention_mask(...). Trigger: FLASH_ATTN on XPU/NPU where supports_piecewise_spans() is False. Unmet requirement: no test reads supports_piecewise_attention (repo-wide only those two production hits). Unchanged-by-diff import in the new test file: test_hidream_o1_attention_mask.py:25-27 from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import ( build_hidream_o1_attention_mask, ). test_mask_capability_flags.py:148-149 assert callable(AttentionBackend.supports_piecewise_spans) / assert AttentionBackend.supports_piecewise_spans() is False never instantiates the HiDream property. Reverting the () at :518 leaves a truthy bound method, all(...) is True, the pipeline skips the mask, and those tests still pass.

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 29 days

@tthakkal this pull request has had no human commit, comment or review since 2026-09-17. Please consider marking this PR as draft until work can resume. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[tthakkal]

Omni ReviewBot review

2 actionable finding(s).

CI at 2ca811e00f16 (2026-10-10T01:42:52.024242+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Full review analysis

Scan:
Category Result
Tests / verification 2 finding(s) below
Security no finding reported
Docs / comments no finding reported
Behavior / compatibility no finding reported
Correctness no finding reported

Validated:

* [resolved] NPU 4D mask → mindiesd: PR table marks unverified after Joshna; HiDream uses causal=False (hidream_o1_image_transformer.py:154) so forward_fa_npu takes the mindiesd path (flash_attn.py:962), not the torch_npu keep-mask invert (916).

* [resolved] split request: net diff matches the stated XPU-only set.

* [resolved] sdpa_fallback None: custom_attention returns at layer.py:584-585 before the reroute; sdpa_fallback is constructed at layer.py:226 when custom is None.

* [resolved] attn_spec dropped in _mask_needs_sdpa — author self-review already named FLASHINFER True/False-with-spec; residual: a future backend that answers False spec-less / True spec-ful would still reroute. Not republished.

* [resolved] Riding-along mask_excludes_tokens / NPU kv-quant raise / sequence_parallel truthiness gate — at this head sequence_parallel.py has no supports_piecewise_spans; grep confirms.

* [claim-verified] 'flag's one consumer' besides layer.py: repo-wide grep of supports_piecewise_spans in vllm_omni/ hits only hidream transformer + flash_attn*.py + abstract.py + layer.py.

2 actionable finding(s).

Verdict: COMMENT

Findings

* **[P2] HiDream's piecewise probe is now a call (`supports_piecewise_spans()` at hidrea…** — `vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py`

Evidence for HiDream's piecewise probe is now a call (supports_piecewise_spans() at hidrea…

HiDream's piecewise probe is now a call (supports_piecewise_spans() at hidream_o1_image_transformer.py:518), which is what pipeline_hidream_o1_image.py:646 (if not self.model.supports_piecewise_attention) uses to build the 4D mask. No test exercises that property: test_hidream_o1_attention_mask.py only imports build_hidream_o1_attention_mask, and test_piecewise_spans_is_a_call_not_a_truthy_attribute only checks AttentionBackend.supports_piecewise_spans. Dropping the () makes all(...) True on a bound method, so the pipeline skips the mask and every added test still passes. Add a CPU test that stubs a transformer whose backends return False and asserts supports_piecewise_attention is False.

Evidence: hidream_o1_image_transformer.py:518 layer.self_attn.attn.attn_backend.supports_piecewise_spans() for layer in self.model.language_model.layers — only production probe for HiDream piecewise. Unchanged by this diff, present in the PR-time tree: pipeline_hidream_o1_image.py:646 if not self.model.supports_piecewise_attention: then :647-650 attention_mask = build_hidream_o1_attention_mask(...). Trigger: FLASH_ATTN on XPU/NPU where supports_piecewise_spans() is False. Unmet requirement: no test reads supports_piecewise_attention (repo-wide only those two production hits). Unchanged-by-diff import in the new test file: test_hidream_o1_attention_mask.py:25-27 from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import ( build_hidream_o1_attention_mask, ). test_mask_capability_flags.py:148-149 assert callable(AttentionBackend.supports_piecewise_spans) / assert AttentionBackend.supports_piecewise_spans() is False never instantiates the HiDream property. Reverting the () at :518 leaves a truthy bound method, all(...) is True, the pipeline skips the mask, and those tests still pass.

🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!

@vllm-omni-review-bot
both findings addressed

nothing read supports_piecewise_attention itself. Added test_supports_piecewise_attention_calls_the_backend_probe in test_hidream_o1_attention_mask.py (575d2f4). It calls the real property on a stub transformer whose layers' backends report spans as all on / all off / mixed, and asserts True / False / False. Dropping the () at hidream_o1_image_transformer.py:518 fails the two False cases.

While auditing the other callers I found the same problem from the latest merge with main:

• CUDNN_ATTN (flag added in [New model] Qwen-Image-2.1 support (rebase + review follow-ups) #8099) and VDNH3_ATTN ([Model] Add VDN-H3 hybrid attention (VDNH3_ATTN) for MiniMax-H3 #8439) still declared supports_piecewise_spans as a bool, so calling it raised TypeError: 'bool' object is not callable.
• Qwen-Image-2.1 ([New model] Qwen-Image-2.1 support (rebase + review follow-ups) #8099) read it without (), so attn_path_hint="auto" picked "piecewise" for every backend, including TORCH_SDPA.

All three are fixed in this PR. A new test also checks every backend in DiffusionAttentionBackendEnum on all five platforms, so a bool flag can't come back. mypy doesn't flag this case, so the test is the guard.

[xuechendi]

Thanks @tthakkal for the thorough analysis here. I re-checked this against current main (2077854). #9123 (merged 2026-10-22) made forward_xpu honor full_attn_spans through piecewise_attn, the same way forward_cuda does, so HiDream-O1 on XPU is already correct on main via the fast flash path. With this PR on top (it still merges cleanly), XPU reports supports_piecewise_spans() == False, HiDream builds the 4D additive mask, and the layer reroutes to SDPA, which is correct but much slower.

Measured on an Arc Pro B70 (vLLM 0.32.0, torch 2.14.0+xpu): real Attention layer on FLASH_ATTN, HiDream's own build_hidream_o1_full_attn_spans / build_hidream_o1_attention_mask (text causal, image full), bf16, H=32, D=128, 20 iters after warm-up, error vs fp32 SDPA with the HiDream mask:

tree	path	max err	S=1024	S=4231 (2048x2048)
main before #9123 (this PR's base)	spans ignored	3.25 (unmasked)	0.30 ms	2.75 ms
main	piecewise flash (#9123)	0.0086	0.50 ms	2.77 ms
main + this PR	4D mask -> SDPA	0.0079	0.62 ms	16.29 ms

The XPU fix is superseded by #9123, and merging this PR would make HiDream-O1 attention ~6x slower per call on XPU. So I'd suggest closing this PR. Two things in it are still valid on main and would be better as small follow-ups:

1. NPU: forward_npu / forward_fa_npu still ignore full_attn_spans while FLASH_ATTN advertises supports_piecewise_spans = True. From code inspection, HiDream-O1 / Qwen-Image-2.1 on NPU still run without the mask. Implementing spans in the NPU forward, as [Bugfix][XPU] Honor full_attn_spans in the XPU FlashAttention path #9123 did for XPU, seems the most consistent fix, and it needs an NPU owner to verify.
2. XPU, >2D masks without spans: e.g. flux/flux2/glm_image/ernie_image passing a 3D/4D mask still hit _forward_varlen_masked's 2D assert on XPU. The SDPA reroute from this PR would fix that as a standalone change, scoped to masks without spans. One caveat: it would also move all-True 4D bool masks off flash (2.04 -> 4.05 ms per call at S=4096/H=24/D=128 on B70), so keeping the dense flash shortcut for that case would be good.