[Feature][Diffusion] Plugin-based sparse attention interface for DiT modules by zzhang-fr · Pull Request #2231 · vllm-project/vllm-omni

zzhang-fr · 2026-03-26T11:38:30Z

Purpose

Build a unified plugin-based sparse attention interface for DiT modules by
extending the existing AttentionBackend infrastructure. Resolves #1568.

Changes in this revision

SparseAttentionBackend now inherits AttentionBackend; _SparseAttentionImpl
inherits AttentionImpl. SPARSE_ATTENTION is a single line in
DiffusionAttentionBackendEnum — the independent registry is deleted.
contrib/ removed. Reference implementations live in external repos.
Sparse is strictly opt-in via --attn-backend sparse_attention or a plugin
name directly. No auto-activation.
Cross-attention isolated via is_self_attention=False in WanCrossAttention;
enable_sparse_attention() hook removed from WanTransformer3DModel.
entry_points fallback in selector.py: --attn-backend spargeattn works
directly once a plugin is installed — no --sparse-attn-backend needed.
DiffusionSparseAttnConfig simplified to backend + params: dict —
vllm-omni does not interpret any kernel-level parameter.

Design

SparseAttentionBackend(AttentionBackend) is registered as SPARSE_ATTENTION
in DiffusionAttentionBackendEnum. All kernel implementations live in external
pip-installable packages — zero vendor backends in vllm-omni core.

Core changes (vllm-omni)

vllm_omni/diffusion/attention/backends/sparse_attention.py — new file:
SparseAttentionBackend(AttentionBackend) + _SparseAttentionImpl(AttentionImpl)
SPARSE_ATTENTION added to DiffusionAttentionBackendEnum
is_self_attention: bool = True param on Attention.__init__ — cross-attention
passes is_self_attention=False, permanently dense, zero overhead
entry_points fallback _get_plugin_backend_cls in selector.py
WanCrossAttention marked is_self_attention=False; hook removed
DiffusionSparseAttnConfig fields: backend: str + params: dict

Plugin interface

Plugins register a function via entry_points — not a class:

[project.entry-points."vllm_omni.sparse_attn"]
spargeattn = "sparge_vllm_omni:sparge_attn_fn"

Function signature: fn(q, k, v, params: dict, is_causal: bool) -> Tensor

The params dict is forwarded verbatim from DiffusionSparseAttnConfig.params.
vllm-omni does not interpret any kernel-level parameter — the plugin reads what it needs.

Configuration

DiffusionSparseAttnConfig accepts flat or nested JSON:

# Flat JSON — unknown keys auto-routed into params:
DiffusionSparseAttnConfig.from_dict({"backend": "spargeattn", "topk": 0.5})
# → DiffusionSparseAttnConfig(backend="spargeattn", params={"topk": 0.5})

# Nested JSON:
DiffusionSparseAttnConfig.from_dict({"backend": "spargeattn", "params": {"topk": 0.5}})

# Programmatic:
DiffusionSparseAttnConfig(backend="spargeattn", params={"topk": 0.5})

Usage

pip install sparge-vllm-omni        # external package, not in this repo
vllm-omni serve Wan2.2 --attn-backend spargeattn

# With custom params via JSON (flat form):
python wan2_2_sparse.py --attn-backend '{"backend":"spargeattn","topk":0.3}'

Test Result

Backend	Total	Per-step	Speedup
Dense	785s	19.6s	1.00×
SpargeAttn topk=0.7	763s	19.1s	1.03×
SpargeAttn topk=0.5	677s	16.9s	1.16×
SpargeAttn topk=0.3	590s	14.8s	1.33×

Visual quality samples (dense vs topk=0.5 vs topk=0.3):
https://github.com/user-attachments/assets/62a7abfd-758f-4ec0-958b-fa3de672f972
https://github.com/user-attachments/assets/016c8463-f0bc-4a69-a891-b07fbc784a68
https://github.com/user-attachments/assets/17d2d464-0205-4d2c-8683-bd2e768c11f3

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 058884a921

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

lishunyang12 · 2026-03-27T05:40:31Z

I think you should wait for feedback from maintainers to disssus about value of this feature and maintainability before opening a PR. This is common practice as attn backend bas been altered. Otherwise, it is quite hard to be accepted.

gcanlin

I don't think we should introduce contrib/ directory and we have implemented the registration for Attention backend. This PR is introducing one independent registry for sparse attention. And it's trying to make sparse default when users have installed it. But in most of scenarios, sparse attention could lead to accuracy regression. So it's better to make it optional instead of default.

Could we make SparseAttentionBackend inherit AttentionBackend?

zzhang-fr · 2026-03-27T21:45:40Z

I don't think we should introduce contrib/ directory and we have implemented the registration for Attention backend. This PR is introducing one independent registry for sparse attention. And it's trying to make sparse default when users have installed it. But in most of scenarios, sparse attention could lead to accuracy regression. So it's better to make it optional instead of default.

Could we make SparseAttentionBackend inherit AttentionBackend?

@gcanlin Thanks for the detailed review — these are all valid points and I want to engage with them carefully.

On inheriting AttentionBackend

This is architecturally sound and I've prototyped it. The inheritance chain works:

SparseAttentionBackend(AttentionBackend)
SparseAttentionImpl(AttentionImpl[SparseAttentionMetadata])
SparseAttentionMetadata(AttentionMetadata)

VIDEO_SPARSE_ATTN becomes one line in DiffusionAttentionBackendEnum, the independent
registry is deleted, and --sparse-attn-backend is replaced by --attn-backend.

However, inheriting introduces a new problem the current design avoids

--attn-backend is a global setting. get_attn_backend() returns the same backend class
for every Attention instance in the model — including WanCrossAttention.attn (line 531
in wan2_2_transformer.py).

Cross-attention (Q from patch tokens, KV from text encoder) must remain fully dense.
Applying sparse attention there is semantically wrong and will degrade quality.

The current per-model enable_sparse_attention() hook exists precisely because it
targets only self-attention layers, leaving cross-attention untouched. With the
inheritance approach, we need an equivalent mechanism to achieve the same protection —
for example, passing a module_type parameter to every Attention(...) constructor call
and checking it against the backend's supported_module_types().

The complexity is equivalent; it just moves from one place to another:

Current PR: one enable_sparse_attention() per model (gcanlin's concern)
Inheritance: one module_type= per Attention() call (same surface, different form)

My proposal

I'm happy to go either direction, but I'd like to resolve this in RFC #2233 before
reworking the implementation. Specifically:

Should we go with inheritance + module_type parameter?
Or keep the current separate interface but address the other concerns
(remove independent registry, remove --sparse-attn-backend, make sparse
strictly opt-in via explicit config)?

On "sparse should not be default"

Completely agreed. The _auto_select() behavior in registry.py is wrong —
it activates sparse just because a plugin is installed. I'll fix this regardless
of which architectural direction we take: sparse attention must require explicit
opt-in via --attn-backend video_sparse_attn (or equivalent), never auto-activated.

On contrib/

Will remove from this PR. Reference implementations will move to separate repos.

Happy to hear from @tzhouam @hsliuustc0106 @SamitHuang on the module_type question
before I revise the implementation.

zzhang-fr · 2026-04-01T00:47:30Z

@gcanlin Thanks for the detailed review. This revision addresses all four points:

contrib/ removed. Reference implementations will live in external repos.
- SparseAttentionBackend now inherits AttentionBackend; _SparseAttentionImpl inherits AttentionImpl. SPARSE_ATTENTION is a single line in DiffusionAttentionBackendEnum — the independent registry is deleted.
- Sparse is strictly opt-in: --attn-backend sparse_attention (or a plugin name directly). No auto-activation.
- Cross-attention is isolated via is_self_attention=False in WanCrossAttention; enable_sparse_attention() hook removed entirely.
  The plugin interface is simplified — plugins register a function fn(q, k, v, topk, is_causal) -> Tensor via vllm_omni.sparse_attn entry_points.

@alex-jw-brooks The new implementation inherits the existing AttentionBackend and AttentionImpl — sparse attention is now just another entry in DiffusionAttentionBackendEnum. The enable_sparse_attention() per-model hook is removed; cross-attention is protected via a single is_self_attention=False parameter on WanCrossAttention (1 line). The broader design question is still open in RFC #2233 — happy to continue the discussion there.

@lishunyang12 This revision was made in direct response to maintainer feedback from this PR. The architecture now follows the existing AttentionBackend pattern as gcanlin suggested. RFC #2233 remains open for the broader design discussion.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 51fc4cec07

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

lishunyang12

Left a few comments. Overall direction looks good

lishunyang12 · 2026-04-02T03:11:49Z

+            for ep in importlib.metadata.entry_points(group="vllm_omni.sparse_attn"):
+                if ep.name.lower() == name.lower():
+                    return ep.load()
+        except Exception:


If ep.load() raises (e.g. plugin exists but has an ImportError), this silently returns None — user sees "falling back to dense" with no clue the plugin was found but broken. Please distinguish not-found vs load-failure, e.g. except Exception as e: logger.warning("Failed to load sparse attn plugin %s: %s", name, e).

Fixed — ep.load() failure now logs logger.warning(..., exc_info=True) with the plugin name and entry_point value, distinct from the "not found" path which returns None silently.

lishunyang12 · 2026-04-02T03:11:49Z

+            from vllm_omni.diffusion.forward_context import get_forward_context
+
+            return get_forward_context().omni_diffusion_config.sparse_attn
+        except Exception:


Bare except Exception masks real bugs (e.g. AttributeError from misconfigured config) the same way it handles the expected "no forward context" case. Narrow the catch or add logger.debug so users know when sparse attention silently fell back to dense.

Fixed — narrowed to except (RuntimeError, LookupError, AssertionError) for the "not yet initialised" case, with a separate except AttributeError: logger.debug(...) for misconfigured config.

lishunyang12 · 2026-04-02T03:11:49Z

+                    normalized_kwargs["sparse_attn"] = merged
+            except json.JSONDecodeError:
+                logger.warning("Invalid --attn-backend JSON: %s", attn_backend_raw)
+        elif attn_backend_raw:


Plain names like flash_attn or sage_attn also hit this branch and get routed into sparse_attn → wrong dispatcher. Known dense backend names should be excluded before this elif.

Fixed — known dense backends are now excluded via _dense_names (built dynamically from DiffusionAttentionBackendEnum minus SPARSE_ATTENTION) before routing a plain name to the sparse_attn dispatcher.

lishunyang12

still waiting on the earlier threads. one more nit.

lishunyang12 · 2026-04-02T15:44:58Z

+        # Handle --attn-backend: accepts simple name or JSON with config.
+        #   --attn-backend spargeattn
+        #   --attn-backend '{"backend":"spargeattn","topk":0.3}'
+        import json


Nit: move import json to the module-level imports.

Fixed — moved to module-level imports (line 12).

zzhang-fr · 2026-04-07T14:52:17Z

@gcanlin @alex-jw-brooks @SamitHuang @lishunyang12 @hsliuustc0106

gcanlin · 2026-04-08T01:18:04Z



+@dataclass
+class DiffusionSparseAttnConfig:


Could we create DiffusionAttentionConfig first? Actually, this PR is doing two things implementing sparse attention and convert the current attention env way to CLI way. So after we add --attn-backend, we need to refactor the original attention backends. It's better to separate them into 2 PRs.

I didn't think so.

DiffusionAttentionConfig would make sense when there are shared fields across multiple attention backends — for example, per-layer backend selection or step-level scheduling. Right now neither exists: FLASH_ATTN, SAGE_ATTN, and TORCH_SDPA have no kernel-level parameters, and DiffusionSparseAttnConfig is the only config type needed.

OmniDiffusionConfig already has attention_backend: str (which backend to use) and sparse_attn: DiffusionSparseAttnConfig (sparse-specific params) as two separate flat fields. Wrapping them into DiffusionAttentionConfig now would produce an empty shell with no new capability.

On splitting into 2 PRs: the CLI changes (--attn-backend, JSON routing in async_omni_engine.py, _dense_names exclusion) are tightly coupled to the sparse attention plugin mechanism — they exist specifically to route plugin names into DiffusionSparseAttnConfig. A PR1 without the sparse attention backend would have no consumers and no testable behavior.

When a second attention type needs config params (e.g. a quantized FlashAttn with a head_dim_override), introducing DiffusionAttentionConfig as a common wrapper would be the right time. Happy to open a follow-up RFC for that if it's the preferred direction.

does vllm upstream provide sparseattnconfig for llm?

@gcanlin

RFC #2632 makes the right call — I'm happy to extend this PR to cover P1–P3.

Concretely, the changes needed are minimal:

is_self_attention: bool → AttentionRole(SELF | CROSS) enum

DiffusionSparseAttnConfig(backend, params) → AttentionSpec(backend, extra)

Add AttentionConfig(default, per_role) to replace the flat
attention_backend + sparse_attn fields on OmniDiffusionConfig

Make get_attn_backend() role-aware (cache key fix included)

Add extra: dict to AttentionMetadata

This PR already covers P4 (first sparse backend as end-to-end validation).
I'll push the P1–P3 changes in a follow-up commit — should keep the diff
to ~150 additional lines across the affected files.

@zzhang-fr Would appreciate it if you could open a new PR to implement #2632. Let this PR #2231 focus on the interface definition of sparse attention backend. What do you think?

Agreed on the split. Will keep this PR focused on the sparse attention plugin interface.

For the attention backend refactoring, I'll study the design and ongoing discussion in #2632 first before opening a new PR.

does vllm upstream provide sparseattnconfig for llm?

@hsliuustc0106 Upstream doesn't have a sparse attention config, but it does have AttentionConfig as a structured config class (vllm/config/attention.py) exposed via EngineArgs.attention_config (L575 in arg_utils.py). It holds backend: AttentionBackendEnum plus backend-specific flags like flash_attn_version and use_cudnn_prefill.
The gap: upstream's AttentionConfig is flat — single global backend, no per-role selection, no opaque extra dict. #2632 follows the same pattern but extends it with per-role (self/cross) backend selection for diffusion.

gcanlin · 2026-04-08T01:23:16Z

+            return None
+
+    @staticmethod
+    def _resolve_plugin(cfg: DiffusionSparseAttnConfig) -> SparseAttnFn | None:


Is it necessary to introduce plugin for attention backend? If developers want to integrate their new sparse attention backend, just adding a new enum in DiffusionAttentionBackendEnum would be fine?

Good question. For dense backends (FLASH_ATTN, SAGE_ATTN, etc.), adding an enum entry is exactly the right approach — they are in-tree and vllm-omni owns their maintenance.

Sparse attention is different: the maintainers explicitly don't want to take on the burden of maintaining sparse kernel backends inside vllm-omni. Each kernel (SpargeAttn, RainFusion, FlashInfer-sparse, MindIE-SD) has its own release cycle, hardware requirements, and API surface. Pulling them in-tree means vllm-omni has to track upstream changes, handle import errors, and gate releases on kernel compatibility.

The plugin system delegates that responsibility to the kernel authors. They ship sparge-vllm-omni, users pip install it, done — no PR to vllm-omni core required.

SPARSE_ATTENTION in the enum is the single in-tree entry point — it's the dispatcher. The actual kernels live outside.

What do you think?

I suggest to provide simplest and easy-to-extend abstraction for sparse attn since hardware differs a lot here

Adds sparse attention support for DiT video models through the existing AttentionBackend plugin system — no new registry needed. Usage: # Via CLI (serve) vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --attn-backend spargeattn # Via CLI with JSON config vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --attn-backend '{"backend":"spargeattn","topk_ratio":0.3}' # Via Python API from vllm_omni.diffusion.data import DiffusionSparseAttnConfig omni = Omni(model=..., sparse_attn=DiffusionSparseAttnConfig( backend="spargeattn", topk_ratio=0.5)) How it works: - SparseAttentionBackend is an AttentionBackend dispatcher that resolves sparse plugins (e.g. SpargeAttn) via entry_points at init time - Self-attention layers delegate to the plugin; cross-attention always uses dense SDPA (controlled by is_self_attention flag) - Sparse attention runs through the Attention wrapper, so SP/ring parallel communication is preserved - Opt-in only: no effect unless user provides sparse_attn config or --attn-backend pointing to a sparse plugin Signed-off-by: Zhen Zhang <zhen.zhang.fr@huawei.com>

chatgpt-codex-connector · 2026-04-16T08:15:57Z

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.

zzhang-fr requested a review from hsliuustc0106 as a code owner March 26, 2026 11:38

zzhang-fr mentioned this pull request Mar 26, 2026

[RFC]: Plugin-based Sparse Attention Interface for DiT Modules #2233

Open

chatgpt-codex-connector Bot reviewed Mar 26, 2026

View reviewed changes

Comment thread vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py Outdated

Comment thread vllm_omni/diffusion/sparse_attn/attention.py Outdated

Comment thread vllm_omni/diffusion/sparse_attn/attention.py Outdated

alex-jw-brooks reviewed Mar 26, 2026

View reviewed changes

Comment thread contrib/flashinfer-vllm-omni/flashinfer_vllm_omni/backend.py Outdated

alex-jw-brooks reviewed Mar 26, 2026

View reviewed changes

Comment thread contrib/sparge-vllm-omni/pyproject.toml Outdated

zzhang-fr changed the title ~~[Feature][Backend] Plugin-based sparse attention for DiT modules~~ [Feature][Diffusion] Plugin-based sparse attention interface for DiT modules Mar 26, 2026

gcanlin requested changes Mar 27, 2026

View reviewed changes

SamitHuang reviewed Mar 27, 2026

View reviewed changes

Comment thread examples/offline_inference/text_to_video/wan2_2_sparse.py Outdated

zzhang-fr marked this pull request as draft March 31, 2026 11:39

zzhang-fr closed this Mar 31, 2026

zzhang-fr force-pushed the main branch from 26ec47b to f8d0bf5 Compare March 31, 2026 23:29

zzhang-fr reopened this Mar 31, 2026

zzhang-fr force-pushed the main branch from 2d132c5 to b801405 Compare April 1, 2026 00:22

zzhang-fr force-pushed the main branch from b801405 to 51fc4ce Compare April 1, 2026 01:07

zzhang-fr marked this pull request as ready for review April 1, 2026 01:11

zzhang-fr requested review from SamitHuang, alex-jw-brooks and gcanlin April 1, 2026 01:11

chatgpt-codex-connector Bot reviewed Apr 1, 2026

View reviewed changes

Comment thread vllm_omni/engine/async_omni_engine.py Outdated

Comment thread vllm_omni/engine/async_omni_engine.py Outdated

zzhang-fr force-pushed the main branch 3 times, most recently from 0511e65 to 12af120 Compare April 1, 2026 22:01

lishunyang12 reviewed Apr 2, 2026

View reviewed changes

zzhang-fr force-pushed the main branch 2 times, most recently from eb92fcc to 49b576c Compare April 2, 2026 12:56

zzhang-fr requested a review from lishunyang12 April 2, 2026 13:19

lishunyang12 reviewed Apr 2, 2026

View reviewed changes

zzhang-fr force-pushed the main branch 4 times, most recently from bd69bd2 to d0a2da0 Compare April 2, 2026 19:24

zzhang-fr requested a review from lishunyang12 April 2, 2026 20:21

zzhang-fr force-pushed the main branch from ff81828 to 028157c Compare April 7, 2026 14:20

gcanlin reviewed Apr 8, 2026

View reviewed changes

zzhang-fr requested a review from gcanlin April 8, 2026 10:10

hsliuustc0106 mentioned this pull request Apr 9, 2026

[RFC]: Per-Role Attention Backend Configuration for Diffusion Models in vllm-omni #2632

Closed

zzhang-fr marked this pull request as draft April 9, 2026 11:01

zzhang-fr force-pushed the main branch from 028157c to 5413c95 Compare April 13, 2026 16:08

zzhang-fr marked this pull request as ready for review April 16, 2026 08:15

zzhang-fr mentioned this pull request May 18, 2026

[RFC]: Pluggable custom attention backends with first-class per-forward state for diffusion #3715

Open

harenome mentioned this pull request Jun 16, 2026

[Feature][Diffusion] Plugin interface for third-party sparse attention for per-forward DiT #4493

Closed

BeatSeat mentioned this pull request Jun 23, 2026

[RFC]: Third-party Sparse Attention Backends Support #4661

Open

1 task done

hsliuustc0106 added diffusion codes related to diffusion models enhancement New feature or request labels Jul 17, 2026 — with ChatGPT Codex Connector

zzhang-fr closed this Jul 27, 2026

BeatSeat mentioned this pull request Sep 14, 2026

[Tracking] vLLM-Omni 阶段性 Milestone 已完成路线图与已解决待清理 Issue 归档清单 BeatSeat/vllm-omni#1

Open

Conversation

zzhang-fr commented Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Changes in this revision

Design

Core changes (vllm-omni)

Plugin interface

Configuration

Usage

Test Result

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lishunyang12 commented Mar 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gcanlin left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zzhang-fr commented Mar 27, 2026

Uh oh!

zzhang-fr commented Apr 1, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

lishunyang12 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lishunyang12 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zzhang-fr commented Apr 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gcanlin Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

zzhang-fr commented Mar 26, 2026 •

edited

Loading

lishunyang12 commented Mar 27, 2026 •

edited

Loading

gcanlin left a comment •

edited

Loading

lishunyang12 left a comment •

edited

Loading

gcanlin Apr 9, 2026 •

edited

Loading