From 9dd36e3afd0940752c2a96f9d39f8d211c2b9195 Mon Sep 17 00:00:00 2001 From: hxhhhlalala Date: Wed, 20 May 2026 14:21:29 +0800 Subject: [PATCH] [NPU][Quant] Add W4A4 MXFP4 online & MXFP4 dual-scale online/offline quantization support for Wan2.2 T2V / I2V inference on Ascend NPU (#3578) Signed-off-by: hyh_hh Co-authored-by: hyh_hh --- docs/user_guide/quantization/mxfp4.md | 479 ++++++++++++ docs/user_guide/quantization/overview.md | 25 +- .../image_to_video/image_to_video.py | 4 +- .../text_to_video/text_to_video.py | 4 +- .../wan2_2/test_create_transformer_quant.py | 383 ++++++++++ .../quantization/test_mxfp4_config.py | 539 +++++++++++++ .../quantization/test_mxfp4_key_remap.py | 399 ++++++++++ .../quantization/test_mxfp8_config.py | 308 ++++++++ .../quantization/test_mxfp8_key_remap.py | 214 ++++++ .../models/wan2_2/pipeline_wan2_2.py | 48 +- vllm_omni/quantization/factory.py | 102 +++ vllm_omni/quantization/mxfp4_config.py | 709 ++++++++++++++++++ vllm_omni/quantization/tools/__init__.py | 2 + .../tools/merge_mxfp4_dualscale_checkpoint.py | 491 ++++++++++++ 14 files changed, 3646 insertions(+), 61 deletions(-) create mode 100644 docs/user_guide/quantization/mxfp4.md create mode 100644 tests/diffusion/models/wan2_2/test_create_transformer_quant.py create mode 100644 tests/diffusion/quantization/test_mxfp4_config.py create mode 100644 tests/diffusion/quantization/test_mxfp4_key_remap.py create mode 100644 tests/diffusion/quantization/test_mxfp8_config.py create mode 100644 tests/diffusion/quantization/test_mxfp8_key_remap.py create mode 100644 vllm_omni/quantization/mxfp4_config.py create mode 100644 vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md new file mode 100644 index 000000000..2fd4a8395 --- /dev/null +++ b/docs/user_guide/quantization/mxfp4.md @@ -0,0 +1,479 @@ +# W4A4 MXFP4 Quantization + +## Overview + +W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4 +(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups +of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale. + +vLLM-Omni provides two quantization methods with different scale structures: + +| Method | Scale structure | Mode | Use case | +|--------|----------------|------|----------| +| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed | +| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended | + +!!! tip "Recommended: `mxfp4_dualscale` offline" + For production deployments, use the `mxfp4_dualscale` offline mode with a + pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load + calibrated `mul_scale` tensors from disk, providing measurably better accuracy + than any online method. The one-time preprocessing cost amortises across all + subsequent inference runs. + + Use `mxfp4` online only for quick experimentation where preprocessing time + is not acceptable and accuracy loss is tolerable. + +!!! warning "Online single-scale ≠ Offline dual-scale" + `mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`: + fine scale (per-32 K), coarse scale (per-512 K), and per-input-channel + `mul_scale` from calibration — all loaded from the checkpoint. + `mxfp4_dualscale` online mode uses `NPUMxfp4DualScaleOnlineLinearMethod`: + dual-level scales computed on the fly from BF16 weights; no calibration + `mul_scale` is available. Loading an offline checkpoint with the online + method (or vice versa) will produce incorrect results or shape errors. + +## Hardware Support + +| Device | Support | +|--------|---------| +| NVIDIA Blackwell GPU (SM 100+) | ⭕ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | +| NVIDIA Ampere GPU (SM 80+) | ⭕ | +| AMD ROCm | ⭕ | +| Intel XPU | ⭕ | +| Ascend NPU (Atlas 950 A5) | ✅ | + +Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. + +## Model Type Support + +### Diffusion Model (Wan2.2) + +| Model | Online | Offline | Notes | +|-------|--------|---------|-------| +| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config | +| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B | +| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss | + +The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about +quantization quality, not model compatibility — both work on cascade (A14B) and +single-transformer models alike, the same as `mxfp8` online: + +- `mxfp4`: single-scale, lower overhead, simpler compute, online only +- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline + +**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced +by the merge script); there is no offline `mxfp4` single-scale format. + +!!! note "Per-layer BF16 fallback in offline cascade models" + The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most + linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers + retain their original BF16 weights and are listed in `ignored_layers` inside + each transformer's `config.json`. The two transformers may have different + `ignored_layers` sets — the pipeline reads each transformer's own `config.json` + and rebuilds the config locally when they differ, so routing is always + per-transformer-accurate. + +!!! warning "TI2V-5B not supported" + Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter + count makes it significantly more sensitive to 4-bit quantization noise, + resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B. + +## Configuration + +### `mxfp4` — Single-Scale Online Mode + +Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to +MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale +(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no +calibration `mul_scale` is available. Applies equally to single-transformer +and cascade (A14B) models — both transformers in a cascade receive the same +quantization config automatically. + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni(model="", quantization="mxfp4") +outputs = omni.generate( + "A cat sitting on a windowsill", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +```bash +# Single-transformer or cascade model — same command +python text_to_video.py --model --quantization mxfp4 + +# Online serving +vllm serve --omni --quantization mxfp4 +``` + +### `mxfp4_dualscale` — DualScale Online Mode + +Online DualScale mode computes both fine and coarse scales on the fly from BF16 +weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to +single-transformer and cascade (A14B) models. Compared to `mxfp4` online, +DualScale provides better quantization accuracy at higher compute cost. + +The default configuration keeps the leading 5 transformer blocks in BF16 +(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this +is sufficient to meet quality requirements and is the recommended setting. + +```python +omni = Omni(model="", quantization="mxfp4_dualscale") +``` + +```bash +python text_to_video.py --model --quantization mxfp4_dualscale +``` + +If accuracy debugging identifies additional precision-sensitive layers, they +can be pinned to BF16 via the Python API: + +```python +omni = Omni( + model="", + quantization={ + "method": "mxfp4_dualscale", + "ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override + }, +) +``` + +BF16 fallback routing in online mode applies two rules in priority order: + +1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix + matches is kept in BF16 regardless of block index. +2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N + transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults + to `5` (recommended). Layers outside `blocks.N.*` + (e.g. `condition_embedder`) are always quantized. + +### `mxfp4_dualscale` — DualScale Offline Mode (Recommended) + +Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A +preprocessing step converts the raw quantized output to the diffusers format +expected by vLLM-Omni and injects the quantization config into each +`transformer/config.json` so that vLLM-Omni auto-detects the offline path +without a `--quantization` flag. + +BF16 fallback layers may be interleaved anywhere in the transformer — they are +not restricted to leading blocks. The merge script detects them from +`quant_model_description.json` and writes their prefixes into `ignored_layers` +inside `config.json`. At runtime, each layer's prefix is matched against +`ignored_layers` to decide BF16 vs. MXFP4 DualScale. + +#### Checkpoint tensor layout + +Each quantized linear layer stores four tensors: + +| Tensor | Shape | dtype | Description | +|--------|-------|-------|-------------| +| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) | +| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) | +| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale | +| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) | + +BF16 fallback layers have no quantization tensors; only the original `weight` +(and optional `bias`) are present, loaded directly from the base checkpoint. + +#### Step 1 — Quantize with msModelSlim + +```bash +msmodelslim quant \ + --model_path /path/to/Wan2.2-T2V-A14B-Diffusers \ + --save_path /path/to/wan2_2_t2v_quantized_raw \ + --device npu \ + --model_type Wan2_2 \ + --config_path /path/to/wan2_2_w4a4_mxfp4_dualscale.yaml \ + --trust_remote_code True +``` + +After this step, `--save_path` contains raw quantized safetensors files, +scale files, and a metadata JSON (`quant_model_description*.json`). + +For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two +subdirectories: `high_noise_model/` (transformer) and `low_noise_model/` +(transformer_2). + +#### Step 2 — Preprocess with merge_mxfp4_dualscale_checkpoint.py + +The script (`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`): + +1. Copies the original diffusers model to `--output-path` (VAE, text encoder, + scheduler, etc. are preserved). +2. Remaps tensor names from msModelSlim convention to diffusers convention and + strips `.linear.` / `.div.` wrappers added by the quantization tool. +3. Overlays MXFP4 tensors (weight, fine/coarse scales, `mul_scale`) onto the + BF16 base checkpoint. Non-quantized layers keep their original BF16 weights. +4. Detects all linear layers that remain in BF16 and writes their prefixes into + `ignored_layers` in `config.json`. +5. Injects `quantization_config` so vLLM-Omni auto-detects offline MXFP4 + DualScale. + +For cascade MoE models, steps 2–5 run separately for each transformer. + +```bash +python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \ + --model-type Wan2.2-T2V-A14B \ + --original-model /path/to/Wan2.2-T2V-A14B-Diffusers \ + --quant-path /path/to/wan2_2_t2v_quantized_raw \ + --output-path /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale +``` + +| Argument | Description | +|----------|-------------| +| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` | +| `--original-model` | Root directory of the original BF16 diffusers model | +| `--quant-path` | Root directory of the msModelSlim quantized output | +| `--output-path` | Output directory for the merged model (created by the script) | + +The script outputs a complete diffusers model directory at `--output-path`, +with each transformer subfolder containing: + +- `diffusion_pytorch_model.safetensors` — MXFP4 weights + scale tensors, with BF16 fallback layers from the base checkpoint +- `config.json` — original transformer config with `quantization_config` injected +- `quant_model_description.json` — quantization metadata (reference only) + +The `quantization_config` injected into `config.json` for each transformer: + +```json +{ + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": true, + "ignored_layers": [ + "blocks.0.attn1.to_qkv", + "blocks.0.attn1.to_out", + "proj_out" + ] +} +``` + +`ignored_layers` lists every linear layer that retains its original BF16 weight, +using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out` +unindexed). The exact entries are determined by the quantization tool (msModelSlim) +and may differ between `transformer` and `transformer_2` in a cascade model. + +#### Step 3 — Serve + +```bash +python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale + +# Online serving +vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni +``` + +```python +omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale") +``` + +!!! note + No `--quantization` flag is needed for offline mode. The preprocessing + script injects `quantization_config` into each `transformer/config.json`, + which vLLM-Omni reads automatically to activate the correct offline path. + +## Parameters + +### `mxfp4` (single-scale, online only) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | — | `"mxfp4"` | +| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 | + +### `mxfp4_dualscale` (dual-scale, online + offline) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | — | `"mxfp4_dualscale"` | +| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script | +| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override | +| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B | + +#### BF16 fallback priority (online mode) + +``` +for each linear layer: + if prefix in ignored_layers → BF16 (explicit override, highest priority) + elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule) + else → MXFP4 DualScale online +``` + +Layers outside `blocks.N.*` (e.g. `condition_embedder.*`) are always quantized +unless they appear in `ignored_layers`. + +## Validation and Notes + +1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using + `npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is + available — all output partitions receive an identity pre-scale. No offline + checkpoint format exists for this method. + +2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)** + quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse + scales computed on the fly). No calibration `mul_scale`; leading blocks or + explicit `ignored_layers` stay in BF16 for accuracy. + +3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** — + **recommended for production** — loads four tensors per quantized layer: FP4 + weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale + (`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback + layers have no quantization tensors and are routed via `ignored_layers`. + +4. **Scale dtype**: fine scales are stored as `uint8` in safetensors (same bit + layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy + float32 round-trip. + +5. **Cascade model config propagation**: in a cascade model (transformer + + transformer_2), vLLM-Omni reads each transformer's own `config.json` and + rebuilds the quant config locally when `ignored_layers` differs between + transformers, ensuring per-layer routing is accurate for each. The first + transformer's config is propagated to `od_config` so the second transformer + can reuse it as a starting point. + +6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a + single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the + fused name (`attn1.to_qkv`), written automatically by the merge script. + +7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The + DualScale offline method mitigates this with calibrated `mul_scale` smooth + quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade + off compression vs. accuracy for precision-sensitive layers. + +## Adapting MXFP4 for a New Model + +This section is aimed at developers who want to add MXFP4 support to a model +other than Wan2.2. The three integration points are: (1) discovering the correct +runtime layer names, (2) wiring `ignored_layers` into the model, and (3) writing +a merge script for offline checkpoints. + +### Step 1 — Discover runtime layer names + +`ignored_layers` entries must match the **runtime parameter names** used inside +vllm-omni, which may differ from the names stored in the diffusers checkpoint. +The canonical source of truth is the model's own `named_parameters()`. + +```python +from vllm_omni import Omni + +# Load the model without quantization to inspect parameter names. +omni = Omni(model="/path/to/your-model") # no --quantization flag +for name, _ in omni.pipeline.transformer.named_parameters(): + if "weight" in name and "scale" not in name: + print(name) +``` + +Compare the printed names against the diffusers checkpoint keys +(`safetensors.safe_open` or `torch.load`) to identify any renames your model +applies. Common patterns that differ in Wan2.2 (and may appear in other +models): + +| Diffusers checkpoint name | vllm-omni runtime name | Reason | +|---------------------------|------------------------|--------| +| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` | +| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores | +| `ffn.net.2` | `ffn.net_2` | Same underscore rule | +| `to_out.0` | `to_out` | Sequential index stripped | + +If your model has different fusion patterns, inspect `packed_modules_mapping` +on the model class — this dict records how checkpoint keys are mapped to +fused runtime parameters. + +!!! warning "Partial QKV fallback is not allowed" + If your model fuses Q, K, V into a single layer, `ignored_layers` must + include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but + `to_k`, `to_v` quantized) cannot be expressed at runtime because they share + one `QKVParallelLinear`. The merge script enforces this and raises an error + if only some of the trio appear as non-quantized. + +### Step 2 — Add ignored_layers to the model + +#### Online mode + +Pass `ignored_layers` directly in the quantization config using the **runtime +names** discovered in Step 1. No code changes to the model are required. + +```python +omni = Omni( + model="/path/to/your-model", + quantization={ + "method": "mxfp4_dualscale", + "ignored_layers": [ + "blocks.0.attn1.to_qkv", # runtime name, not diffusers name + "blocks.0.attn1.to_out", + "blocks.0.ffn.net_0.proj", + ], + }, +) +``` + +```bash +# CLI does not support list-typed ignored_layers directly. +# Use the Python API or set ignored_layers in config.json (offline). +python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale +``` + +The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers +individually: set it to N to keep all linear layers in blocks 0 … N-1 in BF16. +The right value depends on the model's sensitivity; evaluate on a validation +set and pick the smallest N that meets your accuracy target. + +#### Offline mode + +For offline checkpoints, `ignored_layers` is written into each transformer's +`config.json` by the merge script (see Step 3). No manual editing is needed if +the merge script is correct. The injected block: + +```json +{ + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": true, + "ignored_layers": [ + "blocks.0.attn1.to_qkv", + "blocks.0.attn1.to_out" + ] +} +``` + +To add a layer manually (e.g. to pin an additional layer to BF16 without +re-running the merge script), edit `config.json` inside the transformer +subfolder. Use runtime names, not diffusers checkpoint names. + +### Step 3 — Write a merge script for offline mode + +The merge script for a new model mirrors +`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`. The four +things it must do: + +1. **Remap tensor names** from the quantization tool convention to diffusers + convention (strip wrappers like `.linear.`, `.div.`; fix any prefix + differences). + +2. **Collect ignored_layers**: after loading, enumerate all `*.weight` keys that + have no corresponding `*.weight_scale` (i.e. layers the tool left in BF16). + Convert diffusers names to vllm-omni runtime names (fuse QKV, rename FFN + sub-modules, etc.). Write the result to `config.json`. + +3. **Inject `quantization_config`** into `config.json`: + ```python + config["quantization_config"] = { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ignored_layers, # runtime names + } + ``` + +4. **Save** the merged safetensors and the updated `config.json`. + +The key helper to implement is the diffusers-to-runtime name translator +(equivalent to `_diffusers_to_vllm_ignored` in the Wan2.2 merge script). +For each non-quantized diffusers weight key, apply your model's specific +renaming rules and collect the results. + +!!! tip "Validate before serving" + After producing the offline checkpoint, load it without a `--quantization` + flag and verify that vLLM-Omni auto-detects the correct method. Check that + the layer count reported in the startup log matches expectations: quantized + layer count + `ignored_layers` count should equal total linear layer count. + Any mismatch indicates a name-mapping bug in the merge script. diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 39669bf5d..38d027f12 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -9,20 +9,20 @@ type has a different quantization scope. | Mode | Guide | Description | Methods | |------|-------|-------------|---------| -| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8 | +| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 | | Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA | -| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8 | +| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | ## Hardware Support -| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | GGUF | AutoRound | msModelSlim | -|--------|----------|-----------|----------|------------|------|-----------|-------------| -| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ | -| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | -| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | +| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim | +|--------|----------|-----------|----------|------------|------------|------|-----------|-------------| +| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | +| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | +| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. FP8 on Ampere may use a weight-only path where available. @@ -42,6 +42,7 @@ otherwise. | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | | ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints | | MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 | +| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy | | GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists | | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven | | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path | @@ -58,6 +59,7 @@ in BF16 unless the model guide explicitly adds support. | ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | +| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints | | msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | @@ -73,6 +75,7 @@ attached to the intended stage rather than applied globally. | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | | ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation | | MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | +| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | | GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed | | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed | | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation | @@ -100,7 +103,7 @@ config = build_quant_config({ | Component | Default quantized? | Notes | |-----------|--------------------|-------| -| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim | +| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim | | Text encoder | No | Keep BF16 unless a method-specific guide documents support | | VAE | No | Keep BF16; storage-only paths are method-specific | | Scheduler/tokenizer | No | Loaded from the base model repository | diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 9e2521233..e5bbacdca 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -173,8 +173,8 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "int8", "gguf"], - help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) # Distributed and parallel execution diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 9e1775dd1..8bbc871ea 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -163,8 +163,8 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "int8", "gguf"], - help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).", + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) # Distributed and parallel execution diff --git a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py new file mode 100644 index 000000000..bb95363fd --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py @@ -0,0 +1,383 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for transformer quant-config auto-detection. + +The loader path at pipeline_wan2_2.py carries two quantization contracts: + + create_transformer_from_config (~L137) + - Reads quantization_config from config.json (injected by the merge scripts) + - Auto-detects the quant method when no CLI quant_config is provided + - Rejects method mismatches (CLI vs disk) + - Upgrades online → offline when disk marks is_checkpoint_*_serialized=True + - Rebuilds when the active ignored_layers differs from the disk value + + Wan22Pipeline._create_transformer (~L456) + - Passes od_config.quantization_config (set by CLI or externally) to + create_transformer_from_config for each transformer independently. + - When od_config.quantization_config is None, each transformer auto-detects + from its own config.json; od_config is NOT modified by this call. + +All tests are pure-CPU and do not load model weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +# Minimum config that create_transformer_from_config accepts without raising. +_MIN_CFG: dict = {"patch_size": [1, 2, 2]} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_fake_transformer(): + """Return (FakeTransformer class, captured list). + + Each _FakeTransformer.__init__ call appends its **kwargs to captured, + letting tests inspect what quant_config was passed per transformer. + """ + captured: list[dict] = [] + + class _FakeTransformer: + def __init__(self, **kwargs): + captured.append(kwargs) + + return _FakeTransformer, captured + + +class _FakePipeline: + """Minimal stand-in exposing only what _create_transformer needs from self.""" + + def __init__(self, od_config: SimpleNamespace) -> None: + self.od_config = od_config + + # Bind the real unbound method so the tests exercise production code. + _create_transformer = wan22_module.Wan22Pipeline._create_transformer + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — auto-detection +# --------------------------------------------------------------------------- + + +def test_create_transformer_detects_mxfp8_serialized_from_config_json(monkeypatch): + """When config.json carries MXFP8 quant and no CLI quant_config is provided, + the transformer must receive a DiffusionMXFP8Config with serialized=True.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + create_transformer_from_config(config) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config) + assert qc.is_checkpoint_mxfp8_serialized is True + + +def test_create_transformer_detects_mxfp4_dualscale_from_config_json(monkeypatch): + """config.json with mxfp4_dualscale + ignored_layers must produce + a DiffusionMXFP4DualScaleMixedConfig with the correct ignored_layers and + serialized flag.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"], + "is_checkpoint_serialized": True, + }, + } + create_transformer_from_config(config) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig) + assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.0.attn1.to_k"} + assert qc.is_checkpoint_serialized is True + + +def test_create_transformer_without_quantization_config_passes_no_quant(monkeypatch): + """A plain BF16 config.json (no quantization_config key) must result in no + quant_config being passed to WanTransformer3DModel.""" + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + create_transformer_from_config(_MIN_CFG) + + assert "quant_config" not in captured[0] + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — method-mismatch guard +# --------------------------------------------------------------------------- + + +def test_create_transformer_rejects_method_mismatch(monkeypatch): + """Passing a CLI quant_config whose get_name() differs from the config.json + quant_method must raise ValueError immediately (prevents silent weight corruption). + + fp8 (vLLM built-in, get_name()=='fp8') vs disk 'mxfp8' triggers the guard. + These are distinct methods; using the same method for both would not trigger it. + """ + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + FakeTransformer, _ = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + fp8_cli = Fp8Config(is_checkpoint_fp8_serialized=True, activation_scheme="dynamic") + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + with pytest.raises(ValueError, match="quant_method"): + create_transformer_from_config(config, quant_config=fp8_cli) + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — online → offline upgrade +# --------------------------------------------------------------------------- + + +def test_create_transformer_upgrades_to_serialized_when_disk_marks_it(monkeypatch): + """CLI passes online (is_checkpoint_mxfp8_serialized=False) but config.json + marks is_checkpoint_mxfp8_serialized=True → must switch to offline (serialized) + so that pre-quantized FP8 tensors are loaded correctly.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + online_cli = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False) + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + create_transformer_from_config(config, quant_config=online_cli) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config) + assert qc.is_checkpoint_mxfp8_serialized is True + + +# --------------------------------------------------------------------------- +# create_transformer_from_config — ignored_layers rebuild +# --------------------------------------------------------------------------- + + +def test_create_transformer_rebuilds_when_ignored_layers_differ(monkeypatch): + """When the active quant_config has different ignored_layers than config.json, + the config must be rebuilt from disk so per-layer BF16 routing is authoritative + for this specific transformer.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + stale = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + create_transformer_from_config(config, quant_config=stale) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig) + assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.1.attn1.to_q"} + + +def test_create_transformer_does_not_rebuild_when_ignored_layers_match(monkeypatch): + """When the active quant_config already has the same ignored_layers, + the same instance must be passed through unchanged (no unnecessary rebuild).""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + matching = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + create_transformer_from_config(config, quant_config=matching) + + assert captured[0].get("quant_config") is matching + + +# --------------------------------------------------------------------------- +# Wan22Pipeline._create_transformer — od_config passthrough +# --------------------------------------------------------------------------- + + +def test_pipeline_create_transformer_auto_detects_from_config_json(monkeypatch): + """When od_config.quantization_config is None, _create_transformer must + auto-detect the quant method from config.json and pass it to the transformer. + od_config itself must remain unchanged (None).""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + od_config = SimpleNamespace(quantization_config=None) + pipeline = _FakePipeline(od_config) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + pipeline._create_transformer(config) + + qc = captured[0].get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config) + assert qc.is_checkpoint_mxfp8_serialized is True + # od_config must NOT be modified — each transformer auto-detects independently. + assert od_config.quantization_config is None + + +def test_pipeline_create_transformer_does_not_overwrite_existing_od_config(monkeypatch): + """If od_config.quantization_config is already set (e.g. via CLI), _create_transformer + must pass it through to the transformer unchanged and leave od_config unmodified.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, _ = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + existing = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True) + od_config = SimpleNamespace(quantization_config=existing) + pipeline = _FakePipeline(od_config) + + config = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp8", + "is_checkpoint_mxfp8_serialized": True, + }, + } + pipeline._create_transformer(config) + + assert od_config.quantization_config is existing + + +# --------------------------------------------------------------------------- +# Wan22Pipeline._create_transformer — cascade contracts +# --------------------------------------------------------------------------- + + +def test_pipeline_cascade_both_transformers_get_mxfp8_serialized_config(monkeypatch): + """Cascade model (transformer + transformer_2) with MXFP8 checkpoint: + - Each transformer auto-detects from its own config.json independently. + Both must receive is_checkpoint_mxfp8_serialized=True.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + od_config = SimpleNamespace(quantization_config=None) + pipeline = _FakePipeline(od_config) + + mxfp8_qc = {"quant_method": "mxfp8", "is_checkpoint_mxfp8_serialized": True} + pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc}) + pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc}) + + assert len(captured) == 2 + for i, kwargs in enumerate(captured): + qc = kwargs.get("quant_config") + assert isinstance(qc, DiffusionMXFP8Config), f"transformer[{i}]: expected DiffusionMXFP8Config, got {type(qc)}" + assert qc.is_checkpoint_mxfp8_serialized is True, f"transformer[{i}]: expected serialized=True" + + +def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_layers(monkeypatch): + """Cascade with mxfp4_dualscale where transformer and transformer_2 have + different ignored_layers in their config.json. + + Expected outcome: + transformer → ignored_layers=["blocks.0.attn1.to_q"] (auto-detected from its config.json) + transformer_2 → ignored_layers=["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"] + (auto-detected from its own config.json independently) + od_config → quantization_config remains None (not modified by _create_transformer) + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + od_config = SimpleNamespace(quantization_config=None) + pipeline = _FakePipeline(od_config) + + cfg1 = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + cfg2 = { + **_MIN_CFG, + "quantization_config": { + "quant_method": "mxfp4_dualscale", + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"], + "is_checkpoint_serialized": True, + }, + } + + pipeline._create_transformer(cfg1) + pipeline._create_transformer(cfg2) + + assert len(captured) == 2 + qc1 = captured[0].get("quant_config") + qc2 = captured[1].get("quant_config") + + assert isinstance(qc1, DiffusionMXFP4DualScaleMixedConfig) + assert isinstance(qc2, DiffusionMXFP4DualScaleMixedConfig) + assert set(qc1.ignored_layers) == {"blocks.0.attn1.to_q"}, f"transformer expected 1 layer, got {qc1.ignored_layers}" + assert set(qc2.ignored_layers) == { + "blocks.0.attn1.to_q", + "blocks.1.attn1.to_q", + }, f"transformer_2 expected 2 layers, got {qc2.ignored_layers}" + + # od_config must remain unchanged — _create_transformer does not modify it. + assert od_config.quantization_config is None diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py new file mode 100644 index 000000000..af3a9c9c4 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp4_config.py @@ -0,0 +1,539 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config.""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _patch_tp_state(monkeypatch): + """Patch TP rank/world_size so ModelWeightParameter can be instantiated on CPU + without an initialized distributed group. Returns TP=1 rank=0 for all tests.""" + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1) + + +# --------------------------------------------------------------------------- +# DiffusionMXFP4Config +# --------------------------------------------------------------------------- + + +def test_mxfp4_config_get_name(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + assert DiffusionMXFP4Config.get_name() == "mxfp4" + + +def test_mxfp4_config_from_config_defaults(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({}) + assert cfg.is_checkpoint_mxfp4_serialized is False + assert cfg.ignored_layers == [] + + +def test_mxfp4_config_from_config_serialized(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({"is_checkpoint_mxfp4_serialized": True}) + assert cfg.is_checkpoint_mxfp4_serialized is True + + +def test_mxfp4_config_from_config_ignored_layers(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({"ignored_layers": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +def test_mxfp4_config_from_config_modules_to_not_convert_fallback(): + """modules_to_not_convert must be accepted as an alias for ignored_layers.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = DiffusionMXFP4Config.from_config({"modules_to_not_convert": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +# --------------------------------------------------------------------------- +# build_quant_config integration +# --------------------------------------------------------------------------- + + +def test_build_quant_config_mxfp4_string(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = build_quant_config("mxfp4") + assert isinstance(cfg, DiffusionMXFP4Config) + assert cfg.get_name() == "mxfp4" + assert cfg.is_checkpoint_mxfp4_serialized is False + + +def test_build_quant_config_mxfp4_dict(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + cfg = build_quant_config({"method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}) + assert isinstance(cfg, DiffusionMXFP4Config) + assert cfg.is_checkpoint_mxfp4_serialized is True + + +def test_build_quant_config_mxfp4_dualscale_string(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = build_quant_config("mxfp4_dualscale") + assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig) + assert cfg.is_checkpoint_serialized is False + assert cfg.num_bf16_fallback_layers == 5 + assert cfg.ignored_layers == [] + + +def test_build_quant_config_mxfp4_dualscale_dict_offline(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = build_quant_config( + { + "method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"], + } + ) + assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig) + assert cfg.is_checkpoint_serialized is True + assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"] + + +def test_build_quant_config_mxfp4_dualscale_dict_online_custom_fallback(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = build_quant_config({"method": "mxfp4_dualscale", "num_bf16_fallback_layers": 10}) + assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig) + assert cfg.num_bf16_fallback_layers == 10 + + +# --------------------------------------------------------------------------- +# Block-index dispatch (_parse_block_idx) +# --------------------------------------------------------------------------- + + +def test_parse_block_idx_valid(): + from vllm_omni.quantization.mxfp4_config import _parse_block_idx + + assert _parse_block_idx("blocks.0.attn1.to_q") == 0 + assert _parse_block_idx("blocks.5.ffn.net.0.proj") == 5 + assert _parse_block_idx("blocks.40.norm1.weight") == 40 + + +def test_parse_block_idx_non_block_prefixes(): + """Prefixes that do not start with 'blocks.N.' must return None.""" + from vllm_omni.quantization.mxfp4_config import _parse_block_idx + + assert _parse_block_idx("condition_embedder.time_embedder.linear_1") is None + assert _parse_block_idx("proj_out.weight") is None + assert _parse_block_idx("model.layers.0.self_attn.q_proj") is None + assert _parse_block_idx("scale_shift_table") is None + + +# --------------------------------------------------------------------------- +# SUPPORTED_QUANTIZATION_METHODS +# --------------------------------------------------------------------------- + + +def test_supported_methods_include_mxfp4_variants(): + from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS + + assert "mxfp4" in SUPPORTED_QUANTIZATION_METHODS + assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS + assert "mxfp4_dualscale" in SUPPORTED_QUANTIZATION_METHODS + + +# --------------------------------------------------------------------------- +# DiffusionMXFP4DualScaleMixedConfig — config roundtrips +# --------------------------------------------------------------------------- + + +def test_mixed_dualscale_config_get_name(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + assert DiffusionMXFP4DualScaleMixedConfig.get_name() == "mxfp4_dualscale" + + +def test_mixed_dualscale_config_no_args_defaults(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig() + assert cfg.is_checkpoint_serialized is False + assert cfg.ignored_layers == [] + assert cfg.num_bf16_fallback_layers == 5 + + +def test_mixed_dualscale_config_from_config_offline(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig.from_config( + { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ["blocks.0.attn1.to_q", "proj_out"], + } + ) + assert cfg.is_checkpoint_serialized is True + assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "proj_out"] + assert cfg.num_bf16_fallback_layers == 5 # default + + +def test_mixed_dualscale_config_from_config_online_custom_fallback(): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig.from_config({"num_bf16_fallback_layers": 10}) + assert cfg.is_checkpoint_serialized is False + assert cfg.num_bf16_fallback_layers == 10 + + +def test_mixed_dualscale_config_from_config_modules_to_not_convert_fallback(): + """modules_to_not_convert must be accepted as an alias for ignored_layers.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig.from_config( + {"is_checkpoint_serialized": True, "modules_to_not_convert": ["proj_out"]} + ) + assert cfg.ignored_layers == ["proj_out"] + + +# --------------------------------------------------------------------------- +# DiffusionMXFP4DualScaleMixedConfig — get_quant_method dispatch +# --------------------------------------------------------------------------- + + +def test_mixed_dualscale_offline_ignored_layer_returns_unquantized( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Offline: a prefix in ignored_layers must return UnquantizedLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = cfg.get_quant_method(layer, "blocks.0.attn1.to_q") + assert isinstance(method, UnquantizedLinearMethod) + + +def test_mixed_dualscale_offline_non_ignored_returns_mxfp4( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Offline: a prefix NOT in ignored_layers must return NPUMxfp4DualScaleLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + cfg = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=True, + ignored_layers=["blocks.0.attn1.to_q"], + ) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = cfg.get_quant_method(layer, "blocks.1.attn1.to_q") + assert isinstance(method, NPUMxfp4DualScaleLinearMethod) + + +def test_mixed_dualscale_online_fallback_block_returns_unquantized( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: blocks < num_bf16_fallback_layers must return UnquantizedLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + assert isinstance(cfg.get_quant_method(layer, "blocks.0.attn1.to_q"), UnquantizedLinearMethod) + assert isinstance(cfg.get_quant_method(layer, "blocks.4.ffn.net.0.proj"), UnquantizedLinearMethod) + + +def test_mixed_dualscale_online_quantized_block_returns_mxfp4( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: blocks >= num_bf16_fallback_layers must return NPUMxfp4DualScaleOnlineLinearMethod.""" + from vllm.model_executor.layers.linear import LinearBase + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4DualScaleMixedConfig, + NPUMxfp4DualScaleOnlineLinearMethod, + ) + + cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + assert isinstance(cfg.get_quant_method(layer, "blocks.5.attn1.to_q"), NPUMxfp4DualScaleOnlineLinearMethod) + assert isinstance(cfg.get_quant_method(layer, "blocks.40.ffn.net.0.proj"), NPUMxfp4DualScaleOnlineLinearMethod) + + +def test_mixed_dualscale_online_non_block_prefix_returns_mxfp4( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: layers outside 'blocks.N.*' (condition_embedder etc.) always use MXFP4 online.""" + from vllm.model_executor.layers.linear import LinearBase + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4DualScaleMixedConfig, + NPUMxfp4DualScaleOnlineLinearMethod, + ) + + cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = cfg.get_quant_method(layer, "condition_embedder.time_embedder.linear_1") + assert isinstance(method, NPUMxfp4DualScaleOnlineLinearMethod) + + +def test_mixed_dualscale_online_ignored_layers_override( + mocker, + monkeypatch: pytest.MonkeyPatch, +): + """Online: explicit ignored_layers must return UnquantizedLinearMethod regardless of block index. + + A layer that is NOT in the leading-block range (block 10 >= num_bf16_fallback_layers=5) + but IS listed in ignored_layers must still fall back to BF16. This lets power users + pin specific interleaved layers to BF16 during online quantization without needing an + offline checkpoint. + """ + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=False, + num_bf16_fallback_layers=5, + ignored_layers=["blocks.10.attn1.to_q"], + ) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + # block 10 is above the leading-block threshold but is in ignored_layers → BF16 + assert isinstance(cfg.get_quant_method(layer, "blocks.10.attn1.to_q"), UnquantizedLinearMethod) + + +def test_mixed_dualscale_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch): + """Non-LinearBase layers (norms, embeddings) must return None → no quantization.""" + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + cfg = DiffusionMXFP4DualScaleMixedConfig() + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + norm_layer = torch.nn.LayerNorm(64) + assert cfg.get_quant_method(norm_layer, "blocks.0.norm1") is None + + +# --------------------------------------------------------------------------- +# TP=2 create_weights: parameter shapes and input_dim/output_dim +# +# Two scenarios mirror real Wan2.2 A14B linear layer types: +# Column-parallel (to_q, ffn.net_0): output is sharded (N/TP), input is full (K). +# Row-parallel (to_out, ffn.net_2): input is sharded (K/TP), output is full (N). +# +# Tests verify: +# 1. Registered parameter shapes are correct for each partition configuration. +# 2. input_dim/output_dim attributes are set so RowParallelLinear.weight_loader +# can shard scale tensors correctly (the fix for the TP>1 shape-mismatch bug). +# 3. Simulated loader slicing: slicing the full checkpoint tensor along the +# declared input_dim produces the exact shape stored in the parameter — +# proving the dim declaration is consistent with the allocation. +# --------------------------------------------------------------------------- + +# K must be divisible by 32 (fine groups) and 512 (coarse groups). +_TP2_K, _TP2_N, _TP2 = 1024, 512, 2 + + +class _FakeLayer(torch.nn.Module): + """Bare nn.Module that accepts register_parameter without a real weight_loader.""" + + +def _create_weights(method, *, input_size_per_partition, output_partition_sizes): + layer = _FakeLayer() + method.create_weights( + layer=layer, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + input_size=_TP2_K, + output_size=_TP2_N, + params_dtype=torch.bfloat16, + ) + return layer + + +def _shard(tensor, param, rank, tp, dim_attr): + """Slice `tensor` along the dimension given by `param.` for `rank`.""" + dim = getattr(param, dim_attr) + if dim is None: + return tensor # not sharded along this axis + shard_size = param.shape[dim] + slices = [slice(None)] * tensor.ndim + slices[dim] = slice(rank * shard_size, (rank + 1) * shard_size) + return tensor[tuple(slices)] + + +# ── DualScale method ───────────────────────────────────────────────────────── + + +def test_dualscale_column_parallel_tp2_shapes(): + """Column-parallel TP=2: output halved, fine/coarse groups stay full, mul_scale full.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2]) + + assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K) + assert layer.weight_scale.shape == (_TP2_N // _TP2, _TP2_K // 32) + assert layer.weight_dual_scale.shape == (_TP2_N // _TP2, _TP2_K // 512, 1) + assert layer.mul_scale.shape == (_TP2_K,) + + +def test_dualscale_row_parallel_tp2_shapes(): + """Row-parallel TP=2: input halved, fine/coarse groups halved, mul_scale halved.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) + assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32) + assert layer.weight_dual_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 512, 1) + assert layer.mul_scale.shape == (_TP2_K // _TP2,) + + +def test_dualscale_scale_parameter_input_dims(): + """weight_scale/weight_dual_scale must have input_dim=1; mul_scale must have input_dim=0. + + RowParallelLinear.weight_loader only shards a parameter when input_dim is set. + Without these, loading a full checkpoint tensor into a per-rank shape causes a + shape mismatch for TP>1 row-parallel layers (to_out, ffn.net_2). + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N]) + + assert layer.weight_scale.input_dim == 1 + assert layer.weight_scale.output_dim == 0 + assert layer.weight_dual_scale.input_dim == 1 + assert layer.weight_dual_scale.output_dim == 0 + assert layer.mul_scale.input_dim == 0 + assert layer.mul_scale.output_dim is None + + +def test_dualscale_row_parallel_tp2_loader_simulation(): + """Slicing full checkpoint tensors along input_dim must match row-parallel parameter shapes. + + Simulates what RowParallelLinear.weight_loader does: for each scale parameter, + take the slice at rank*shard_size:(rank+1)*shard_size along input_dim. + The resulting shape must equal the per-rank parameter shape allocated by create_weights. + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + # Full checkpoint tensors (what the loader reads from disk). + ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32) + ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1) + ckpt_mul_scale = torch.zeros(_TP2_K) + + for rank in range(_TP2): + assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape + assert ( + _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "input_dim").shape + == layer.weight_dual_scale.shape + ) + assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "input_dim").shape == layer.mul_scale.shape + + +def test_dualscale_column_parallel_tp2_loader_simulation(): + """Slicing full checkpoint tensors along output_dim must match column-parallel parameter shapes. + + For column-parallel layers, the loader shards along output_dim (rows). + mul_scale has output_dim=None → not sharded (full tensor, same for all ranks). + """ + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod + + method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2]) + + ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32) + ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1) + ckpt_mul_scale = torch.zeros(_TP2_K) + + for rank in range(_TP2): + assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "output_dim").shape == layer.weight_scale.shape + assert ( + _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "output_dim").shape + == layer.weight_dual_scale.shape + ) + # mul_scale: output_dim=None → no sharding → full tensor fits the column-parallel parameter + assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "output_dim").shape == layer.mul_scale.shape + + +# ── Single-scale method ─────────────────────────────────────────────────────── + + +def test_single_scale_row_parallel_tp2_shapes(): + """Row-parallel TP=2: input halved → weight_scale groups halved.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod + + method = NPUMxfp4LinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) + assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32) + + +def test_single_scale_scale_parameter_input_dims(): + """Single-scale weight_scale must have input_dim=1 for RowParallel TP sharding.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod + + method = NPUMxfp4LinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N]) + + assert layer.weight_scale.input_dim == 1 + assert layer.weight_scale.output_dim == 0 + + +def test_single_scale_row_parallel_tp2_loader_simulation(): + """Slicing full checkpoint weight_scale along input_dim matches row-parallel parameter shape.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod + + method = NPUMxfp4LinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32) + + for rank in range(_TP2): + assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape diff --git a/tests/diffusion/quantization/test_mxfp4_key_remap.py b/tests/diffusion/quantization/test_mxfp4_key_remap.py new file mode 100644 index 000000000..e0cbc84df --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp4_key_remap.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for merge_mxfp4_dualscale_checkpoint.py key-remapping helpers. + +These are pure-Python unit tests that exercise the transformation functions +without loading any actual checkpoint files or requiring NPU hardware. +""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +# --------------------------------------------------------------------------- +# SUPPORTED_MODEL_TYPES +# --------------------------------------------------------------------------- + + +def test_supported_model_types_includes_a14b(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES + assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES + + +def test_supported_model_types_excludes_ti2v_5b(): + """TI2V-5B is explicitly NOT supported under W4A4 quantization.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-TI2V-5B" not in SUPPORTED_MODEL_TYPES + + +# --------------------------------------------------------------------------- +# _apply_rename_dict +# --------------------------------------------------------------------------- + + +def test_apply_rename_dict_self_attn_qkvo(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.0.self_attn.q.weight") == "blocks.0.attn1.to_q.weight" + assert _apply_rename_dict("blocks.0.self_attn.k.weight") == "blocks.0.attn1.to_k.weight" + assert _apply_rename_dict("blocks.0.self_attn.v.weight") == "blocks.0.attn1.to_v.weight" + assert _apply_rename_dict("blocks.0.self_attn.o.weight") == "blocks.0.attn1.to_out.0.weight" + + +def test_apply_rename_dict_ffn(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.1.ffn.0.weight") == "blocks.1.ffn.net.0.proj.weight" + assert _apply_rename_dict("blocks.1.ffn.2.weight") == "blocks.1.ffn.net.2.weight" + + +def test_apply_rename_dict_norm_order_swap(): + """norm2↔norm3 swap: quant tool uses norm1/norm3/norm2 order, + Diffusers uses norm1/norm2/norm3.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.0.norm2.weight") == "blocks.0.norm3.weight" + assert _apply_rename_dict("blocks.0.norm3.weight") == "blocks.0.norm2.weight" + + +def test_apply_rename_dict_cross_attn(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("blocks.0.cross_attn.q.weight") == "blocks.0.attn2.to_q.weight" + assert _apply_rename_dict("blocks.0.cross_attn.k.weight") == "blocks.0.attn2.to_k.weight" + assert _apply_rename_dict("blocks.0.cross_attn.v.weight") == "blocks.0.attn2.to_v.weight" + assert _apply_rename_dict("blocks.0.cross_attn.o.weight") == "blocks.0.attn2.to_out.0.weight" + + +def test_apply_rename_dict_head_and_modulation(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict + + assert _apply_rename_dict("head.head.weight") == "proj_out.weight" + + +# --------------------------------------------------------------------------- +# _strip_mxfp4_wrapper +# --------------------------------------------------------------------------- + + +def test_strip_mxfp4_wrapper_linear_weight(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight") == "blocks.0.attn1.to_q.weight" + + +def test_strip_mxfp4_wrapper_linear_weight_scale(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_scale") == "blocks.0.attn1.to_q.weight_scale" + + +def test_strip_mxfp4_wrapper_linear_weight_dual_scale(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert ( + _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_dual_scale") == "blocks.0.attn1.to_q.weight_dual_scale" + ) + + +def test_strip_mxfp4_wrapper_linear_bias(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.bias") == "blocks.0.attn1.to_q.bias" + + +def test_strip_mxfp4_wrapper_div_mul_scale(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.div.mul_scale") == "blocks.0.attn1.to_q.mul_scale" + + +def test_strip_mxfp4_wrapper_noop_for_plain_weight(): + """MXFP8 / FLOAT tensors have no wrapper — must be returned unchanged.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper + + assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.weight") == "blocks.0.attn1.to_q.weight" + assert _strip_mxfp4_wrapper("blocks.0.norm_q.weight") == "blocks.0.norm_q.weight" + assert _strip_mxfp4_wrapper("condition_embedder.time_embedder.linear_1.weight") == ( + "condition_embedder.time_embedder.linear_1.weight" + ) + + +# --------------------------------------------------------------------------- +# _classify_blocks +# --------------------------------------------------------------------------- + + +def test_classify_blocks_mixed_mxfp8_and_mxfp4(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks + + quant_meta = { + "blocks.0.attn1.to_q.weight": "W8A8_MXFP8", + "blocks.1.attn1.to_q.weight": "W8A8_MXFP8", + "blocks.2.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE", + "blocks.3.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE", + "condition_embedder.time_embedder.linear_1.weight": "FLOAT", + } + block_types = _classify_blocks(quant_meta) + assert block_types[0] == "mxfp8" + assert block_types[1] == "mxfp8" + assert block_types[2] == "mxfp4_dualscale" + assert block_types[3] == "mxfp4_dualscale" + # Non-block key must not produce an entry + assert None not in block_types + + +def test_classify_blocks_float_entries_skipped(): + """FLOAT-typed tensors must not contribute to block classification.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks + + quant_meta = { + "blocks.0.bias": "FLOAT", + "blocks.1.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE", + } + block_types = _classify_blocks(quant_meta) + assert 0 not in block_types + assert block_types[1] == "mxfp4_dualscale" + + +def test_classify_blocks_empty(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks + + assert _classify_blocks({}) == {} + + +# --------------------------------------------------------------------------- +# _is_mxfp4_tensor +# --------------------------------------------------------------------------- + + +def test_is_mxfp4_tensor_quantized_weight(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight", "W4A4_MXFP4_DUALSCALE") is True + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_scale", "W4A4_MXFP4_DUALSCALE") is True + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_dual_scale", "W4A4_MXFP4_DUALSCALE") is True + + +def test_is_mxfp4_tensor_companion_bias_and_mul_scale(): + """Companion tensors (bias, mul_scale) are FLOAT but belong to MXFP4 layers + and must be included in the merged checkpoint.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + # bias inside .linear. wrapper → must be included + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.bias", "FLOAT") is True + # mul_scale inside .div. wrapper → must be included + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.div.mul_scale", "FLOAT") is True + + +def test_is_mxfp4_tensor_bf16_fallback_weight(): + """BF16 fallback linear layers have plain .weight keys (no wrapper) → must NOT be included.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + assert _is_mxfp4_tensor("blocks.0.attn1.to_q.weight", "FLOAT") is False + assert _is_mxfp4_tensor("condition_embedder.time_embedder.linear_1.weight", "FLOAT") is False + + +def test_is_mxfp4_tensor_norm_weight(): + """Norm layers are always FLOAT with plain .weight → not MXFP4 related.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor + + assert _is_mxfp4_tensor("blocks.0.norm1.weight", "FLOAT") is False + + +# --------------------------------------------------------------------------- +# _collect_ignored_layers +# --------------------------------------------------------------------------- + + +def test_collect_ignored_layers_basic(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers + + # block 0: all three of Q/K/V are BF16 → fused as to_qkv in ignored_layers + # block 1: to_qkv is MXFP4 → not in ignored_layers + merged = { + "blocks.0.attn1.to_q.weight": torch.zeros(4, 4), + "blocks.0.attn1.to_k.weight": torch.zeros(4, 4), + "blocks.0.attn1.to_v.weight": torch.zeros(4, 4), + "blocks.1.attn1.to_q.weight": torch.zeros(4, 4), + "blocks.1.attn1.to_q.weight_scale": torch.zeros(4, 1), + } + mxfp4_prefixes = {"blocks.1.attn1.to_q"} + + ignored = _collect_ignored_layers(merged, mxfp4_prefixes) + assert "blocks.0.attn1.to_qkv" in ignored # fused vllm-omni name + assert "blocks.0.attn1.to_q" not in ignored # diffusers name must not appear + assert "blocks.0.attn1.to_k" not in ignored + assert "blocks.1.attn1.to_qkv" not in ignored # MXFP4, not ignored + + +def test_collect_ignored_layers_empty_mxfp4(): + """When no layers are MXFP4 (all BF16), all .weight prefixes become ignored_layers.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers + + # Use non-attn layers to avoid the Q/K/V completeness requirement. + merged = { + "blocks.0.ffn.net.0.proj.weight": torch.zeros(4, 4), + "proj_out.weight": torch.zeros(4, 4), + } + ignored = _collect_ignored_layers(merged, set()) + assert sorted(ignored) == ["blocks.0.ffn.net_0.proj", "proj_out"] + + +def test_collect_ignored_layers_returns_sorted(): + """ignored_layers must be sorted for deterministic config.json output.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers + + # Use FFN layers (no Q/K/V fusion concern) to isolate the sort guarantee. + merged = { + "blocks.2.ffn.net.0.proj.weight": torch.zeros(1), + "blocks.0.ffn.net.0.proj.weight": torch.zeros(1), + "blocks.1.ffn.net.0.proj.weight": torch.zeros(1), + } + ignored = _collect_ignored_layers(merged, set()) + assert ignored == sorted(ignored) + + +# --------------------------------------------------------------------------- +# _build_quant_config (new mxfp4_dualscale format) +# --------------------------------------------------------------------------- + + +def test_build_quant_config_new_format(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config + + config = _build_quant_config(["blocks.0.attn1.to_q", "proj_out"]) + assert config["quant_method"] == "mxfp4_dualscale" + assert config["is_checkpoint_serialized"] is True + assert config["ignored_layers"] == ["blocks.0.attn1.to_q", "proj_out"] + assert "num_mxfp8_blocks" not in config + + +def test_build_quant_config_empty_ignored_layers(): + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config + + config = _build_quant_config([]) + assert config["ignored_layers"] == [] + + +# --------------------------------------------------------------------------- +# _diffusers_to_vllm_ignored +# --------------------------------------------------------------------------- + + +def test_remap_self_attn_qkv_fusion(): + """All three of attn1.to_q/k/v in BF16 → fused attn1.to_qkv.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_k", "blocks.0.attn1.to_q", "blocks.0.attn1.to_v"]) + assert result == ["blocks.0.attn1.to_qkv"] + + +def test_remap_self_attn_qkv_partial_raises(): + """Partial Q/K/V in ignored_layers is invalid — must raise ValueError. + + Self-attention Q/K/V are fused into a single to_qkv layer at runtime; + partial precision (some BF16, some MXFP4) cannot be expressed and must + be caught early in the merge script. + """ + import pytest + + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + with pytest.raises(ValueError, match="Partial BF16 fallback"): + _diffusers_to_vllm_ignored(["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"]) + + +def test_remap_cross_attn_kept_separate(): + """Cross-attention (attn2) to_q/k/v must NOT be fused.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.attn2.to_k", "blocks.0.attn2.to_q", "blocks.0.attn2.to_v"]) + assert "blocks.0.attn2.to_qkv" not in result + assert "blocks.0.attn2.to_q" in result + assert "blocks.0.attn2.to_k" in result + assert "blocks.0.attn2.to_v" in result + + +def test_remap_ffn_dot_to_underscore(): + """ffn.net.0.proj → ffn.net_0.proj; ffn.net.2 → ffn.net_2.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.ffn.net.0.proj", "blocks.0.ffn.net.2"]) + assert "blocks.0.ffn.net_0.proj" in result + assert "blocks.0.ffn.net_2" in result + assert "blocks.0.ffn.net.0.proj" not in result + assert "blocks.0.ffn.net.2" not in result + + +def test_remap_to_out_strips_index(): + """attn1.to_out.0 and attn2.to_out.0 must have the .0 suffix removed.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_out.0", "blocks.0.attn2.to_out.0"]) + assert "blocks.0.attn1.to_out" in result + assert "blocks.0.attn2.to_out" in result + assert "blocks.0.attn1.to_out.0" not in result + assert "blocks.0.attn2.to_out.0" not in result + + +def test_remap_noop_for_non_block_layers(): + """Layers outside 'blocks.N.*' (condition_embedder, proj_out) pass through unchanged.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + inputs = [ + "condition_embedder.time_embedder.linear_1", + "condition_embedder.text_embedder.linear_1", + "proj_out", + ] + result = _diffusers_to_vllm_ignored(inputs) + assert result == sorted(inputs) + + +def test_remap_mixed_block_layers(): + """Typical mixed block: QKV fused, to_out stripped, FFN renamed, cross-attn separate.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + inputs = [ + "blocks.0.attn1.to_k", + "blocks.0.attn1.to_out.0", + "blocks.0.attn1.to_q", + "blocks.0.attn1.to_v", + "blocks.0.attn2.to_k", + "blocks.0.attn2.to_out.0", + "blocks.0.attn2.to_q", + "blocks.0.attn2.to_v", + "blocks.0.ffn.net.0.proj", + "blocks.0.ffn.net.2", + ] + result = _diffusers_to_vllm_ignored(inputs) + assert "blocks.0.attn1.to_qkv" in result + assert "blocks.0.attn1.to_out" in result + assert "blocks.0.attn2.to_q" in result + assert "blocks.0.attn2.to_k" in result + assert "blocks.0.attn2.to_v" in result + assert "blocks.0.attn2.to_out" in result + assert "blocks.0.ffn.net_0.proj" in result + assert "blocks.0.ffn.net_2" in result + # Old names must be gone + assert "blocks.0.attn1.to_q" not in result + assert "blocks.0.ffn.net.0.proj" not in result + + +def test_remap_returns_sorted(): + """Output of _diffusers_to_vllm_ignored must always be sorted.""" + from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored + + inputs = [ + "blocks.2.ffn.net.0.proj", + "blocks.0.attn1.to_q", + "blocks.0.attn1.to_k", + "blocks.0.attn1.to_v", + "blocks.1.attn2.to_out.0", + ] + result = _diffusers_to_vllm_ignored(inputs) + assert result == sorted(result) diff --git a/tests/diffusion/quantization/test_mxfp8_config.py b/tests/diffusion/quantization/test_mxfp8_config.py new file mode 100644 index 000000000..d831b5d47 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp8_config.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP8 quantization config and linear method dispatch. + +Coverage: +- DiffusionMXFP8Config.from_config roundtrips (CPU, no NPU required) +- get_quant_method dispatch (mocked platform) +- MXFPLinearMethodBase.apply() reshape skeleton (CPU) +- Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU) +- build_quant_config integration +- MXFP8_QUANT_CONFIG structure as the auto-detection contract +""" + +import pytest +import torch +from pytest_mock import MockerFixture +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import build_quant_config + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available") + + +# --------------------------------------------------------------------------- +# DiffusionMXFP8Config — from_config roundtrips +# --------------------------------------------------------------------------- + + +def test_mxfp8_config_get_name(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + assert DiffusionMXFP8Config.get_name() == "mxfp8" + + +def test_mxfp8_config_from_config_defaults(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({}) + assert cfg.is_checkpoint_mxfp8_serialized is False + assert cfg.ignored_layers == [] + + +def test_mxfp8_config_from_config_serialized(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({"is_checkpoint_mxfp8_serialized": True}) + assert cfg.is_checkpoint_mxfp8_serialized is True + + +def test_mxfp8_config_from_config_ignored_layers(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({"ignored_layers": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +def test_mxfp8_config_from_config_modules_to_not_convert_fallback(): + """modules_to_not_convert must be accepted as an alias for ignored_layers.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = DiffusionMXFP8Config.from_config({"modules_to_not_convert": ["proj_out"]}) + assert cfg.ignored_layers == ["proj_out"] + + +# --------------------------------------------------------------------------- +# build_quant_config integration +# --------------------------------------------------------------------------- + + +def test_build_quant_config_mxfp8_string(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = build_quant_config("mxfp8") + assert isinstance(cfg, DiffusionMXFP8Config) + assert cfg.get_name() == "mxfp8" + assert cfg.is_checkpoint_mxfp8_serialized is False + + +def test_build_quant_config_mxfp8_dict(): + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + cfg = build_quant_config({"method": "mxfp8", "is_checkpoint_mxfp8_serialized": True}) + assert isinstance(cfg, DiffusionMXFP8Config) + assert cfg.is_checkpoint_mxfp8_serialized is True + + +def test_build_quant_config_mxfp8_config_json_format(): + """Verify that the exact quantization_config injected by merge_mxfp8_checkpoint.py + is accepted by build_quant_config and selects the offline (serialized) path. + + This is the critical auto-detection contract: TransformerConfig.from_dict() + reads quant_method + is_checkpoint_mxfp8_serialized to pick NPUMxfp8LinearMethod. + """ + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG + + cfg = build_quant_config(MXFP8_QUANT_CONFIG) + assert isinstance(cfg, DiffusionMXFP8Config) + assert cfg.is_checkpoint_mxfp8_serialized is True + + +def test_mxfp8_quant_config_structure(): + """MXFP8_QUANT_CONFIG must contain exactly the keys that auto-detection reads.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG + + assert MXFP8_QUANT_CONFIG.get("quant_method") == "mxfp8" + assert MXFP8_QUANT_CONFIG.get("is_checkpoint_mxfp8_serialized") is True + + +# --------------------------------------------------------------------------- +# get_quant_method dispatch +# --------------------------------------------------------------------------- + + +def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Offline (serialized) path on NPU must return NPUMxfp8LinearMethod.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod + + config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = config.get_quant_method(layer, "blocks.0.attn1.to_q") + assert isinstance(method, NPUMxfp8LinearMethod) + + +def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod + + config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = config.get_quant_method(layer, "blocks.0.attn1.to_q") + assert isinstance(method, NPUMxfp8OnlineLinearMethod) + + +def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Non-NPU platform must raise NotImplementedError (CUDA not yet supported).""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config() + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + + with pytest.raises(NotImplementedError): + config.get_quant_method(layer, "blocks.0.attn1.to_q") + + +def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """A prefix in ignored_layers must bypass quantization → UnquantizedLinearMethod.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config(ignored_layers=["proj_out"]) + layer = mocker.Mock(spec=LinearBase) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + method = config.get_quant_method(layer, "proj_out") + assert isinstance(method, UnquantizedLinearMethod) + + +def test_get_quant_method_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch): + """Non-LinearBase layers (norms, embeddings) must get None → no quantization.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config() + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + + norm_layer = torch.nn.LayerNorm(64) + assert config.get_quant_method(norm_layer, "blocks.0.norm1") is None + + +# --------------------------------------------------------------------------- +# MXFPLinearMethodBase.apply() — reshape skeleton (CPU, no NPU) +# --------------------------------------------------------------------------- + + +def test_apply_reshape_skeleton(): + """apply() must flatten batch dims → _apply_inner → restore original leading dims.""" + from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase + + OUT_FEATURES = 4 + + class _StubMethod(MXFPLinearMethodBase): + def create_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ): + pass + + def _quantize_activation(self, x): + return x, None + + def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype): + return torch.zeros(x_q.shape[0], OUT_FEATURES, dtype=ori_dtype) + + method = _StubMethod() + x = torch.randn(2, 3, 8) # (batch=2, seq=3, K=8) + out = method.apply(None, x) + assert out.shape == (2, 3, OUT_FEATURES) + + +def test_apply_reshape_with_bias(): + """apply() must pass bias through to _apply_inner unchanged.""" + from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase + + received_bias = [] + + class _StubMethod(MXFPLinearMethodBase): + def create_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ): + pass + + def _quantize_activation(self, x): + return x, None + + def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype): + received_bias.append(bias) + return torch.zeros(x_q.shape[0], 4, dtype=ori_dtype) + + method = _StubMethod() + bias = torch.zeros(4) + method.apply(None, torch.randn(2, 8), bias=bias) + assert received_bias[0] is bias + + +# --------------------------------------------------------------------------- +# process_weights_after_loading shape arithmetic (pure torch, no NPU ops) +# +# These tests replicate the CPU-safe portions of process_weights_after_loading +# to guard the key layout contract: (N,K) weight → (K,N) and (N,S) scale → +# (S/2,N,2). They do NOT call NPU ops; they test only the math. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "n, k", + [(64, 128), (32, 64), (16, 96)], + ids=["64x128", "32x64", "16x96"], +) +def test_weight_transpose_contract(n: int, k: int): + """Weight must be transposed from (N, K) to (K, N) and be contiguous.""" + w = torch.zeros(n, k, dtype=torch.uint8) + w = w.transpose(0, 1).contiguous() + assert w.shape == (k, n) + assert w.is_contiguous() + + +@pytest.mark.parametrize( + "n, k_groups, expected_groups", + [ + (64, 4, 2), # even K_groups — no padding + (64, 3, 2), # odd K_groups — padded to 4 + (32, 1, 1), # odd K_groups — padded to 2 + ], + ids=["even", "odd-to-4", "odd-to-2"], +) +def test_weight_scale_reshape_contract(n: int, k_groups: int, expected_groups: int): + """Scale must be reshaped from (N, K_groups) to (K_groups_even//2, N, 2). + + Odd K_groups must be padded to even before the reshape. + """ + s = torch.zeros(n, k_groups, dtype=torch.uint8) + if k_groups % 2 == 1: + s = torch.cat([s, torch.zeros(n, 1, dtype=s.dtype)], dim=1) + k_groups += 1 + s = s.reshape(n, k_groups // 2, 2).transpose(0, 1).contiguous() + assert s.shape == (expected_groups, n, 2) + assert s.is_contiguous() + + +def test_num_groups_formula(): + """K_groups formula: ceil(K / 32) — spot-check boundary values.""" + assert (31 + 31) // 32 == 1 # K=31 → 1 group + assert (32 + 31) // 32 == 1 # K=32 → 1 group + assert (33 + 31) // 32 == 2 # K=33 → 2 groups + assert (128 + 31) // 32 == 4 # K=128 → 4 groups (even) + assert (96 + 31) // 32 == 3 # K=96 → 3 groups (odd → needs padding) + + +# --------------------------------------------------------------------------- +# SUPPORTED_QUANTIZATION_METHODS +# --------------------------------------------------------------------------- + + +def test_supported_methods_include_mxfp8(): + from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS + + assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS diff --git a/tests/diffusion/quantization/test_mxfp8_key_remap.py b/tests/diffusion/quantization/test_mxfp8_key_remap.py new file mode 100644 index 000000000..f6eea04a0 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp8_key_remap.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for merge_mxfp8_checkpoint.py key-remapping helpers and model metadata. + +These are pure-Python unit tests that exercise the transformation functions +without loading any actual checkpoint files or requiring NPU hardware. + +Key contracts verified: +- SUPPORTED_MODEL_TYPES includes Wan2.2-TI2V-5B (MXFP8 supports it; MXFP4 does not) +- MXFP8_QUANT_CONFIG structure matches what TransformerConfig.from_dict() reads +- _remap_keys correctly translates msModelSlim naming → Diffusers naming +- _get_transformer_dirs routes cascade vs single-transformer models +- _get_quant_subdir maps high/low noise subdirs for cascade models +""" + +import pathlib + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- +# SUPPORTED_MODEL_TYPES and CASCADE_MODEL_TYPES +# --------------------------------------------------------------------------- + + +def test_supported_model_types_includes_all_wan22(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES + assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES + + +def test_supported_model_types_includes_ti2v_5b(): + """MXFP8 supports TI2V-5B (contrast: MXFP4 explicitly excludes it).""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES + + assert "Wan2.2-TI2V-5B" in SUPPORTED_MODEL_TYPES + + +def test_cascade_model_types_excludes_ti2v(): + """TI2V-5B is a single-transformer model — must not be in CASCADE_MODEL_TYPES.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import CASCADE_MODEL_TYPES + + assert "Wan2.2-TI2V-5B" not in CASCADE_MODEL_TYPES + assert "Wan2.2-T2V-A14B" in CASCADE_MODEL_TYPES + assert "Wan2.2-I2V-A14B" in CASCADE_MODEL_TYPES + + +# --------------------------------------------------------------------------- +# MXFP8_QUANT_CONFIG — auto-detection contract +# --------------------------------------------------------------------------- + + +def test_mxfp8_quant_config_has_required_keys(): + """MXFP8_QUANT_CONFIG must carry exactly the keys that auto-detection reads: + quant_method (selects DiffusionMXFP8Config) and is_checkpoint_mxfp8_serialized + (selects NPUMxfp8LinearMethod over the online path).""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG + + assert MXFP8_QUANT_CONFIG["quant_method"] == "mxfp8" + assert MXFP8_QUANT_CONFIG["is_checkpoint_mxfp8_serialized"] is True + + +# --------------------------------------------------------------------------- +# _get_transformer_dirs +# --------------------------------------------------------------------------- + + +def test_get_transformer_dirs_cascade(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs + + assert _get_transformer_dirs("Wan2.2-T2V-A14B") == ["transformer", "transformer_2"] + assert _get_transformer_dirs("Wan2.2-I2V-A14B") == ["transformer", "transformer_2"] + + +def test_get_transformer_dirs_single(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs + + assert _get_transformer_dirs("Wan2.2-TI2V-5B") == ["transformer"] + + +# --------------------------------------------------------------------------- +# _get_quant_subdir +# --------------------------------------------------------------------------- + + +def test_get_quant_subdir_cascade_high_noise(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir + + base = pathlib.Path("/quant") + result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer") + assert result == base / "high_noise_model" + + +def test_get_quant_subdir_cascade_low_noise(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir + + base = pathlib.Path("/quant") + result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer_2") + assert result == base / "low_noise_model" + + +def test_get_quant_subdir_non_cascade(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir + + base = pathlib.Path("/quant") + result = _get_quant_subdir("Wan2.2-TI2V-5B", base, "transformer") + assert result == base + + +# --------------------------------------------------------------------------- +# _remap_keys — msModelSlim naming → Diffusers naming +# --------------------------------------------------------------------------- + + +def test_remap_keys_self_attn_q(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"blocks.0.self_attn.q.weight": torch.zeros(1)} + meta = {"blocks.0.self_attn.q.weight": "W8A8_MXFP8"} + new_state, new_meta = _remap_keys(state, meta) + assert "blocks.0.attn1.to_q.weight" in new_state + assert new_meta.get("blocks.0.attn1.to_q.weight") == "W8A8_MXFP8" + + +def test_remap_keys_self_attn_all_heads(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + pairs = { + "self_attn.q": "attn1.to_q", + "self_attn.k": "attn1.to_k", + "self_attn.v": "attn1.to_v", + "self_attn.o": "attn1.to_out.0", + } + for src_part, dst_part in pairs.items(): + src_key = f"blocks.0.{src_part}.weight" + state = {src_key: torch.zeros(1)} + new_state, _ = _remap_keys(state, {}) + expected = f"blocks.0.{dst_part}.weight" + assert expected in new_state, f"{src_key} → expected {expected}, got {list(new_state)}" + + +def test_remap_keys_ffn(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = { + "blocks.1.ffn.0.weight": torch.zeros(1), + "blocks.1.ffn.2.weight": torch.zeros(1), + } + new_state, _ = _remap_keys(state, {}) + assert "blocks.1.ffn.net.0.proj.weight" in new_state + assert "blocks.1.ffn.net.2.weight" in new_state + + +def test_remap_keys_norm_order_swap(): + """norm2↔norm3 swap: msModelSlim uses norm1/norm3/norm2 order, + Diffusers uses norm1/norm2/norm3.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = { + "blocks.0.norm2.weight": torch.zeros(1), + "blocks.0.norm3.weight": torch.zeros(1), + } + new_state, _ = _remap_keys(state, {}) + # norm2 → norm3 and norm3 → norm2 + assert "blocks.0.norm3.weight" in new_state + assert "blocks.0.norm2.weight" in new_state + # Both must be present (swap, not collapse) + assert len([k for k in new_state if "norm" in k and "norm_q" not in k and "norm_k" not in k]) == 2 + + +def test_remap_keys_head(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"head.head.weight": torch.zeros(1)} + new_state, _ = _remap_keys(state, {}) + assert "proj_out.weight" in new_state + + +def test_remap_keys_cross_attn(): + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"blocks.0.cross_attn.q.weight": torch.zeros(1)} + new_state, _ = _remap_keys(state, {}) + assert "blocks.0.attn2.to_q.weight" in new_state + + +def test_remap_keys_meta_only_mapped_for_existing_state_keys(): + """quant_meta entries are only emitted for keys present in state_dict.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + state = {"blocks.0.self_attn.q.weight": torch.zeros(1)} + # meta has an extra key not in state_dict + meta = { + "blocks.0.self_attn.q.weight": "W8A8_MXFP8", + "blocks.0.self_attn.q.weight_scale": "W8A8_MXFP8", + } + _, new_meta = _remap_keys(state, meta) + assert "blocks.0.attn1.to_q.weight" in new_meta + # weight_scale was in meta but not state_dict → must be absent + assert "blocks.0.attn1.to_q.weight_scale" not in new_meta + + +def test_remap_keys_preserves_tensors(): + """Tensor values must survive the key rename unchanged.""" + from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys + + t = torch.randn(4, 8) + state = {"blocks.0.self_attn.q.weight": t} + new_state, _ = _remap_keys(state, {}) + assert torch.equal(new_state["blocks.0.attn1.to_q.weight"], t) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index b4438c795..3ee46ffb0 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -159,54 +159,10 @@ def create_transformer_from_config( if "pos_embed_seq_len" in config: kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"] - # Auto-detect quantization from transformer's config.json when not explicitly provided. - # merge_mxfp8_checkpoint.py injects quantization_config into config.json so that - # offline quantized checkpoints are recognized here without a CLI flag. if "quantization_config" in config: - from vllm_omni.quantization.factory import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk - disk_qc = config["quantization_config"] - if isinstance(disk_qc, dict) and "quant_method" in disk_qc: - qc_method = disk_qc["quant_method"] - qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"} - if quant_config is None: - # No CLI flag: full auto-detection. - quant_config = build_quant_config(qc_method, **qc_kwargs) - logger.info( - "Auto-detected quantization from transformer config.json: method=%s kwargs=%s", - qc_method, - qc_kwargs, - ) - elif quant_config.get_name() != qc_method: - # The caller supplied a quant_config of a different method than - # what the checkpoint was built with. Loading serialized tensors - # (e.g. MXFP8 weight scales) with the wrong linear method would - # produce corrupt output or a shape mismatch crash. Reject early - # so the user gets a clear message instead of a silent failure. - raise ValueError( - f"Checkpoint config.json declares quant_method={qc_method!r} but the " - f"active quantization config is {quant_config.get_name()!r}. " - "Pass a matching --quantization flag or omit it for auto-detection." - ) - elif ( - qc_kwargs.get("is_checkpoint_mxfp8_serialized", False) - and hasattr(quant_config, "is_checkpoint_mxfp8_serialized") - and not quant_config.is_checkpoint_mxfp8_serialized - ): - # Same method: CLI provided online mode but config.json marks this - # as a pre-quantized offline checkpoint. Switch to offline mode so - # users can pass --quantization mxfp8 without knowing the - # online/offline distinction. - quant_config = build_quant_config(qc_method, **qc_kwargs) - logger.info( - "config.json marks checkpoint as serialized; switching from online to offline MXFP8 mode.", - ) - elif isinstance(disk_qc, str) and quant_config is None: - quant_config = build_quant_config(disk_qc) - logger.info( - "Auto-detected quantization from transformer config.json: method=%s", - disk_qc, - ) + quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"]) if quant_config is not None: kwargs["quant_config"] = quant_config diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index ea7b8a7fe..955f97cef 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -48,6 +48,29 @@ def _build_mxfp8(**kw: Any) -> QuantizationConfig: return DiffusionMXFP8Config(**kw) +def _build_mxfp4(**kw: Any) -> QuantizationConfig: + """Lazy import for W4A4 MXFP4 diffusion config (NPU only).""" + from .mxfp4_config import DiffusionMXFP4Config + + return DiffusionMXFP4Config(**kw) + + +def _build_mxfp4_dualscale(**kw: Any) -> QuantizationConfig: + """Lazy import for MXFP4 DualScale + BF16 mixed diffusion config (NPU only). + + Offline mode (is_checkpoint_serialized=True): + ignored_layers from config.json marks interleaved BF16 fallback layers. + All other linear layers use W4A4 MXFP4 DualScale. + + Online mode (is_checkpoint_serialized=False): + num_bf16_fallback_layers leading transformer blocks use BF16 original weights + (default 5 when not specified). Remaining blocks use W4A4 MXFP4 DualScale online. + """ + from .mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + return DiffusionMXFP4DualScaleMixedConfig(**kw) + + def _build_inc(**kw: Any) -> QuantizationConfig: """Lazy import for INC/AutoRound config with checkpoint kwarg normalization.""" from .inc_config import OmniINCConfig @@ -66,6 +89,8 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "gguf": _build_gguf, "int8": _build_int8, "mxfp8": _build_mxfp8, + "mxfp4": _build_mxfp4, + "mxfp4_dualscale": _build_mxfp4_dualscale, "inc": _build_inc, "auto-round": _build_inc, "auto_round": _build_inc, @@ -255,3 +280,80 @@ def build_quant_config( return _build_from_method_and_config(method, merged) raise TypeError(f"quantization config must be str, dict, QuantizationConfig, or None, got {type(spec).__name__}") + + +def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> bool: + """Return True when config.json says serialized but the active quant_config does not. + + Matches any flag following the is_checkpoint_*_serialized naming convention, + so new quant methods don't require updating an explicit allowlist. + """ + for key, val in qc_kwargs.items(): + if key.startswith("is_checkpoint_") and key.endswith("_serialized"): + if val and hasattr(quant_config, key) and not getattr(quant_config, key): + return True + return False + + +def resolve_quant_config_from_disk( + quant_config: QuantizationConfig | None, + disk_qc: dict[str, Any] | str | None, +) -> QuantizationConfig | None: + """Reconcile an active quant_config against quantization_config from a transformer's config.json. + + Used when loading individual transformer blocks that each have their own config.json + (e.g. cascade models with separate transformer and transformer_2 directories). + + Rules: + - disk_qc is None: return quant_config unchanged. + - quant_config is None: auto-detect from disk_qc (full build). + - Methods mismatch: raise ValueError — prevents silent weight corruption. + - Disk marks serialized but quant_config is online: rebuild from disk. + - ignored_layers differ: rebuild from disk (per-transformer BF16 routing). + """ + if disk_qc is None: + return quant_config + + if isinstance(disk_qc, str): + if quant_config is None: + logger.info("Auto-detected quantization from config.json: method=%s", disk_qc) + return build_quant_config(disk_qc) + return quant_config + + if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc: + return quant_config + + qc_method: str = disk_qc["quant_method"] + qc_kwargs: dict[str, Any] = {k: v for k, v in disk_qc.items() if k != "quant_method"} + + if quant_config is None: + logger.info( + "Auto-detected quantization from config.json: method=%s kwargs=%s", + qc_method, + qc_kwargs, + ) + return build_quant_config(qc_method, **qc_kwargs) + + if quant_config.get_name() != qc_method: + raise ValueError( + f"Checkpoint config.json declares quant_method={qc_method!r} but the " + f"active quantization config is {quant_config.get_name()!r}. " + "Pass a matching --quantization flag or omit it for auto-detection." + ) + + if _disk_marks_serialized(qc_kwargs, quant_config): + logger.info( + "config.json marks checkpoint as serialized; switching to offline %s mode.", + qc_method, + ) + return build_quant_config(qc_method, **qc_kwargs) + + if ( + "ignored_layers" in qc_kwargs + and hasattr(quant_config, "ignored_layers") + and set(qc_kwargs.get("ignored_layers") or []) != set(quant_config.ignored_layers or []) + ): + logger.info("config.json ignored_layers differs from active config; rebuilding quant_config.") + return build_quant_config(qc_method, **qc_kwargs) + + return quant_config diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py new file mode 100644 index 000000000..6a02d5edb --- /dev/null +++ b/vllm_omni/quantization/mxfp4_config.py @@ -0,0 +1,709 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers. + +Architecture mirrors mxfp8_config.py: + + MXFPLinearMethodBase – platform-agnostic skeleton (imported from mxfp8_config) + NPUMxfp4LinearMethod – NPU single-scale offline (W4A4 MXFP4) + NPUMxfp4OnlineLinearMethod – NPU single-scale online (BF16 → FP4) + NPUMxfp4DualScaleLinearMethod – NPU dual-scale offline (W4A4 MXFP4 DualScale) + NPUMxfp4DualScaleOnlineLinearMethod – NPU dual-scale online (BF16 → FP4) + +Quantization configs: + + DiffusionMXFP4Config – single-scale online/offline (quant_method="mxfp4") + DiffusionMXFP4DualScaleMixedConfig – dual-scale + per-layer BF16 fallback (quant_method="mxfp4_dualscale") + Offline: ignored_layers from config.json routes interleaved BF16 layers + Online: num_bf16_fallback_layers leading blocks stay in BF16 (default 5) + +Key differences from MXFP8: + + 1. Precision: float4_e2m1fn_x2 (FP4 packed, 2 values per element). + npu_dynamic_mx_quant(x) without dst_type defaults to float4_e2m1fn_x2. + + 2. Weight layout: stored as (N, K) — NOT pre-transposed. + FP4 uses a packed format; transposing a packed tensor is not safely contiguous. + Transpose is done inline in _quant_matmul via layer.weight.transpose(0, 1). + + 3. GEMM signature: npu_quant_matmul requires explicit + x1_dtype=float4_e2m1fn_x2 and x2_dtype=float4_e2m1fn_x2. + + Scale layout: (N, S/2, 2) — same reshape as MXFP8, also NOT pre-transposed; + transposed inline in _quant_matmul. + +Dual-scale (W4A4_MXFP4_DUALSCALE): + + Two-level quantization: fine scale (per-32 K) + coarse scale (per-512 K) + per-channel + activation pre-scale (mul_scale from calibration). Uses npu_dynamic_dual_level_mx_quant + and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast. + + Checkpoint tensor shapes for dual-scale: + weight : (N, K) float8_e4m3fn – FP4 packed + weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits) + weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert) + mul_scale : (K,) float32 – per-input-channel activation pre-scale + +Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/quantization/layer.py). +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Any + +import torch +from torch.nn import Module +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped +from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.model_executor.utils import replace_parameter + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.mxfp8_config import ( + MXFPLinearMethodBase, + _LazyWeightMixin, +) + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +class DiffusionMXFP4Config(QuantizationConfig): + """W4A4 MXFP4 quantization config for diffusion transformers. + + Supports both online (BF16 checkpoint → quantize at load time) and offline + (pre-quantized MXFP4 checkpoint) modes, mirroring DiffusionMXFP8Config. + + MX (microscaling) format: groups of 32 K-dimension elements share one + float8_e8m0fnu exponent scale. Weight and activation are float4_e2m1fn_x2. + """ + + def __init__( + self, + is_checkpoint_mxfp4_serialized: bool = False, + ignored_layers: list[str] | None = None, + ) -> None: + super().__init__() + self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized + self.ignored_layers = ignored_layers or [] + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "mxfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + if self.ignored_layers: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> DiffusionMXFP4Config: + is_serialized = cls.get_from_keys_or(config, ["is_checkpoint_mxfp4_serialized"], False) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + return cls( + is_checkpoint_mxfp4_serialized=is_serialized, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> QuantizeMethodBase | None: + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if current_omni_platform.is_npu(): + if self.is_checkpoint_mxfp4_serialized: + return NPUMxfp4LinearMethod(self) + return NPUMxfp4OnlineLinearMethod(self) + raise NotImplementedError( + "DiffusionMXFP4Config (W4A4 MXFP4) is currently only supported " + "on NPU (Ascend) platforms. CUDA support is not yet implemented." + ) + return None + + +# --------------------------------------------------------------------------- +# NPU MXFP4 single-scale offline method (pre-quantized checkpoint) +# --------------------------------------------------------------------------- + + +class NPUMxfp4LinearMethod(MXFPLinearMethodBase): + """NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints. + + Weight canonical layout after process_weights_after_loading: + weight : (N, K) in float4_e2m1fn_x2 — NOT pre-transposed (FP4 packed) + weight_scale: (N, S/2, 2) in float8_e8m0fnu — NOT pre-transposed + + Both are transposed inline in _quant_matmul, unlike MXFP8 which pre-transposes. + NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared. + """ + + def __init__(self, quant_config: DiffusionMXFP4Config) -> None: + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights. + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # Scale stored as uint8 in safetensors (float8_e8m0fnu is same bit width). + # Using uint8 avoids a lossy float32 round-trip when loading the checkpoint. + num_groups = (input_size_per_partition + 31) // 32 + layer.register_parameter( + "weight_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups, dtype=torch.uint8), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + # NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose. + w = layer.weight + if w.dtype != torch_npu.float4_e2m1fn_x2: + w = torch_npu.npu_dtype_cast(w.npu(), torch_npu.float4_e2m1fn_x2) + + # Scale: checkpoint stores uint8 bytes that ARE float8_e8m0fnu bits. + # Only convert if neither uint8 nor the target NPU dtype already. + # (N, K_groups) → (N, K_groups/2, 2). Not pre-transposed; done inline in _quant_matmul. + s = layer.weight_scale.data + if s.dtype not in (torch.uint8, torch_npu.float8_e8m0fnu): + s = s.to(torch_npu.float8_e8m0fnu) + N, K_groups = s.shape + if K_groups % 2 == 1: + s = torch.cat([s, torch.zeros(N, 1, dtype=s.dtype, device=s.device)], dim=1) + s = s.reshape(N, -1, 2).contiguous() + + replace_parameter(layer, "weight", w) + replace_parameter(layer, "weight_scale", s) + layer._already_called_process_weights_after_loading = True + + # --- NPU MXFP4 ops — shared with online path via inheritance --- + + def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + import torch_npu + + # No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2. + return torch_npu.npu_dynamic_mx_quant(x) + + def _quant_matmul( + self, + x_q: torch.Tensor, + x_scale: torch.Tensor, + layer: torch.nn.Module, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + import torch_npu + + if bias is not None and bias.dtype != torch.float32: + bias = bias.to(torch.float32) + # FP4 differences vs FP8: + # weight (N,K) transposed inline → (K,N); scale (N,S/2,2) transposed inline → (S/2,N,2). + # x1_dtype / x2_dtype required — FP4 dtype not inferred from tensor dtype. + return torch_npu.npu_quant_matmul( + x_q, + layer.weight.transpose(0, 1), # (K, N) inline + layer.weight_scale.transpose(0, 1), # (S/2, N, 2) inline + scale_dtype=torch_npu.float8_e8m0fnu, + x1_dtype=torch_npu.float4_e2m1fn_x2, + x2_dtype=torch_npu.float4_e2m1fn_x2, + pertoken_scale=x_scale, + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, + bias=bias, + output_dtype=ori_dtype, + group_sizes=[1, 1, 32], + ) + + +# --------------------------------------------------------------------------- +# NPU MXFP4 single-scale online method (BF16 checkpoint → quantize at load time) +# --------------------------------------------------------------------------- + + +class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): + """NPU W4A4 MXFP4 online linear method. + + MRO: NPUMxfp4OnlineLinearMethod → _LazyWeightMixin → NPUMxfp4LinearMethod + → MXFPLinearMethodBase → LinearMethodBase + + create_weights : _LazyWeightMixin (meta device + patched loader) + process_weights : NPUMxfp4OnlineLinearMethod (BF16 → FP4 + normalize) + apply / ops : NPUMxfp4LinearMethod / MXFPLinearMethodBase (shared) + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + # NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2. + weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight) + + # Weight stays (N, K) — no pre-transpose for FP4 packed format. + # Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul. + weight_scale = weight_scale_raw.reshape(weight_scale_raw.shape[0], -1, 2).contiguous() + + replace_parameter(layer, "weight", weight_fp4) + replace_parameter(layer, "weight_scale", weight_scale) + layer._already_called_process_weights_after_loading = True + + +# --------------------------------------------------------------------------- +# NPU MXFP4 dual-scale offline method (W4A4_MXFP4_DUALSCALE checkpoint) +# --------------------------------------------------------------------------- + + +class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): + """NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints. + + Checkpoint tensors and their canonical post-load shapes: + weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte) + weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2) + weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N) + mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration) + + Forward pass: + x_q, l0, l1 = npu_dynamic_dual_level_mx_quant(x, smooth_scale=mul_scale) + out = npu_dual_level_quant_matmul(x_q, weight, l0, weight_dual_scale, l1, weight_scale) + + Reference: MindIE-SD W4A4MXFP4DualQuantLinear. + """ + + def __init__(self, quant_config: Any) -> None: + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16 + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # Fine scale: one uint8 exponent (float8_e8m0fnu bit pattern) per group of 32 K elements. + # input_dim=1: RowParallelLinear weight_loader shards along the K-group dimension (dim 1) + # so each rank receives only its (K/TP)//32 groups; ColumnParallelLinear only uses + # output_dim=0 for sharding and leaves dim 1 intact. + num_groups_fine = (input_size_per_partition + 31) // 32 + layer.register_parameter( + "weight_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups_fine, dtype=torch.uint8), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # Coarse scale: one float32 per group of 512 K elements. + # Shape (N, K_coarse, 1) matches checkpoint layout exactly, avoiding the + # shape-mismatch assert in linear.py:1344. + # input_dim=1: same TP sharding rationale as weight_scale above. + num_groups_coarse = (input_size_per_partition + 511) // 512 + layer.register_parameter( + "weight_dual_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups_coarse, 1, dtype=torch.float32), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + # mul_scale is a float32 calibration tensor; register as float32 + # to avoid precision loss from an implicit BF16 cast during weight loading. + # input_dim=0: RowParallelLinear shards the 1-D per-input-channel tensor along dim 0 + # so each rank receives only its K/TP channels; ColumnParallelLinear leaves it intact. + layer.register_parameter( + "mul_scale", + ModelWeightParameter( + data=torch.empty(input_size_per_partition, dtype=torch.float32), + input_dim=0, + output_dim=None, + weight_loader=weight_loader, + ), + ) + setattr(layer.mul_scale, "ignore_warning", True) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + # float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29). + # NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices. + w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2) + w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8) + + # Fine scale: (N, K//32) uint8 → cast to float8_e8m0fnu → (N, K//64, 2). + # The cast reinterprets the stored uint8 bit-patterns as float8_e8m0fnu exponents, + # which is required for npu_dual_level_quant_matmul to compute correct scale values. + s = layer.weight_scale.data.npu() + s = s.reshape(s.shape[0], -1, 2).contiguous() + + # Coarse scale: (N, K//512, 1) → squeeze → (N, K//512) → transpose → (K//512, N). + ds = layer.weight_dual_scale.data.squeeze(-1).transpose(0, 1).npu().contiguous() + + ms = layer.mul_scale.to(torch.bfloat16).data.view(-1).npu().contiguous() + + replace_parameter(layer, "weight", w) + replace_parameter(layer, "weight_scale", s) + replace_parameter(layer, "weight_dual_scale", ds) + replace_parameter(layer, "mul_scale", ms) + layer._already_called_process_weights_after_loading = True + + def _apply_inner( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + """Dual-scale inner loop: dtype guard → 3-tuple quantize (mul_scale as smooth) → matmul. + + Overrides the default single-scale _apply_inner from MXFPLinearMethodBase. + apply() in the base class handles reshape/unreshape; this method is not + responsible for that. + """ + if ori_dtype not in (torch.bfloat16, torch.float16): + x = x.to(torch.bfloat16) + x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale) + return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype) + + def _quantize_activation( # type: ignore[override] + self, + x: torch.Tensor, + smooth_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + import torch_npu + + return torch_npu.npu_dynamic_dual_level_mx_quant(x, smooth_scale=smooth_scale) + + def _quant_matmul( # type: ignore[override] + self, + x_q: torch.Tensor, + l0_scale: torch.Tensor, + l1_scale: torch.Tensor, + layer: torch.nn.Module, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + import torch_npu + + if bias is not None and bias.dtype != torch.float32: + bias = bias.to(torch.float32) + # weight_scale is (N, K//64, 2) float8_e8m0fnu — operator expects output-major layout. + # weight_dual_scale is (K//512, N) float32 — transposed to K-major in process_weights. + return torch_npu.npu_dual_level_quant_matmul( + x_q, + layer.weight, + l0_scale, + layer.weight_dual_scale, + l1_scale, + layer.weight_scale, + bias=bias, + output_dtype=ori_dtype, + ) + + +# --------------------------------------------------------------------------- +# NPU MXFP4 dual-scale online method (BF16 checkpoint → quantize at load time) +# --------------------------------------------------------------------------- + + +class NPUMxfp4DualScaleOnlineLinearMethod(_LazyWeightMixin, NPUMxfp4DualScaleLinearMethod): + """NPU W4A4 MXFP4 dual-scale online method: quantises BF16 weights at load time. + + MRO: NPUMxfp4DualScaleOnlineLinearMethod → _LazyWeightMixin + → NPUMxfp4DualScaleLinearMethod → MXFPLinearMethodBase + + create_weights : _LazyWeightMixin (meta device + patched loader) + process_weights : NPUMxfp4DualScaleOnlineLinearMethod (BF16 → FP4 + dual scales) + apply / _quant_matmul : NPUMxfp4DualScaleLinearMethod (shared with offline path) + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import torch_npu + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + # Quantize BF16 weight → FP4 + dual-level scales (no smooth pre-scale for online). + # Returns: (weight_fp4, l0_scale[coarse per-512], l1_scale[fine per-32]) + weight_fp4, w_l0_scale, w_l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant( + layer.weight.data.npu(), smooth_scale=None + ) + + # NZ hardware format for the FP4 weight (same as offline path). + w = torch_npu.npu_format_cast(weight_fp4.view(torch.int8), 29, customize_dtype=torch.int8) + + # Fine scale (l1): (N, K//32) → (N, K//64, 2). Dtype from op output (float8_e8m0fnu). + s = w_l1_scale.reshape(w_l1_scale.shape[0], -1, 2).contiguous() + + # Coarse scale (l0): (N, K_coarse) → (K_coarse, N). Dtype from op output. + ds = w_l0_scale.reshape(w_l0_scale.shape[0], -1).transpose(0, 1).contiguous() + + # No calibration available: identity pre-scale (no smooth quantization effect). + ms = torch.ones(layer.input_size_per_partition, dtype=torch.bfloat16, device="npu") + + replace_parameter(layer, "weight", w) + replace_parameter(layer, "weight_scale", s) + replace_parameter(layer, "weight_dual_scale", ds) + replace_parameter(layer, "mul_scale", ms) + layer._already_called_process_weights_after_loading = True + + +# --------------------------------------------------------------------------- +# Block-index helper (shared by DiffusionMXFP4DualScaleMixedConfig) +# --------------------------------------------------------------------------- + +_BLOCK_IDX_RE = re.compile(r"^blocks\.(\d+)\.") + + +def _parse_block_idx(prefix: str) -> int | None: + """Extract block index from prefix like 'blocks.5.attn1.to_q'.""" + m = _BLOCK_IDX_RE.match(prefix) + return int(m.group(1)) if m else None + + +# --------------------------------------------------------------------------- +# Config: MXFP4 DualScale + per-layer BF16 fallback +# --------------------------------------------------------------------------- + + +class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig): + """W4A4 MXFP4 DualScale with per-layer BF16 fallback for diffusion transformers. + + Sensitive layers fall back to BF16 (original weights) while all other linear + layers use W4A4 MXFP4 DualScale. BF16 fallback layers may be interleaved + anywhere in the transformer. + + Offline mode (is_checkpoint_serialized=True): + Layers whose prefix appears in ignored_layers → UnquantizedLinearMethod (BF16) + All other linear layers → NPUMxfp4DualScaleLinearMethod + + ignored_layers is injected into transformer/config.json by the merge script + and contains the prefixes of all non-MXFP4 linear layers. + + Online mode (is_checkpoint_serialized=False): + Layer routing applies two rules in priority order: + 1. ignored_layers (explicit per-layer BF16 override, user-supplied) → BF16 + 2. Blocks 0 .. num_bf16_fallback_layers-1 (coarse leading-block rule) → BF16 + 3. All other linear layers → NPUMxfp4DualScaleOnlineLinearMethod + + num_bf16_fallback_layers defaults to 5 when not specified. + Set ignored_layers to pin arbitrary interleaved layers to BF16 without + needing an offline checkpoint (useful for accuracy debugging). + Layers outside "blocks.N.*" (condition_embedder etc.) always use online MXFP4 + unless they appear in ignored_layers. + + Config injected by merge_mxfp4_dualscale_checkpoint.py: + { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": true, + "ignored_layers": ["blocks.0.attn1.to_q", ...] + } + """ + + def __init__( + self, + is_checkpoint_serialized: bool = False, + ignored_layers: list[str] | None = None, + num_bf16_fallback_layers: int = 5, + ) -> None: + super().__init__() + self.is_checkpoint_serialized = is_checkpoint_serialized + self.ignored_layers = ignored_layers or [] + self.num_bf16_fallback_layers = num_bf16_fallback_layers + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "mxfp4_dualscale" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + if self.ignored_layers: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> DiffusionMXFP4DualScaleMixedConfig: + is_serialized = cls.get_from_keys_or(config, ["is_checkpoint_serialized"], False) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + num_bf16_fallback_layers = cls.get_from_keys_or(config, ["num_bf16_fallback_layers"], 5) + return cls( + is_checkpoint_serialized=is_serialized, + ignored_layers=ignored_layers, + num_bf16_fallback_layers=num_bf16_fallback_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> QuantizeMethodBase | None: + if not isinstance(layer, LinearBase): + return None + + if self.is_checkpoint_serialized: + # Offline: ignored_layers lists interleaved BF16 fallback layer prefixes. + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if not current_omni_platform.is_npu(): + raise NotImplementedError( + "DiffusionMXFP4DualScaleMixedConfig is currently only supported on NPU (Ascend) platforms." + ) + return NPUMxfp4DualScaleLinearMethod(self) + + # Online: explicit ignored_layers take priority (user-specified per-layer BF16 override), + # then fall back to the coarse leading-block rule. + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + block_idx = _parse_block_idx(prefix) + if block_idx is not None and block_idx < self.num_bf16_fallback_layers: + return UnquantizedLinearMethod() + + if not current_omni_platform.is_npu(): + raise NotImplementedError( + "DiffusionMXFP4DualScaleMixedConfig is currently only supported on NPU (Ascend) platforms." + ) + return NPUMxfp4DualScaleOnlineLinearMethod(self) diff --git a/vllm_omni/quantization/tools/__init__.py b/vllm_omni/quantization/tools/__init__.py index e69de29bb..208f01a7c 100644 --- a/vllm_omni/quantization/tools/__init__.py +++ b/vllm_omni/quantization/tools/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py b/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py new file mode 100644 index 000000000..0786d4607 --- /dev/null +++ b/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Merge W4A4_MXFP4_DUALSCALE quantized Wan2.2 weights into HF Diffusers format. + +msModelSlim produces a checkpoint where each linear layer is either: + - W4A4_MXFP4_DUALSCALE: quantized weights with fine/coarse scales and mul_scale + - FLOAT (BF16 fallback): kept as original BF16 weights for precision-sensitive layers + +BF16 fallback layers may be interleaved anywhere in the transformer (not just leading +blocks). The merge script detects them from quant_model_description.json and writes +their prefixes into config.json as ignored_layers so the runtime dynamically routes +each layer to the correct quantization method at weight-loading time. + +MXFP4_DUALSCALE key structure per linear layer (msModelSlim naming) +--------------------------------------------------------------------- + blocks.N.X.linear.weight W4A4_MXFP4_DUALSCALE – int8 (FP4 packed) + blocks.N.X.linear.weight_scale W4A4_MXFP4_DUALSCALE – uint8 (float8_e8m0fnu fine scale, per-32K) + blocks.N.X.linear.weight_dual_scale W4A4_MXFP4_DUALSCALE – float32 (coarse scale, per-512K) + blocks.N.X.linear.bias FLOAT – bias (if present) + blocks.N.X.div.mul_scale FLOAT – float32 per-input-channel activation pre-scale + +BF16 fallback key structure (no wrappers, plain linear weights): + blocks.N.X.weight FLOAT + condition_embedder.*.weight FLOAT (always BF16 — not quantized) + +Self-attention QKV notes +------------------------ +Self-attention Q/K/V weights are separate in the checkpoint (self_attn.q/k/v) but fused +into a single to_qkv layer in vllm-omni. The transformer's load_weights() handles this via +stacked_params_mapping. This script keeps Q/K/V keys separate — do NOT pre-fuse them here. + +For mul_scale specifically: even though Q/K/V process the same input (same mul_scale value), +they are kept separate. load_weights() routes all three to the same to_qkv.mul_scale parameter +and each overwrites the previous. Since Q=K=V for mul_scale, the final value is correct. + +NOTE: Pre-fusing as to_qkv.mul_scale would BREAK loading because ".attn1.to_q" is a +substring of ".attn1.to_qkv", causing load_weights() stacked_params_mapping to produce a +garbage key ("to_qkvkv") that is not in params_dict, triggering a break that skips the +direct-load else branch entirely. + +Supported model types: + - Wan2.2-T2V-A14B (MoE cascade: transformer + transformer_2) + - Wan2.2-I2V-A14B (MoE cascade: transformer + transformer_2) + +Note: Wan2.2-TI2V-5B is NOT supported. Its smaller parameter count causes +unacceptable accuracy loss under W4A4 quantization. Use MXFP8 for TI2V-5B. + +Usage: + python merge_mxfp4_dualscale_checkpoint.py \\ + --model-type Wan2.2-T2V-A14B \\ + --original-model /path/to/Wan2.2-T2V-A14B-Diffusers \\ + --quant-path /path/to/msmodelslim-output \\ + --output-path /path/to/merged-output +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import re +import shutil +from typing import Any + +import torch +from safetensors.torch import load_file, save_file + +# --------------------------------------------------------------------------- +# Key rename: msModelSlim naming → Diffusers / vllm-omni naming +# Identical to merge_mxfp8_checkpoint.py; applied to all blocks uniformly. +# --------------------------------------------------------------------------- + +TRANSFORMER_KEYS_RENAME_DICT: dict[str, str] = { + "time_embedding.0": "condition_embedder.time_embedder.linear_1", + "time_embedding.2": "condition_embedder.time_embedder.linear_2", + "text_embedding.0": "condition_embedder.text_embedder.linear_1", + "text_embedding.2": "condition_embedder.text_embedder.linear_2", + "time_projection.1": "condition_embedder.time_proj", + "head.modulation": "scale_shift_table", + "head.head": "proj_out", + "modulation": "scale_shift_table", + "ffn.0": "ffn.net.0.proj", + "ffn.2": "ffn.net.2", + # Norm order swap (quant tool: norm1, norm3, norm2 → diffusers: norm1, norm2, norm3) + "norm2": "norm__placeholder", + "norm3": "norm2", + "norm__placeholder": "norm3", + # Self-attention + "self_attn.q": "attn1.to_q", + "self_attn.k": "attn1.to_k", + "self_attn.v": "attn1.to_v", + "self_attn.o": "attn1.to_out.0", + "self_attn.norm_q": "attn1.norm_q", + "self_attn.norm_k": "attn1.norm_k", + # Cross-attention + "cross_attn.q": "attn2.to_q", + "cross_attn.k": "attn2.to_k", + "cross_attn.v": "attn2.to_v", + "cross_attn.o": "attn2.to_out.0", + "cross_attn.norm_q": "attn2.norm_q", + "cross_attn.norm_k": "attn2.norm_k", + "attn2.to_k_img": "attn2.add_k_proj", + "attn2.to_v_img": "attn2.add_v_proj", + "attn2.norm_k_img": "attn2.norm_added_k", + # I2V image embedder + "img_emb.proj.0": "condition_embedder.image_embedder.norm1", + "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", + "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", + "img_emb.proj.4": "condition_embedder.image_embedder.norm2", + "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed", +} + +SUPPORTED_MODEL_TYPES = ["Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"] +CASCADE_MODEL_TYPES = {"Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"} + +# Suffixes that appear inside .linear.* wrapper for MXFP4 tensors. +_LINEAR_ATTRS = ("weight_dual_scale", "weight_scale", "weight", "bias") + +_BLOCK_IDX_RE = re.compile(r"^blocks\.(\d+)\.") + + +# --------------------------------------------------------------------------- +# Key transformation helpers +# --------------------------------------------------------------------------- + + +def _parse_block_idx(key: str) -> int | None: + """Extract block index from a key like 'blocks.5.attn1.to_q.weight'.""" + m = _BLOCK_IDX_RE.match(key) + return int(m.group(1)) if m else None + + +def _apply_rename_dict(key: str) -> str: + for src, dst in TRANSFORMER_KEYS_RENAME_DICT.items(): + key = key.replace(src, dst) + return key + + +def _strip_mxfp4_wrapper(key: str) -> str: + """Strip .linear.ATTR or .div.mul_scale wrappers added by msModelSlim. + + MXFP4 tensors are wrapped in sub-modules: + X.linear.weight / weight_scale / weight_dual_scale / bias + X.div.mul_scale + + MXFP8 and FLOAT tensors have no wrappers — this function is a no-op for them. + + Examples after apply_rename_dict: + attn1.to_q.linear.weight → attn1.to_q.weight + attn1.to_q.linear.weight_scale → attn1.to_q.weight_scale + attn1.to_q.div.mul_scale → attn1.to_q.mul_scale + attn1.norm_q.weight → attn1.norm_q.weight (unchanged) + """ + # Check longest attribute names first to avoid partial suffix matches. + for attr in _LINEAR_ATTRS: + suffix = f".linear.{attr}" + if key.endswith(suffix): + return key[: -len(suffix)] + f".{attr}" + if key.endswith(".div.mul_scale"): + return key[: -len(".div.mul_scale")] + ".mul_scale" + return key + + +# --------------------------------------------------------------------------- +# Quantization metadata helpers +# --------------------------------------------------------------------------- + +# Known quantized weight types (FLOAT tensors don't determine block type). +_MXFP8_TYPE = "mxfp8" +_MXFP4_DUALSCALE_TYPE = "mxfp4_dualscale" + + +def _classify_blocks(quant_meta: dict[str, str]) -> dict[int, str]: + """Classify each transformer block by quantization type from quant_meta. + + Returns a dict mapping block_idx → 'mxfp8' | 'mxfp4_dualscale'. + A block's type is determined by the first quantized (non-FLOAT) tensor found for it. + """ + block_types: dict[int, str] = {} + for key, qtype in quant_meta.items(): + idx = _parse_block_idx(key) + if idx is None or idx in block_types: + continue + if qtype.startswith("W8A8_MXFP8"): + block_types[idx] = _MXFP8_TYPE + elif qtype.startswith("W4A4_MXFP4_DUALSCALE"): + block_types[idx] = _MXFP4_DUALSCALE_TYPE + return block_types + + +def _print_block_summary(block_types: dict[int, str]) -> None: + """Print a compact run-length summary of the block layout.""" + if not block_types: + print(" Block layout: (empty)") + return + + sorted_indices = sorted(block_types) + runs: list[tuple[int, int, str]] = [] + run_start = sorted_indices[0] + run_type = block_types[run_start] + for idx in sorted_indices[1:]: + if block_types[idx] != run_type: + runs.append((run_start, idx - 1, run_type)) + run_start = idx + run_type = block_types[idx] + runs.append((run_start, sorted_indices[-1], run_type)) + + print(f" Block layout ({len(sorted_indices)} blocks classified):") + for start, end, btype in runs: + count = end - start + 1 + range_str = f"{start}" if start == end else f"{start}–{end}" + print(f" blocks {range_str:>8}: {btype} ({count} block{'s' if count > 1 else ''})") + + +# --------------------------------------------------------------------------- +# Safetensors I/O +# --------------------------------------------------------------------------- + + +def _load_safetensors_dir(directory: pathlib.Path, glob: str = "*.safetensors") -> dict[str, torch.Tensor]: + candidates = sorted(directory.glob(glob)) + if not candidates: + raise FileNotFoundError(f"No safetensors matching '{glob}' found in {directory}") + state: dict[str, torch.Tensor] = {} + for f in candidates: + state.update(load_file(str(f))) + return state + + +def _load_quant_safetensors(directory: pathlib.Path) -> dict[str, torch.Tensor]: + try: + return _load_safetensors_dir(directory, "quant_model_weight*.safetensors") + except FileNotFoundError: + return _load_safetensors_dir(directory) + + +def _load_quant_meta(directory: pathlib.Path) -> dict[str, str]: + candidates = sorted(directory.glob("quant_model_description*.json")) + if not candidates: + print(f" WARNING: No quant_model_description*.json in {directory}; treating all tensors as FLOAT.") + return {} + with open(candidates[0]) as f: + return json.load(f) + + +# --------------------------------------------------------------------------- +# Per-transformer conversion +# --------------------------------------------------------------------------- + + +def _is_mxfp4_tensor(key: str, qtype: str) -> bool: + """Return True for MXFP4 quantized tensors and their companion tensors. + + MXFP4 DualScale layers produce four tensor types: + .linear.weight / .linear.weight_scale / .linear.weight_dual_scale (W4A4_MXFP4_DUALSCALE) + .linear.bias / .div.mul_scale (FLOAT but companion to a quantized layer) + + BF16 fallback layers have plain .weight / .bias without .linear. / .div. wrappers. + """ + if qtype.startswith("W4A4_MXFP4_DUALSCALE"): + return True + # Companion tensors: bias and mul_scale that belong to an MXFP4 layer. + return ".linear." in key or ".div.mul_scale" in key + + +def _diffusers_to_vllm_ignored(diffusers_ignored: list[str]) -> list[str]: + """Translate diffusers checkpoint prefixes to vllm-omni model parameter names. + + Three transformations align the naming conventions: + + 1. Self-attention QKV fusion (attn1 only, not attn2): + attn1.to_q + attn1.to_k + attn1.to_v → attn1.to_qkv + Cross-attention (attn2) to_q/k/v are not fused and remain separate. + + 2. FFN naming: + ffn.net.0.proj → ffn.net_0.proj + ffn.net.2 → ffn.net_2 + + 3. to_out index: + attn1.to_out.0 / attn2.to_out.0 → attn1.to_out / attn2.to_out + """ + ignored_set = set(diffusers_ignored) + result: set[str] = set() + + for name in diffusers_ignored: + m = re.match(r"^(.*\.attn1)\.to_([qkv])$", name) + if m: + prefix = m.group(1) + if all(f"{prefix}.to_{c}" in ignored_set for c in ("q", "k", "v")): + result.add(f"{prefix}.to_qkv") + else: + present = [f"to_{c}" for c in ("q", "k", "v") if f"{prefix}.to_{c}" in ignored_set] + missing = [f"to_{c}" for c in ("q", "k", "v") if f"{prefix}.to_{c}" not in ignored_set] + raise ValueError( + f"Partial BF16 fallback for '{prefix}': " + f"{', '.join(present)} in ignored_layers but {', '.join(missing)} is not. " + f"Self-attention Q/K/V are fused into a single to_qkv layer at runtime; " + f"all three must share the same precision. " + f"Either quantize all of to_q/to_k/to_v or keep all three in BF16." + ) + continue + + name = re.sub(r"\.ffn\.net\.0\.proj$", ".ffn.net_0.proj", name) + name = re.sub(r"\.ffn\.net\.2$", ".ffn.net_2", name) + name = re.sub(r"\.to_out\.0$", ".to_out", name) + result.add(name) + + return sorted(result) + + +def _collect_ignored_layers(merged: dict[str, Any], mxfp4_layer_prefixes: set[str]) -> list[str]: + """Collect vllm-omni parameter name prefixes for BF16 fallback layers. + + A layer is BF16 if it has a .weight tensor but its prefix is not in + mxfp4_layer_prefixes. Prefixes are returned in vllm-omni parameter naming + (QKV-fused, FFN underscored, to_out unindexed) so the list can be written + directly into ignored_layers in config.json without further translation. + """ + all_weight_prefixes = {key[: -len(".weight")] for key in merged if key.endswith(".weight")} + return _diffusers_to_vllm_ignored(sorted(all_weight_prefixes - mxfp4_layer_prefixes)) + + +def _convert_transformer( + quant_subdir: pathlib.Path, + output_dir: pathlib.Path, + original_transformer_dir: pathlib.Path, +) -> None: + """Convert one transformer directory to the mxfp4_dualscale + BF16 mixed format.""" + output_dir.mkdir(parents=True, exist_ok=True) + + # BF16 base: provides the scaffold for non-MXFP4 tensors (norms, embeddings, + # BF16 fallback linear layers that msModelSlim may omit or keep as FLOAT). + print(f" Loading BF16 base from {original_transformer_dir} …") + base_state = _load_safetensors_dir(original_transformer_dir) + print(f" {len(base_state)} BF16 tensors loaded") + + print(f" Loading quantized weights from {quant_subdir} …") + quant_state = _load_quant_safetensors(quant_subdir) + quant_meta = _load_quant_meta(quant_subdir) + print(f" {len(quant_state)} quant tensors, {len(quant_meta)} meta entries") + + # Classify blocks for a compact summary (informational only in the new scheme). + block_types = _classify_blocks(quant_meta) + _print_block_summary(block_types) + + # Remap MXFP4 quantized tensors only. + # BF16 fallback tensors (FLOAT in quant_meta, no .linear./.div. wrapper) are + # skipped here; the base_state provides them unchanged. + remapped: dict[str, torch.Tensor] = {} + remapped_meta: dict[str, str] = {} + mxfp4_layer_prefixes: set[str] = set() + skipped: list[str] = [] + + for key, tensor in quant_state.items(): + qtype = quant_meta.get(key, "FLOAT") + + if not _is_mxfp4_tensor(key, qtype): + continue # BF16 fallback — base_state already covers this tensor + + renamed = _apply_rename_dict(key) + final_key = _strip_mxfp4_wrapper(renamed) + + if final_key.endswith(".quant_type"): + skipped.append(key) + continue + + remapped[final_key] = tensor + remapped_meta[final_key] = qtype + + # Track MXFP4 layer prefixes (via their .weight keys) for ignored_layers. + if final_key.endswith(".weight") and qtype.startswith("W4A4_MXFP4_DUALSCALE"): + mxfp4_layer_prefixes.add(final_key[: -len(".weight")]) + + if skipped: + print(f" Skipped {len(skipped)} metadata keys (quant_type markers): {skipped[:5]}") + print(f" {len(remapped)} MXFP4 tensors remapped, {len(mxfp4_layer_prefixes)} quantized layers") + + # Merge: base_state provides BF16 scaffold; MXFP4 tensors override their BF16 counterparts + # and add the new scale tensors (weight_scale, weight_dual_scale, mul_scale). + merged = {**base_state, **remapped} + + # Determine ignored_layers: BF16 fallback layer prefixes in vllm-omni parameter naming. + ignored_layers = _collect_ignored_layers(merged, mxfp4_layer_prefixes) + print(f" {len(ignored_layers)} BF16 fallback layers → ignored_layers in config.json") + + # Save weights. + out_weights = output_dir / "diffusion_pytorch_model.safetensors" + save_file(merged, str(out_weights)) + print(f" Saved {len(merged)} tensors → {out_weights}") + + # Save remapped quant metadata (for inspection / debugging). + out_meta_path = output_dir / "quant_model_description.json" + with open(out_meta_path, "w") as f: + json.dump(remapped_meta, f, indent=2) + + # Inject quantization_config into config.json. + src_config = original_transformer_dir / "config.json" + if src_config.is_file(): + with open(src_config) as f: + config = json.load(f) + config["quantization_config"] = _build_quant_config(ignored_layers) + out_config = output_dir / "config.json" + with open(out_config, "w") as f: + json.dump(config, f, indent=2) + print(f" Injected quantization_config (mxfp4_dualscale, {len(ignored_layers)} ignored_layers) → {out_config}") + else: + print(f" WARNING: No config.json at {src_config}; quantization_config not injected.") + + +def _build_quant_config(ignored_layers: list[str]) -> dict[str, Any]: + return { + "quant_method": "mxfp4_dualscale", + "is_checkpoint_serialized": True, + "ignored_layers": ignored_layers, + } + + +# --------------------------------------------------------------------------- +# Model-type helpers +# --------------------------------------------------------------------------- + + +def _get_transformer_dirs(model_type: str) -> list[str]: + return ["transformer", "transformer_2"] if model_type in CASCADE_MODEL_TYPES else ["transformer"] + + +def _get_quant_subdir(model_type: str, quant_path: pathlib.Path, transformer_dir: str) -> pathlib.Path: + if model_type in CASCADE_MODEL_TYPES: + sub = "high_noise_model" if transformer_dir == "transformer" else "low_noise_model" + return quant_path / sub + return quant_path + + +# --------------------------------------------------------------------------- +# Main repack +# --------------------------------------------------------------------------- + + +def repack( + model_type: str, + original_model_path: pathlib.Path, + quant_path: pathlib.Path, + output_path: pathlib.Path, +) -> None: + transformer_dirs = _get_transformer_dirs(model_type) + + print(f"Copying original model to {output_path} (skipping {transformer_dirs}) …") + shutil.copytree( + str(original_model_path), + str(output_path), + ignore=shutil.ignore_patterns(*transformer_dirs), + ) + + for tdir in transformer_dirs: + q_subdir = _get_quant_subdir(model_type, quant_path, tdir) + out_tdir = output_path / tdir + orig_tdir = original_model_path / tdir + print(f"\nConverting {tdir} (quant source: {q_subdir.name}) …") + _convert_transformer(q_subdir, out_tdir, orig_tdir) + + print(f"\nDone. Merged model → {output_path}") + print("\nRun inference (quantization auto-detected from config.json):") + print(f" python text_to_video.py --model {output_path}") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--model-type", required=True, choices=SUPPORTED_MODEL_TYPES, help="Model variant.") + parser.add_argument("--original-model", required=True, help="Original HF Diffusers model directory (BF16).") + parser.add_argument("--quant-path", required=True, help="msModelSlim quantized weights directory.") + parser.add_argument("--output-path", required=True, help="Output directory for merged model.") + args = parser.parse_args() + + repack( + model_type=args.model_type, + original_model_path=pathlib.Path(args.original_model), + quant_path=pathlib.Path(args.quant_path), + output_path=pathlib.Path(args.output_path), + ) + + +if __name__ == "__main__": + main()