Repository navigation
Description
[Perf] Adopt flashdreams.accelerated across the public demos
Relationship to existing issues
- [Perf]
flashdreams.acceleratedLibrary #482[Perf] flashdreams.accelerated Library— builds the library. This
issue is about consuming it, and depends on it. - [Perf] Use
flashdreams.acceleratedin New API #536[Perf] Use flashdreams.accelerated in New API— closely related and
overlapping. [Perf] Useflashdreams.acceleratedin New API #536 is scoped to wiring the library into the new API's runner
configs for headless, WebRTC, and native-window output. This issue is scoped
to the per-model attention adapters that have to exist before any runner
config can select them. If we would rather keep one issue, this can be folded
into [Perf] Useflashdreams.acceleratedin New API #536 as its implementation plan. - [Acceleration API] cuDNN FP8 SDPA fails on Windows because cuDNN Frontend looks for Linux CUDA runtime libraries #540
[Acceleration API] cuDNN FP8 SDPA fails on Windows— constrains
which schedules are selectable per platform; not a blocker for adapter work.
Summary
OptimizedMultiHeadAttention and the quantization toolkit are used by exactly
one public integration: integrations/omnidreams, and only in its OPTIMIZED
backend configs. Every other public demo runs attention through
F.scaled_dot_product_attention (via ContextParallelAttention), a bundled
Triton sparse kernel, or its own custom math.
The reason this is worth doing as one piece of work rather than twelve: the
public demos do not have twelve attention implementations. Most of them
compose a shared recipe, so a small number of adapters covers most of the tree.
The single highest-leverage fact: flashdreams.accelerated already ships a
documented Wan 2.1 attention config in
docs/source/developer_guides/flashdreams_accelerated.md —
query_dim=1536, context_dim=1536, n_heads=12, head_dim=128,
QKNormScope.INNER, RoPEStyle.INTERLEAVED, RoPEScope.BEFORE_KV_CACHE.
That is exactly the geometry of the Wan recipe (dim=1536, num_heads=12 in
flashdreams/flashdreams/recipes/wan/transformer/impl/network.py, with
apply_rope_before_kvcache: bool = True). The library documents how to
configure Wan attention, and no Wan integration uses it.
Survey of the public demos
| Integration | Transformer | Attention today | Uses accelerated |
|---|---|---|---|
wan21, wan22 |
recipes/wan |
cuDNN SDPA | No |
causal_forcing, self_forcing |
recipes/wan |
cuDNN SDPA | No |
fastvideo_causal_wan22 |
recipes/wan |
cuDNN SDPA | No |
lingbot |
recipes/wan + camera block |
cuDNN SDPA | No |
hy_worldplay |
recipes/wan + dual PRoPE self-attn |
cuDNN SDPA ×2 branches | No |
flashvsr |
recipes/wan + sparse self-attn |
Triton block-sparse (self), cuDNN SDPA (cross) | No |
omnidreams |
own transformer fork | optimized MHA, or cuDNN SDPA, or native CUDA | Partial |
sana |
custom stage1_model |
GDN scan + SDPA | No |
waypoint |
custom Waypoint DiT | SDPA with GQA + flex_attention |
No |
apps/cam2v, apps/t2v, and the integrations_v2/* wrappers are application
shells with no attention code — they inherit whatever the integration does.
color_fade, red_screen, null_model, and slangpy_ui_demo have no model.
Collapsing by implementation rather than by integration, the Wan recipe is
one adapter that covers six named integrations plus the cross-attention of
two more.
Why the Wan family is low-friction
Beyond the geometry already being expressible:
- The KV cache is already the right class. The recipes build
flashdreams.core.attention.BlockKVCache, which is exactly what
OptimizedMultiHeadAttention.allocate_kv_cachereturns. Both name the
sequence axis with heads and features trailing, so the layout matches too. - The self/cross split already matches the library's two adapter shapes.
The recipes have distinctSelfAttentionandCrossAttentionmodules, with
rolling-window self-attention and a static cross-attention cache — the same
division the library assumes. - The attention is unmasked. The library's SDPA backends are the non-causal
kernels; the Wan path gets its causality from the rolling cache window rather
than a mask, so nothing needs a mask to be expressed. - There is a working reference to copy. The
omnidreamsOPTIMIZED
adapters show the whole pattern working on a real integration — the six
required interface properties, checkpoint-native attribute naming, and the
_initialize_derived_weights()hook that rebuilds fused and quantized weights
after checkpoint load. Tier 1 is following a proven path rather than inventing
one, even though the geometry it configures is different.
Known gaps that limit the long tail
These are library capabilities, not integration problems, and each blocks a
specific demo:
- No grouped-query attention.
AttentionConfig.n_headsis one count for
query, key, and value heads.waypointis GQA (32 query heads, 16 K/V
heads) and cannot be expressed at all. Pre-expanding K/V would forfeit the
memory saving that motivates GQA. This is the only public GQA model today,
but it is a hard blocker for it. - No attention mask.
forward(x, kv_cache, rope_freqs)takes no mask and no
is_causal.flashvsrself-attention is a Triton block-sparse kernel driven
by a sparsity mask, so it cannot move without mask support — and it is
plausibly already faster than a dense optimized kernel, so it may not want to. sanais a different operation. Most of its blocks are a bidirectional
GDN scan rather than attention, it has no streaming KV cache, and it carries
its own quantization recipes. Only its every-fourth-block SDPA and its
cross-attention are candidates.
Proposed plan
Tier 1 — one Wan adapter, six integrations. Write self- and
cross-attention adapters for the Wan recipe using the already-documented config,
select them in the recipe's block construction, and expose the schedule on the
network config rather than hard-coding it. This is the bulk of the value and
should land before anything else.
Tier 2 — hy_worldplay. Wan-based, so it inherits the Tier 1 adapter for
cross-attention, but its dual PRoPE self-attention branches and prefilled memory
cache need their own handling.
Tier 3 — quantization, independently of the MHA work.
QuantizedNonPersistentLinear is a drop-in nn.Linear replacement at inference
and needs none of the above. It can be applied to projections and MLPs in any
integration on its own schedule, including ones whose attention cannot move.
Deferred — flashvsr self-attention, sana, waypoint, pending mask and
GQA support in the library, and a decision on whether they should move at all.
Schedules do not transfer between platforms
OptimizedImplConfig is a scheduling space, not a setting. The benchmark-
selected policies documented for omnidreams already show self- and
cross-attention preferring different fusion options, SDPA backends, TMA
settings, and precisions on the same GPU, and different choices across GPUs — to
the point where its RTX PRO 6000 variant keeps the original implementation for
cross-attention. Every adapter should therefore expose separate self- and
cross-attention schedule fields on the integration's network config, so
platform-specific runner configs can override them after benchmarking. Adopting
the library should not mean baking one hardware-specific result into a model.
Acceptance criteria
- The Wan recipe runs its attention through
flashdreams.accelerated, with
self- and cross-attention schedules selectable from the network config, and
all six Wan-family integrations inherit it. - No regression in the benchmark suite's existing quality gates for any migrated
scenario, and a measured before/after per integration on at least one
platform, recording commit, config, GPU, and warmup policy. - Correctness tests comparing each new adapter against the current
ContextParallelAttentionoutput as the oracle, following the pattern in
integrations/omnidreams/tests/test_transformer_attention_backend.py. - Integrations that cannot migrate are documented as such, with the blocking
library gap named.
Activity