Skip to content
BranchesTags

Repository files navigation

SparkDiffusion

A DiT video-generation acceleration framework for 265× faster inference

English | 简体中文

Blog Hugging Face License

News

Overview

SparkDiffusion is a video-generation acceleration framework for Diffusion Transformer (DiT) models. It combines sparse low-rank attention (RoLa), few-step distillation (CrossDistill), and custom high-performance operators to deliver 200×+ end-to-end inference speedups over the dense multi-step baseline, while preserving generation quality.

The framework targets the Wan 2.1 and Wan 2.2 video diffusion models and provides an end-to-end pipeline: sparse-attention finetuning, few-step distillation, and optimized single-case T2V/I2V inference. It also ships optional weight-activation quantization and self-developed inference operators that require no external sparse-attention checkout.

Demo

demo.mp4

Resources

Model Weights

All released checkpoints are available in the SparkDiffusion Hugging Face collection:

Model Name Model Link Resolution Sparsity Steps
SparkWan2.1-T2V-1.3B-480P-0.90Sparsity Huggingface Model 480P 0.90 4
SparkWan2.1-T2V-14B-480P-0.90Sparsity Huggingface Model 480P 0.90 4
SparkWan2.1-T2V-14B-720P-0.97Sparsity Huggingface Model 720P 0.97 4
SparkWan2.1-T2V-14B-720P-0.95Sparsity-3Step Huggingface Model 720P 0.95 3
SparkWan2.1-I2V-14B-720P-0.97Sparsity Huggingface Model 720P 0.97 4
SparkWan2.2-T2V-14B-480P-0.95Sparsity Huggingface Model 480P 0.95 4

Highlights

  • 200×+ inference acceleration through joint sparse attention, low-rank factorization, and few-step distillation.
  • RoLa sparse low-rank attention — an efficient attention design usable for both training (finetuning/distillation) and inference.
  • CrossDistill few-step distillation — trajectory-level hybrid distillation that balances generation quality and diversity.
  • Custom high-performance operators under sparkdiffusion/ops/, self-developed and dependency-free at inference time.
  • Wan 2.1 & Wan 2.2 support for both T2V and I2V, with dense/sparse/distilled checkpoints comparable through the same inference wrappers.
  • Pluggable sparse-attention registry so custom attention variants integrate without touching the core training/distillation/inference code.

Repository Layout

Path Purpose
sparkdiffusion/ Model, sampler, dataset, checkpoint, inference, and operator code
imaginaire/ Training framework and configuration utilities
scripts/sparse_finetune/ Wan 2.1/2.2 sparse finetuning launchers
scripts/distill/ Distillation launchers for supported configurations
scripts/inference/ Single-case shell wrappers for inference
datasets/distill/ Local distillation dataset mount point
datasets/rola/ Local sparse-finetuning dataset mount point
pretrain_weights/ Local pretrained model mount point
outputs/distill/ Distillation outputs
outputs/rola/ Sparse-finetuning outputs

Weights, datasets, checkpoints, and generated videos are intentionally not included in the repository.

Requirements

  • Linux with a CUDA-capable GPU
  • Python 3.10 or newer
  • A CUDA-compatible PyTorch installation
  • Triton, flash-attn, and the packages listed in requirements.txt

Install the Python dependencies after installing the CUDA-matched PyTorch:

pip install -r requirements.txt
source scripts/env.sh

scripts/env.sh adds the repository to PYTHONPATH, enables offline defaults for Hugging Face and W&B, and defines repository-relative data/output roots. Override any root when local storage is elsewhere:

PRETRAIN_ROOT=/path/to/pretrain_weights \
DISTILL_DATA_ROOT=/path/to/distill_data \
ROLA_DATA_ROOT=/path/to/rola_data \
DISTILL_OUTPUT_ROOT=/path/to/distill_outputs \
ROLA_OUTPUT_ROOT=/path/to/rola_outputs \
source scripts/env.sh

Quick Start

The inference examples below use Wan 2.1 T2V at 480p. You need:

  • A native Wan 2.1 T2V model repository containing the VAE, T5 encoder, tokenizer, and base DiT assets.
  • A compatible SparkDiffusion checkpoint from the SparkDiffusion Hugging Face collection, such as SparkWan2.1-T2V-14B-480P-0.90Sparsity.
  • A CUDA-capable GPU. Use fp8 on GPUs with FP8 Tensor Core support; use bf16 otherwise.

Clone the repository and install a CUDA-matched PyTorch and a compatible flash-attn build before installing the remaining dependencies:

git clone git@github.com:AlibabaResearch/SparkDiffusion.git
cd SparkDiffusion
pip install -r requirements.txt
source scripts/env.sh

Place the native Wan assets under pretrain_weights/, for example:

pretrain_weights/
└── Wan2.1-T2V-14B/
    ├── Wan2.1_VAE.pth
    ├── models_t5_umt5-xxl-enc-bf16.pth
    ├── google/umt5-xxl/
    └── diffusion_pytorch_model-*.safetensors

Download the SparkDiffusion checkpoint separately and set DIT_PATH to its checkpoint file or directory. Then run three sequential samples in one process to see the warmup and steady-state timings:

NUM_SAMPLES=3 SEED=0 \
DIT_PATH=/path/to/SparkWan2.1-T2V-14B-480P-0.90Sparsity \
bash scripts/inference/eval_student_2pt1_distilled.sh \
  pretrain_weights/Wan2.1-T2V-14B \
  outputs/inference/quickstart \
  4 fp8 14B_rola 0.1 "" \
  "A cat playing in the garden under the sun."

The first sample is labeled warmup and may include compilation or kernel autotuning. Later samples are labeled after warmup and are the appropriate ones for steady-state latency comparisons. Generated videos are saved under outputs/inference/quickstart; with NUM_SAMPLES=3, each sample is saved separately with a _sample_<index>_seed_<seed> suffix.

Data and Weights

Use the following layout convention:

pretrain_weights/
  Wan2.1-T2V-14B/
  Wan2.1-I2V-14B-480P/
  Wan2.2-T2V-A14B/
datasets/
  distill/
  rola/
outputs/
  distill/
  rola/

Wan 2.1 native repositories keep the DiT safetensors, Wan2.1_VAE.pth, models_t5_umt5-xxl-enc-bf16.pth, and google/umt5-xxl directly under the model root; I2V repositories additionally contain the native CLIP .pth. Wan 2.2 native repositories keep shared assets at the root and the two DiT experts under high_noise_model/ and low_noise_model/.

The exact dataset shard names are experiment-specific. Set DATASET for sparse finetuning or DATASET_ROOT for distillation when using a different layout.

Sparse Finetuning

Wan 2.1:

SLA_SRC=path/to/SLA MAX_ITER=250 SAVE_ITER=50 \
  bash scripts/sparse_finetune/run_finetune_2pt1.sh

Wan 2.2 high-noise and low-noise experts with a native model repository:

export SLA_SRC=path/to/SLA

MAX_ITER=250 SAVE_ITER=50 EXPERT=high bash scripts/sparse_finetune/run_finetune_2pt2.sh
MAX_ITER=250 SAVE_ITER=50 EXPERT=low bash scripts/sparse_finetune/run_finetune_2pt2.sh

Use EXPERT=joint for the joint two-expert training path. Use EXPERT=both to launch high-noise and low-noise training sequentially. RoLa training requires SLA_SRC to point to the external SLA checkout before launch; every training launcher validates it before starting torchrun. Important overrides include PRETRAINED_CKPT, DATASET, NUM_GPUS, MAX_ITER, SAVE_ITER, BATCH_SIZE, LR, and EXPERIMENT. MAX_ITER sets the total number of optimizer update steps, and SAVE_ITER sets the checkpoint interval in optimizer steps. Override these variables to choose your own training length and save interval.

Pretrained checkpoint (PRETRAINED_CKPT)

The finetuning loader auto-detects the checkpoint format and adapts the state dict. Point PRETRAINED_CKPT at the path required by your format:

Format Required path
Native Wan 2.1 (default) The model repository directory, e.g. pretrain_weights/Wan2.1-T2V-1.3B, containing diffusion_pytorch_model.safetensors (optionally sharded with a *.index.json). You may also pass the .safetensors file directly.
Native Wan 2.2 The required expert directory, e.g. pretrain_weights/Wan2.2-T2V-A14B/high_noise_model or low_noise_model, each containing native sharded safetensors.
.pth / .pt A Wan-official or SparkDiffusion training checkpoint file.
DCP A distributed-checkpoint directory containing *.distcp shards.

By default the launchers load the native Wan repository directory. If the path or format is wrong, loading fails fast: when a checkpoint matches zero backbone parameters the loader raises an error (instead of silently training from random weights), and a partial match logs a warning.

Note: RoLa sparse parameters (proj_q, proj_k, gate_proj, gate_bias) are newly added and are expected to be missing from a stock checkpoint; they start at random init and are trained during finetuning. Only missing backbone weights indicate a wrong path/format.

Distillation

The supported distillation launchers are grouped under scripts/distill/:

export SLA_SRC=path/to/SLA
bash scripts/distill/wan2.1_14b_t2v_480p.sh
bash scripts/distill/wan2.1_14b_i2v_480p.sh

STUDENT_CKPT=path/to/high_noise_student.pth \
STUDENT_CKPT_LOW=path/to/low_noise_student.pth \
DATASET_ROOT=path/to/distillation_dataset \
  bash scripts/distill/wan2.2_a14b_t2v_480p_joint.sh

Each launcher uses repository-relative defaults. Override WAN_REPO, STUDENT_CKPT, TEACHER_CKPT, DATASET_ROOT, NEG_EMBED, and OUTPUT_ROOT for a different local layout. The Wan 2.2 joint launcher loads both noise experts in one process (also TEACHER_CKPT_LOW / STUDENT_CKPT_LOW).

The VAE, text encoder, tokenizer, and DiT paths can be overridden independently:

  • VAE_PATH, T5_PATH, TOKENIZER_PATH (and CLIP_ENCODER for Wan 2.1 I2V).
  • Wan 2.1 uses the native repository layout: ${WAN_REPO}/Wan2.1_VAE.pth, ${WAN_REPO}/models_t5_umt5-xxl-enc-bf16.pth, ${WAN_REPO}/google/umt5-xxl, and native DiT safetensors at the repository root.
  • For a native Wan 2.2 repository, set WAN_REPO=pretrain_weights/Wan2.2-T2V-A14B, TEACHER_CKPT=${WAN_REPO}/high_noise_model, TEACHER_CKPT_LOW=${WAN_REPO}/low_noise_model, and use the shared root assets ${WAN_REPO}/Wan2.1_VAE.pth, ${WAN_REPO}/models_t5_umt5-xxl-enc-bf16.pth, and ${WAN_REPO}/google/umt5-xxl through VAE_PATH, T5_PATH, and TOKENIZER_PATH.

RoLa distillation requires SLA_SRC even when the teacher is dense because the student sparse-attention path needs the external backward kernel. The launchers validate it and every other required path before starting, then abort with a clear message if one is missing.

Inference

Inference wrappers run one prompt per process. Use --prompt in the Python entrypoint or pass the prompt as the final positional argument to a shell wrapper. PROMPT_FILE is not used by the public inference path. The first positional argument is the native asset root for the VAE, text encoder, tokenizer, and optional CLIP encoder; set DIT_PATH to the distilled student checkpoint, and set CKPT_LOW for the Wan 2.2 low-noise student.

Set NUM_SAMPLES=3 (Python: --num_samples 3) to generate three videos sequentially in one process, reusing the loaded models with batch size 1. The default is 1; values must be positive integers. Seeds are SEED, SEED+1, and SEED+2. Each sample logs its seed and CUDA-synchronized denoising time: the first is labeled warmup (may include compilation and autotuning), and later samples are labeled after warmup. Compare the later times to observe warmed-up generation; decoding and video writing are excluded, while expert transfers during denoising are included. All samples, including the warmup sample, are saved separately. With multiple samples, filenames gain _sample_00_seed_0, etc.; a single sample retains the requested video filename. This replaces the previous batched meaning of num_samples.

Wan 2.1 distilled T2V:

NUM_SAMPLES=3 SEED=0 DIT_PATH=path/to/distill_model.pt \
  bash scripts/inference/eval_student_2pt1_distilled.sh \
  pretrain_weights/Wan2.1-T2V-14B \
  outputs/inference/wan21_t2v \
  4 fp8 14B_rola 0.1 "" \
  "A playful raccoon is seen playing an electronic guitar, strumming the strings with its front paws. The raccoon has distinctive black facial markings and a bushy tail. It sits comfortably on a small stool, its body slightly tilted as it focuses intently on the instrument. The setting is a cozy, dimly lit room with vintage posters on the walls, adding a retro vibe. The raccoon's expressive eyes convey a sense of joy and concentration. Medium close-up shot, focusing on the raccoon's face and hands interacting with the guitar."

Wan 2.1 distilled I2V:

NUM_SAMPLES=3 SEED=0 DIT_PATH=path/to/distill_model.pt \
  bash scripts/inference/eval_student_2pt1_distilled.sh \
  pretrain_weights/Wan2.1-I2V-14B-480P \
  outputs/inference/wan21_i2v \
  4 fp8 14B_rola 0.05 examples/i2v_input_1.jpg \
  "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."

Wan 2.2 distilled T2V:

NUM_SAMPLES=3 SEED=0 DIT_PATH=path/to/distill_high_noise_model.pt \
CKPT_LOW=path/to/distill_low_noise_model.pt \
  bash scripts/inference/eval_student_2pt2_distilled.sh \
  pretrain_weights/Wan2.2-T2V-A14B \
  outputs/inference/wan22_t2v \
  4 fp8 A14B_rola 0.1 \
  "A playful raccoon is seen playing an electronic guitar, strumming the strings with its front paws. The raccoon has distinctive black facial markings and a bushy tail. It sits comfortably on a small stool, its body slightly tilted as it focuses intently on the instrument. The setting is a cozy, dimly lit room with vintage posters on the walls, adding a retro vibe. The raccoon's expressive eyes convey a sense of joy and concentration. Medium close-up shot, focusing on the raccoon's face and hands interacting with the guitar."

The corresponding *_diffusion.sh wrappers run the original multi-step CFG sampler for teacher/reference comparisons. The final positional argument is always the text prompt. Use CKPT_LOW when Wan 2.2 high- and low-noise checkpoints are stored separately.

Common environment variables are NUM_FRAMES, RESOLUTION, ASPECT_RATIO, SEED, NUM_SAMPLES, OUT_ROOT, and FIXED_RESOLUTION for I2V. The topk argument is a keep ratio: 0.1, 0.05, and 0.03 correspond to 90%, 95%, and 97% sparsity.

Operators

The fused operators in sparkdiffusion/ops/fused_kernel/ are SparkDiffusion-specific, self-developed operators and carry Alibaba copyright headers. RoLa sparse finetuning and distillation require the external backward-compatible training kernel selected by the explicitly configured SLA_SRC environment variable. Standard dense and fused RoLa inference use the repository's inference operators and do not require SLA_SRC.

Checkpoints

Training checkpoints may use the repository's distributed checkpoint format. Inference loading supports the checkpoint layouts handled by sparkdiffusion.utils.model_utils.load_checkpoint_auto and the inference wrappers: native Wan repository directories, native safetensors, supported .pth / .pt checkpoints, and DCP directories. Checkpoint conversion is not required for standard launcher usage.

License Agreement

This repository is released under the Apache License 2.0. Model weights, datasets, upstream dependencies, and generated content may have separate licenses and usage restrictions.

Acknowledgments

We learned the design and reused or adapted code from the following projects:

We thank the authors and contributors of these projects for making their work available to the community. Relevant source files retain local attribution comments where an adaptation is implementation-specific. Please review the upstream licenses before redistributing derived artifacts.

Citation

If you use this code or find our work valuable, please cite:

@misc{liu2026sparkdiffusionmitigatinghighsparsitytrap,
  title={SparkDiffusion: Mitigating the High-Sparsity Trap --- A Unified Framework for up to $265\times$ Single-GPU Acceleration of Visual Generation},
  author={Yuxi Liu and Haoyu Li and Zekun Zhang and Tengxu Sun and Yixiang Cai and Jiayong Li and Yifei Xia and Tianle Liu and Baole Ai and Ang Wang and Jiamang Wang and Lin Qu and Kai Zhang and Kun Yuan and Bin Cui},
  year={2026},
  eprint={2609.23153},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2609.23153},
}

@misc{zhang2026rolarotarypositionedlowranklinear,
  title={RoLA: Rotary-Positioned Low-Rank Linear Attention for Efficient Diffusion Transformers},
  author={Zekun Zhang and Yixiang Cai and Yuxi Liu and Tengxu Sun and Tianle Liu and Zhoutong Wu and Haoyu Li and Baole Ai and Ang Wang and Jiamang Wang and Lin Qu and Kun Yuan},
  year={2026},
  eprint={2609.06712},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2609.06712},
}

@misc{liu2026crossdistillbalancingqualitydiversity,
  title={CrossDistill: Balancing Quality and Diversity via Trajectory-Level Hybrid Few-Step Distillation},
  author={Yuxi Liu and Haoyu Li and Yixiang Cai and Tengxu Sun and Zekun Zhang and Baole Ai and Ang Wang and Jiamang Wang and Lin Qu and Kun Yuan and Kai Zhang},
  year={2026},
  eprint={2609.14725},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2609.14725},
}

@misc{liu2026ropeslr3dropedrivensparselowrank,
  title={RoPeSLR: 3D RoPE-driven Sparse-LowRank Attention for Efficient Diffusion Transformers},
  author={Yuxi Liu and Zekun Zhang and Yixiang Cai and Renjia Deng and Yutong He and Kun Yuan},
  year={2026},
  eprint={2605.20659},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2605.20659},
}

About

No description, website, or topics provided.

Resources

Readme
Apache-2.0 license
Activity
Custom properties

Stars

196 stars

Watchers

17 watching

Forks

31 forks
Report repository

Releases

Packages

Contributors

Languages