（Phase 2）Add ModelOpt mixed FP8/NVFP4 support for image generation (#3570)

Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com>

docs/user_guide/quantization/modelopt.md

@@ -6,34 +6,46 @@ ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
 quantized weights and scale tensors are generated before serving, so inference
 does not run online calibration or convert a BF16 checkpoint at startup.
 
-vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion
-transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs
-and keeps non-transformer components, such as the tokenizer, scheduler, text
+vLLM-Omni validates ModelOpt FP8, ModelOpt NVFP4, and ModelOpt mixed
+FP8/NVFP4 checkpoint loading for diffusion transformer stages. The loader
+auto-detects supported ModelOpt checkpoint configs and keeps non-transformer
+components, such as the tokenizer, scheduler, text encoder, vision/audio
 encoder, and VAE, on the base checkpoint unless a model-specific guide says
 otherwise.
 
 !!! note
-    `--force-cutlass-fp8` is an explicit runtime override for diffusion
-    checkpoints that already carry a supported ModelOpt FP8 config. It does not
-    quantize BF16 checkpoints and it does not apply to online `--quantization
-    fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA
-    SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the
-    normal vLLM kernel selection.
+    ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
+    `--quantization fp8` for these checkpoints. The checkpoint
+    `quantization_config` selects the ModelOpt path.
+
+!!! note
+    `--force-cutlass-fp8`, `--linear-backend cutlass`, and
+    `--moe-backend cutlass` are runtime backend selections for checkpoints that
+    already carry supported ModelOpt quantized weights and scales. They do not
+    quantize BF16 checkpoints at startup.
 
 ## Supported ModelOpt Checkpoint Formats
 
 vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
 checkpoint config must identify ModelOpt as the quantization method or producer,
-and the quantization algorithm must be one of the validated FP8 algorithms.
+and the quantization algorithm must be one of the validated algorithms below.
 
 | Checkpoint field | Supported value |
 |------------------|-----------------|
-| `method` / `quant_method` | `modelopt` |
+| `method` / `quant_method` | `modelopt`, `modelopt_fp4`, `modelopt_mixed` |
 | `producer.name` | `modelopt` |
-| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` |
+| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `NVFP4`, `MIXED_PRECISION` |
 
-Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion
-FP8 path until they have separate model and quality validation.
+| `quant_algo` | Runtime method | Typical use |
+|--------------|----------------|-------------|
+| `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` | `modelopt` | FP8 diffusion transformer checkpoints |
+| `NVFP4` | `modelopt_fp4` | NVFP4 diffusion transformer checkpoints |
+| `MIXED_PRECISION` | `modelopt_mixed` | Mixed FP8/NVFP4 checkpoints with a ModelOpt per-layer policy |
+
+For multi-component diffusion or omni models, only the checkpoint component
+that contains ModelOpt quantized weights should use the ModelOpt quantization
+method. Encoders, decoders, tokenizers, schedulers, and other BF16 components
+stay unquantized unless the model-specific recipe validates otherwise.
 
 ## Hardware Support
 
@@ -47,7 +59,10 @@ FP8 path until they have separate model and quality validation.
 | Ascend NPU | ❌ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
-The optional CUTLASS FP8 runtime override requires CUDA SM89+.
+The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
+and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
+CUDA systems in the recipes below; other CUDA generations require separate
+backend and quality validation.
 
 ## Model Type Support
 
@@ -56,12 +71,19 @@ The optional CUTLASS FP8 runtime override requires CUDA SM89+.
 | Model | HF checkpoint | Scope | Status |
 |-------|---------------|-------|--------|
 | Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy` | Diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
 | Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
 | HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
 | Wan2.2 | Not available | Diffusion transformer | Not validated |
 
+For full serving commands and benchmark context, see
+[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
+and
+[`recipes/Tencent/HunyuanImage-3.0-Instruct.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md).
+
 ### Multi-Stage Omni/TTS Model
 
 | Model | Scope | Status |
@@ -80,15 +102,28 @@ ModelOpt targets yet.
 
 ## Configuration
 
-For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is
-needed. The checkpoint config selects the ModelOpt path.
+For pre-quantized ModelOpt checkpoints, no `--quantization fp8` flag is needed.
+The checkpoint config selects the ModelOpt path.
 
 Online serving:
 
 ```bash
-vllm serve <modelopt-fp8-checkpoint> \
+vllm serve <modelopt-checkpoint> \
   --omni \
   --tensor-parallel-size <N> \
+  --linear-backend cutlass \
+  --force-cutlass-fp8
+```
+
+For mixed FP8/NVFP4 MoE checkpoints, also select the validated MoE backend:
+
+```bash
+vllm serve <modelopt-mixed-moe-checkpoint> \
+  --omni \
+  --tensor-parallel-size <N> \
+  --enable-expert-parallel \
+  --linear-backend cutlass \
+  --moe-backend cutlass \
   --force-cutlass-fp8
 ```
 
@@ -96,14 +131,14 @@ Offline inference:
 
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
-  --model <modelopt-fp8-checkpoint> \
+  --model <modelopt-checkpoint> \
   --tensor-parallel-size <N> \
   --prompt "a red ceramic teapot on a wooden table" \
   --height 1024 \
   --width 1024 \
   --num-inference-steps 20 \
   --seed 42 \
-  --output outputs/modelopt_fp8.png
+  --output outputs/modelopt.png
 ```
 
 Python API:
@@ -112,7 +147,7 @@ Python API:
 from vllm_omni import Omni
 
 omni = Omni(
-    model="<modelopt-fp8-checkpoint>",
+    model="<modelopt-checkpoint>",
     tensor_parallel_size=2,
     force_cutlass_fp8=True,
 )
@@ -123,16 +158,34 @@ omni = Omni(
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
+| `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages |
+| `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints |
 
 ## Validation and Notes
 
-1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same
+1. Compare the ModelOpt checkpoint against the BF16 baseline with the same
    prompt, resolution, seed, and inference steps.
 2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
    `VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
    paths.
-3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a
-   new validated ModelOpt diffusion checkpoint.
-4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
-   ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights
-   and scales.
+3. For HunyuanImage-3.0 quantized DiT checkpoints, the opt-in accuracy check is:
+
+   ```bash
+   CUDA_VISIBLE_DEVICES=2,3 \
+   HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY=1 \
+   HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
+   HUNYUAN_IMAGE3_QUANT_TP=2 \
+   HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
+   HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
+   HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+   PYTHONPATH=/path/to/vllm-omni:${PYTHONPATH:-} \
+   python -m pytest -s -v \
+     tests/e2e/accuracy/test_hunyuan_image3.py \
+     -k quantized_dit_matches_bf16_accuracy
+   ```
+
+4. Report CLIP score deltas, SSIM, PSNR, throughput, latency, and peak memory
+   when adding a new validated ModelOpt diffusion checkpoint.
+5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
+   ModelOpt path only when the checkpoint already contains ModelOpt quantized
+   weights and scales.

mkdocs.yml

@@ -120,7 +120,9 @@ plugins:
           inventories:
           - https://docs.python.org/3/objects.inv
           - https://typing-extensions.readthedocs.io/en/latest/objects.inv
-          - https://docs.aiohttp.org/en/stable/objects.inv
+          # Disabled because Read the Docs strict builds fail if this
+          # external inventory intermittently resets the connection.
+          # - https://docs.aiohttp.org/en/stable/objects.inv
           - https://pillow.readthedocs.io/en/stable/objects.inv
           - https://numpy.org/doc/stable/objects.inv
           # Temporarily disabled due to decompression errors

recipes/Qwen/Qwen-Image.md

@@ -130,3 +130,108 @@ and compare the output JSON or terminal metrics.
 - Key flags: `--step-execution` enables the step-wise runtime; `--max-num-seqs` controls how many compatible requests may stay active together.
 - Keep `--max-num-seqs 1` when you want the more conservative path, when traffic is mostly single-request, or when you are debugging correctness before measuring throughput.
 - Current batching is still shape-sensitive: different step progress can co-batch, but different resolutions do not yet co-batch.
+
+### 2x B200 ModelOpt mixed FP8/NVFP4
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs
+- vLLM version: Match the vLLM-Omni checkout used for deployment
+- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed
+  FP8/NVFP4 checkpoint loading for diffusion models
+
+#### Command
+
+This configuration uses a pre-quantized ModelOpt mixed checkpoint for
+Qwen-Image-2512. The validated checkpoint keeps boundary, modulation,
+normalization, and output layers in BF16, uses FP8 for sensitive attention
+projections, and uses NVFP4 for heavier MLP / attention-output linear layers.
+
+Validated checkpoint:
+
+```text
+feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy
+```
+
+Start the server with the CUTLASS quantized linear backend:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+vllm serve feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091 \
+  --trust-remote-code \
+  --tensor-parallel-size 2 \
+  --linear-backend cutlass \
+  --force-cutlass-fp8
+```
+
+#### Verification
+
+Run a direct image generation request after the server is ready:
+
+```bash
+curl http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy",
+    "prompt": "A cinematic Warhammer 40,000 Ultramarine warrior in blue power armor, armor stained with alien ichor, swinging a chainsword at a Tyranid creature, gothic battlefield, dramatic lighting, high detail, no text, no watermark",
+    "negative_prompt": "blurry, low quality, text, watermark, logo, malformed armor, bad hands",
+    "size": "1024x1024",
+    "num_inference_steps": 20,
+    "true_cfg_scale": 4.0,
+    "guidance_scale": 4.0,
+    "seed": 48
+  }'
+```
+
+Check that:
+
+- The server responds on `http://localhost:8091/health`.
+- The response contains one generated image.
+- Logs show CUTLASS selection for ModelOpt FP8 / mixed-precision linear
+  layers instead of falling back to the default automatic backend choice.
+
+#### Benchmark
+
+The numbers below are loaded-once request-level measurements, not an HTTP
+concurrency benchmark. Each case loads the model once, runs warmup requests,
+then measures 100 sequential requests.
+
+```text
+Hardware: 2x B200
+Tensor parallel size: 2
+Requests: 100
+Concurrency: 1
+Resolution: 1024x1024
+Denoising steps: 20
+```
+
+| Config | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 2.696s | 2.700s | 2.707s | 99,000 MiB |
+| ModelOpt FP8 CUTLASS | 2.663s | 2.671s | 2.676s | 85,760 MiB |
+| Mixed FP8/NVFP4 CUTLASS | 2.584s | 2.590s | 2.594s | 84,412 MiB |
+
+The same 2-GPU setup was also used to sweep NVFP4 linear backends. Direct
+CUTLASS was the fastest validated backend for this checkpoint.
+
+| Backend | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| auto | 2.927s | 2.948s | 3.535s | 85,340 MiB |
+| flashinfer-cudnn | 3.523s | 3.528s | 3.766s | 84,500 MiB |
+| flashinfer-trtllm | 2.711s | 2.727s | 3.010s | 84,996 MiB |
+| flashinfer-cutlass | 2.935s | 2.950s | 3.011s | 85,340 MiB |
+| cutlass | 2.578s | 2.584s | 2.592s | 84,412 MiB |
+
+#### Notes
+
+- Use `--linear-backend cutlass` to select the validated quantized linear
+  backend from the CLI.
+- `--force-cutlass-fp8` keeps ModelOpt FP8 layers on the CUTLASS FP8 path.
+- The benchmark table intentionally reports Qwen-Image-2512 as
+  concurrency-1 request-level data; do not compare it directly with an online
+  HTTP concurrency benchmark.

recipes/README.md

@@ -33,13 +33,14 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
+| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
 | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
 | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
 | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
-| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark | 4x H100/H800 80GB |
+| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -1,6 +1,8 @@
 # HunyuanImage-3.0-Instruct
 
-> DiT-only text-to-image serving and benchmark
+> DiT-only text-to-image serving and benchmark with FP8, tensor parallelism,
+> sequence parallelism, CFG parallelism, and ModelOpt mixed FP8/NVFP4
+> checkpoints.
 
 ## Summary
 
@@ -16,13 +18,15 @@ Use this recipe when you want to run or benchmark the HunyuanImage-3.0 DiT
 stage directly. This is the recommended first setup when validating DiT
 throughput, memory, FP8 kernels, sequence parallelism, or CFG parallelism.
 
-The recipe covers three 4-GPU configurations:
+The recipe covers three 4-GPU FP8 configurations and one 2-GPU ModelOpt mixed
+FP8/NVFP4 configuration:
 
 | Configuration | Parallelism | Notes |
 | --- | --- | --- |
 | `tp4_fp8` | TP=4 | Lowest per-GPU memory, higher communication overhead |
 | `tp2_fp8_sp2` | TP=2, SP=2, Ulysses=2 | Splits sequence work across two GPUs per TP group |
 | `tp2_fp8_cfgp2` | TP=2, CFG=2 | Runs CFG branches in parallel; fastest validated DiT setup |
+| `tp2_mixed_fp8_nvfp4` | TP=2, EP enabled | Uses FP8 dense layers and NVFP4 routed experts on B200 |
 
 ## References
 
@@ -222,3 +226,127 @@ of:
 - If you see OOM on 80GB GPUs, reduce image size, request concurrency, or use
   the TP=4 configuration before increasing batch size. GPU memory utilization
   is not a useful primary tuning knob for this DiT-only recipe.
+
+### 2x B200 ModelOpt mixed FP8/NVFP4
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs
+- vLLM version: Match the vLLM-Omni checkout used for deployment
+- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed
+  FP8/NVFP4 checkpoint loading for diffusion models
+
+#### Command
+
+This configuration uses a pre-quantized ModelOpt mixed checkpoint for the
+HunyuanImage-3.0 DiT stage. Dense attention and shared dense projections use
+FP8, while routed MoE expert projections use NVFP4. Embeddings, norms,
+routers, output layers, and non-DiT components stay in BF16.
+
+Validated checkpoint:
+
+```text
+feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8
+```
+
+Start the server with CUTLASS selected from the CLI for quantized linear and
+MoE kernels:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+vllm serve feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8091 \
+  --trust-remote-code \
+  --tensor-parallel-size 2 \
+  --enable-expert-parallel \
+  --linear-backend cutlass \
+  --moe-backend cutlass \
+  --force-cutlass-fp8
+```
+
+#### Verification
+
+Run a direct image generation request after the server is ready:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8",
+    "messages": [
+      {"role": "user", "content": "A cinematic portrait of an Adeptus Custodes warrior standing on Terra before the Imperial Palace, ornate golden armor, guardian spear, crimson cloak, marble steps, gothic architecture, holy golden light, highly detailed, no text, no watermark"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 20,
+      "guidance_scale": 4.0,
+      "true_cfg_scale": 4.0,
+      "seed": 48
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' \
+     | cut -d',' -f2- \
+     | base64 -d > hunyuan_image3_mixed_nvfp4.png
+```
+
+Check that:
+
+- The server responds on `http://localhost:8091/health`.
+- The response contains one generated image.
+- Logs show CUTLASS selection for ModelOpt FP8 dense layers and ModelOpt
+  NVFP4 routed expert layers.
+
+#### Benchmark
+
+The numbers below are online HTTP measurements with 16 concurrent requests.
+Each case serves 100 image-generation requests on the same 2-GPU B200 setup.
+
+```text
+Hardware: 2x B200
+Tensor parallel size: 2
+Requests: 100
+Concurrency: 16
+Resolution: 1024x1024
+Denoising steps: 20
+```
+
+| Config | Mean | P90 | P99 | Peak VRAM |
+| --- | ---: | ---: | ---: | ---: |
+| BF16 | 70.729s | 77.716s | 77.758s | 191,920 MiB |
+| ModelOpt FP8 | 58.603s | 64.435s | 64.469s | 115,816 MiB |
+| Mixed FP8/NVFP4 | 54.739s | 60.244s | 60.317s | 83,560 MiB |
+
+#### Accuracy
+
+Run the HunyuanImage-3.0 quantized DiT accuracy check with the BF16, FP8, and
+mixed FP8/NVFP4 checkpoints:
+
+```bash
+CUDA_VISIBLE_DEVICES=2,3 \
+HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
+HUNYUAN_IMAGE3_QUANT_TP=2 \
+HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
+HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
+HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
+PYTHONPATH=/path/to/vllm-omni:/path/to/vllm:${PYTHONPATH:-} \
+python -m pytest -s -v \
+  tests/e2e/accuracy/test_hunyuan_image3.py \
+  -k quantized_dit_matches_bf16_accuracy
+```
+
+The test compares quantized images with the BF16 image using CLIP-based image
+and text scores, plus structural SSIM / PSNR checks. The current structural
+thresholds are `SSIM >= 0.20` and `PSNR >= 10.0`.
+
+#### Notes
+
+- Use `--linear-backend cutlass` to select the validated quantized linear
+  backend from the CLI.
+- Use `--moe-backend cutlass` for the routed expert path.
+- `--force-cutlass-fp8` keeps the FP8 dense layers on the CUTLASS FP8 path.
+- The benchmark table is an online HTTP concurrency-16 result; do not compare
+  it directly with loaded-once request-level benchmarks.

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -6,9 +6,11 @@ from __future__ import annotations
 import copy
 import gc
 import importlib
+import json
 import os
 import tempfile
 import time
+from dataclasses import dataclass
 from pathlib import Path
 
 import pytest
@@ -19,16 +21,17 @@ from PIL import Image
 from tests.e2e.accuracy.helpers import (
     CLIPScorer,
     SemanticSimilarityScorer,
+    assert_similarity,
     compute_image_ssim_psnr,
     download_images,
     model_output_dir,
 )
-from tests.helpers.mark import hardware_marks
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
 
 os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
 
-pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, *hardware_marks(res={"cuda": "H100"}, num_cards=8)]
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
 
 # ============================================================================
 # Configurable Parameters
@@ -76,6 +79,18 @@ THRESHOLDS = {
     "ssim": 0.20,  # Structural similarity
     "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
 }
+
+QUANT_PROMPT = "A brown and white dog is running on the grass."
+QUANT_HEIGHT, QUANT_WIDTH = 1024, 1024
+QUANT_PSNR_THRESHOLD = 10.0
+QUANT_SSIM_THRESHOLD = 0.20
+QUANT_CLIP_SCORE_THRESHOLD = 20.0
+QUANT_CLIP_SCORE_DROP_THRESHOLD = float(os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0"))
+QUANT_RUN_ENV = "HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY"
+QUANT_BF16_ENV = "HUNYUAN_IMAGE3_BF16_MODEL"
+QUANT_FP8_ENV = "HUNYUAN_IMAGE3_FP8_MODEL"
+QUANT_NVFP4_ENV = "HUNYUAN_IMAGE3_NVFP4_MODEL"
+_TRUE_ENV_VALUES = {"1", "true", "yes", "on"}
 # fmt: off
 _BASE_CONFIG = {
     "stage_args": [
@@ -124,6 +139,64 @@ _BASE_CONFIG = {
 }
 # fmt: on
 
+_QUANT_DIT_CONFIG = {
+    "stage_args": [
+        {
+            "stage_id": 0,
+            "stage_type": "diffusion",
+            "runtime": {
+                "devices": "0,1",
+                "max_batch_size": 1,
+            },
+            "engine_args": {
+                "model_stage": "dit",
+                "enforce_eager": True,
+                "trust_remote_code": True,
+                "distributed_executor_backend": "mp",
+                "force_cutlass_fp8": True,
+                "moe_backend": "cutlass",
+                "parallel_config": {
+                    "tensor_parallel_size": 2,
+                    "enable_expert_parallel": True,
+                },
+                "omni_kv_config": {"need_recv_cache": True},
+            },
+            "final_output": True,
+            "final_output_type": "image",
+            "is_comprehension": False,
+            "default_sampling_params": {"seed": SEED},
+        }
+    ],
+    "runtime": {"enabled": True},
+}
+
+
+@dataclass(frozen=True)
+class _QuantAccuracyCase:
+    name: str
+    model_env: str
+    nvfp4_backend: str | None = None
+
+
+def _quant_accuracy_cases() -> list[pytest.ParameterSet]:
+    cases = [
+        _QuantAccuracyCase(name="fp8", model_env=QUANT_FP8_ENV),
+        _QuantAccuracyCase(name="mixed_nvfp4", model_env=QUANT_NVFP4_ENV, nvfp4_backend="cutlass"),
+    ]
+    params: list[pytest.ParameterSet] = []
+    run_quant_accuracy = os.environ.get(QUANT_RUN_ENV, "").lower() in _TRUE_ENV_VALUES
+    for case in cases:
+        marks = []
+        if not run_quant_accuracy:
+            marks.append(pytest.mark.skip(reason=f"Set {QUANT_RUN_ENV}=1 to run HunyuanImage3 quant accuracy."))
+        if not os.environ.get(QUANT_BF16_ENV):
+            marks.append(pytest.mark.skip(reason=f"Set {QUANT_BF16_ENV} to run HunyuanImage3 quant accuracy."))
+        if not os.environ.get(case.model_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.model_env} to run this quant accuracy case."))
+        params.append(pytest.param(case, id=case.name, marks=marks))
+    return params
+
+
 # fmt: off
 COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
            "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
@@ -144,6 +217,21 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None:
     path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
 
 
+def _quant_devices() -> str:
+    return os.environ.get("HUNYUAN_IMAGE3_QUANT_DEVICES", "0,1")
+
+
+def _quant_tensor_parallel_size() -> int:
+    return int(os.environ.get("HUNYUAN_IMAGE3_QUANT_TP", str(len(_quant_devices().split(",")))))
+
+
+def _make_quant_dit_config(path: Path) -> None:
+    config = copy.deepcopy(_QUANT_DIT_CONFIG)
+    config["stage_args"][0]["runtime"]["devices"] = _quant_devices()
+    config["stage_args"][0]["engine_args"]["parallel_config"]["tensor_parallel_size"] = _quant_tensor_parallel_size()
+    path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+
+
 def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from transformers import AutoTokenizer
 
@@ -220,6 +308,65 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
     return image, cot_text, elapsed
 
 
+def _extract_image(outputs) -> Image.Image:
+    assert outputs, "Pipeline produced no outputs"
+    for output in outputs:
+        images = getattr(output, "images", None)
+        request_output = getattr(output, "request_output", None)
+        if not images and request_output is not None:
+            images = getattr(request_output, "images", None)
+        if images:
+            image = images[0].convert("RGB")
+            image.load()
+            return image
+    raise AssertionError("Pipeline output had no images")
+
+
+def _run_dit_model(
+    model: str,
+    stage_config_path: str,
+    output_path: Path,
+    *,
+    nvfp4_backend: str | None = None,
+) -> tuple[Image.Image, float]:
+    from tests.helpers.runtime import OmniRunner
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    old_backend = os.environ.get("VLLM_NVFP4_GEMM_BACKEND")
+    if nvfp4_backend is not None:
+        os.environ["VLLM_NVFP4_GEMM_BACKEND"] = nvfp4_backend
+
+    try:
+        with OmniRunner(model, stage_configs_path=stage_config_path, mode="text-to-image") as runner:
+            generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
+            params = OmniDiffusionSamplingParams(
+                height=QUANT_HEIGHT,
+                width=QUANT_WIDTH,
+                seed=SEED,
+                generator=generator,
+                num_inference_steps=20,
+                guidance_scale=4.0,
+                guidance_scale_provided=True,
+            )
+            t0 = time.perf_counter()
+            outputs = runner.omni.generate({"prompt": QUANT_PROMPT}, params)
+            elapsed = time.perf_counter() - t0
+            image = _extract_image(outputs)
+            image.save(output_path)
+            return image, elapsed
+    finally:
+        if nvfp4_backend is not None:
+            if old_backend is None:
+                os.environ.pop("VLLM_NVFP4_GEMM_BACKEND", None)
+            else:
+                os.environ["VLLM_NVFP4_GEMM_BACKEND"] = old_backend
+        gc.collect()
+        if torch.accelerator.is_available():
+            torch.accelerator.empty_cache()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=8)
 @pytest.mark.skipif(
     torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
     reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
@@ -265,3 +412,88 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
     )
     assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
     assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"
+
+
+@pytest.mark.parametrize("case", _quant_accuracy_cases())
+@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Needs 2+ GPUs for HunyuanImage3 DiT")
+def test_quantized_dit_matches_bf16_accuracy(
+    case: _QuantAccuracyCase,
+    accuracy_artifact_root: Path,
+) -> None:
+    """Quantized DiT checkpoints should preserve prompt-aligned image quality."""
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-quant")
+    bf16_model = os.environ[QUANT_BF16_ENV]
+    quant_model = os.environ[case.model_env]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        stage_config_path = Path(tmpdir) / "hunyuan_image3_quant_dit.yaml"
+        _make_quant_dit_config(stage_config_path)
+
+        bf16_image, bf16_time = _run_dit_model(
+            bf16_model,
+            str(stage_config_path),
+            output_dir / "bf16.png",
+        )
+        quant_image, quant_time = _run_dit_model(
+            quant_model,
+            str(stage_config_path),
+            output_dir / f"{case.name}.png",
+            nvfp4_backend=case.nvfp4_backend,
+        )
+
+    ssim_score, psnr_score = compute_image_ssim_psnr(
+        prediction=quant_image,
+        reference=bf16_image,
+    )
+    assert_similarity(
+        model_name=f"{MODEL_NAME} {case.name} vs bf16",
+        vllm_image=quant_image,
+        diffusers_image=bf16_image,
+        ssim_threshold=QUANT_SSIM_THRESHOLD,
+        psnr_threshold=QUANT_PSNR_THRESHOLD,
+        width=QUANT_WIDTH,
+        height=QUANT_HEIGHT,
+    )
+
+    clip_scorer = CLIPScorer()
+    bf16_clip_score = clip_scorer.score(bf16_image, QUANT_PROMPT)
+    quant_clip_score = clip_scorer.score(quant_image, QUANT_PROMPT)
+    clip_score_drop = bf16_clip_score - quant_clip_score
+
+    metrics = {
+        "case": case.name,
+        "bf16_model": bf16_model,
+        "quant_model": quant_model,
+        "prompt": QUANT_PROMPT,
+        "seed": SEED,
+        "height": QUANT_HEIGHT,
+        "width": QUANT_WIDTH,
+        "num_inference_steps": 20,
+        "guidance_scale": 4.0,
+        "bf16_elapsed_s": bf16_time,
+        "quant_elapsed_s": quant_time,
+        "ssim": ssim_score,
+        "psnr": psnr_score,
+        "bf16_clip_score": bf16_clip_score,
+        "quant_clip_score": quant_clip_score,
+        "clip_score_drop": clip_score_drop,
+    }
+    metrics_path = output_dir / f"{case.name}_metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+
+    print(f"\nHunyuanImage3 quant accuracy ({case.name})")
+    print(f"  bf16 model:       {bf16_model}")
+    print(f"  quant model:      {quant_model}")
+    print(f"  BF16 CLIP score:  {bf16_clip_score:.4f}")
+    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={QUANT_CLIP_SCORE_THRESHOLD:.4f}")
+    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}")
+    print(f"  metrics:          {metrics_path}")
+
+    assert quant_clip_score >= QUANT_CLIP_SCORE_THRESHOLD, (
+        f"{case.name} CLIP score below threshold: got {quant_clip_score:.4f}, "
+        f"expected >= {QUANT_CLIP_SCORE_THRESHOLD:.4f}"
+    )
+    assert clip_score_drop <= QUANT_CLIP_SCORE_DROP_THRESHOLD, (
+        f"{case.name} CLIP score drop too large: got {clip_score_drop:.4f}, "
+        f"expected <= {QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}"
+    )

tests/model_executor/models/test_encoder_quant_config.py

@@ -58,7 +58,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 @pytest.mark.parametrize("method", sorted(PRE_QUANTIZED_METHODS))
 def test_pre_quantized_returns_none(method: str) -> None:
     """visual_quant_config and audio_quant_config must be None for
-    pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8)."""
+    pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8,
+    modelopt_mixed)."""
     mock_config = MagicMock()
     mock_config.get_name.return_value = method
 
@@ -100,5 +101,5 @@ def test_component_config_passed_through() -> None:
 
 def test_pre_quantized_methods_contains_expected() -> None:
     """Guard against accidental removal of a known pre-quantized method."""
-    expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8"}
+    expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}
     assert PRE_QUANTIZED_METHODS == expected

vllm_omni/diffusion/data.py

@@ -745,8 +745,11 @@ class OmniDiffusionConfig:
             return
 
         is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False))
-        should_use_checkpoint_config = self.quantization_config is None or (
-            is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config)
+        is_checkpoint_nvfp4 = bool(getattr(tf_config.quant_config, "is_checkpoint_nvfp4_serialized", False))
+        should_use_checkpoint_config = (
+            self.quantization_config is None
+            or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config))
+            or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
         )
         if should_use_checkpoint_config:
             self.quantization_config = tf_config.quant_config
@@ -766,6 +769,17 @@ class OmniDiffusionConfig:
             return quant_config.get_name() == "fp8"
         return False
 
+    @staticmethod
+    def _is_generic_nvfp4_quant_config(quant_config: object) -> bool:
+        if isinstance(quant_config, str):
+            return quant_config.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
+        if isinstance(quant_config, Mapping):
+            method = quant_config.get("method", quant_config.get("quant_method"))
+            return isinstance(method, str) and method.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
+        if hasattr(quant_config, "get_name"):
+            return quant_config.get_name() == "modelopt_fp4"
+        return False
+
     def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
         """Assign `tf_model_config` and propagate quantization if detected.
 

vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py

@@ -2,7 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from torch import nn
 
-from .modelopt import ModelOptFp8CheckpointAdapter
+from .modelopt import (
+    ModelOptFp8CheckpointAdapter,
+    ModelOptMixedPrecisionCheckpointAdapter,
+    ModelOptNvFp4CheckpointAdapter,
+)
 
 
 def get_checkpoint_adapter(
@@ -10,13 +14,19 @@ def get_checkpoint_adapter(
     source: object,
     quant_config: object | None,
     use_safetensors: bool,
-) -> ModelOptFp8CheckpointAdapter | None:
+) -> ModelOptFp8CheckpointAdapter | ModelOptNvFp4CheckpointAdapter | ModelOptMixedPrecisionCheckpointAdapter | None:
     if ModelOptFp8CheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
         return ModelOptFp8CheckpointAdapter(model, source)
+    if ModelOptNvFp4CheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
+        return ModelOptNvFp4CheckpointAdapter(model, source)
+    if ModelOptMixedPrecisionCheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
+        return ModelOptMixedPrecisionCheckpointAdapter(model, source)
     return None
 
 
 __all__ = [
     "ModelOptFp8CheckpointAdapter",
+    "ModelOptMixedPrecisionCheckpointAdapter",
+    "ModelOptNvFp4CheckpointAdapter",
     "get_checkpoint_adapter",
 ]

vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py

@@ -11,7 +11,7 @@ from vllm.model_executor.utils import get_packed_modules_mapping
 
 logger = init_logger(__name__)
 
-MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_inv")
+MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_2", ".weight_scale_inv")
 DEFAULT_PACKED_MODULES_MAPPING = {
     "to_qkv": ("to_q", "to_k", "to_v"),
     "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
@@ -249,3 +249,24 @@ class ModelOptFp8CheckpointAdapter:
 
         self._check_pending_weights(state)
         self._log_adaptation_summary(state)
+
+
+class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
+    @staticmethod
+    def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
+        return (
+            quant_config is not None
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "modelopt_fp4"
+            and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False))
+        )
+
+
+class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter):
+    @staticmethod
+    def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
+        return (
+            quant_config is not None
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "modelopt_mixed"
+        )

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -3,6 +3,7 @@
 # Adapted from Helios (https://github.com/BestWishYsh/Helios)
 
 import math
+from collections import OrderedDict
 from collections.abc import Iterable
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any
@@ -447,21 +448,32 @@ class HeliosCrossAttention(nn.Module):
             causal=False,
         )
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-    ) -> torch.Tensor:
-        query = self.to_q(hidden_states)
-        query = self.norm_q(query)
-
+    def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         key = self.to_k(encoder_hidden_states)
         value = self.to_v(encoder_hidden_states)
         key = self.norm_k(key)
 
-        query = query.unflatten(2, (self.num_heads, self.head_dim))
         key = key.unflatten(2, (self.num_heads, self.head_dim))
         value = value.unflatten(2, (self.num_heads, self.head_dim))
+        return key, value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        query = self.to_q(hidden_states)
+        query = self.norm_q(query)
+
+        if encoder_key_value is None:
+            if encoder_hidden_states is None:
+                raise ValueError("encoder_hidden_states is required when encoder_key_value is not provided.")
+            key, value = self.project_kv(encoder_hidden_states)
+        else:
+            key, value = encoder_key_value
+
+        query = query.unflatten(2, (self.num_heads, self.head_dim))
 
         hidden_states = self.attn(query, key, value)
         hidden_states = hidden_states.flatten(2, 3)
@@ -526,6 +538,7 @@ class HeliosTransformerBlock(nn.Module):
         temb: torch.Tensor,
         rotary_emb: torch.Tensor,
         original_context_length: int | None = None,
+        cross_attn_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         if temb.ndim == 4:
             shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
@@ -556,12 +569,20 @@ class HeliosTransformerBlock(nn.Module):
                 hidden_states[:, history_seq_len:],
             )
             norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states)
-            attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states,
+                encoder_key_value=cross_attn_key_value,
+            )
             current_hidden_states = current_hidden_states + attn_output
             hidden_states = torch.cat([history_hidden_states, current_hidden_states], dim=1)
         else:
             norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states)
-            attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states,
+                encoder_key_value=cross_attn_key_value,
+            )
             hidden_states = hidden_states + attn_output
 
         # 3. Feed-forward
@@ -723,11 +744,80 @@ class HeliosTransformer3DModel(nn.Module):
         # 5. Output norm & projection
         self.norm_out = HeliosOutputNorm(inner_dim, eps)
         self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
+        self._projected_encoder_cache: OrderedDict[tuple, torch.Tensor] = OrderedDict()
+        self._cross_attn_kv_cache: OrderedDict[tuple, list[tuple[torch.Tensor, torch.Tensor]]] = OrderedDict()
+        self._cross_attn_cache_size = 2
 
     @property
     def dtype(self) -> torch.dtype:
         return next(self.parameters()).dtype
 
+    @staticmethod
+    def _tensor_cache_key(tensor: torch.Tensor) -> tuple:
+        try:
+            version = tensor._version
+        except RuntimeError:
+            version = None
+
+        return (
+            tensor.data_ptr(),
+            tuple(tensor.shape),
+            tuple(tensor.stride()),
+            tensor.dtype,
+            tensor.device.type,
+            tensor.device.index,
+            version,
+        )
+
+    @staticmethod
+    def _get_from_lru(cache: OrderedDict, key: tuple):
+        value = cache.get(key)
+        if value is not None:
+            cache.move_to_end(key)
+        return value
+
+    def _put_lru(self, cache: OrderedDict, key: tuple, value) -> None:
+        cache[key] = value
+        cache.move_to_end(key)
+        while len(cache) > self._cross_attn_cache_size:
+            cache.popitem(last=False)
+
+    def clear_cross_attention_cache(self) -> None:
+        self._projected_encoder_cache.clear()
+        self._cross_attn_kv_cache.clear()
+
+    def _cache_enabled(self) -> bool:
+        return not self.training and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+
+    def _project_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        if not self._cache_enabled():
+            return self.condition_embedder.text_embedder(encoder_hidden_states)
+
+        cache_key = self._tensor_cache_key(encoder_hidden_states)
+        cached = self._get_from_lru(self._projected_encoder_cache, cache_key)
+        if cached is not None:
+            return cached
+
+        projected = self.condition_embedder.text_embedder(encoder_hidden_states)
+        self._put_lru(self._projected_encoder_cache, cache_key, projected)
+        return projected
+
+    def _get_cross_attn_key_values(
+        self,
+        encoder_hidden_states: torch.Tensor,
+    ) -> list[tuple[torch.Tensor, torch.Tensor]] | None:
+        if not self._cache_enabled():
+            return None
+
+        cache_key = self._tensor_cache_key(encoder_hidden_states)
+        cached = self._get_from_lru(self._cross_attn_kv_cache, cache_key)
+        if cached is not None:
+            return cached
+
+        key_values = [block.attn2.project_kv(encoder_hidden_states) for block in self.blocks]
+        self._put_lru(self._cross_attn_kv_cache, cache_key, key_values)
+        return key_values
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -836,7 +926,12 @@ class HeliosTransformer3DModel(nn.Module):
                 .expand(batch_size, -1, history_context_length, -1)
             )
 
-        temb, timestep_proj, encoder_hidden_states = self.condition_embedder(timestep, encoder_hidden_states)
+        temb, timestep_proj, _ = self.condition_embedder(
+            timestep,
+            encoder_hidden_states,
+            is_return_encoder_hidden_states=False,
+        )
+        encoder_hidden_states = self._project_encoder_hidden_states(encoder_hidden_states)
         timestep_proj = timestep_proj.unflatten(-1, (6, -1))
 
         if indices_hidden_states is not None and not self.zero_history_timestep:
@@ -857,14 +952,17 @@ class HeliosTransformer3DModel(nn.Module):
         hidden_states = hidden_states.contiguous()
         encoder_hidden_states = encoder_hidden_states.contiguous()
         rotary_emb = rotary_emb.contiguous()
+        cross_attn_key_values = self._get_cross_attn_key_values(encoder_hidden_states)
 
-        for block in self.blocks:
+        for block_idx, block in enumerate(self.blocks):
+            cross_attn_key_value = None if cross_attn_key_values is None else cross_attn_key_values[block_idx]
             hidden_states = block(
                 hidden_states,
                 encoder_hidden_states,
                 timestep_proj,
                 rotary_emb,
                 original_context_length,
+                cross_attn_key_value,
             )
 
         # 7. Output normalization

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2073,7 +2073,7 @@ class HunyuanImage3Model(nn.Module):
             return False
 
         def is_scalar_quant_scale(name: str, tensor: torch.Tensor) -> bool:
-            return tensor.numel() == 1 and name.endswith((".input_scale", ".weight_scale"))
+            return tensor.numel() == 1 and name.endswith((".input_scale", ".weight_scale", ".weight_scale_2"))
 
         def load_split_param(
             name: str,

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -49,6 +49,10 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 logger = init_logger(__name__)
 
 
+def _join_prefix(prefix: str, suffix: str) -> str:
+    return f"{prefix}.{suffix}" if prefix else suffix
+
+
 class ImageRopePrepare(nn.Module):
     """Prepares image hidden_states and RoPE embeddings for sequence parallel.
 
@@ -471,7 +475,12 @@ class FeedForward(nn.Module):
 
         layers: list[nn.Module] = [
             ColumnParallelApproxGELU(
-                dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config, prefix=prefix
+                dim,
+                inner_dim,
+                approximate="tanh",
+                bias=bias,
+                quant_config=quant_config,
+                prefix=_join_prefix(prefix, "net.0.proj"),
             ),
             nn.Identity(),  # placeholder for weight loading
             RowParallelLinear(
@@ -480,7 +489,7 @@ class FeedForward(nn.Module):
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
-                prefix=prefix,
+                prefix=_join_prefix(prefix, "net.2"),
             ),
         ]
 
@@ -507,6 +516,7 @@ class QwenImageCrossAttention(nn.Module):
         context_pre_only: bool = False,
         out_dim: int | None = None,
         quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         assert dim % num_heads == 0
@@ -523,7 +533,7 @@ class QwenImageCrossAttention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=num_heads,
             quant_config=quant_config,
-            prefix="to_qkv",
+            prefix=_join_prefix(prefix, "to_qkv"),
         )
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -539,7 +549,7 @@ class QwenImageCrossAttention(nn.Module):
             head_size=head_dim,
             total_num_heads=num_heads,
             quant_config=quant_config,
-            prefix="add_kv_proj",
+            prefix=_join_prefix(prefix, "add_kv_proj"),
         )
         self.add_query_num_heads = self.add_kv_proj.num_heads
         self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
@@ -552,7 +562,7 @@ class QwenImageCrossAttention(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix="to_add_out",
+            prefix=_join_prefix(prefix, "to_add_out"),
         )
 
         assert not pre_only
@@ -563,7 +573,7 @@ class QwenImageCrossAttention(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix="to_out",
+            prefix=_join_prefix(prefix, "to_out"),
         )
 
         self.norm_added_q = RMSNorm(head_dim, eps=eps)
@@ -697,6 +707,7 @@ class QwenImageTransformerBlock(nn.Module):
         eps: float = 1e-6,
         zero_cond_t: bool = False,
         quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -716,7 +727,7 @@ class QwenImageTransformerBlock(nn.Module):
                 bias=True,
                 return_bias=False,
                 quant_config=None,
-                prefix="img_mod.1",
+                prefix=_join_prefix(prefix, "img_mod.1"),
             ),
         )
         self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -727,9 +738,15 @@ class QwenImageTransformerBlock(nn.Module):
             context_pre_only=False,
             head_dim=attention_head_dim,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
         )
         self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.img_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="img_mlp")
+        self.img_mlp = FeedForward(
+            dim=dim,
+            dim_out=dim,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "img_mlp"),
+        )
 
         # Text processing modules.
         self.txt_mod = nn.Sequential(
@@ -740,13 +757,18 @@ class QwenImageTransformerBlock(nn.Module):
                 bias=True,
                 return_bias=False,
                 quant_config=None,
-                prefix="txt_mod.1",
+                prefix=_join_prefix(prefix, "txt_mod.1"),
             ),
         )
         self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
         # Text doesn't need separate attention - it's handled by img_attn joint computation
         self.txt_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="txt_mlp")
+        self.txt_mlp = FeedForward(
+            dim=dim,
+            dim_out=dim,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "txt_mlp"),
+        )
 
         self.zero_cond_t = zero_cond_t
 
@@ -996,8 +1018,9 @@ class QwenImageTransformer2DModel(CachedTransformer):
                     attention_head_dim=attention_head_dim,
                     zero_cond_t=zero_cond_t,
                     quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -84,7 +84,7 @@ from vllm.multimodal.processing import (
     PromptUpdateDetails,
 )
 from vllm.sequence import IntermediateTensors
-from vllm.transformers_utils.tokenizer import get_tokenizer
+from vllm.tokenizers import get_tokenizer
 from vllm.utils.tensor_schema import TensorSchema
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata

vllm_omni/quantization/component_config.py

@@ -23,10 +23,10 @@ if TYPE_CHECKING:
     )
 
 
-# Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8) only quantize the
+# Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8/mixed) only quantize the
 # Thinker LM.  Vision and audio encoder weights remain in BF16 with no
 # corresponding scale tensors in the checkpoint.
-PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8"})
+PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"})
 
 
 def resolve_encoder_quant_config(

vllm_omni/quantization/factory.py

@@ -138,11 +138,19 @@ def _normalize_quant_method_alias(method: str | None) -> str | None:
 
 _MODEL_OPT_METHODS = {
     "modelopt",
+    "modelopt_fp4",
+    "modelopt_mixed",
 }
 _MODEL_OPT_FP8_ALGOS = {
     "FP8",
     "FP8_PER_CHANNEL_PER_TOKEN",
 }
+_MODEL_OPT_NVFP4_ALGOS = {
+    "NVFP4",
+}
+_MODEL_OPT_MIXED_ALGOS = {
+    "MIXED_PRECISION",
+}
 
 
 def _normalize_method_name(method: Any) -> str:
@@ -170,6 +178,10 @@ def _detect_modelopt_method(config: Mapping[str, Any]) -> str | None:
     if quant_algo:
         if quant_algo in _MODEL_OPT_FP8_ALGOS:
             return "modelopt"
+        if quant_algo in _MODEL_OPT_NVFP4_ALGOS:
+            return "modelopt_fp4"
+        if quant_algo in _MODEL_OPT_MIXED_ALGOS:
+            return "modelopt_mixed"
         return None
 
     if method is not None: