Optimizing block-sparse INT8 attention and multi-expert step caching for Wan 2.2

Self-attentionEvery token attends to every other token, so cost grows with the square of the token count. That is why it dominates per-step compute here. accounts for approximately $73\%$ of per-step computational overhead in the Wan 2.2 video generation transformer. Execution operates across a 3D latent token gridThe compressed 3D array the model denoises: 21 time positions by 30 by 52 in space, 32,760 tokens flattened for attention. containing 32,760 tokens (21 frames × 30 × 52). Because quadratic sequence length scaling makes attention the primary bottleneck, acceleration strategies focus on two mechanisms: intra-step block-sparse attention pruning and inter-step residual feature caching.

On NVIDIA A100 GPUs, hardware quantization constraints cap dense INT8 attention gains to $1.33\times$ due to missing FP8 Tensor Core support for value matrix projections ($PV$). Furthermore, existing attention caches fail on Wan 2.2 because the pipeline swaps between two distinct 14B expert networksWan 2.2 swaps between two separate 14B networks by denoising stage: a high-noise layout expert early, a low-noise detail expert late. mid-generation.

By replacing published block selection gates with a lightweight mean-pooling selector and calibrating step caching independently for each expert network, we achieve 1.30–1.76× speedups in attention kernel execution and 1.78× speedups in step caching at near-lossless quality (65.7 dB PSNR).

Hardware Architecture and Precision Constraints on NVIDIA A100

Self-attention computes two primary matrix multiplications per head: score calculation $QK^T$ and value aggregation $PV$. On NVIDIA A100 hardware, INT8 Tensor Cores provide $2\times$ the instruction throughput of standard FP16 operations. However, because the A100 lacks FP8 hardware units, only $QK^T$ can be quantized to INT8 while $PV$ must execute in FP16 to preserve precision.

When half the FLOP workload runs at $2\times$ speed while the remaining half runs at $1\times$ speed, the maximum theoretical speedup for dense INT8 attention is bounded at $1.33\times$:

$$\text{Speedup}_{\text{dense}} = 1.0 / (0.5 \cdot 0.5 + 0.5 \cdot 1.0) = 1.33$$

Consequently, exceeding $1.33\times$ speedup on A100 GPUs requires skipping sequence block evaluations via sparse indexing.

FlashAttention-2 (FP16, dense baseline)1.00× · Baseline reference

INT8 QK^T + FP16 PV (A100 hardware limit)1.33× · Max theoretical speedup for dense INT8

SageAttention (Consumer FP8/INT8 units)2.1–3.0× · Hardware features absent on A100

Mathematical derivation of the 1.33× A100 execution limit

For $N$ tokens and head dimension $d$, $QK^T$ requires $2 N^2 d\text{ FLOPs}$ and $PV$ requires $2 N^2 d\text{ FLOPs}$. On A100 hardware, FP16 peak performance is $312\text{ TFLOPs}$ and INT8 peak performance is $624\text{ TOPs}$. Quantizing $QK^T$ reduces execution time from $T$ to $0.5 T$, while $PV$ remains at $0.5 T$. Total execution time is $0.75 T$, capping speedup at $1 / 0.75 = 1.33\times$.

Evaluation of Dense INT8 Attention Kernels

Custom dense INT8 Triton kernels written specifically for $QK^T$ score calculation achieved only $0.77\times$ the speed of FlashAttention-2The production attention kernel that computes exact attention without materializing the full score matrix. It is the dense baseline every lever must clear. due to kernel launch and quantization overheads.

Similarly, evaluating published sparse INT8 kernels such as SpargeAttnA published INT8 block-sparse attention kernel, with a tuned CUDA path and its own block selector. yielded performance near or below the dense baseline:

FlashAttention-2 (dense baseline)1.00× · Reference threshold

Custom Triton INT8 (dense)0.77× · Quantization overhead regression

SpargeAttn (out-of-the-box, worst layer)0.81× · Fallback to dense computation

SpargeAttn (out-of-the-box, typical layer)1.02× · Marginal gain

SpargeAttn (out-of-the-box, best layer)1.10× · Peak un-modified result

Profiling revealed that SpargeAttn’s default block selector—which evaluates block cosine similarity—marked 65%–99.6% of video latent blocks as “unpredictable”, forcing them to fall back to dense execution.

Passing the resulting sparse block mask into SpargeAttn’s execution kernel via its mask_id interface yielded 1.30–1.76× speedups across the 40 layers of Wan 2.2 A14B.

Spatial wrapping of 3D video latents in 128-token contiguous blocks

Video latents are structured as a 3D grid (21 × 30 × 52). Standard block-sparse kernels flatten this grid into a 1D sequence and divide it into contiguous 128-token blocks. Because a single 30 × 52 frame row contains 52 tokens, a 128-token block wraps across ~2.4 spatial rows. As a result, tokens within a single 128-token block span disjoint spatial regions, causing cosine similarity selectors to misclassify the block as unpredictable.

Decoupling Block Selection from Sparse CUDA Kernels

To resolve block selector misclassification without rewriting lower-level CUDA routines, we decoupled the block selection heuristic from SpargeAttn’s execution kernel.

The replacement block selector uses mean-pooling across query and key blocks:

1. Mean-pool each 128-token query and key block into single summary vectors $\bar{q}_i$ and $\bar{k}_j$.
2. Compute block pair interaction scores via dot products: $S_{i,j} = \bar{q}_i \cdot \bar{k}_j^T$.
3. Sort block scores per query block and retain key blocks based on a cumulative distribution threshold (CDF).

Passing the resulting sparse block mask into SpargeAttn’s execution kernel via its mask_id interface yielded $1.30\text{–}1.76\times$ speedups across the 40 layers of Wan 2.2 A14B.

FlashAttention-2 (dense baseline)1.00× · Reference threshold

Layer 39 (low-noise, dense fallback)1.22× · rel-L1 0.150 (exceeds quality gate)

Layer 10 (high-noise expert)1.48× · rel-L1 0.075

Layer 20 (high-noise expert)1.72× · rel-L1 0.089

Layer 20 (low-noise expert)1.76× · rel-L1 0.086

Layer & Expert Stage	FlashAttention-2 Latency	Mean-Pool Latency	Speedup	Relative-L1 Error
Layer 20 (low-noise expert)	117.8 ms	67.1 ms	1.76×	0.086
Layer 20 (high-noise expert)	118.3 ms	68.9 ms	1.72×	0.089
Layer 10 (high-noise expert)	118.2 ms	80.0 ms	1.48×	0.075

Middle transformer layers achieve 1.72–1.76× speedups while maintaining rel-L1 ≤ 0.09. Dense fallback logic preserves full precision on highly sensitive layers.

Selection Heuristic Ablation and Recall Analysis

To measure selection accuracy independently of execution speed, we compared the block selection masks generated by the cosine-gram gate and the mean-pool selector against an exact dense-attention ground truth.

Evaluation Probe	Cosine-Gram Gate Recall	Mean-Pool Selector Recall	Recall 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

Mean-pool selection matches cosine-gram recall to three decimal places while eliminating complex internal similarity transformations.

Ablation conclusion on selector complexity

Because the mean-pool selector reproduces the exact block recall of the cosine-gram gate ($0.921, 0.960, 0.865$), the additional computational complexity of cosine-gram similarity metrics adds no predictive value for 3D video latent sequences.

Multi-Expert Step Caching in Wan 2.2

Diffusion step cachingReusing a previous diffusion step's network output when the step-to-step change is small, so the step skips a full network evaluation. avoids redundant network evaluations by reusing latent residual updates from previous timesteps. Wan 2.2 A14B uses a Mixture-of-Experts (MoE) architecture containing two 14B sub-networks:

• High-Noise Layout Expert: Operates during early denoising steps (steps 1–13) to establish global scene structure.
• Low-Noise Detail Expert: Operates during late denoising steps (steps 14–40) to refine fine spatial details.

high-noise layout expertThreshold τ = 0.05steps 1–13

low-noise detail expertThreshold τ = 0.20steps 14–40

Existing step caching implementations apply a uniform distance threshold ($\tau$) across all timesteps. However, applying a single threshold fails when transitioning between distinct expert networks.

To address this, we implement a two-expert step cache:

1. Assign a conservative distance threshold ($\tau = 0.05$) to the high-noise layout expert to preserve global structure.
2. Assign an aggressive distance threshold ($\tau = 0.20$) to the low-noise detail expert, which exhibits high step-to-step similarity.
3. Enforce a hard cache flush at step 13 when switching network weights.

Caching Configuration	Generation Latency	Net Speedup	Output PSNR vs Un-cached Reference
No Cache (Baseline)	184.5 s	1.00×	Reference (Infinity)
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 Cache (τ_high=0.05, τ_low=0.20)	103.4 s	1.78×	65.7 dB

Per-expert thresholding and hard boundary resets achieve higher speedup (1.78× vs 1.60×) and superior fidelity (65.7 dB vs 64.6 dB) compared to uniform sweeps.

• uniform sweep
• two-expert

Mathematical formulation of residual step caching

At step $t$, given block input $x_t$ and previously evaluated network output $\hat{\epsilon}_{t-1}$, the cache computes relative L1 change: $\Delta = \|x_t - x_{t-1}\|_1 / \|x_{t-1}\|_1$. If accumulated drift $\sum \Delta < \tau_{\text{expert}}$, the model reuses $\hat{\epsilon}_{t-1}$. When transitioning across expert boundaries, accumulated drift is reset to $\infty$.

Composition of Step Caching and Block-Sparse Attention

Combining step caching ($1.78\times$) with block-sparse attention ($1.40\times$ average) yields a net composed speedup of approximately 2.1–2.3×.

The speedups do not multiply multiplicatively ($1.78 \times 1.40 = 2.49\times$) because step caching selectively skips late-stage timesteps—the exact steps where block-sparse attention achieves its highest sparsity rates.

Dynamic Sparsity Budgeting via Softmax Normalizer Registers

Instead of relying on pre-computed offline calibration tables to set block sparsity limits, dynamic sparsity budgets can be extracted directly from online FlashAttention register states.

FlashAttention tracks the softmaxThe function that turns attention scores into weights that sum to one. Its largest weight measures how peaked a row is. normalization factor $l_i = \sum_j \exp(s_{i,j} - m_i)$ in register memory for each query row $i$. The reciprocal of this factor, $\text{max\_prob}_i = 1 / l_i$, measures row attention peakedness:

• High $\text{max\_prob}$ ($\approx 1.0$): Attention mass is concentrated on a small subset of key blocks (high sparsity potential).
• Low $\text{max\_prob}$ ($\approx 1/N$): Attention mass is uniformly distributed (low sparsity potential).

Timestep Proxy (Offline Table)0.49 · Requires manual calibration

max_prob / Peakedness (Register State)0.75 · Zero-overhead online signal

• high-noise layout
• low-noise detail

Zero-overhead extraction of max_prob

Because $l_i$ is computed during the forward softmax pass of FlashAttention, evaluating $\text{max\_prob}_i = 1 / l_i$ requires zero additional memory reads or FLOP computations. $\text{max\_prob}$ correlates with achievable block sparsity at Spearman $\rho = 0.75$, outperforming offline timestep lookup tables ($\rho = 0.49$).

Implementation Considerations and Future Research

1. Classifier-Free Guidance (CFG) Caching: Production video generation uses CFG, evaluating conditional and unconditional network passes at each step. Cache implementations must maintain independent residual buffers for both conditional and unconditional paths.
2. Perceptual Video Evaluation: Quality verification for step caching and block-sparse attention should incorporate full-sequence perceptual metrics (FVD / LPIPS) across generated MP4 outputs alongside layer-wise rel-L1 gates.

Conclusion

Accelerating attention in Wan 2.2 requires aligning kernel heuristics with model architecture. Replacing complex cosine-similarity selectors with a mean-pooled block selector yields 1.30–1.76× speedups in INT8 attention. Structuring diffusion step caching around Wan 2.2’s dual-expert MoE boundary delivers an additional $1.78\times$ speedup at 65.7 dB PSNR.

Glossary

Glossary — every term, defined

Self-attention

Every token attends to every other token. Cost grows with the square of the token count, which is why it dominates per-step compute in this model.

Latent token grid

The compressed 3D array the model denoises, here 21 × 30 × 52 = 32,760 tokens (time by height by width), flattened into one sequence for attention.

FlashAttention (FA2)

The production attention kernel that computes exact attention without materializing the full score matrix. It is the dense baseline every lever must clear.

INT8 quantization

Running a matmul in 8-bit integers instead of 16-bit floats. On the A100 it accelerates the score matmul QK^T at ~2×, but not the value matmul PV.

Block-sparse attention

Attention that skips whole query-block by key-block pairs judged unimportant, computing only the blocks that carry attention mass.

SpargeAttn

A published state-of-the-art INT8 block-sparse attention kernel, with a tuned CUDA path, a self-similarity block gate, and a mask_id hook that accepts an external block mask.

Selector

The part of a sparse-attention method that decides which blocks to keep. Here a cheap mean-pool selector replaces SpargeAttn cosine-gram gate through the same kernel.

relative-L1 (rel-L1)

Mean absolute error of an approximate attention output against an FP32 reference, divided by the reference scale. The quality gate for the sparse lever is 0.09.

PSNR

Peak signal-to-noise ratio against a reference output, in dB. Higher is closer; above about 60 dB the difference is near-lossless. Used here to score the cache against the no-cache output.

Step cache

Reusing a previous diffusion step network output when the step-to-step change is small, skipping a full 14B network evaluation for that step.

Two-expert / mixture of experts

Wan 2.2 A14B carries two separate 14B networks routed by denoising stage: a high-noise layout expert early and a low-noise detail expert late, ~27B parameters total, 14B active per step.

Softmax

The function that turns attention scores into weights that sum to one. Its largest weight, read free as 1 / l_i, measures how peaked a query row is.

max_prob

The largest softmax weight in a query row, equal to 1 / l_i, read free from flash attention registers. High means peaked and sparsifiable, low means diffuse.

Classifier-free guidance (CFG)

A sampling method that evaluates the network twice per step, conditional and unconditional. A deployable cache needs a separate slot for each pass.