Sglang Diffusion Modelopt Quant

SGLang Diffusion ModelOpt Quant

Overview

Use this skill when the task is to take a diffusion transformer through the full ModelOpt workflow:

• quantize it with NVIDIA ModelOpt
• adapt the exported checkpoint to SGLang Diffusion
• verify that quality holds up
• benchmark whether the quantized checkpoint is actually faster

This skill owns the ModelOpt-to-SGLang bridge. It is not a generic kernel-tuning skill.

Core Rules

• Use ModelOpt's official quantize.py as the PTQ source of truth.
• Keep the workflow generic. Put model-specific fallback logic in small isolated branches, not in the main conversion path.
• Benchmark only when BF16 and quantized commands are identical except for the checkpoint override being tested.
• For diffusion FP8, keep dit_cpu_offload=false. dit_layerwise_offload=true is valid on the fixed path when you want lower DiT residency.
• For multi-transformer pipelines, use per-component overrides when different components need different checkpoints.
• For B200 NVFP4 validation, keep backend-sensitive environment variables explicit. The current default is FlashInfer TensorRT-LLM (flashinfer_trtllm); high-resolution Qwen Image can favor SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=cutlass, while 1024x1024 can remain BF16-faster. Benchmark the exact shape instead of assuming one backend or quantized checkpoint wins.
• When a branch is missing the validated helper tools, refresh python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py, python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py, and python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py instead of inventing one-off scripts elsewhere.
• After validating a new ModelOpt quant path, update the ModelOpt support matrix in docs/docs/sglang-diffusion/quantization.mdx before closing the task.

Read First

Read these sources before changing code:

• NVIDIA ModelOpt diffusers guide: examples/diffusers/README.md
• ModelOpt quantization entrypoint: examples/diffusers/quantization/quantize.py
• ModelOpt diffusers quant presets: examples/diffusers/quantization/config.py
• SGLang diffusion quant runtime:
  • python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py
  • python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
  • python/sglang/multimodal_gen/runtime/utils/quantization_utils.py
  • python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
• Helper tools in this repo:
  • python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py
  • python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py
  • python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py

If you are working on a new model family, inspect the transformer's config and tensor naming before changing the generic converter.

What SGLang Supports Here

This repo now contains:

• flat quant_method=modelopt plus quant_algo=FP8/NVFP4 resolution
• diffusion-side ModelOpt FP8 linear loading
• diffusion-side NVFP4 loading from ModelOpt exports
• FLUX.2 packed-QKV detection that distinguishes packed NVFP4 checkpoints from standard diffusers exports
• automatic protection against incompatible FP8 CPU offload while keeping layerwise DiT offload available
• separate online diffusion quantization paths such as --quantization fp8 / mxfp4; keep those out of this ModelOpt PTQ/export workflow unless the user explicitly asks for runtime quantization
• FP8 transformer build: python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py
• NVFP4 mixed transformer build: python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py
• trajectory similarity validation: python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py

Validated documentation and CI coverage currently center on these ModelOpt diffusion transformer override families:

• FP8: FLUX.1-dev, FLUX.2-dev, Wan2.2, HunyuanVideo, Qwen Image, Qwen Image Edit
• NVFP4: FLUX.1-dev, FLUX.2-dev, Wan2.2, Qwen Image, Qwen Image 2512, Qwen Image Edit, Qwen Image Edit 2511

Treat a new family, a new precision, or a new checkpoint layout as unsupported until it has a documented matrix row and a matching validation story. Current B200 CI also contains an Ideogram4 NVFP4 native load case (ideogram4_nvfp4_t2i via Comfy-Org/Ideogram-4). Treat that as source evidence for an existing NVFP4 path, but do not expand the ModelOpt support matrix to Ideogram4 unless docs/docs/sglang-diffusion/quantization.mdx is updated with the exact checkpoint, loader path, quality check, and benchmark scope. Before writing CLI examples, re-read the active branch's docs/docs/sglang-diffusion/quantization.mdx: FLUX.2 NVFP4 is an official black-forest-labs/* repo rather than a lmsys/* converted repo, and its preferred flag depends on the current documented loader flow. Use --transformer-path for a component override directory with config.json; use --transformer-weights-path when the repo or path should be probed as raw weights.

B200 CI coverage can include loose BF16-vs-quantized quality checks. Inspect the active branch's run_suite.py before assuming they are part of the suite; mainline and feature branches may differ. Those checks are intended to catch blank, corrupted, or obviously divergent images, not exact image parity.

Mainline documentation now tracks thirteen published ModelOpt checkpoints. Twelve live under lmsys/*; the FLUX.2 NVFP4 raw export remains black-forest-labs/FLUX.2-dev-NVFP4. Do not use older BBuf/* examples unless you are explicitly testing a historical branch.

MiniMax-H3 boundary

MiniMax-H3 is current-main evidence for the separate online FP8 path, not a validated ModelOpt PTQ/export family. Its verified B200/B300 serving recipe loads the unquantized root checkpoint with --quantization fp8 and preserves the video/audio patch projections, timestep MLP, and final video/audio heads in FP32. Do not add H3 to the ModelOpt support matrix or run the generic ModelOpt converter until an exact H3 export, loader mapping, accuracy check, and benchmark scope have been validated.

If the user asks for current H3 online quantization, route the command and quality caveats through sglang-diffusion-performance and the MiniMax-H3 cookbook. Online FP8 is approximate and must be compared against eager BF16/FP32 for both video and audio; combining it with Cache-DiT compounds two approximations.

Related PR Watchlist

These related SGLang PRs are useful as ModelOpt diffusion support history. Re-check the PR state and the active source tree before treating any item as current behavior, and keep the docs/CI matrix as the support boundary.

• #23155 added Qwen Image ModelOpt FP8 support.
• #23199 adds HunyuanVideo ModelOpt FP8 support.
• #23373 adds a runtime quantization flag; keep PTQ/export workflows separate from runtime quant examples until the CLI behavior is merged.
• #24024 adds transformer FP8-cast compatibility mode.
• #24186 re-enables B200 multimodal CI with NVFP4 fixes for FLUX.2 and Wan2.2.

Do not expand the validated matrix beyond the documented rows solely because a related PR exists. Add a row only after the exact checkpoint, loader path, accuracy check, and benchmark scope are validated on the active branch.

Documentation Maintenance

• Keep the validated ModelOpt support matrix in docs/docs/sglang-diffusion/quantization.mdx.
• Each row should record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats.
• If the quantized DiT weights are not published yet, write unpublished explicitly instead of leaving the field blank.

FP8 Vs NVFP4

FP8 and NVFP4 are not wired into SGLang in exactly the same way.

FP8:

• the validated ModelOpt diffusers FP8 export still needs an extra SGLang-side conversion step
• SGLang expects explicit weight_scale and input_scale
• the validated path also materializes SGLang-native float8_e4m3fn weights from backbone.pt

NVFP4:

• the official diffusers export often already contains packed FP4 weights, scale tensors, and enough safetensors metadata for SGLang to rebuild the quant config
• in that case SGLang mainly needs to detect the checkpoint family and rearrange tensors into the runtime layout
• this is why NVFP4 often does not need an extra offline conversion pass like FP8 does
• backend choice matters on B200; record whether the run used the default CUTLASS path or a cuDNN-backed FlashInfer FP4 GEMM path

Important caveat:

• "often" does not mean "always"
• the exact load path still depends on the checkpoint family, especially whether a model uses a packed-QKV layout

Generic Workflow

1. Verify The BF16 Baseline First

Before quantizing anything:

• run the original BF16 model in SGLang
• fix the prompt, seed, size, step count, and GPU topology
• save the output and perf.json

Do not start quantization work until the BF16 path is already healthy.

2. Quantize With Official ModelOpt

Use ModelOpt's official script. Generic template:

python quantize.py \
  --model <model-name> \
  --override-model-path <hf-repo-or-local-model> \
  --model-dtype <Half|BFloat16> \
  --format <fp8|fp4> \
  --batch-size 1 \
  --calib-size <calib-size> \
  --n-steps <calib-steps> \
  --quantize-mha \
  --prompts-file <prompt-file> \
  --quantized-torch-ckpt-save-path <out>/ckpt \
  --hf-ckpt-dir <out>/hf

For current ModelOpt diffusion examples, use --format fp4 for NVFP4 exports. Do not assume the checked-out ModelOpt version accepts a literal nvfp4 format string unless you verified it locally.

For multi-transformer models:

• quantize each backbone deliberately
• keep each output directory separate
• save both backbone.pt and the matching hf/<component> export