ASTRAEA: A GPU-Oriented Token-wise Acceleration Framework for Video Diffusion Transformers

ICLR 2026

Haosong Liu^1*, Yuge Cheng^1*, Wenxuan Miao¹, Zihan Liu^1,2, Aiyue Chen³, Jing Lin³, Yiwu Yao³

Chen Chen¹, Jingwen Leng^1,2, Minyi Guo^1,2 Yu Feng^1,2,#,

¹Shanghai Jiao Tong University ²Shanghai Qizhi Institute ³Huawei Technologies Co.,Ltd

^*Equal Contribution ^#Corresponding Author

Comparison of video generation speeds between original method and ours. We test on Wan with videos of 2s 480p resolution. Baseline and ours use 1 A100 GPU.

Abstract

We introduce ASTRAEA, an automatic framework that searches for near-optimal configurations for vDiT-based video generation.At its core, ASTRAEA proposes a lightweight token selection mechanism and a memory-efficient, GPU-parallel sparse attention strategy, enabling linear reductions in execution time with minimal impact on generation quality.To determine optimal token reduction for different timesteps, we further design a search framework that leverages a classic evolutionary algorithm to automatically determine the distribution of the token budget effectively.Together, ASTRAEA achieves up to 2.4x inference speedup on a single GPU with great scalability (up to 13.2x speedup on 8 GPUs) while retaining better video quality compared to the state-of-the-art methods (<0.5% loss on the VBench score compared to the baseline vDiT models).

-->

Quantitative evaluation of our method against the state-of-the-arts PAB, ToCa and Delta-Dit on two vDiT models: Wan v2.1 1.3B and OpenSora v1.2.

The speedup of AsTRAEA against the baseline models across various numbers of GPUs.

Video Comparisons Wan

Origin

Ours 0.7

PSNR: 29.21 Speedup: 1.36x

Ours 0.5

PSNR: 22.45 Speedup: 1.86x

Ours 0.4

PSNR: 22.35 Speedup: 2.29x

PAB 26

PSNR: 18.21 Speedup: 1.37x

TOCA 0.8

PSNR: 19.77 Speedup: 1.60x

Prompt: a dog running happily

Origin

Ours 0.7

PSNR: 29.37 Speedup: 1.36x

Ours 0.5

PSNR: 24.95 Speedup: 1.86x

Ours 0.4

PSNR: 24.29 Speedup: 2.29x

PAB 26

PSNR: 18.55 Speedup: 1.37x

TOCA 0.8

PSNR: 17.26 Speedup: 1.60x

Prompt: A jellyfish floating through the ocean, with bioluminescent tentacles

Origin

Ours 0.7

PSNR: 27.58 Speedup: 1.36x

Ours 0.5

PSNR: 23.51 Speedup: 1.86x

Ours 0.4

PSNR: 22.33 Speedup: 2.29x

PAB 26

PSNR: 16.37 Speedup: 1.37x

TOCA 0.8

PSNR: 18.87 Speedup: 1.60x

Prompt: A robot DJ is playing the turntable, in heavy raining futuristic tokyo rooftop cyberpunk night, sci-fi, fantasy

Origin

Ours 0.7

PSNR: 35.00 Speedup: 1.36x

Ours 0.5

PSNR: 25.76 Speedup: 1.86x

Ours 0.4

PSNR: 25.28 Speedup: 2.29x

PAB 26

PSNR: 18.28 Speedup: 1.37x

TOCA 0.8

PSNR: 24.08 Speedup: 1.60x

Prompt: A raccoon dressed in suit playing the trumpet, stage background

Video Comparisons HunYuan

Origin

Ours 0.7

PSNR: 29.67 Speedup: 1.39x

Ours 0.5

PSNR: 25.33 Speedup: 1.93x

Ours 0.4

PSNR: 24.14 Speedup: 2.38x

PAB 26

PSNR: 23.19 Speedup: 1.17x

SVG 0.7

PSNR: 19.33 Speedup: 1.69x

Prompt: The bund Shanghai, tilt down

Origin

Ours 0.7

PSNR: 30.69 Speedup: 1.39x

Ours 0.5

PSNR: 25.65 Speedup: 1.93x

Ours 0.4

PSNR: 24.14 Speedup: 2.38x

PAB 26

PSNR: 23.67 Speedup: 1.17x

SVG 0.7

PSNR: 22.06 Speedup: 1.69x

Prompt: Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, tilt down