[Research] Explore linear attention for the streaming video models

[jarcherNV]

[Research] Explore linear attention for the streaming video models

What this is

An exploration issue, not an implementation plan. The goal is a written
recommendation, backed by measurements, on whether linear attention is worth
adopting in any of our models — and if so, which one, in what form, and at what
quality cost. It is expected to involve post-training, which means it cannot be
answered inside this repo alone (see Scope below).

Why, and the honest version of the argument

The obvious pitch for linear attention is that it replaces the quadratic
attention cost with something linear. For our streaming models that pitch is
mostly wrong, and it is worth saying so up front so we evaluate this for the
right reason.

Our autoregressive models already bound per-step attention cost with a short
rolling K/V window plus a small attention sink — the LingBot runner configs, for
example, are built around a 15-frame window with a 3-frame sink. Per-step
attention over a 15-frame window is already cheap, and it does not grow as the
rollout gets longer. Swapping in linear attention would not make that step
dramatically faster.

The actual prize is different: a recurrent state carries unbounded history at
fixed per-step cost and fixed memory. Today the window is the model's
memory, so anything older than ~15 frames is simply gone, and no amount of
kernel optimization changes that. Linear attention would let a model condition
on the whole rollout without the cost growing, which targets long-horizon drift
and temporal consistency — the class of problem a short window cannot solve at
any speed. Secondary benefits are O(1) state instead of an O(window) K/V cache,
and no sink/window tuning per model.

So the question to evaluate is not "is it faster?" but "does it buy us
long-horizon coherence we cannot otherwise get, at an acceptable quality and
complexity cost?"

We already ship a hybrid linear-attention model

Worth knowing before anyone treats this as greenfield: the SANA-WM Stage-1 DiT
in integrations/sana is already a hybrid. Most of its blocks run a gated
delta-net linear attention with a matrix-valued per-head recurrent state, and
every fourth block runs ordinary softmax attention — the interleave is literally
a spec field, softmax_every_n = 4.

That gives us two useful things: an existence proof that a hybrid works in one
of our models, and a design template for the interleave ratio. The important
caveat is that its scan is bidirectional, which suits a bidirectional DiT
but is not the form the streaming models need. The causal, forward-only variant
is the one where the O(1) state actually pays off at inference time, and it is
the one that would need new kernel work.

Related axes already in the tree

Linearization is not the only way at this problem, and we are already exploring
the neighbours. FlashVSR runs a Triton block-sparse self-attention, and Waypoint
mixes locally-cached attention with flex_attention global layers. Sparsity and
linearization are complementary, and part of the value of this issue is deciding
deliberately which we pursue where rather than accumulating one of each.

Candidate models, in rough priority order

1. The Wan-family streaming models (wan21, wan22, self_forcing,
   causal_forcing, fastvideo_causal_wan22, lingbot). Best first target for
   two reasons: they share one transformer recipe, so a single architectural
   change reaches six integrations, and the benchmark suite already has
   one-minute scenarios for most of them, so a long-horizon eval exists.
2. omnidreams, particularly the multiview configurations, where sequence
   length scales with the number of views and the window limit bites hardest.
3. Not SANA-WM — it is already hybrid, so it serves as the reference
   implementation rather than a target.
4. Deprioritize flashvsr and waypoint for now. Both already have bespoke
   sparse or hybrid schemes that would have to be unpicked first, and neither is
   obviously limited by window length.

Suggested phases

Phase 0 — quantify what we would be buying. Before any post-training, use
the existing models to measure how per-step time and memory scale with window
size, and how quality degrades as the rollout lengthens at a fixed window. If
long-rollout quality is already flat out to a minute, the case for this weakens
considerably. If drift is measurable, we have a target number to beat. This
phase is cheap, needs no new weights, and should gate the rest.

Phase 1 — choose a form and prototype inference-side. Pick the linear
attention variant (gated delta net as in SANA-WM, or an alternative) and the
hybrid ratio, then build it against frozen weights to validate shapes, the
streaming state lifecycle, and per-step cost. Quality will be meaningless at
this stage; the goal is a working causal scan and a real performance curve.

Phase 2 — conversion and post-training. Linearizing pretrained weights
rather than training from scratch is the tractable path: attention-transfer
style distillation from the softmax teacher, then finetuning. This is the
expensive, uncertain part, and it happens outside this repo.

Phase 3 — evaluate against the long-horizon scenarios. Judge on the existing
one-minute scenarios and the long-form quality metrics we already compute,
alongside per-step cost and memory versus the softmax baseline. Short clips will
not show the effect this is meant to buy, and are also where linear attention is
most likely to look worse.

Scope, and what does not belong here

FlashDreams is an inference and serving library — there is no training code in
it. This issue should therefore own:

• the inference-side architecture and streaming state handling,
• the kernel work for the causal scan,
• the evaluation used to make the call.

Weight conversion, distillation, and finetuning need a training environment
elsewhere. This issue should reference that work rather than pretend to contain
it, and should not be closed on the basis of a prototype that never received
converted weights.

Open questions and risks

• Quality regression is the main risk. Linear attention generally gives up
  precise retrieval, which in video is likely to show up as texture, identity,
  or fine-detail drift rather than as a metric cliff. Human review on long
  rollouts probably matters more than the aggregate scores here.
• What is the right hybrid ratio, and which layers keep softmax? The
  every-fourth-block choice in SANA-WM is one data point, not a law, and the
  answer may differ for causal models.
• Kernel cost. A streaming causal gated scan needs an efficient kernel; the
  existing scan is written for the bidirectional bulk case. It also has to stay
  compatible with torch.compile and CUDA graph capture, which the current
  streaming paths depend on.
• Composition with the rest of the acceleration work. Does a linear
  attention block coexist with the optimized attention library and the
  quantization toolkit, or does it fork the execution path? Worth answering
  before committing, so this does not become an island.
• Is the underlying complaint actually long-horizon drift? If nobody is
  hitting the window limit in practice, this is interesting research with no
  user attached. Phase 0 should force that question.

What "done" looks like for this issue

A recommendation document with measurements behind it: whether to pursue linear
attention, on which model, in which form, with the observed cost and quality
trade-off, and an explicit answer on whether the long-horizon benefit showed up.
A well-evidenced "not worth it" closes this issue successfully.