[RFC]: NVIDIA ModelOpt Quantized Diffusion & Video Model Deployment in vLLM-OMNI

Motivation.

Diffusion and video-generation models (FLUX, Wan 2.2, LTX-2, HunyuanVideo, HunyuanImage-3) are compute- and memory-bound, which makes quantization a requirement for practical deployment rather than a nice-to-have. NVIDIA ModelOpt produces high-quality PTQ checkpoints in FP8 and NVFP4, with quantization_config metadata (quant_algo, quant_method: "modelopt") embedded in config.json. Today, vLLM-OMNI has no path to consume those checkpoints for diffusion models:

• vllm_omni/diffusion/ hosts ~26 transformer backbones, but only 8 thread quant_config through their linear layers (Flux, Flux2-Klein, HunyuanImage-3, Helios, Bagel, Qwen-Image, Z-Image). Wan 2.2 and LTX-2 still use plain nn.Linear.
• The ModelOpt integration in vllm/model_executor/layers/quantization/modelopt.py (ModelOptFp8Config, ModelOptNvFp4Config, ModelOptMxFp8Config, ModelOptMixedPrecisionConfig) is wired for LLMs only — diffusion pipelines don't hit that code path.
• Stage-config YAMLs (vllm_omni/model_executor/stage_configs/*.yaml) accept quantization: "fp8" as a string, but there's no auto-detection of ModelOpt metadata in diffusers checkpoints, and no NVFP4 path at all on the diffusion side.

The goal of this RFC is to close that gap: auto-detect ModelOpt checkpoints for diffusion/video models, reuse the existing FP8 / NVFP4 kernel infrastructure from the LLM runtime, and extend coverage to the major video backbones.

Proposed Change.

Architecture

Reuse vLLM's existing quantization plugin pattern; do not fork a parallel stack for diffusion.

ModelOpt Checkpoint (config.json: quantization_config)
        │
        ▼
┌─────────────────────────────────────────┐
│  Quantization Registry                   │
│  vllm/model_executor/layers/             │
│    quantization/__init__.py              │
│  "modelopt"      → ModelOptFp8Config     │
│  "modelopt_fp4"  → ModelOptNvFp4Config   │
└──────────────┬──────────────────────────┘
               │
        ┌──────┴──────┐
        ▼             ▼
┌──────────────┐ ┌──────────────────┐
│ FP8 Linear   │ │ NVFP4 Linear     │
│ Method       │ │ Method           │
└──────┬───────┘ └───────┬──────────┘
       │                 │
       ▼                 ▼
┌─────────────────────────────────────────┐
│  Shared Kernel Infrastructure            │
│  • utils/fp8_utils.py (CUTLASS W8A8)     │
│  • utils/nvfp4_utils.py                  │
│  • utils/marlin_utils_fp4.py             │
│  • flashinfer FP4/FP8 backends           │
└─────────────────────────────────────────┘

Key design decisions

1. Config-driven auto-detection. DiffusersPipelineLoader (vllm_omni/diffusion/model_loader/diffusers_loader.py) reads quantization_config from the checkpoint's config.json and dispatches to the matching ModelOpt*Config. No new CLI flag — quantization: "fp8" in stage YAMLs continues to work, and checkpoints with embedded ModelOpt metadata Just Work.
2. Ignore-list routing. ModelOpt checkpoints mark norm layers, embedders, and adapters as unquantized. Reuse the existing glob-based ignore matching in modelopt.py to route those modules to UnquantizedLinearMethod.
3. Auto-dequantization for diffusers-style modules. Blocks like AdaLayerNormZero hold linears that are quantized in the checkpoint but declared as nn.Linear in the module tree. Extend the diffusers loader to detect FP8/FP4 → bf16 dtype mismatches and dequantize using weight_scale on load.
4. Kernel reuse, no new kernels. Call into apply_fp8_linear and the NVFP4 utilities already used by the LLM runtime. HSDP sharding (vllm_omni/diffusion/distributed/hsdp.py) already supports mixed-precision policies; quantized params plug into that.
5. Thread quant_config through new backbones. Wan 2.2 and LTX-2 transformers (vllm_omni/diffusion/models/wan2_2/, vllm_omni/diffusion/models/ltx2/) need their linears replaced with ColumnParallelLinear / RowParallelLinear / QKVParallelLinear that accept quant_config, following the pattern in flux/flux_transformer.py.

Components per quantization method

Component	Location	Purpose
ModelOpt<Fmt>Config	vllm/model_executor/layers/quantization/modelopt.py (exists)	Parse checkpoint config, ignore-list matching
Diffusion loader hook	vllm_omni/diffusion/model_loader/diffusers_loader.py	Detect quantization_config, handle dtype mismatches
Stage-config wiring	vllm_omni/model_executor/stage_configs/*.yaml	Add fp8 / nvfp4 variants for new models
Backbone linears	vllm_omni/diffusion/models/{wan2_2,ltx2}/…	Thread quant_config through Column/Row/QKVParallelLinear

Roadmap

Phase 1 — FP8 image generation (ModelOpt auto-detect).

• Land auto-detection of quantization_config in diffusers_loader.py.
• Verify on already-quantized-capable backbones: Flux, Flux2-Klein, HunyuanImage-3, Qwen-Image, Z-Image.
• Add FP8 stage configs mirroring hunyuan_image3_moe_dit_2gpu_fp8.yaml.

Phase 2 — NVFP4 image generation.

• Route modelopt_fp4 through ModelOptNvFp4Config on the diffusion side.
• Use FlashInfer / CUTLASS FP4 backends via nvfp4_utils.py; group_size=16 block scaling.
• Target: Ada / Hopper / Blackwell.

Phase 3 — Video backbones.

• Wan 2.2: convert wan2_2_transformer.py linears to parallel-linear + quant_config; validate FP8 and NVFP4.
• LTX-2: same treatment on ltx2_transformer.py.
• Flux 2 completion (currently partial quant support).

Phase 4 — Hardening.

• Mixed-precision (FP8 attention + NVFP4 FFN) via ModelOptMixedPrecisionConfig.
• Unit tests under @diffusion / @cuda markers for each config + loader path.
• Benchmarks and user docs.

Execution plan

Phase	Deliverable	Dependencies
1	FP8 auto-detect in diffusion loader + Flux/Flux2-Klein/HunyuanImage-3 validation	—
1	Unit tests for ModelOpt FP8 config + loader	Phase 1
2	NVFP4 routing for FLUX / Qwen-Image / Z-Image	Phase 1, NVFP4 backend availability
3a	FP8 / NVFP4 Wan 2.2 (backbone rewrite)	Phase 1–2
3b	FP8 / NVFP4 LTX-2 (backbone rewrite)	Phase 1–2
4	Mixed-precision policies, benchmarks, docs	Phase 1–3

Open questions

1. NVFP4 backend default. The LLM runtime has multiple FP4 paths (FlashInfer cuDNN, CUTLASS, TRT-LLM, Marlin). Which is the sensible default for diffusion, and do DiT blocks need tuning distinct from LLM FFNs?
2. TP and HSDP interaction. Large video models need tensor parallelism. How should quantized weight sharding interact with the existing HSDP policies in vllm_omni/diffusion/distributed/hsdp.py?
3. Per-layer format selection. Should we expose a stage-config syntax for mixing FP8 attention with NVFP4 FFN, and what's the right YAML schema — extending the existing quantization: "fp8" string, or a structured block?
4. Alignment with RFC_unified_quantization.md. That RFC proposes a ComponentQuantizationConfig router for multi-stage models. Should ModelOpt dispatch register as a component there from day one, or land standalone and migrate later?

Feedback Period.

One week.

Any Other Things.

• A prototype for Phase 1 (FP8 auto-detect + Flux2-Klein validation) is straightforward to stand up against the existing ModelOptFp8Config; I can open a tracking PR once there's directional agreement.
• Debug notes from NVFP4 Flux 2 work (orange-square failure mode with otherwise normal alpha values) are relevant to Phase 2 and will be linked there.