[NumberWan]

Fixes #9092.

1. What outcome do you want to achieve?

Put Qwen-Image dense attention back on the pre-#7735 FA3 path when fa3-fwd is installed, so single-device 512 latency returns to that band. Keep fa3-fwd optional. Do not revert #7735.

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

Whether CUDA dense attention has flash_attn_func. After #7735 that binding is None, so FlashAttentionImpl.forward_cuda emulates the dense call with vLLM flash_attn_varlen_func. #9100 only passes fa_version on that varlen call. On this L20X that is still much slower than fa3_fwd_interface.flash_attn_func.

3. How will this work move that variable?

After the vLLM varlen candidate is bound, and only on non-Blackwell, import fa3_fwd_interface if it is present and bind both flash_attn_func and flash_attn_varlen_func from it. If the import fails, leave the vLLM varlen + #9100 fa_version path unchanged. requirements/cuda.txt stays without fa3-fwd.

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

People running Qwen-Image (and any other dense diffusion attention on the same backend) on Hopper. #9092 reports the perf case about 60% slower than baseline after the image moved off dense FA3.

Purpose

Before #7735 (ebe2c1827), CUDA candidate 2 imported fa3_fwd_interface and set both flash_attn_func and flash_attn_varlen_func. Dense, unmasked attention in FlashAttentionImpl.forward_cuda called flash_attn_func directly.

#7735 removed fa3-fwd from requirements/cuda.txt and replaced that candidate with vLLM's flash_attn_varlen_func only. flash_attn_func stayed None, so dense attention always went through _forward_varlen_dense. vLLM's dispatcher defaults to FA2 unless fa_version is passed. #9100 (vllm_varlen_version_kwargs) fixes that version on the varlen call site (FA3 on SM90). It does not restore a dense FA3 kernel. On Qwen-Image the varlen FA3 path is still far from the old dense FA3 number.

This PR inserts one optional bind, _bind_optional_fa3_fwd, after the vLLM varlen candidate and before standalone flash_attn_interface / FA2, and only when the device is not Blackwell:

• flash_attn_func already set (FA4, AITER, MUSA, or an earlier candidate) is left alone.
• fa3_fwd_interface importable: both dense and varlen callables come from that package, matching pre-[Kernel] Remove fa3-fwd dependency #7735 candidate 2.
• import fails: bindings are unchanged, so [Bugfix][Diffusion] Pass the resolved FA version into dense varlen attention #9100 varlen FA3 remains the path.
• fa3-fwd is not added back as a hard dependency.

No user guide or example changes. Backend selection is unchanged. The only new behavior is "use the optional dense FA3 package when it is already installed."

CI images that do not install fa3-fwd will keep the post-#9100 varlen FA3 path. This PR does not by itself move an H100 nightly that lacks the package back to the old JSON baseline.

What we measured before the fix

Same host (NVIDIA L20X, SM90), same probe test_qwen_image_single_device / 512x512_steps20 (512x512, 20 steps, 10 prompts, concurrency 1), model Qwen/Qwen-Image from the local hub snapshot.

tree	flash_attn_func	path	latency_mean
#7735 parent 4ec9e5a31	fa3_fwd_interface	dense FA3	2.2832 s
#7735 head ebe2c1827	None	vLLM varlen, default FA2	3.3190 s
#9100 head 6839c4d16	None	vLLM varlen, fa_version=3	3.5379 s
main 15bce3458 with fa3-fwd==0.0.3 installed but not selected	None	vLLM varlen, fa_version=3	3.4534 s

Installing the wheel is not enough. Main never imports it, so the package sitting in the environment does not change the path.

Test Plan

vLLM Version: 0.32.0 (torch 2.13.0+cu130, CUDA 13.0)

vLLM-Omni Commit: 0997e4b08 on base 15bce3458

Hardware: 1x NVIDIA L20X (SM90). Optional package: fa3-fwd==0.0.3.

• pre-commit on the two changed files (ruff, mypy 3.10, SPDX, forbidden imports, test marks)
• CPU contract tests for _bind_optional_fa3_fwd (package present, package absent, existing dense callable left alone)
• CPU tests/diffusion/attention/test_flash_attn.py -m "core_model and cpu" on this commit
• L20X test_qwen_image_single_device / 512x512_steps20 on this commit with fa3-fwd installed
• Same probe on unmodified 15bce3458 with fa3-fwd installed, to show the wheel alone does not recover latency
• Same probe on this commit after uninstalling fa3-fwd, to show the varlen FA3 fallback still runs

Test Result

Passed: optional bind

With fa3-fwd==0.0.3 importable, this commit binds:

• flash_attn_func = fa3_fwd_interface.flash_attn_func
• flash_attn_varlen_func = fa3_fwd_interface.flash_attn_varlen_func

With the package uninstalled, flash_attn_func stays None, varlen stays vllm.vllm_flash_attn, and vllm_varlen_version_kwargs returns {"fa_version": 3} on this SM90 host.

requirements/cuda.txt does not mention fa3-fwd.

Passed: CPU tests

pytest tests/diffusion/attention/test_flash_attn.py -m "core_model and cpu" -o addopts= -p no:xdist

58 passed, 7 deselected in 0.97s. The three new tests are in that run:

• test_bind_optional_fa3_fwd_sets_dense_when_package_present
• test_bind_optional_fa3_fwd_keeps_vllm_varlen_when_package_absent
• test_bind_optional_fa3_fwd_does_not_override_existing_dense

Passed: L20X Qwen-Image 512, same machine

Command shape (working tree, no checkout of the dirty fix during the "with fix" run):

pytest -s -o addopts= -p no:xdist \
  tests/dfx/perf/scripts/run_diffusion_benchmark.py \
  --test-config-file /tmp/9092_bisect/single_device_512.json \
  -k test_qwen_image_single_device

Config is the upstream test_qwen_image_single_device entry restricted to 512x512_steps20, model path pointed at the local Qwen-Image snapshot. num-prompts=10, max-concurrency=1, 20 steps, negative prompt enabled.

label	commit	fa3-fwd	latency_mean	throughput_qps	QwenImagePipeline.diffuse
#7735 parent (reference)	4ec9e5a31	selected by old candidate 2	2.2832 s	0.4379	2.0902 s
main, package installed, code does not bind it	15bce3458	installed, unused	3.4534 s	0.2896	3.2609 s
this PR, package installed	0997e4b08	bound	2.2444 s	0.4455	2.0545 s
this PR, package uninstalled	0997e4b08	absent	3.4906 s	0.2865	3.2898 s

With the package, latency is 35% lower than the same-tree control (3.4534 s -> 2.2444 s) and sits next to the #7735 parent (2.2832 s). Without the package the server still comes up and the probe finishes on varlen FA3, in the same slow band as current main.

Not run

• H100 / B200 JSON baseline in test_qwen_image_vllm_omni.json. No H100 or B200 on this host. An image that does not install fa3-fwd will not pick up the dense path.
• 1536x1536_steps35, 4-GPU, and cache_dit. [Bug][Perf]: qwen-image, performance metrics regressed by more than 10% compared to the baseline in some scenarios #9092 CI moved those the same direction as single-device 512; this validation stays on the single-device probe that isolated [Kernel] Remove fa3-fwd dependency #7735.
• Pixel / image-quality diff against the pre-[Kernel] Remove fa3-fwd dependency #7735 outputs. The check here is latency and that the fallback does not crash.

[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-21T06:38:38.706199Z We couldn’t complete this review. Comment @codex review to try again.	0997e4b	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.

[NumberWan]

@amy-why-3459 PTAL

[NumberWan]

@natureofnature PTAL

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

@NumberWan, 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.

[amy-why-3459]

Thanks for the careful bisect, the before/after table makes the regression and the fix easy to follow. Keeping fa3-fwd optional and leaving the #9100 fallback intact makes sense.

Two things I'd like to see tightened before this lands, both about the scope of the new bind:

1. The bind runs on every non-Blackwell CUDA device, but fa3_fwd_interface only ships Hopper (SM90) kernels. On SM80/SM89 with the wheel present (e.g. a shared image), we'd now route dense attention to a kernel that can't launch, instead of vLLM varlen FA2 as today.
2. It also swaps flash_attn_varlen_func, which moves the masked, packed and piecewise paths (and the FP8 path, via flash_attn_func) off vLLM's dispatcher. The measurements only cover the dense unmasked Qwen-Image case.

Details inline. Minor: an info log when fa3_fwd_interface gets bound (as the FA4 branch does) would help, since the PR body itself shows that whether the wheel is actually used isn't obvious from the outside.

[natureofnature]

Reviewed both changed files, provider discovery, dense/masked/packed/piecewise and FP8 call sites, and the reported current-head tests. The latency recovery is useful evidence for the measured Qwen-Image BF16 workload. I left one FP16 compatibility finding inline and followed up in the existing provider-validation discussion; wheel/build checks passed, but numerical coverage of the newly selected routes is still needed.

【Conclusion: Found blocking code issues and validation gaps in commit 0997e4b.】
【Action recommended: Preserve working FP16 attention with the wheel installed, and verify BF16 dense/masked/packed outputs plus Hopper FP8 accuracy before merging.】

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit ffd87f4d9a20 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

@NumberWan this PR reads as a performance or value claim:

• claim: With the package, latency is 35% lower than the same-tree control (3.4534 s -> 2.2444 s) and sits next to the [Kernel] Remove fa3-fwd dependency #7735 parent (2.2832 s).

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.

[NumberWan]

Author self-review

Checked on ffd87f4d9:

• Optional fa3-fwd is Hopper-only (_is_hopper), dense-only (optional_fa3_fwd_dense_func), and FP16-safe (falls back to vLLM varlen).
• Masked / packed / piecewise / FP8 still use flash_attn_varlen_func from vLLM, so [Bugfix][Diffusion] Pass the resolved FA version into dense varlen attention #9100 fa_version routing is unchanged for those paths.
• requirements/cuda.txt still has no fa3-fwd hard dependency.
• CPU tests/diffusion/attention/test_flash_attn.py -m "core_model and cpu": 63 passed.
• Re-running the L20X 512x512_steps20 probe on this narrowed commit to confirm latency recovery still holds.

[NumberWan]

Follow-up: L20X test_qwen_image_single_device / 512x512_steps20 on ffd87f4d9 with fa3-fwd==0.0.3 still recovers latency:

label	commit	latency_mean	throughput_qps
narrowed optional dense	ffd87f4d9	2.1625 s	0.4624
previous (broader bind)	0997e4b08	2.2444 s	0.4455
#7735 parent (ref)	4ec9e5a31	2.2832 s	0.4379

Binding check on that run: optional_fa3_fwd_dense_func from fa3_fwd_interface, flash_attn_func is None, varlen still vllm.vllm_flash_attn, BF16 selects optional dense, FP16 selects None (varlen fallback).

[NumberWan]

@natureofnature @amy-why-3459 PTAL

[natureofnature]

Rechecked all three changed files and the dense/masked/packed/piecewise/FP8 call paths. The FP16 finding is fixed: the optional callable is separate, and FP16 falls back to the existing implementation. The broader provider changes are also removed, so my previous masked/packed/FP8 validation request no longer applies to this revision. The reported 63 CPU passes and current-head L20X latency run support the routing fix and performance recovery.

I suggest verifying BF16 unmasked dense outputs against SDPA on this head with fa3-fwd==0.0.3 installed on Hopper before merging, including the Qwen-Image attention geometry and confirming that optional_fa3_fwd_dense_func is actually selected. Outputs should be finite and agree within the existing attention tests' BF16 tolerances. This is the remaining provider change; the latency run does not check numerical equivalence, and the direct vLLM adapter test bypasses this selector.

【Conclusion: Static review found no blocking code issues in commit ffd87f4.】
【Action recommended: Verify the selected Hopper BF16 dense path against SDPA within the existing BF16 tolerances before merging.】

[NumberWan]

@natureofnature Verified the selected Hopper BF16 unmasked dense path against SDPA on ffd87f4d9 before merge.

Host: NVIDIA L20X, SM90 (9, 0), fa3-fwd==0.0.3, vLLM 0.32.0.
Path: FlashAttentionImpl.forward with attn_metadata=None (unmasked dense). A wrapper on optional_fa3_fwd_dense_func counted one call per shape, so the selector actually ran. flash_attn_func stayed None; varlen stayed vllm.vllm_flash_attn.

Qwen-Image geometry from transformer/config.json: 24 heads, head dim 128. 512px image tokens are (64/2)*(64/2)=1024; joint adds a representative 256 text tokens.

Tolerance is the existing BF16 bound in test_vllm_dense_adapter_cuda_matches_sdpa_and_lse: rtol=1e-2, atol=1e-2. Outputs were finite.

shape	B, S, H, D	max abs	mean abs	assert_close
512 image tokens	1, 1024, 24, 128	0.001953	0.000027	pass
512 joint (image+text)	1, 1280, 24, 128	0.001953	0.000023	pass

Both also sit inside the tighter CUDA bounds used by test_fa_vs_sdpa (max < 0.01, mean < 0.001).