intel-ai/Wan2.2-I2V-A14B-Lightning-Diffusers-MXFP8-AutoRound

Wan2.2-I2V-A14B-Lightning · MXFP8 (AutoRound OptRTN)

This repository contains an MXFP8 W8A8 quantized version of Wan2.2-I2V-A14B-Lightning, produced with AutoRound OptRTN and calibration for NVIDIA Blackwell (sm_120) + vLLM-Omni inference.

Model Summary

Item	Value
Base model	Wan-AI/Wan2.2-I2V-A14B
Acceleration LoRA	lightx2v/Wan2.2-Lightning, I2V 4-step, rank64 Seko V1
Conversion	LightX2V converter.py (LoRA merge) + vLLM-Omni assemble_wan22_i2v_diffusers.py
Quantization	AutoRound 0.17.0
Scheme	MXFP8, OptRTN + calibration
Calibration data	opens2v, 8 samples × 4 denoise steps
Model size	39 GB
Inference	vLLM-Omni 0.29, FlashInfer FlashInferCutlassMxfp8LinearKernel (W8A8)

Both diffusion transformers (transformer and transformer_2) are quantized. The text encoder, VAE, tokenizer, and scheduler are unchanged.

Quantization Details

transformer/quantization_config.json and transformer_2/quantization_config.json:

{
  "bits": 8,
  "act_bits": 8,
  "data_type": "mx_fp",
  "act_data_type": "mx_fp",
  "group_size": 32,
  "act_group_size": 32,
  "sym": true,
  "act_sym": true,
  "act_dynamic": true,
  "block_name_to_quantize": "blocks",
  "quant_method": "auto-round",
  "packing_format": "auto_round:llm_compressor"
}

Coverage per transformer:

• Quantized: 400 linear layers inside the 40 transformer blocks (q/k/v/out for attn1 and attn2; ffn.net.0.proj and ffn.net.2). Weights are stored as FP8 E4M3 with per-32 E8M0 scales; activations are dynamically quantized to MXFP8 at runtime.
• Kept at higher precision:
  • F32: time_embedder, norm2, scale_shift_table;
  • BF16: text_embedder, time_proj, patch_embedding, proj_out, norm_q / norm_k;
  • F16: biases of the 400 quantized linear layers.

The module-count log reports 401/407; the six unquantized linear layers are the four condition_embedder linears, time_proj, and proj_out, plus patch_embedding. The quantized weights are about 13.99 GiB per transformer.

Inference

No --quantization flag is needed; vLLM-Omni detects the format from quantization_config.json.

Start the server from this model directory:

export MODEL_DIR=$(pwd)
export CUDA_VISIBLE_DEVICES=0,1

python -m vllm_omni.entrypoints.cli.main serve "$MODEL_DIR" \
  --omni --host 127.0.0.1 --port 8000 \
  --tensor-parallel-size 2 --vae-patch-parallel-size 2 \
  --boundary-ratio 0.9 --flow-shift 1.0 --enforce-eager

Adjust --tensor-parallel-size and --vae-patch-parallel-size to your GPU count.

The startup log should contain:

Using FlashInferCutlassMxfp8LinearKernel for MXFP8 GEMM

If FlashInfer JIT fails with a missing cuda-13.0 path, set CUDA_HOME=/usr/local/cuda and clear the stale FlashInfer JIT cache.

Running the Official Example

Official example input:

• Image: examples/i2v_input.JPG (832×1104, 3:4)
• Prompt: examples/prompt.txt
• 81 frames, 16 FPS, 4 steps, guidance 1.0, boundary 0.9, seed 0

About 832×1104: for I2V, the official --size 1280*720 is passed as max_area=1280*720=921600; the output aspect ratio follows the input image. For this 3:4 image, WanI2V.generate() computes 832×1104 (about 720P area). Forcing 1280×720 would change the composition.

With the server running, from this model directory (and with requests installed), execute the following in another terminal:

python - <<'PY'
import json, time, requests

base = "http://127.0.0.1:8000"
prompt = open("examples/prompt.txt", encoding="utf-8").read().strip()
with open("examples/i2v_input.JPG", "rb") as f:
    files = {"input_reference": ("i2v_input.JPG", f, "image/jpeg")}
    data = {
        "prompt": prompt,
        "width": "832", "height": "1104",
        "num_frames": "81", "fps": "16",
        "num_inference_steps": "4",
        "guidance_scale": "1.0",
        "boundary_ratio": "0.9",
        "flow_shift": "1.0",
        "seed": "0",
        "extra_params": json.dumps({"sample_solver": "euler"}),
    }
    job = requests.post(f"{base}/v1/videos", data=data, files=files).json()

while job["status"] not in ("completed", "failed"):
    time.sleep(2)
    job = requests.get(f"{base}/v1/videos/{job['id']}").json()
assert job["status"] == "completed", job

video = requests.get(f"{base}/v1/videos/{job['id']}/content")
open("example.mp4", "wb").write(video.content)
print("saved example.mp4")
PY

For the other attention backends, append one of the following to the server command:

--diffusion-attention-backend SAGE_ATTN_1   # Sage v1
--diffusion-attention-backend SAGE_ATTN     # Sage v2
--diffusion-attention-backend SAGE_ATTN_3   # Sage v3

Measured on this model:

Attention backend	Resolution	Generation time	motion	Sample
FlashInfer (CUDNN_ATTN)	832×1104	99.3 s	38.2	mxfp8_flashinfer.mp4
Sage v1 (SAGE_ATTN_1)	832×1104	93.3 s	37.7	mxfp8_sage_v1.mp4
Sage v2 (SAGE_ATTN)	832×1104	81.5 s	38.2	mxfp8_sage_v2.mp4
Sage v3 (SAGE_ATTN_3)	832×1104	78.3 s	30.5	mxfp8_sage_v3.mp4

VBench-I2V 50-pair

imaging_quality on a 960×960, 81-frame, 4-step, 50-pair VBench-I2V subset:

Attention backend	imaging_quality	Δ vs BF16
BF16 CUDNN_ATTN (reference)	0.736986	—
FlashInfer	0.735057	-0.26%
Sage v1	0.735935	-0.14%
Sage v2	0.735655	-0.18%
Sage v3	0.732355	-0.63%

• v1 and v2 are within noise of both FlashInfer and BF16; v3 is the lowest.
• This is a 50-prompt subset, not the full VBench. Calibration used an 8×4 budget (the model card uses 32×8), so these numbers are only comparable within this local evaluation.

Performance Profile

Kernel time for one 960×960 / 81-frame request (torch profiler, TP=2):

Part	FlashInfer	Sage v1	Sage v2	Sage v3
attention	45.8 s	38.5 s	24.6 s	19.8 s
linear (MXFP8 W8A8)	13.9 s	14.0 s	13.8 s	14.6 s
total kernel	93.5 s	86.7 s	72.9 s	71.2 s

The W8A8 linear layers remain the main non-attention cost; Sage v2/v3 mainly reduce attention time.

Reproduction

Conversion: merge the I2V 4-step LoRA with the official LightX2V converter.py, then assemble the model with vLLM-Omni's assemble_wan22_i2v_diffusers.py and the official Diffusers skeleton. Before quantization, rename blocks.N.modulation to blocks.N.scale_shift_table, which diffusers expects.

Quantization:

export DIFFUSERS_KEYS_DIR=/path/to/Wan2.2-I2V-A14B-Lightning-DiffusersKeys
export OUTPUT_DIR=/path/to/Wan2.2-I2V-A14B-Lightning-Diffusers-MXFP8-calib

CUDA_VISIBLE_DEVICES=0,1 auto_round \
  --model "$DIFFUSERS_KEYS_DIR" \
  --scheme MXFP8 \
  --dataset opens2v --nsamples 8 --calib_num_inference_steps 4 \
  --iters 0 \
  --disable_model_free --device_map 0,1 --low_gpu_mem_usage \
  --diffusion_tuning_cache_size auto --batch_size 1 \
  --output_dir "$OUTPUT_DIR"

Measured time: about 89 minutes (calibration cache + OptRTN over 80 blocks).

Limitations

• Calibration used only 8 samples × 4 steps, below the model card's 32×8; rerun with a larger budget for closer accuracy alignment.
• VBench results are a 50-prompt subset with a local 960×960 protocol, not the full VBench, and are not directly comparable to the model card.
• All numbers were measured on 8× RTX PRO 5000 72GB (sm_120) + vLLM-Omni; other hardware and runtimes may differ.
• For research and local evaluation only. Follow the licenses of the base model and the LoRA.

Signed-off-by: Kaokao, Lv kaokao.lv@intel.com