Attention is about 73% of Wan 2.2’s per-step compute at 32,760 tokens, so caching and sparsity both point at the same place. Two obstacles sit between that observation and an actual speedup. The A100 the model runs on has INT8 tensor cores but no FP8, which caps INT8 attention hard before any sparsity is even considered. And nearly every published cache or sparsity method assumes one fixed network runs the whole generation, an assumption Wan 2.2 breaks by swapping its entire 14B network partway through.
What follows is the full path across three attacks on that per-step cost. First, a published state-of-the-art sparse-attention kernel that did nothing on this workload out of the box, and the one-line selector swap that recovered 1.3–1.76× from it. Second, a step cache rebuilt around the two experts, 1.78× at near-lossless quality. Third, a free signal already sitting in the attention kernel’s registers that predicts achievable sparsity better than the calibration tables the field ships. The first two are measured speedups. The third is a validated signal with no wall-clock, and it is presented without one. None of it shipped yet.
First, the hardware ceiling
Before optimizing attention on this card, the real question is how much INT8 can buy at all. Attention runs two matmuls of roughly equal cost: the score matmul QKT and the value matmul PV. On the A100, INT8 tensor cores run at about 2× the FP16 rate, but the card has no FP8 units, so only QKT can be quantized to INT8 and accelerated. PV has to stay in FP16. Speed up half the work by 2× and leave the other half unchanged, and the blended ceiling is 1.33×, before a single attention block is skipped.
Where the 1.33× ceiling comes from, exactly
Attention’s two matmuls have equal FLOP counts: QKT is N × N × d and PV is N × N × d for
N tokens and head dimension d. Split the work evenly, put the score matmul on INT8 at 2× and leave
the value matmul on FP16 at 1×, and the total time scales as 0.5/2 + 0.5/1 = 0.75 of the dense FP16
time. That is a 1/0.75 = 1.33× speedup, and it is a hard ceiling: no amount of INT8 tuning moves it,
because PV is pinned to FP16 by the missing FP8 units. Everything above 1.33× on this card has to come
from skipping blocks, not from precision.
Measure the published kernel before building one
The obvious move is to write a video-tuned sparse-attention kernel. The disciplined move is to first measure the best shipped tool on the real workload, because “state of the art” is relative to the workload a method was tuned on. SpargeAttn is a strong INT8 + block-sparse attention kernel with a tuned CUDA path. Run on Wan 2.2’s real attention op, out of the box, it lands at roughly 1×.
Every bar sits on top of the baseline, and the worst layer is below it. As shipped, the tool is a dead-end on this workload, sometimes an active regression. The diagnosis is specific: SpargeAttn’s own self-similarity gate marks 65–99.6% of video-latent blocks “unpredictable” and forces them down the dense path. With almost everything dense, there is no sparsity to convert into speed, and the INT8 path’s own overhead is what pushes some layers below 1.0×.
SpargeAttn's cosine-gram gate, and why video trips it
SpargeAttn decides, per block, whether attention there is predictable enough to approximate. It builds a compressed gram-like summary of each block and measures self-similarity by cosine distance; a block only takes the sparse path if that similarity clears a confidence bar, otherwise it falls back to dense to stay safe. That gate was tuned on workloads where blocks are internally coherent. Wan’s video latents are the opposite case: attention over 32,760 spatio-temporal tokens is diffuse, and per-block internal similarity is low, so the gate marks 65–99.6% of blocks “unpredictable” and routes them dense. The gate is doing exactly what it was designed to do; it was designed for a different distribution.
Keep the kernel, swap only the selector
That out-of-the-box result reads as a dead-end, but it also pins the problem to one place. SpargeAttn is two separable pieces. One is a tuned CUDA kernel that runs the INT8 score matmul and the block-sparse value matmul. The other is a selector, the cosine-gram gate, that decides which blocks the kernel is allowed to skip. The kernel is generic arithmetic. A matrix multiply does not know whether its tokens came from text, an image, or a video latent; hand it a block mask and it computes that mask fast on any of them. The selector is the only piece that carries an assumption about the data, and video latents are the distribution that assumption breaks on.
That splits the failure cleanly. A 65–99.6% dense rate is a selection decision, not a kernel result. The kernel never ran sparse on this workload because the gate never handed it a sparse mask to run, so its sparse path, the part that would produce the speedup, went unexercised here entirely. The natural hypothesis follows. The kernel is fine and only the selector is wrong, and if that holds, feeding the same kernel a better mask should recover the speedup the gate was suppressing.
That is a causal claim, so testing it means changing exactly one variable. Forking the kernel would
change both the math and the selection at once and prove nothing. Instead the kernel is held
byte-for-byte fixed and only the selector is replaced, injected through SpargeAttn’s own mask_id hook.
Any change in speed or quality is then attributable to the selector alone.
Why inject through mask_id instead of forking the kernel
The one-variable test needs an injection point that leaves the kernel untouched, and SpargeAttn already
exposes one. Its mask_id parameter accepts an
externally-computed block mask and runs its identical tuned path on it. Feeding a different selector
through that hook keeps the kernel bit-identical and varies only the selection heuristic. A rewritten
kernel that happened to be faster would leave the original question open; the same kernel with a new
selector answers it directly.
The replacement selector is deliberately cheap: mean-pool each query and key block to a single vector, score block pairs by the pooled dot product, and keep the top key blocks per query block by a cumulative threshold. No cosine-gram summary, no per-timestep calibration table. Feeding that into SpargeAttn’s own kernel recovers 1.3–1.76× at matched quality on the real 40-layer 14B target.
The mean-pool + per-q-block-CDF selector, step by step
Split queries and keys into blocks. Mean-pool each block to one vector, then score every query-block/key-block pair by the dot product of the two pooled vectors, a cheap estimate of where the attention mass sits. For each query block, sort its key-block scores, walk the cumulative distribution, and keep key blocks until a fixed fraction of the estimated mass is covered. Everything below that per-query-block CDF threshold is dropped before the kernel runs. The whole selector is one pooling pass and one sort per query block, with no offline artifact to calibrate.
| Layer / expert | FA2 dense | Selector-swap | Speedup | rel-L1 |
|---|---|---|---|---|
| Layer 20, low-noise | 117.8 ms | 67.1 ms | 1.76× | 0.086 |
| Layer 20, high-noise | 118.3 ms | 68.9 ms | 1.72× | 0.089 |
| Layer 10, high-noise | 118.2 ms | 80.0 ms | 1.48× | 0.075 |
Middle layers reach 1.72–1.76× at rel-L1 ≤ 0.09. Layer 39, the densest, barely sparsifies and stays ~1.22×, reported, not hidden. Some layers simply do not sparsify.
The expensive gate was dead weight
The speedup is the smaller of the two findings here. The sharper one is what the ablation says about the cosine-gram gate. Compare the blocks the cheap mean-pool selector keeps against the blocks SpargeAttn’s expensive gate keeps, scored by recall against an exact dense-attention oracle, and the two selectors agree to three decimal places at every point tested.
| Probe | Cosine-gram gate | Mean-pool selector | Difference |
|---|---|---|---|
| Probe 1 | 0.921 | 0.921 | 0.000 |
| Probe 2 | 0.960 | 0.960 | 0.000 |
| Probe 3 | 0.865 | 0.865 | 0.000 |
Identical to three decimals at every point tested. The expensive gate contributes no selection quality on this workload; the cheap selector reaches the same choice for less compute.
What 'recall' means here
Take exact dense attention as the oracle and rank key blocks by the true attention mass each contributes to a given query block. A selector’s recall is the fraction of that truly-important set it keeps. A selector can be cheap and still have high recall if it happens to keep the same blocks the oracle would. Across the probes the cosine-gram gate and the mean-pool selector keep the same blocks to three decimals, so they land at the same quality. The cheap one is then strictly better, because it computes an identical selection for less work.
A cache that knows there are two experts
Sparsity exploited the structure inside a single attention op. A step cache exploits a different Wan-specific fact: the model swaps its entire 14B network partway through generation. Wan 2.2 A14B is a dual-expert diffusion transformer, roughly 27B parameters total with 14B active per step. A high-noise “layout” expert runs the early steps, a low-noise “detail” expert runs the late ones, switching at step 26 of the 40-step native schedule.
That structure is exactly what the field’s cache methods assume away. TeaCache, FasterCache, AdaCache, BWCache, TaoCache, SpargeAttn, SVG, and DFSAttn all assume one fixed network runs the whole trajectory. Step caching skips recomputing a step when it is similar enough to the previous one, gated by a single similarity threshold. Calibrated once, that threshold is implicitly tuned for whichever expert dominated calibration, and wrong for the other.
The swap makes the fix structural, not a matter of tuning. Two networks generate the trajectory, so there are two step-to-step similarity regimes, and no single threshold can sit correctly in both at once. Calibrating each expert on its own is the direct consequence. The high-noise layout expert gets a conservative threshold (0.05), because its consecutive steps still differ a lot; the low-noise detail expert gets an aggressive one (0.20), because its steps barely move. The hard barrier at step 26 follows from the same fact. A cached update produced by the layout expert describes a network that no longer exists once the switch happens, so the detail expert can never reuse it, and the cache resets across the boundary instead of skipping through it. All of it Pareto-dominates a single global threshold, faster and higher quality at once.
| Cache config | Time | Speedup | PSNR |
|---|---|---|---|
| No cache | 184.5 s | 1.00× | ref |
| Uniform threshold 0.06 | 131.5 s | 1.40× | 64.5 dB |
| Uniform threshold 0.10 | 115.4 s | 1.60× | 64.6 dB |
| Two-expert (hi 0.05 / lo 0.20) | 103.4 s | 1.78× | 65.7 dB |
Conservative on the high-noise layout expert, aggressive on the low-noise detail expert, hard barrier at the switch. Both PSNRs are near-lossless; the per-expert version is both faster and slightly cleaner.
- uniform sweep
- two-expert
How step caching decides to skip
The cache stores the previous step’s transformer update and, at the next step, estimates how much the signal changed using a relative-L1 distance between the current and cached activations. If that distance is below a threshold τ, the step reuses the cached update instead of running the 14B network. A larger τ skips more steps: faster, but riskier, because a skipped step that actually mattered leaves visible error in the frame. The right τ is the largest one whose skips all stay under the quality bar.
Why the detail expert tolerates a higher threshold
The high-noise layout expert runs the early steps, where the sample is still resolving global structure and consecutive steps differ a lot, so nearly every step carries new information and skipping corrupts the layout. The low-noise detail expert runs the late steps, where the sample is nearly settled and each step makes a small refinement, so consecutive steps are highly similar and many can be reused safely. A single global threshold cannot express both regimes: set it for the layout expert and the detail steps are needlessly recomputed; set it for the detail expert and the layout steps get skipped and break. Per-expert τ (0.05 high-noise, 0.20 low-noise) matches each regime to its own similarity profile.
A sparsity budget the kernel already computes for free
Every sparse-attention method needs a budget: how many blocks to keep per layer and per step. And each ships its own offline artifact to supply it, a fitted polynomial, a codebook, or a per-timestep table. The alternative is to look for a quantity the kernel already computes. Flash attention keeps a softmax normalizer in registers for every query row, and its reciprocal is the peakedness of that row’s attention. Peaked rows concentrate on a few keys and sparsify well; diffuse rows spread their mass and do not. The signal that predicts the budget is already sitting there.
max_prob = 1 / l_i, and why it is free
Flash attention never materializes the full score matrix. It walks key blocks while keeping two running
numbers per query row: the max score m_i and the normalizer l_i, the sum of exp(s_ij - m_i) over
keys j. The softmax weight of key j is exp(s_ij - m_i) / l_i, so the largest weight in the row is
exp(m_i - m_i) / l_i = 1 / l_i. That reciprocal, max_prob, is the row’s peakedness: near 1 when one
key dominates (sparsifiable), near 1/N when the mass is spread flat (diffuse, not sparsifiable). It is
already in registers at the end of each row’s accumulation, so reading it costs no compute and no extra
pass over memory.
Measured across a full 40-step generation, 45 layer-and-step points in all, max_prob predicts the
achievable per-layer block-sparsity at Spearman 0.75, against the timestep proxy the field currently
uses at 0.49. And the expert switch shows up in it directly: attention density jumps roughly 2.2–2.5× at
the boundary, so the signal detects the network change on its own, with no hardcoded step number.
Where this goes
- The levers are small and self-contained. These results came from short-lived rented GPU pods that
are now gone, so unlike the ollama VRAM fix that was upstreamed, this work has not been pushed yet. That
the wins are a selector swapped through SpargeAttn’s own
mask_idhook and a pair of per-expert thresholds, not a new kernel, is what keeps the path from here short. Re-run on persistent hardware, they land as a patch, not a rebuild. - The cache’s next measurement is under full guidance. It was measured at 20 steps with classifier-free guidance off, to isolate a clean step-to-step similarity signal. Production T2V runs under CFG, which evaluates the network twice per step, so the deployable version carries separate conditional and unconditional cache slots. The per-expert structure carries straight into that run; it is the next measurement, not a redesign.
- The next quality gate is the assembled clip. Quality here is attention-output relative-L1 against an FP32 reference, and PSNR against the no-cache output for the cache. That is the right gate for an op and a cache in isolation, and the next one is a full-clip perceptual metric that scores the finished video rather than the op alone. It raises the bar from a faithful op to a faithful clip.
- These sit at a different operating point than distillation. At the four distilled steps from the Wan distillation work there is little step-to-step redundancy left for a cache to exploit and little cross-step compute to sparsify away, so the 1.3–1.8× here does not multiply the 16.8×. The two are separate levers on the same model. Distillation cuts how many steps run; these cut what a full schedule of steps costs, and which one applies is a choice of operating point.
Three levers, one method. Establish the hardware ceiling before claiming a speedup, so 1.33× is not mistaken for a paper’s 2.1–3×. Measure the field’s best shipped tool on the real workload, so a “SOTA” kernel that regresses to below 1× is caught before a new one is built. Then exploit the model’s actual structure, the two experts that generic tooling ignores. Two of the three wins came from deleting machinery and changing a threshold, not from writing a new kernel.
The three findings also compose into the next build. Lever one swapped SpargeAttn’s expensive gate for a
cheap selector on a fixed keep-fraction, and lever three found the signal that fraction should track,
max_prob, already sitting free in the kernel’s registers and already following the expert switch on its
own. Feeding that diffuseness signal into the selector as an adaptive per-layer budget turns three
separate results into one kernel, a domain-correct selector spending a free, self-calibrating budget on
the tool’s own tuned math. Re-measure the cache under full guidance, and that composed kernel is a patch
away, not a rewrite.