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

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

This repository contains an MXFP4 W4A4 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	MXFP4, OptRTN + calibration
Calibration data	opens2v, 8 samples × 4 denoise steps
Model size	26 GB
Inference	vLLM-Omni 0.29, FlashInfer FlashInferMxFp4LinearKernel + cuDNN blockScaleMatMul (W4A4)

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": 4,
  "act_bits": 4,
  "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 packed FP4 E2M1 (two FP4 values per byte) with per-32 E8M0 scales; activations are dynamically quantized to MXFP4 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 7.45 GiB per transformer.

Inference

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

Required patch: the stock vLLM MXFP4 FlashInfer wrapper hardcodes the GEMM backend to cute-dsl, but FlashInfer 0.6.18 does not support MXFP4 with cute-dsl on sm120. Change the GEMM backend to cudnn:

VLLM_ROOT=$(python -c "import vllm, pathlib; print(pathlib.Path(vllm.__file__).parent.parent)")
patch -p1 -d "$VLLM_ROOT" < examples/mxfp4-flashinfer-cudnn.patch

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 FlashInferMxFp4LinearKernel for MXFP4 GEMM
cudnn_handle created for device_id = 0

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	90.3 s	29.7	mxfp4_flashinfer.mp4
Sage v1 (SAGE_ATTN_1)	832×1104	84.3 s	29.4	mxfp4_sage_v1.mp4
Sage v2 (SAGE_ATTN)	832×1104	69.3 s	27.1	mxfp4_sage_v2.mp4
Sage v3 (SAGE_ATTN_3)	832×1104	72.6 s	25.4	mxfp4_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.726135	-1.47%
Sage v1	0.725050	-1.62%
Sage v2	0.727190	-1.33%
Sage v3	0.722222	-2.00%

• W4A4 loses noticeably more quality than MXFP8 (about -1.3% to -2.0%);
• v2 is the best and v3 is the lowest; v3 follows the same ordering as in BF16, but W4A4 + FP4 attention compounds the loss;
• 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.9 s	38.6 s	24.6 s	19.4 s
linear (MXFP4 W4A4, cuDNN)	5.3 s	5.3 s	5.3 s	6.0 s
total kernel	84.9 s	78.2 s	64.2 s	62.3 s

MXFP4 linear layers are about 2.6× faster than MXFP8, so attention accounts for a larger share of DiT time (63.7% with FlashInfer).

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-MXFP4-calib

CUDA_VISIBLE_DEVICES=0,1 auto_round \
  --model "$DIFFUSERS_KEYS_DIR" \
  --scheme MXFP4 \
  --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 88 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.
• MXFP4 is W4A4 and has lower quality than MXFP8; use MXFP8 or BF16 when quality is critical.
• The MXFP4 cudnn patch modifies site-packages and is lost when vLLM is reinstalled or upgraded; reapply it after upgrades.
• 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