[Feature][Diffusion] Plugin interface for third-party sparse attention for per-forward DiT by harenome · Pull Request #4493 · vllm-project/vllm-omni

harenome · 2026-06-16T15:02:19Z

Purpose

Refactors and supersedes (#2231).
Implements a light plugin mechanism for sparse attention as proposed in RFCs #2233 and #3715.
It includes two parts:

a generic per-forward state channel (to be forwarded to sparse attention kernels)
a sparse attention dispatcher

This mechanism reuses and extends the ideas of per-role attention config (#2681) and ForwardContext.denoise_step_idx (#2640).

The dispatcher is built to make shipping a third-party sparse kernel as simple as possible. A plugin only has to expose one callable in its entry point and there is no base class to subclass or boilerplate: a minimal plugin can be just a function. Optional additional capabilities are advertised in a simple capabilities dict.

Design

Per-forward state (generic)

New field AttentionMetadata.per_forward holds a PerForwardState object for per-forward metadata.
It holds common metadata such as framework state or geometry.
Framework state (such as denoising steps) is read from ForwardContext by Attention layers.
Geometry is model-specific and set by the transformer forward of supported models.

Sparse attention dispatcher

The sparse attention dispatcher lives in SparseAttentionBackend and selected via DiffusionAttentionBackendEnum.SPARSE_ATTN.
Plugins register a callable under the vllm_omni.sparse_attn entry-point group.

A minimal plugin can be as simple as the following snippet (on top of setting the entry-point):

# A minimal plugin: a callable, plus an optional capabilities dict.
def my_attn(q, k, v, params, is_causal, softmax_scale) -> Tensor:
    ...  # return the attention output in BSND layout

# Optional: advertise framework inputs the kernel actually consumes.
my_attn.capabilities = {"supports_attention_mask": True}

It is enabled on per role with a short form: for instance, per_role.self.backend=<plugin> or per_role.cross.backend=<plugin> will be routed to the dispatcher automatically.
Multiple sparse plugins can coexist.
Built-in backends (other DiffusionAttentionBackendEnum values) have priority over third-party plugins.

The callable signature is fn(query, key, value, params, is_causal, softmax_scale) -> Tensor, with q/k/v in BSND layout [batch, seq, heads, head_dim].

params is a flat dict (RFC #3715, Proposal B), merged from three sources in increasing priority:

static config — the kernel's own extra.params;
PerForwardState — the per-forward fields the framework fills (step, geometry);
AttentionMetadata.extra — dynamic per-forward values.

On a key collision the later source wins, and None values are dropped, so a kernel can rely on params.get(key, default).

Most params entries are optional and may be silently ignored by kernels.
Entries that can not be ignored, if present, are guarded by capabilities.
Currently, two entries must be handled if provided:

attn_mask (for padding/validity, capability name: supports_attention_mask)
full_attn_spans (for mixed causal/full attention, capability name: supports_full_attn_spans).

Additional capabilities can be advertised by setting an additional attribute capabilities to the callable.

Ready for future evolutions

Current demo examples are simple examples where attn_mask and full_attn_spans are None.
The plugin mechanism is already ready for models and usecases that require it.
New entries can be added in params as new usecases emerge and new entries will not break compatibility with older plugins
PerForwardState uses per-sample tensors for denoising steps: ready for future work such as
RFC [RFC]: Refactor engine/runner/pipeline to support step-wise and continuos batching #874 and feat: Support dynamic batching for heterogeneous diffusion requests #4024.

Tests and docs

tests/diffusion/attention/test_per_forward_state.py: the state dataclass and the layer helper (with a mocked SDPA backend).
tests/diffusion/sparse_attn/test_sparse_attn.py: plugin lookup (mocked entry points), the three-layer merge order and None-drop, capability checks, and the short-form routing.

User docs: docs/user_guide/diffusion/attention_backends.md (the Sparse attention (plugin) section: the contract, params, capabilities, and a small example plugin).

Demo

Two external plugins exercise the dispatcher end-to-end, each a standalone package with no dependency on vllm-omni: spargeattn (SpargeAttn) and lvsa_sparse (LVSA, with SDPA or FlashInfer kernels). See vllm-omni-plugin-demo for setup and run commands.

chatgpt-codex-connector · 2026-06-16T15:02:29Z

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.

Gaohan123 · 2026-06-22T15:31:00Z

@ZJY0516 PTAL

hsliuustc0106 · 2026-06-26T13:13:46Z

+        attn_metadata = self._with_per_forward_state(attn_metadata, query)

        # 2. Kernel Execution (Computation)
        if self.use_ring and strategy is not self._no_parallel_strategy:


This bypasses the selected SPARSE_ATTN backend when ring_degree > 1: RingParallelAttention only runs its own SDPA/flash-style kernels, so a configured sparse plugin is silently ignored. Please either fail fast for sparse plugins with ring SP or route the ring path through the plugin, with a regression test.

hsliuustc0106 · 2026-06-26T13:13:59Z

        query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata)

        attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+        attn_metadata = self._with_per_forward_state(attn_metadata, query)


This now builds per-forward tensors for every attention layer whenever step/geometry exists, even for dense backends that ignore per_forward. Wan sets this state unconditionally, so default inference can pick up many extra GPU allocations per denoise step. Please gate this to backends that consume it, cache the state per forward, or include before/after latency evidence showing the default path is unaffected.

…e for DiT Add a plugin-based sparse-attention path for diffusion (DiT) models on top of the per-role attention config (vllm-project#2681), plus a generic per-forward state channel that step-aware attention backends can read. Implements RFC vllm-project#3715: a dedicated SPARSE_ATTN slot (Proposal A, alpha) and Proposal B's param-merge semantics. Per-forward state (generic, no sparse code): - ForwardContext.total_denoise_steps + setter; the Wan2.2 pipelines set it once before the denoise loop (alongside the existing denoise_step_idx). - PerForwardState: a frozen dataclass carried by composition on AttentionMetadata.per_forward, with per-sample [num_samples] tensor fields denoise_step_idx and total_denoise_steps. A source-isolated bridge in the Attention layer fills it from ForwardContext (device-agnostic), so any backend can read step state without reaching into global context. Sparse attention dispatcher: - SPARSE_ATTN member in DiffusionAttentionBackendEnum; SparseAttentionBackend(AttentionBackend) + _SparseAttentionImpl(AttentionImpl) inherit the conservative ABC defaults. - Resolves an external vllm_omni.sparse_attn callable plugin fn(query, key, value, params, is_causal, softmax_scale) -> Tensor; a missing or broken plugin fails loudly at model init, never a silent fall back to dense. - Per-role config with a shorthand so backend=<plugin> routes to the dispatcher; built-in dense names are excluded. Cross-attention stays dense by role. - The dispatcher merges params in three tiers (static < per-forward < extra, None dropped) into the plugin's flat params dict. Per-forward video geometry: - Extend the per-forward state with optional video grid geometry so structure-aware sparse kernels (e.g. windowed / keyframe attention like LVSA) can map the flat token sequence back to (frame, patch) coordinates instead of relying on env-var or call-counting heuristics. - ForwardContext: total_latent_frames / patches_per_frame fields + set_forward_context_geometry() setter. - PerForwardState: matching per-sample [num_samples] tensor fields, carried and flattened through the existing to_dict() path (so kernels read params["total_latent_frames"] / params["patches_per_frame"]). - Attention bridge: fills the geometry tensors from ForwardContext alongside the step state. - Wan2.2 and HunyuanVideo transformers publish their patch grid once per forward (before attention), reusing the post-patch grid they already compute. - Geometry is optional; backends ignore fields they don't use. seq_len == total_latent_frames * patches_per_frame. Gated attention metadata: - Forward the correctness-critical AttentionMetadata fields attn_mask (padding / validity) and full_attn_spans (mixed causal/full topology) into params, only when set, so a kernel that needs them is not silently starved of them. - A plugin advertises which it consumes via an optional capabilities dict on the callable (read once at resolution, malformed -> empty). When such a field is present but unadvertised, the dispatcher raises instead of letting the kernel ignore it. The dispatcher forward runs under torch.compile, so the gating uses direct attribute access with capability flags precomputed in __init__. Includes CPU-runnable tests for the per-forward state and the dispatcher (param merge, capability self-check, malformed capabilities) plus a GPU-gated on-CUDA dispatch test, and a user-guide section documenting the plugin contract. Co-authored-by: Zhen Zhang <zhen.zhang.fr@huawei.com> Signed-off-by: Harenome Ranaivoarivony-Razanajato <harenome.razanajato1@huawei.com> Signed-off-by: Zhen Zhang <zhen.zhang.fr@huawei.com>

…ia a transformer pre-hook Per-forward video geometry was published per-model: every video DiT had to call set_forward_context_geometry(...) inside its forward (Wan2.2, HunyuanVideo) — an O(models) cost coupling core model code to the sparse-attention feature. Replace it with a framework-registered transformer forward pre-hook that reads the (B,C,T,H,W) latent shape plus the model's patch config and publishes RAW, model-agnostic primitives (latent_shape, patch_size) onto ForwardContext. The SPARSE_ATTN dispatcher's default resolver derives total_latent_frames / patches_per_frame from them; a plugin geometry_fn may override. New video DiTs need no geometry edits. - NEW video_geometry.py: _resolve_patch_size (handles Wan 3-tuple + Hunyuan scalar/patch_size_t), best-effort _video_geometry_pre_hook (no-op on non-5D input / no patch size / no active context), idempotent register_video_geometry_hooks. - forward_context: ForwardContext carries raw latent_shape/patch_size; setter renamed set_forward_context_geometry -> set_forward_context_video_geometry. - diffusion_model_runner: register hooks at end of load_model (after compile/offload). - attention/layer bridge: surface raw primitives as per-sample [n,3] tensors. - abstract PerForwardState: add raw latent_shape/patch_size; derived fields kept for back-compat, now resolver-filled, not bridge-filled. - sparse_attn: _apply_default_geometry resolver + optional plugin geometry_fn. - Remove per-model geometry blocks from wan2_2 / hunyuan_video_15 transformers (post_patch_* locals retained for RoPE/unpatchify). Address review feedback: - Gate per_forward construction to backends that consume it (AttentionImpl .consumes_per_forward; only the SPARSE_ATTN dispatcher opts in), so the default dense path builds no per-layer/per-step step/geometry tensors. - Fail fast when SPARSE_ATTN is combined with ring sequence parallelism: ring routes through RingParallelAttention's dense kernels and would silently bypass the sparse plugin (Attention._assert_ring_backend_supported). - Tests: new test_video_geometry.py; update test_per_forward_state, test_sparse_attn. Docs: attention_backends.md. Signed-off-by: Zhen Zhang <zhen.zhang.fr@huawei.com>

Gaohan123 · 2026-08-29T11:42:38Z

@harenome Please resolve conlicts. Thanks

vllm-omni-review-bot · 2026-09-08T12:22:51Z

Omni ReviewBot: no human activity for 10 days

@harenome this pull request has had no human commit, comment or review since 2026-08-29. Per repository policy it may be closed if it stays inactive.

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 · 2026-09-23T16:10:36Z

Omni ReviewBot: no human activity for 25 days

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

vllm-omni-review-bot · 2026-10-02T13:44:14Z

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 · 2026-10-09T10:48:29Z

Omni ReviewBot attempt record

Review attempt ended as timeout (Strict attempt outlived its budget; falling back to direct/cursor/auto).

vllm-omni-review-bot · 2026-10-09T10:48:30Z

Omni ReviewBot attempt record

Review attempt ended as timeout (Strict attempt outlived its budget; falling back to direct/cursor/auto).

vllm-omni-review-bot · 2026-10-09T10:48:30Z

Omni ReviewBot attempt record

Review attempt ended as timeout (Strict attempt outlived its budget; falling back to direct/cursor/auto).

vllm-omni-review-bot · 2026-10-09T10:48:30Z

Omni ReviewBot attempt record

Review attempt ended as timeout (Strict attempt outlived its budget; falling back to direct/cursor/auto).

harenome · 2026-10-09T11:44:34Z

Can not further work on this at the moment.

harenome requested review from Gaohan123, Isotr0py, RuixiangMa, SamitHuang, ZJY0516, david6666666, hsliuustc0106, princepride, wtomin and yenuo26 as code owners June 16, 2026 15:02

harenome changed the title ~~[Feature][Diffusion] Plugin-based sparse attention for per-forward DiT~~ [Feature][Diffusion] Plugin interface for third-party sparse attention for per-forward DiT Jun 16, 2026

harenome force-pushed the sparse-attn-per-forward branch from 405ecd3 to eb391d7 Compare June 16, 2026 17:39

hsliuustc0106 added diffusion codes related to diffusion models core related to core module: cache, scheduler, engine, worker, modelrunner labels Jun 17, 2026

Gaohan123 added this to the v0.24.0 milestone Jun 17, 2026

BeatSeat mentioned this pull request Jun 23, 2026

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

Open

1 task done

hsliuustc0106 reviewed Jun 26, 2026

View reviewed changes

zzhang-fr force-pushed the sparse-attn-per-forward branch 2 times, most recently from b3066d5 to 462a1f6 Compare June 26, 2026 17:09

zzhang-fr force-pushed the sparse-attn-per-forward branch from 462a1f6 to 1151ff2 Compare July 5, 2026 16:50

zzhang-fr requested a review from xuechendi as a code owner July 5, 2026 16:50

zzhang-fr force-pushed the sparse-attn-per-forward branch from 1151ff2 to 5b7f5f2 Compare July 5, 2026 21:04

Gaohan123 modified the milestones: v0.24.0, v0.26.0 Jul 7, 2026

hsliuustc0106 added the enhancement New feature or request label Jul 17, 2026 — with ChatGPT Codex Connector

Gaohan123 modified the milestones: v0.26.0, v0.28.0 Aug 4, 2026

david6666666 mentioned this pull request Aug 5, 2026

[Roadmap] MiniMax-H3 follow-up: accuracy/performance CI, quantization, and optimizations #5700

Closed

Gaohan123 removed this from the v0.28.0 milestone Aug 29, 2026

Dong1017 mentioned this pull request Sep 4, 2026

[Survey] Kernel 优化全仓梳理：通用 vs 模型绑定 + 归档方案（2026-09-04） Dong1017/vllm-omni#1

Open

Dong1017 mentioned this pull request Sep 10, 2026

[RFC]: Kernel Reuse and Hardware Integration in vLLM-Omni #7382

Open

1 of 16 tasks

blian6 mentioned this pull request Sep 17, 2026

[RFC]: Make attention backends own their cross-rank communication #7727

Open

1 task

hsliuustc0106 added high priority high priority issue, needs to be done asap stale No new commits for at least 60 days labels Sep 30, 2026 — with ChatGPT Codex Connector

rahul-steiger-nv mentioned this pull request Oct 1, 2026

[RFC]: Shared block-sparse execution across attention backends #8396

Open

1 task done

harenome closed this Oct 9, 2026

Conversation

harenome commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Design

Per-forward state (generic)

Sparse attention dispatcher

Ready for future evolutions

Tests and docs

Demo

Uh oh!

chatgpt-codex-connector Bot commented Jun 16, 2026

Uh oh!

Gaohan123 commented Jun 22, 2026

Uh oh!

hsliuustc0106 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Gaohan123 commented Aug 29, 2026

Uh oh!

vllm-omni-review-bot commented Sep 8, 2026

Omni ReviewBot: no human activity for 10 days

Uh oh!

vllm-omni-review-bot commented Sep 23, 2026

Omni ReviewBot: no human activity for 25 days

Uh oh!

vllm-omni-review-bot commented Oct 2, 2026

Uh oh!

vllm-omni-review-bot commented Oct 9, 2026

Uh oh!

vllm-omni-review-bot commented Oct 9, 2026

Uh oh!

vllm-omni-review-bot commented Oct 9, 2026

Uh oh!

vllm-omni-review-bot commented Oct 9, 2026

Uh oh!

harenome commented Oct 9, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

harenome commented Jun 16, 2026 •

edited

Loading