Skip to content

[Perf] Adopt flashdreams.accelerated across the public demos #567

Open

Description

@jarcherNV

[Perf] Adopt flashdreams.accelerated across the public demos

Relationship to existing issues

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_cache returns. 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 distinct SelfAttention and CrossAttention modules, 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 omnidreams OPTIMIZED
    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_heads is one count for
    query, key, and value heads. waypoint is 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. flashvsr self-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.
  • sana is 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
    ContextParallelAttention output 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions