CLiS: Plug-and-Play Attention Acceleration for Video Diffusion on TPUs via Certified Stabilization

Elisa Tsai; Shao-yu Huang; Daniel Mandragona; Wei-Yu Lin; WenXin Dong; Jevin Jiang; Chandra Devarakonda; Mehdy Bohlool

cs.LGOct 30, 2026

CLiS: Plug-and-Play Attention Acceleration for Video Diffusion on TPUs via Certified Stabilization

Authors: Elisa Tsai, Shao-yu Huang, Daniel Mandragona, Wei-Yu Lin, WenXin Dong, Jevin Jiang, Chandra Devarakonda, Mehdy Bohlool,

Organizations: Google

Abstract

High-resolution video diffusion relies on attention over long token sequences, making attention efficiency critical to generation latency. Our analysis identifies repeated maximum reductions and accumulator rescaling in online softmax as a bottleneck in long-context TPU attention. We propose CLiS (Certified Loop-invariant Stabilization), a training-free, plug-and-play optimization for video diffusion inference. CLiS fixes per-row stabilizers before the KV loop and certifies numerical safety per head, falling back to online softmax when certification fails. It integrates into existing fused-attention kernels without changing model weights or architecture. On Wan 2.2 and LTX-2 across TPU v6e and TPU7x, CLiS reduces diffusion transformer (DiT) step latency by up to 19.6% relative to tuned online softmax without measurable quality loss. Controlled TPU adaptations of prior GPU stabilization methods show performance or numerical-safety limitations. Preliminary causal LLM experiments also show up to 13.5% higher training throughput in continued pretraining and 7.7% in supervised fine-tuning.

Figures & tables

Figure 2: (1) Each unit’s modeled share of cycles per attention score (work divided by peak throughput). At

d{=}128

(dotted line, used by Wan 2.2 and LTX-2), exponentiation alone matches the matmuls. On TPU7x, vector-unit max reduction and rescaling push softmax past them. (2) Exposed softmax in the compiled VLIW bundle schedule for Wan 2.2 DiT self-attention (one KV-block iteration) on TPU7x.

Figure 3: Approximate numerical limits of

\operatorname{exp2}(x)=2^{x}

for FP16, BF16, and FP32 ( Kalamkar et al., 2019 ) .

Figure 4: Simplified TPU7x chiplet architecture and dataflow for attention ( Google Cloud, 2025 ; JAX Team, 2026b ) . Arrows show logical dataflow. Each TPU7x chip contains two chiplets.

Figure 5: Online-stabilization overhead for one compiled KV-block iteration. Wan 2.2 on TPU7x, generating a 5 s, 720p video (81 frames), at identical tile sizes. (a) A loop-invariant stabilizer

\mu_{i}

removes the reduction and rescaling chain. (b–c) Compiled VLIW schedules. Fixed stabilization reduces VPU-only issue bundles and shortens the schedule by

24.3\%

. (d–e) Scheduled bundles, VPU vmax reductions, VPU-only issue bundles, and VREG spill stores. Counts describe instruction issue, not elapsed unit-active cycles.

Panel of Figure 6 — Figure 6: Loop-invariant stabilization reduces spill/fill traffic across tile sizes and sequence lengths. One ring attention step of WAN 2.2 on TPU7x (20 heads, $d=128$ , BF16). Times are medians of 20 warm calls. Spill/fill traffic is derived from static LLO counts.

1	$\bar{k}\leftarrow N_{k}^{-1}\sum_{r=1}^{N_{k}}{\color[rgb]{0,0.4453,0.6992}k_{r}}$ .
2	$\widetilde{K}\leftarrow K-\mathbf{1}\bar{k}^{\top}$ .
3	$\kappa\leftarrow$ a conservative upper bound on $\max_{r}\\|{\color[rgb]{0,0.4453,0.6992}\widetilde{k}_{r}}\\|_{2}$ .
4	$B_{h}\leftarrow\max_{1\leq r\leq N_{q}}\textsc{BoundUp}({\color[rgb]{0.9023,0.625,0}q_{r}},\kappa)$ .
5	$G\leftarrow\lceil\log_{2}\max(1,\max_{j,r}\|{\color[rgb]{0,0.4453,0.6992}V_{jr}}\|)\rceil$ .
6	$A\leftarrow\lceil\log_{2}N_{k}\rceil+G$ .
7	$C\leftarrow\min\bigl(\lfloor U_{\exp}\rfloor,\lceil U_{\mathrm{acc}}-A\rceil-1\bigr)$ .

Algorithm 1 CLiS forward pass for one dense attention head.

Panel of Figure 7 — Figure 7: Normalized attention-kernel latency for CLiS and baselines, relative to online softmax (1.0 $\times$ ). Lower is better. All runs use eight TensorCores with CFG $=$ 2 $\times$ Ring $=$ 4 (Ulysses $=$ 1). Panels (a1)–(d2) use focused axes. Out-of-range bars are clipped and labeled with their measured ratios. Crosses mark corrupted outputs.

Panel of Figure 8 — Figure 8: Normalized DiT-step latency (continued). Panels (e1)–(h2) follow the same order and encoding as panels (a1)–(d2).

	Certification		Attention output		$\Delta$ VBench vs. online ( $\rightarrow 0$ )
Method	Valid	Safe-cond. rate	Max abs.	Rel. $\ell_{2}$	Imaging	Subject cons.	Temp. flicker	Motion smooth.	Dynamic degree
Online softmax	10/10	–	0	0	0	0	0	0	0
VFA-style	1/10	40.1	$3.1\!\times\!10^{-2}$	$2.53\!\times\!10^{-3}$	–	–	–	–	–
FD++-style	10/10	91.0	$3.1\!\times\!10^{-2}$	$2.73\!\times\!10^{-3}$	$+0.0011$	$-0.0005$	$+0.0012$	$-0.0001$	$-0.1000$
FA4-style (tile)	10/10	–	$3.1\!\times\!10^{-2}$	$2.05\!\times\!10^{-3}$	$+0.0032$	$+0.0013$	$-0.0001$	$+0.0010$	0
CLiS	10/10	97.9–98.3	$3.1\!\times\!10^{-2}$	$2.7\!\times\!10^{-3}$	$+0.0110$	$+0.0015$	$+0.0004$	$+0.0002$	0

Table 1: Safety and output quality on Wan 2.2. Certification and attention error cover ten TPU7x workloads (17K–184K tokens). Valid counts workloads without NaN/Inf or all-zero outputs. Safe-condition rates use different units and are not directly comparable. The CLiS rate spans the five 720p workloads, from the same run as Table 2 and Figure 9 . Errors and VBench deltas are relative to online softmax. VBench uses 110 prompts with matched seeds at 720p

\times

81f. VFA videos have all-zero frames and are unscored.

Table 2: Certified-head eligibility and the cost of the fallback path. Wan 2.2 T2V,

1280\!\times\!720

, 81 frames, 75,600 tokens.

Panel of Figure 9 — Figure 9: Runtime fallback rate of CLiS on Wan 2.2 (TPU7x) as a function of sequence length (left) and denoising step (right).

Model (Dataset)	Task	Seq.	Fwd. Attn.	Train Speedup	Loss (Base / CLiS)	$\Delta$ Loss	Cert.
Qwen3-8B (C4)	Pretrain	8K	$-37.9\%$	$+7.8\%$	3.182 / 3.185	$+0.10\%$	100%
Qwen3-8B (ClimbMix)	Pretrain	8K	$-37.9\%$	$+8.9\%$	2.678 / 2.673	$-0.18\%$	100%
Qwen3-8B (UltraChat)	SFT	8K	$-36.2\%$	$+7.7\%$	0.807 / 0.806	$-0.03\%$	95.8%
Qwen3-8B (PG-19)	Continual PT	256K	$-35.6\%$	$+12.4\%$	2.182 / 2.182	$-0.02\%$	95.5%
Llama-3.1-8B (UltraChat)	SFT	8K	$-37.8\%$	$+6.7\%$	0.752 / 0.752	$-0.07\%$	100%
Llama-3.1-8B (PG-19)	Continual PT	128K	$-43.3\%$	$+11.6\%$	1.992 / 1.992	$-0.03\%$	99.8%

Table 3: Representative causal LLM pretraining, continual pretraining, and SFT results on TPU7x (64 chips, or 512 chips for DSv3-671B). Forward attention kernel latency change, training speedup (step throughput or wall-clock time-to-train, TTT), and final-50 loss are paired against online softmax.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Dimension	Online softmax	VFA-style	FD++-style	FA4-style (tile)	FA4-style (row)	CLiS
Imaging quality	0.6658	– †	0.6668	0.6690	0.6672	0.6768
Subject consistency	0.9237	– †	0.9231	0.9250	0.9224	0.9252
Temporal flickering	0.9693	– †	0.9705	0.9692	0.9710	0.9697
Motion smoothness	0.9454	– †	0.9453	0.9465	0.9452	0.9456
Dynamic degree	1.0000	– †	0.9000	1.0000	0.9000	1.0000
Aesthetic quality	0.6226	– †	0.6240	0.6223	0.6216	0.6287

Appendix

Table 5: Absolute VBench scores for Wan 2.2 T2V at 720p

\times

81f on 110 prompts, using VBench’s per-dimension aggregates (higher is better). The first five dimensions appear in Table 1 . VFA-style, FD++-style, and FA4-style (row) use TPU v6e. Other methods use TPU7x. † All 110 VFA-style videos have all-zero frames and are unscored.

Dimension	FD++-style	FA4-style (tile)	FA4-style (row)	CLiS	$n$
Imaging quality	$+0.0011$ [ $-0.0158$ , $+0.0198$ ]	$+0.0032$ [ $-0.0061$ , $+0.0149$ ]	$+0.0014$ [ $-0.0124$ , $+0.0181$ ]	$+0.0110$ [ $+0.0015$ , $+0.0239$ ]	10
Subject consistency	$-0.0005$ [ $-0.0106$ , $+0.0096$ ]	$+0.0013$ [ $-0.0029$ , $+0.0061$ ]	$-0.0013$ [ $-0.0082$ , $+0.0068$ ]	$+0.0015$ [ $-0.0043$ , $+0.0076$ ]	10
Temporal flickering	$+0.0012$ [ $-0.0005$ , $+0.0035$ ]	$-0.0001$ [ $-0.0004$ , $+0.0002$ ]	$+0.0017$ [ $-0.0001$ , $+0.0041$ ]	$+0.0004$ [ $+0.0000$ , $+0.0010$ ]	10
Motion smoothness	$-0.0001$ [ $-0.0028$ , $+0.0026$ ]	$+0.0010$ [ $-0.0006$ , $+0.0032$ ]	$-0.0003$ [ $-0.0031$ , $+0.0032$ ]	$+0.0002$ [ $-0.0010$ , $+0.0016$ ]	10
Dynamic degree	$-0.1000$ [ $-0.3000$ , $+0.0000$ ]	$+0.0000$ [ $+0.0000$ , $+0.0000$ ]	$-0.1000$ [ $-0.3000$ , $+0.0000$ ]	$+0.0000$ [ $+0.0000$ , $+0.0000$ ]	10

Appendix

Table 6: Paired VBench differences from online softmax with 95% percentile-bootstrap confidence intervals over prompts.

n

counts prompts scored for both methods.

Attention Kernel Microbenchmark (ms)
Parallelism	CP=1, DP=8	Ulysses=8, Ring=1	Ulysses=4, Ring=2	Ulysses=2, Ring=4	Ulysses=1, Ring=8
Kernel family	Splash	Splash + A2A	2D Ring	2D Ring	Pure Ring
Local shard $(H_{\mathrm{loc}},S_{\mathrm{shard}})$	$(40,\,75{,}600)$	$(5,\,75{,}600)$	$(10,\,37{,}800)$	$(20,\,18{,}900)$	$(40,\,9{,}450)$
Online softmax	185.85	31.62	32.95	28.26	27.67
CLiS	173.32	29.54	29.74	25.06	24.66
Latency change	$\mathbf{-6.7\%}$	$\mathbf{-6.6\%}$	$\mathbf{-9.8\%}$	$\mathbf{-11.3\%}$	$\mathbf{-10.9\%}$

Appendix

Table 7: Parallelism and attention-kernel ablation on TPU7x-8 (Wan 2.2, 720p

\times

81f,

S=75{,}600,H=40,D=128

). Online softmax and CLiS use independently tuned tiles for each setting. All settings use the same eight TensorCores (four chips). The four context-parallel columns disable CFG and use eight-way context parallelism (CFG

=

1, Ulysses

\times

Ring

=8

), whereas Figure 7 and Tables 8 – 9 use CFG

=

\times

Ring

=

4. Absolute latencies are therefore not comparable across them.

Figure 12: Visual examples of video generation on Wan 2.2 (720p $\times$ 81f, 40 denoising steps, with frames 0, 16, 32, 48, 64, 80 shown) across VBench prompts comparing tuned online softmax and CLiS.

Wan 2.2, TPU v6e, 480p
Frames	Tokens	Online	FD++	VFA	FA4 tile	FA4 row	CLiS
41	17,160	4.089	4.809	6.931 †	7.624	4.086	3.344
81	32,760	14.156	18.320	12.926 †	21.098	14.152	11.012
121	48,360	30.666	37.374	26.123 †	47.123	30.645	24.169
161	63,960	52.009	61.793	43.164 †	64.258	52.045	40.841
201	79,560	80.151	97.233	64.993	97.864	80.280	61.110

Appendix

Table 8: Absolute attention-kernel latency (ms) for the 40 workloads in Figure 7 . Values are means. All runs use eight TensorCores with CFG

=

\times

Ring

=

4 (Ulysses

=

1). FD++ and FA4 denote the style baselines. † Corrupted output.

Wan 2.2, TPU v6e, 480p
Frames	Tokens	Online	FD++	VFA	FA4 tile	FA4 row	CLiS
41	17,160	0.606	0.635	0.768 †	0.754	0.606	0.574
81	32,760	1.159	1.325	1.233 †	1.445	1.159	1.039
121	48,360	2.055	2.306	1.873 †	2.696	2.053	1.797
161	63,960	3.159	3.527	2.807 †	3.643	3.160	2.719
201	79,560	4.550	5.203	3.935	5.254	4.555	3.799

Appendix

Table 9: Absolute DiT-step latency (s) for the 40 workloads in Figure 7 . Values are means. All runs use eight TensorCores with CFG

=

\times

Ring

=

4 (Ulysses

=

1). FD++ and FA4 denote the style baselines. † Corrupted output.

Efficient Attention

Efficient attention reduces the compute, memory traffic, and cache costs of Transformer attention while preserving useful long-range interactions.

See all 146 papers

AdaCluster: Adaptive Query-Key Clustering for Sparse Attention in Video Generation
Apr 20, 2026Haoyue Tan, Shengnan Wang, Yulin Qiao +5Video Diffusion Models Diffusion Transformer
University of Science and Technology of China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center · 3Independent Researcher +2
HyperVAttention: Efficient Sparse Attention with Spatio-Temporal Clustering for Video Diffusion
Jul 3, 2026Dongyeun Lee, Amir Zandieh, Vahab Mirrokni +2Video Diffusion Models Diffusion Transformer
KAIST · Google Research
TriTile: Three-Tier Attention for Video Diffusion Transformers
Oct 26, 2026Yeonjae Kim, Omin Kwon, Seong Hoon Seo +1Diffusion Model Inference Acceleration Video Diffusion Models
Department of Computer Science and Engineering, Seoul National University, Seoul, Republic of Korea.

Apr 20, 2026J/K move · Enter open · S save

AdaCluster: Adaptive Query-Key Clustering for Sparse Attention in Video Generation

Haoyue Tan, Shengnan Wang, Yulin Qiao, Juncheng Zhang+4

University of Science and Technology of China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center · 3Independent Researcher+2

Video Diffusion Models Diffusion Transformer Sparse Attention Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Video diffusion transformers (DiTs) suffer from prohibitive inference latency due to quadratic attention complexity. Existing sparse attention methods either overlook semantic similarity or fail to adapt to heterogeneous token distributions across layers, leading to model performance degradation. We propose AdaCluster, a training-free adaptive clustering framework that accelerates the generation of DiTs while preserving accuracy. AdaCluster applies an angle-similarity-preserving clustering method to query vectors for higher compression, and designs a euclidean-similarity-preserving clustering method for keys, covering cluster number assignment, threshold-wise adaptive clustering, and efficient critical cluster selection. Experiments on CogVideoX-2B, HunyuanVideo, and Wan-2.1 on one A40 GPU demonstrate up to 1.67-4.31x speedup with negligible quality degradation.

Abstract

Efficient Attention

Explore similar work

AdaCluster: Adaptive Query-Key Clustering for Sparse Attention in Video Generation

HyperVAttention: Efficient Sparse Attention with Spatio-Temporal Clustering for Video Diffusion

TriTile: Three-Tier Attention for Video Diffusion Transformers