[dougbtv]

Purpose

Remove the fa3-fwd package dependency and use the FlashAttention implementation already maintained and shipped by vLLM.

This is the full-removal alternative to #7671, following the direction explored in #3340. If this approach lands, the packaging-only fallback in #7671 is unnecessary.

vLLM-Omni needed the separate package because its ring-attention path requires each block's log-sum-exp (LSE) for numerically stable accumulation and called fa3_fwd_interface._flash_attn_forward directly. Core vLLM does not need that package: its maintained vllm.vllm_flash_attn dispatcher already ships compatible FA2/FA3/FA4 kernels and can return LSE from its public varlen wrapper.

Changes

• remove fa3-fwd==0.0.3 from the CUDA requirements;
• remove all fa3_fwd_interface discovery and architecture-version handling;
• prefer vllm.vllm_flash_attn for ordinary CUDA FlashAttention dispatch;
• add a dense-to-varlen adapter that preserves ring attention's (B, H, S) LSE contract;
• prefer that vLLM-backed path for CUDA ring attention; and
• retain source-built FA3, standalone FA2, FA4, AITER, and SDPA as platform-specific fallbacks.

The adapter flattens each dense batch into packed sequences, constructs cumulative sequence lengths, calls vLLM with return_softmax_lse=True, and restores both output and LSE layouts before ring accumulation.

Test Plan

vLLM Version: v0.29.0 (the CUDA image base)

vLLM-Omni base commit: eb1abbadf832c625a9560cb6e0142d2cec4998ce

• Add an L1 CPU test for dense-to-varlen metadata and output/LSE layout restoration.
• Add an L1 CPU test proving ring dispatch prefers the vLLM-backed path over standalone FA fallbacks.
• Run the repository's changed-file pre-commit suite and git diff --check.
• Run CUDA ring-attention coverage on Ampere/Hopper/Blackwell in CI or maintainer infrastructure.

Test Result

• Changed-file pre-commit suite passed, including Ruff, formatting, mypy, SPDX, forbidden-import, test-marker, and new-torch.cuda checks.
• git diff --check passed.
• Python byte-compilation passed for the changed Python modules and tests.
• The new L1 tests were not executed locally because this workspace does not have PyTorch or vLLM installed.
• CUDA ring execution and performance have not been measured locally; those remain explicit CI/reviewer validation items.

AI assistance: Codex helped trace the dependency and ring-LSE contracts, implement the vLLM-backed adapter and selection path, add focused tests, run static validation, and draft this description. I reviewed the resulting change and evidence.

[dougbtv]

Author self-review:

• Confirmed the branch is one signed commit on current upstream main (eb1abbadf832c625a9560cb6e0142d2cec4998ce).
• The dependency and every fa3_fwd_interface import are removed; optional source-built FA3 and other platform fallbacks remain.
• Ring attention no longer reaches into a private FA3 symbol. Its vLLM adapter uses the maintained varlen API with LSE enabled and restores the existing (B, H, S) accumulation contract.
• Added CPU contract tests for packed metadata/layout restoration and vLLM-first ring selection.
• Changed-file pre-commit, mypy, Ruff, SPDX, marker checks, byte-compilation, and git diff --check pass.
• I could not execute pytest or CUDA ring kernels locally because this host lacks PyTorch/vLLM and a GPU environment. CUDA correctness and performance remain the main review/CI risk.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/index.md.

Module owners: @princepride @david6666666 @Isotr0py

Routing: @princepride via module of the changed files, semantic router, CODEOWNERS; @david6666666 via module of the changed files, semantic router; @Isotr0py via module of the changed files, semantic router

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

[dougbtv]

Rebased onto current main (62b15142bdab8c80eb83473f65e51774d1a2521b). The new PR head is 601a32a486f3acf9f0aca0afdc9891dfaa7e0dda; the stable patch ID is unchanged from the previously reviewed/tested diff, and the changed-file pre-commit suite passes.

I also exercised the previous exact head (5ed7f6636a7275881110605bb33581aba055d02a) as a built CUDA image using the repository's docker/Dockerfile.ci, then ran the two new dependency-removal contract tests inside that image:

• test_vllm_dense_adapter_builds_varlen_metadata_and_restores_lse
• test_ring_prefers_vllm_flash_attention

Result: 2 passed. The image build/publish completed successfully, and the test run produced JUnit results successfully.

• Image build: https://github.com/neuralmagic/nm-cicd/actions/runs/35257024140
• Focused test run: https://github.com/neuralmagic/nm-cicd/actions/runs/35618690303

This proves that the source image builds without the removed requirement and that the new adapter-layout and backend-selection contracts execute successfully in the packaged dependency environment. These two tests are CPU/monkeypatch contract tests; they do not yet claim real CUDA-kernel or distributed-ring execution. I am treating real FA2/FA3 dispatcher execution plus ring/LSE numerical correctness as the remaining mechanical validation, not performance benchmarking.

[dougbtv]

Mechanical CUDA validation is complete on the current PR head a7137796.

We built a CUDA image directly from this PR's source and ran the focused attention tests below against that image.

Single-GPU validation

The same focused suite passed on both A100 (SM80 / FA2) and H100 (SM90 / FA3):

• test_vllm_dense_adapter_cuda_matches_sdpa_and_lse
• test_padding_equivalence
• test_fa_vs_sdpa
• test_cross_attn_key_padding_vs_sdpa (both parameterizations)

Results:

• A100: vLLM selected FA2; 5 passed.
• H100: vLLM selected FA3; 5 passed.

This exercises the real vLLM FlashAttention dispatcher and checks the new dense-to-varlen adapter's output and LSE against the PyTorch reference, alongside the existing CUDA attention parity cases.

Distributed ring validation

On 2× H100 with NCCL, test_vllm_flash_attention_ring_cuda_matches_full_attention passed. This executes the real two-rank RingComm path with AttnType.VLLM_FA and compares each rank's output and accumulated (B, H, S) LSE with full-attention reference results.

The first version of this new test had a test-only layout mistake: its reference was transposed to (B, S, H), while ring_flash_attn_forward intentionally preserves the established (B, H, S) contract. The current head corrects that expectation; the rebuilt image and full focused matrix passed afterward.

Dependency-removal check

In every test environment we also verified that:

• the fa3-fwd distribution is not installed; and
• neither fa3_fwd nor fa3_fwd_interface is importable.

So the scoped mechanical removal is covered: the old package is absent, real FA2 and FA3 dispatch work, dense output/LSE parity passes, and two-GPU ring output/LSE parity passes.

For Test Approve, the hardware-specific cases to watch are the CUDA adapter/parity tests on A100 and H100, plus tests/diffusion/attention/test_ring_vllm_flash_attn.py on a two-GPU H100 runner. Performance comparison is intentionally outside this change's scope.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 7 days

@dougbtv this pull request has had no human commit, comment or review since 2026-09-21. 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.

[Dong1017]

I reviewed head a713779 against base 62b1514.

The dense vLLM FlashAttention adapter and ring-attention LSE accumulation passed my targeted checks: the dense output/LSE comparison passed on one GPU, and the two-GPU ring output comparison passed. I also ran a fixed-head A/B/A Qwen-Image request on two NVIDIA L20X GPUs (1024×1024, 8 steps, ring degree 2). The median was 1.578 s on this head versus 1.365 s and 1.397 s for the two base runs. This workload therefore did not show a request-level speedup; the head was about 14% slower than the mean of the two base medians.

These measurements are specific to this device and workload and do not invalidate the reported A100/H100 results. Please keep the performance claim scoped to the workloads where it was measured. GitHub currently reports this PR as CONFLICTING; please update the branch before merge so the final diff and required checks apply to the mergeable version.

[dougbtv]

Thanks for the focused validation and the device/workload-specific result. I have rebased the PR onto current main (head b1c608df), so GitHub now reports it mergeable and the required checks are running again.\n\nThe PR is intentionally scoped to removing the fa3-fwd dependency and private interface while preserving the attention/ring correctness contract; it is not presented as a performance improvement. I will keep any performance discussion limited to the measured workloads.\n\nThanks again for reviewing it—once the refreshed checks settle, I’d appreciate another look toward approval.

[Dong1017]

Reviewed b1c608df561684fbeea18a0fb97bbd9c364487fa. I found no blocking issue in the dependency removal, vLLM FlashAttention dispatch, or dense-to-varlen LSE adapter. The dense adapter passed on H200, and a two-rank H200/NCCL check of the production ring path matched the output and LSE references. The original ring pytest could not collect with the available vLLM 0.30.0 environment because its helper imports vllm.v1.kv_hints; the production path was verified separately with a saved harness.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on zcode (GLM-5.3-Flash) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 5a1071e7-0d2e-4533-9cbb-fe1102492902) — check zcode login and the CLI version; retrying strict/zcode/GLM-5.3-Flash in 120s (try ).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 9cbe961d-11d6-4d98-ad50-598b571ec307) — check zcode login and the CLI version; retrying strict/zcode/GLM-5.3-Flash in 600s (try ).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: bd886b8a-b34f-45fa-9a87-cab3f9342fed) — check zcode login and the CLI version; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

2 actionable finding(s).

CI at b1c608df5616 (2026-10-10T08:03:51.510085+00:00): verification incomplete; required-check status is unknown. Observed Buildkite: buildkite/vllm-omni (passed), buildkite/vllm-omni-amd-ci (failed), buildkite/vllm-omni-npu-ci (failed), and 1 more.

Note: The assigned review arm strict/zcode/GLM-5.3-Flash could not complete this review, so it was produced by the fallback arm direct/cursor/auto. It is excluded from the routing experiment.

Full review analysis

PR description

This change removes the separate fa3-fwd CUDA requirement and uses the FlashAttention dispatcher already shipped in vLLM. CUDA discovery now binds vLLM's varlen kernel before source-built FA3 or standalone FA2, while optional FA4, AITER, and SDPA stay as platform fallbacks. Ring attention prefers that path: each dense block is packed into varlen sequences, return_softmax_lse=True is set, and the log-sum-exp is restored to (batch, heads, sequence) before the existing ring accumulation.

Change flow

flowchart LR
  A["[REMOVED] fa3-fwd CUDA pin"]:::removed --> B["[CHANGED] CUDA FA discovery"]:::changed
  B --> C["[NEW] dense-to-varlen LSE adapter"]:::new
  C --> D["[CHANGED] ring prefers VLLM_FA"]:::changed
  D --> E["[EXISTING] ring (B, H, S) accumulation"]:::existing
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

See inline comments below.

---

_{🤖 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]

[P2] Default CUDA dense attention drops fa_deterministic

Evidence and suggested fix

On a normal CUDA build, discovery now imports vllm.vllm_flash_attn.flash_attn_varlen_func and stops. It never assigns flash_attn_func, and the later source-FA3/FA2 candidates run only when that varlen symbol is still missing. FlashAttentionImpl.forward_cuda passes deterministic=True only when flash_attn_func is set; otherwise unmasked dense attention falls through to _forward_varlen_dense, which omits deterministic and emits fa_deterministic=True is ignored on the dense-varlen-fallback path. Previously the pinned fa3-fwd package filled flash_attn_func, so this opt-in reached the dense kernel. The new vllm_flash_attn_varlen_with_lse helper already forwards deterministic, but the ordinary backend does not call it. Forward self.fa_deterministic on the dense varlen call, or bind a dense wrapper that still honors the flag.

[vllm-omni-review-bot]

[P2] L1 layout test is not selected by the CPU diffusion job

Evidence and suggested fix

test_vllm_dense_adapter_builds_varlen_metadata_and_restores_lse has no cpu marker, unlike the version-selection tests above it. The ready-pipeline step Simple · Diffusion Test runs pytest -sv tests/diffusion -m 'core_model and cpu', so this layout contract is not collected there. The single-GPU numerical test is marked for one CUDA card and can be selected by Diffusion · FlashAttention Test (tests/diffusion/attention/ -m 'core_model and cuda and L4 and cards_1'). The two-GPU ring test requests num_cards=2; that same attention step requires cards_1, and the materialized cards_2 steps name test_ulysses_uaa.py, test_attention_sp.py, or tests/diffusion/distributed/ rather than test_ring_vllm_flash_attn.py. tests/helpers/mark.py is not in this checkout, so the exact generated markers cannot be confirmed. Mark the CPU test cpu, and point a two-card attention step at the ring parity test.