From becfefd4402da180d8af5121cef474fee40883af Mon Sep 17 00:00:00 2001 From: HAAZZZEEEE <49724853+HAAZZZEEEE@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:11:40 +0800 Subject: [PATCH] [Hardware][Ascend] Add Wan2.2 MXFP4 UOS, Smooth and W4A8 fallback (#7210) Signed-off-by: HAAZZZEEEE --- docs/user_guide/quantization/mxfp4.md | 414 +++++++++---- .../text_to_video/text_to_video.py | 11 +- tests/config/test_omni_config.py | 30 +- .../wan2_2/test_create_transformer_quant.py | 213 ++++++- .../wan2_2/test_wan22_pipeline_diffuse.py | 98 +++- .../quantization/test_mxfp4_checkpoint.py | 294 ++++++++++ .../quantization/test_mxfp4_config.py | 555 +++++++++++++++++- .../quantization/test_wan_mxfp4_loading.py | 237 ++++++++ .../test_diffusion_config_propagation.py | 6 + .../test_text_to_video_prompts.py | 49 +- .../npu/quant/test_mxfp4_step_fallback.py | 168 ++++++ vllm_omni/config/omni_config.py | 4 + vllm_omni/diffusion/data.py | 4 + .../model_loader/diffusers_loader.py | 37 +- .../models/wan2_2/pipeline_wan2_2.py | 51 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 26 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 11 +- .../models/wan2_2/wan2_2_transformer.py | 106 +++- vllm_omni/quantization/factory.py | 77 ++- vllm_omni/quantization/mxfp4_config.py | 228 ++++++- 20 files changed, 2435 insertions(+), 184 deletions(-) create mode 100644 tests/diffusion/quantization/test_mxfp4_checkpoint.py create mode 100644 tests/diffusion/quantization/test_wan_mxfp4_loading.py create mode 100644 tests/platforms/npu/quant/test_mxfp4_step_fallback.py diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md index b18b3e7b3..00ff97c01 100644 --- a/docs/user_guide/quantization/mxfp4.md +++ b/docs/user_guide/quantization/mxfp4.md @@ -6,22 +6,25 @@ 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. +On Ascend, single-scale MXFP4 also supports Universal Optimal Scaling +(UOS, `Qmax=7.25`), described in the single-scale scaling section below. + 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 | +| -------- | ---------------- | ------ | ---------- | +| `mxfp4` | Single-scale (per-32 fine only) | Online + pre-quantized single-scale weights | Online baseline, or an existing single-scale checkpoint in the layout below | +| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | W4A4 with dual-level scales; offline mode loads calibrated Smooth tensors | -!!! tip "Recommended: `mxfp4_dualscale` offline" - For production deployments, use the `mxfp4_dualscale` offline mode with a +!!! tip "Existing W4A4-only DualScale deployments" + For W4A4-only deployments, the existing `mxfp4_dualscale` offline mode uses 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. + calibrated `mul_scale` tensors from disk. Quality depends on the model, + calibration recipe and BF16 fallback policy; evaluate the checkpoint against + your quality target before deployment. - Use `mxfp4` online only for quick experimentation where preprocessing time - is not acceptable and accuracy loss is tolerable. + W4A8 layer/step fallback requires single-scale `mxfp4`, online or offline. + DualScale supports W4A4 and BF16 routing only. !!! warning "Online single-scale ≠ Offline dual-scale" `mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`: @@ -34,61 +37,244 @@ vLLM-Omni provides two quantization methods with different scale structures: ## Hardware Support -| Device | Support | -|--------|---------| -| NVIDIA Blackwell GPU (SM 100+) | ⭕ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | -| NVIDIA Ampere GPU (SM 80+) | ⭕ | -| AMD ROCm (gfx950 / MI355X) | ✅ | -| Intel XPU | ⭕ | -| Ascend NPU (Atlas 950 A5) | ✅ | +The following matrix describes these vLLM-Omni implementations, rather than +hardware support for the MXFP4 format in general. -Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. +| Device | Single-scale online OCP | Single-scale offline | UOS / W4A8 fallback | DualScale | +| -------- | -------- | -------- | -------- | -------- | +| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ✅ | ✅ | +| AMD ROCm (gfx950 / MI355X) | ✅ | ❌ | ❌ | ❌ | +| Other AMD architectures | ❌ | ❌ | ❌ | ❌ | +| NVIDIA GPU | ❌ | ❌ | ❌ | ❌ | +| Intel XPU | ❌ | ❌ | ❌ | ❌ | + +Legend: `✅` supported by the implementation, `❌` unsupported. +ROCm uses the existing AITER online path and requires `mxfp4_scale_alg=0`. ## 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 | +| ------- | -------- | --------- | ------- | +| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Cascade experts resolve checkpoint metadata independently; explicit global runtime policy applies to both | +| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4` / `mxfp4_dualscale` | Per-expert configuration uses the same interface as T2V; validate I2V checkpoints separately | +| Wan2.2-TI2V-5B | Requires validation | Requires validation | This guide does not establish checkpoint compatibility or generation quality for TI2V-5B | -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: +Both methods can be integrated with single-transformer and cascade models. +This alone does not establish support or quality for every model variant: -- `mxfp4`: single-scale, lower overhead, simpler compute, online only -- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline +- `mxfp4`: single-scale; online and pre-quantized weights, with optional W4A8 fallback on NPU +- `mxfp4_dualscale`: dual-scale; online and offline, with BF16 fallback -**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced -by the merge script); there is no offline `mxfp4` single-scale format. +The existing `merge_mxfp4_dualscale_checkpoint.py` produces **offline +`mxfp4_dualscale`** checkpoints. Single-scale weights must satisfy the separate +layout described below; changing a DualScale checkpoint's method name to +`mxfp4` does not convert its weights. Prepare single-scale checkpoints with +external tooling. Real exported checkpoint compatibility and model quality +require validation of that checkpoint; the synthetic tests do not establish +UOS export or video acceptance. !!! note "Per-layer BF16 fallback in offline cascade models" - The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most + An A14B offline DualScale 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. + and resolves the storage format and BF16 layer policy independently. ## Configuration +### W4A8 fallback at selected layers and denoising steps (Ascend NPU) + +For single-scale `mxfp4`, set `w4a8_fallback_steps` to a list of zero-based +request denoising loop indices (not scheduler timestep values). +Set `w4a8_fallback_layers` to exact runtime Linear paths, for example +`["blocks.10.attn1.to_qkv", "blocks.12.ffn.net_2"]`. A selected layer uses W4A8 +at every step; at a selected step, all other quantized layers also use W4A8. +The rules are combined with **OR**. W4A8 uses MXFP8 activations with MXFP4 +weights. Omitting both lists or passing `[]` keeps the existing W4A4 behavior. Existing +`ignored_layers` take priority: **BF16 > (layer OR step) W4A8 > W4A4**. +DualScale's existing leading BF16 block rule is unchanged; it does not enable +W4A8. + +```python +from vllm_omni import Omni + +omni = Omni( + model="", + quantization_config={ + "method": "mxfp4", + "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"], + "w4a8_fallback_steps": [0, 1, 38, 39], + }, +) +``` + +Both policies also work with offline single-scale `mxfp4`. Any nonempty +W4A8 list with `mxfp4_dualscale`, online or offline, raises an explicit error +at configuration construction. Empty lists preserve its existing W4A4 path. +Each expert's checkpoint still determines its serialized +format and BF16 layer routing. A serialized checkpoint config is the complete +storage contract: `ignored_layers`, or its `modules_to_not_convert` alias, +selects the BF16 layers, and omitting both means no explicit BF16 layers. A +method-only checkpoint value such as `"mxfp4"` or +`{"quant_method": "mxfp4"}` carries no storage policy and leaves a caller's +online `ignored_layers` unchanged. An explicitly supplied quantization config +retains its W4A8 layer and step lists when the loader rebuilds either expert's +config from disk, including empty lists that disable checkpoint policies. +Within an explicit runtime config, an omitted W4A8 list defaults to `[]`; +individual omitted fields do not inherit saved policies. Without an explicit +runtime config, auto-detection uses the saved checkpoint policies. + +Layer names are relative to each expert (`blocks.10...`, without a +`transformer.` or `transformer_2.` prefix). A global config applies the same +list to both experts; per-component `transformer`/`transformer_2` configs can +select different lists. Wan selects these expert names exactly: a `transformer` +entry does not match `transformer_2`; an omitted expert uses `default` if set. +A component mapped to `null` (or an unmatched component +without a default) stays unquantized and requires a BF16 checkpoint; loading a +quantized checkpoint for that disabled component raises an error. A top-level +`quantization_config=None` retains checkpoint auto-detection. +Names are exact, with no regex, wildcard, block expansion, +or checkpoint-name remapping. Wan validates them against its runtime Linear +modules and logs the selected paths/count and BF16 overrides per PP rank. +Use `attn1.to_qkv` for fused self-attention; `attn1.to_q`, `to_k`, or `to_v` +checkpoint names raise an error. Cross-attention `attn2.to_q/to_k/to_v` remains +separate. Layer-only fallback does not require a denoise ForwardContext. + +The text-to-video example accepts the JSON directly: + +```bash +ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 \ +python examples/offline_inference/text_to_video/text_to_video.py \ + --model /path/to/Wan2.2-T2V-A14B-Diffusers \ + --quantization-config '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}' \ + --tensor-parallel-size 1 --ulysses-degree 2 --cfg-parallel-size 2 \ + --enforce-eager --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 40 --seed 42 --output wan22-w4a8.mp4 +``` + +Use either `--quantization` or `--quantization-config` in this example. +Serving accepts the same JSON through `--diffusion-quantization-config`. + +The list uses one request's loop index across both high- and low-noise experts, +and both CFG branches see the same index. A new request starts at step zero. +Indices outside the request's step range are never selected. A pipeline must +publish `ForwardContext.denoise_step_idx`; requesting step fallback without +that context raises an error instead of silently ignoring the list. + +Single-scale W4A4 and W4A8 reuse the same prepared `weight` and `weight_scale`. +Online loading quantizes the original BF16/FP16 weight once; offline loading +packs the numeric FP4 values and reshapes the saved E8M0 scale without +requantizing the weight. Only activation precision changes during forward. +There is no DualScale decoding, second W4 representation or `w4a8_weight` / +`w4a8_weight_scale` cache. No weight quantization runs inside the denoising loop. + +The mixed FP8/FP4 GEMM uses BF16 output and a matching BF16 row bias for +compatibility with the target A5 runtime. An FP16 caller receives that result +converted back to FP16; its output shape and dtype are preserved. + +The single-scale fallback validation target is eager Wan2.2 T2V A14B on A5. +Quality scores, optimal fallback lists, real UOS exports and additional +compile/offload combinations require separate evaluation. Existing W4A4-only +DualScale still needs complete 512-channel input groups; use TP=1 with +sequence/CFG parallelism for its initial multi-card run. + +### `mxfp4` — Pre-Quantized Single-Scale Weights (Ascend NPU) + +Set the following fields in each transformer's `config.json` when a +single-scale checkpoint uses calibrated Smooth scaling: + +```json +{ + "quantization_config": { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": true, + "require_smooth_scale": true, + "ignored_layers": [] + } +} +``` + +Each quantized Linear must provide tensors under its model parameter name: + +| Tensor | Checkpoint layout | Meaning | +| -------- | ------------------- | --------- | +| `weight` | `(N, K)` numeric FP4 values | Packed to FP4 after loading | +| `weight_scale` | `(N, ceil(K / 32))`, `uint8` | Raw E8M0 exponent bytes, required | +| `mul_scale` | `(K,)`, FP32 | Smooth pre-scale, required when declared above | + +Both W4A4 and W4A8 quantize `x * mul_scale`. The Smooth tensor is converted +to the layer's activation dtype after loading. Row-parallel layers load the +corresponding input-channel shard. Pre-fused QKV weights must use one shared +Smooth tensor for the fused projection. + +For backward compatibility, `require_smooth_scale` defaults to `false`: +an absent tensor uses identity scaling, while a present tensor is still loaded. +Without a declaration, an absent Smooth tensor cannot be distinguished from a +legacy unsmoothed checkpoint. Calibrated checkpoints should declare +`require_smooth_scale: true` so missing or misnamed Smooth tensors fail loading; +incorrect tensor shapes, non-floating tensors, NaN/Inf and nonpositive values +also fail. Values that overflow or underflow to zero in the activation dtype +are rejected during weight processing. The flag requires serialized single-scale +weights and is rejected for online quantization. It does not make DualScale +scales optional or change missing-parameter policies for other backends. +Either the checkpoint or an explicit caller can require Smooth: a `true` +declaration is retained when per-expert metadata rebuilds the configuration, +including when the other source omits the flag or sets it to `false`. + +Prepare the canonical numeric FP4 checkpoint before inference using external +tooling. Runtime loading does not parse producer-specific export formats. +Retain the producer recipe to identify the weight-generation algorithm and +Smooth calibration; tensor storage alone does not establish UOS provenance. + +### `mxfp4` — Single-Scale Block Scaling (OCP / UOS) + +`mxfp4_scale_alg` selects the MXFP4 block-scaling method on NPU. +UOS is a scaling rule, not a separate activation dtype or an activation-only +algorithm. Both choices use FP4 E2M1 values with shared E8M0 scales. + +| Mode | Weight preparation | A4 activation quantization | W4A8 fallback | +| --- | --- | --- | --- | +| Online | Selected scaling method, once at loading | Selected scaling method | Same W4; OCP MXFP8 A8 | +| Offline | Load existing W4 and scales unchanged | Selected scaling method | Same W4; OCP MXFP8 A8 | + +The available settings are: + +| Setting | MXFP4 scaling method | NPU parameters | +| --- | --- | --- | +| `0` (default) | OCP MX | `scale_alg=0` | +| `2` | UOS (`Qmax=7.25`) | `scale_alg=2, dst_type_max=7.25` | + +Universal Optimal Scaling (UOS, `Qmax=7.25`) is derived in +[MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attention](https://arxiv.org/abs/2607.24377), +Section 4.1. The value `7.25` is the boundary used to select the shared +power-of-two scale; the maximum finite E2M1 value remains `6`. +This option applies UOS to Linear quantization. It does not implement the +paper's full MXAttention pipeline or Pre-Normalization Quantization (PNQ). + +Both use FP4 E2M1 with E8M0 scales, `axis=-1`, `block_size=32`, +`round_mode="rint"`. W4A8 always quantizes the activation to +`float8_e4m3fn` with `scale_alg=0` and the same grouping/rounding, without +UOS parameters. A4 and A8 share the prepared W4 and scale. UOS is currently +NPU-only; ROCm retains its existing AITER online algorithm. + +An explicit UOS online configuration is +`{"method": "mxfp4", "mxfp4_scale_alg": 2}`. For offline checkpoints the +setting affects **only A4 activations**. Existing weights are never requantized: +a UOS checkpoint can run UOS W4 + UOS A4 or UOS W4 + OCP A4, and both fall back to +UOS W4 + OCP A8. Record the weight-generation algorithm separately from the +runtime activation setting. Smooth remains active with either setting. The +UOS configuration explicitly selects `2`; attention computation is unchanged. + ### `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. +Online mode quantizes BF16/FP16 weights once at loading and computes activation +scales during forward. Each block of 32 K elements shares one E8M0 scale; +no calibrated `mul_scale` is available. An explicit global quantization config +applies to both cascade experts. Per-component configs can select different +policies; checkpoint auto-detection resolves each expert independently. ```python from vllm_omni import Omni @@ -103,7 +289,7 @@ outputs = omni.generate( ```bash # Single-transformer or cascade model — same command -python text_to_video.py --model --quantization mxfp4 +python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4 # Online serving vllm serve --omni --quantization mxfp4 @@ -113,19 +299,19 @@ vllm serve --omni --quantization mxfp4 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. +single-transformer and cascade (A14B) models with compatible Linear layers. +Compare quality and performance using the intended model and workload. 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. +(`num_bf16_fallback_layers=5`). This is an implementation default; select and +validate the BF16 policy against your own quality target. ```python omni = Omni(model="", quantization="mxfp4_dualscale") ``` ```bash -python text_to_video.py --model --quantization mxfp4_dualscale +python examples/offline_inference/text_to_video/text_to_video.py --model --quantization mxfp4_dualscale ``` If accuracy debugging identifies additional precision-sensitive layers, they @@ -136,21 +322,22 @@ omni = Omni( model="", quantization_config={ "method": "mxfp4_dualscale", - "ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override + "ignored_layers": ["blocks.10.attn1.to_qkv"], # 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. +1. **`ignored_layers`** (explicit per-layer override): a listed runtime Linear + module is kept in BF16 regardless of block index. Use exact module paths, + not arbitrary block prefixes or names ending in `.weight`. 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. + to `5`. Layers outside `blocks.N.*` are quantized unless explicitly listed + in `ignored_layers`. -### `mxfp4_dualscale` — DualScale Offline Mode (Recommended) +### `mxfp4_dualscale` — DualScale Offline Mode Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A preprocessing step converts the raw quantized output to the diffusers format @@ -169,8 +356,8 @@ inside `config.json`. At runtime, each layer's prefix is matched against Each quantized linear layer stores four tensors: | Tensor | Shape | dtype | Description | -|--------|-------|-------|-------------| -| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) | +| -------- | ------- | ------- | ------------- | +| `weight` | `(N, K)` | float8_e4m3fn | Numeric FP4 values; packed into two values per byte after loading | | `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) | @@ -223,7 +410,7 @@ python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \ ``` | 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 | @@ -251,14 +438,14 @@ The `quantization_config` injected into `config.json` for each transformer: ``` `ignored_layers` lists every linear layer that retains its original BF16 weight, -using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out` +using vllm-omni runtime module paths (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 +python examples/offline_inference/text_to_video/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 @@ -275,25 +462,32 @@ omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale") ## Parameters -### `mxfp4` (single-scale, online only) +### `mxfp4` (single-scale, online + offline) | Parameter | Type | Default | Description | -|-----------|------|---------|-------------| +| ----------- | ------ | --------- | ------------- | | `method` | str | — | `"mxfp4"` | -| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 | +| `mxfp4_scale_alg` | int | `0` | MXFP4 block scaling: `0` OCP MX, `2` UOS (`Qmax=7.25`, NPU only). Online: W4 and A4; offline: A4 only. A8 remains OCP MXFP8 | +| `is_checkpoint_mxfp4_serialized` | bool | `false` | Load pre-quantized single-scale weights and scales; declared in checkpoint `config.json` | +| `require_smooth_scale` | bool | `false` | Require calibrated Smooth tensors for offline single-scale weights | +| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths to keep in BF16 | +| `w4a8_fallback_layers` | list[str] | `[]` | Exact runtime Linear paths that always use A8; NPU only | +| `w4a8_fallback_steps` | list[int] | `[]` | Zero-based denoise steps that use A8; NPU only | ### `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 | +| `ignored_layers` | list[str] | `[]` | Exact runtime Linear module paths 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. Validate the default of `5` against the intended model and quality target | +| `w4a8_fallback_layers` | list[str] | `[]` | Must remain empty; DualScale does not support W4A8 fallback | +| `w4a8_fallback_steps` | list[int] | `[]` | Must remain empty; DualScale does not support W4A8 fallback | #### BF16 fallback priority (online mode) -``` +```text 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) @@ -307,8 +501,9 @@ unless they appear in `ignored_layers`. 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. + available — all output partitions receive an identity pre-scale. Offline + single-scale checkpoints use the canonical numeric FP4 layout + described above. 2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)** quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse @@ -316,7 +511,7 @@ unless they appear in `ignored_layers`. 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 + 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`. @@ -325,21 +520,21 @@ unless they appear in `ignored_layers`. 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. +5. **Cascade model configuration**: checkpoint auto-detection resolves each + expert's own `config.json` independently; the first expert's detected config + is not reused as the second expert's policy. An explicit global runtime + config applies to both experts, while per-component configs select policies + by the exact `transformer` / `transformer_2` names. Serialized checkpoint + metadata still determines each expert's storage format and BF16 layers. 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. +7. FP4 has fewer representable values than FP8, so evaluate quantization + error and generation quality for the chosen recipe. Offline Smooth scaling + and BF16 fallback can help with sensitive layers. + `num_bf16_fallback_layers` applies only to online DualScale. ## Adapting MXFP4 for a New Model @@ -350,27 +545,26 @@ 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()`. +`ignored_layers` entries use exact **runtime Linear module paths**, without +`.weight`. Inspect the actual transformer instance in model initialization or +a developer harness; `Omni` does not expose an `omni.pipeline.transformer` API. +For a transformer instance named `transformer`: ```python -from vllm_omni import Omni +from vllm.model_executor.layers.linear import LinearBase -# 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: +for name, module in transformer.named_modules(): + if isinstance(module, LinearBase): print(name) ``` -Compare the printed names against the diffusers checkpoint keys +Compare these module paths against the corresponding 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 | @@ -382,8 +576,9 @@ 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 + name the fused runtime module (for Wan, `attn1.to_qkv`). The source + checkpoint must keep **all three or none** in BF16. 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. @@ -392,12 +587,13 @@ fused runtime parameters. #### 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. +names** discovered in Step 1. This assumes the model already passes its +quantization config and module prefixes to compatible vLLM Linear layers. ```python omni = Omni( model="/path/to/your-model", - quantization={ + quantization_config={ "method": "mxfp4_dualscale", "ignored_layers": [ "blocks.0.attn1.to_qkv", # runtime name, not diffusers name @@ -409,9 +605,14 @@ omni = Omni( ``` ```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 text-to-video example accepts list-valued fields through JSON. +python examples/offline_inference/text_to_video/text_to_video.py \ + --model /path/to/your-model \ + --quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}' + +# Serving accepts the same configuration. +vllm serve /path/to/your-model --omni \ + --diffusion-quantization-config '{"method":"mxfp4_dualscale","ignored_layers":["blocks.0.attn1.to_qkv"]}' ``` The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers @@ -436,9 +637,13 @@ the merge script is correct. The injected block: } ``` -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. +Adding an offline BF16 fallback layer requires both the original BF16 weight +from the base model and a matching `ignored_layers` entry. Regenerate the +checkpoint with the revised storage policy, restoring the original weight +(and bias, if present) and removing that layer's quantization tensors. +Editing `config.json` alone does not reconstruct BF16 weights from stored FP4 +values. For fused QKV, restore all three projections together. Use runtime +module paths in `ignored_layers`. ### Step 3 — Write a merge script for offline mode @@ -474,6 +679,7 @@ 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. + each intended runtime Linear module resolves to the expected quantized or + BF16 method and that every BF16 entry has original BF16 weights. Account + for fused modules and, in online DualScale, the leading-block fallback rule; + the raw length of `ignored_layers` alone is not a reliable layer count. 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 06179ef9e..7b7bcb890 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -384,13 +384,20 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) - parser.add_argument( + quantization_group = parser.add_mutually_exclusive_group() + quantization_group.add_argument( "--quantization", type=str, default=None, choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"], 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).", ) + quantization_group.add_argument( + "--quantization-config", + type=json.loads, + default=None, + help='Quantization JSON, e.g. \'{"method":"mxfp4","w4a8_fallback_steps":[0,2]}\'.', + ) # Distributed and parallel execution parser.add_argument( @@ -572,6 +579,8 @@ def main(): omni_kwargs["flow_shift"] = args.flow_shift if args.quantization is not None: omni_kwargs["quantization"] = args.quantization + if args.quantization_config is not None: + omni_kwargs["quantization_config"] = args.quantization_config if args.cache_backend is not None: omni_kwargs["cache_backend"] = args.cache_backend omni_kwargs["cache_config"] = cache_config diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index e0a5ad9b0..1ed62f7f1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1480,7 +1480,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config - engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model=str(tmp_path)) od_config = OmniDiffusionConfig.from_kwargs(**engine_args) assert od_config.kv_transfer_config.engine_id == "dit-engine-1" @@ -1739,3 +1739,31 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder): builder(pipeline, deploy) else: builder(pipeline, user_deploy_config=deploy) + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit): + from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig + from vllm_omni.quantization import build_quant_config + + checkpoint = TransformerConfig.from_dict( + { + "quantization_config": { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "w4a8_fallback_steps": [37], + } + } + ) + requested = build_quant_config("mxfp4", w4a8_fallback_steps=[]) if explicit else None + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + tf_model_config=checkpoint, + quantization_config=requested, + ) + assert cfg.quantization_config_is_auto_detected is not explicit + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda self: self.set_tf_model_config(checkpoint)) + cfg.enrich_config() + restored = ForkingPickler.loads(ForkingPickler.dumps(cfg)) + assert restored.quantization_config_is_auto_detected is not explicit + assert restored.quantization_config.w4a8_fallback_steps == ([] if explicit else [37]) diff --git a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py index bb95363fd..b01f275eb 100644 --- a/tests/diffusion/models/wan2_2/test_create_transformer_quant.py +++ b/tests/diffusion/models/wan2_2/test_create_transformer_quant.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for transformer quant-config auto-detection. The loader path at pipeline_wan2_2.py carries two quantization contracts: @@ -27,6 +27,7 @@ from types import SimpleNamespace import pytest import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module +from vllm_omni.diffusion.data import OmniDiffusionConfig 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] @@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_ # od_config must remain unchanged — _create_transformer does not modify it. assert od_config.quantization_config is None + + +def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch): + import json + + import torch + + import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module + import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module + + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True) + for component in ("transformer", "transformer_2"): + directory = tmp_path / component + directory.mkdir() + config = dict(_MIN_CFG) + expert_quant_config = ( + disk_quant_config[0 if component == "transformer" else 1] + if isinstance(disk_quant_config, tuple) + else disk_quant_config + ) + if expert_quant_config is not None: + config["quantization_config"] = expert_quant_config + (directory / "config.json").write_text(json.dumps(config)) + (tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]})) + + class FakeComponent(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8) + + class FakeTransformer(FakeComponent): + def __init__(self, quant_config=None, **kwargs): + super().__init__() + self.received_quant_config = quant_config + + for module in (wan22_module, i2v_module): + monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent()) + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer) + pipeline_cls = { + "t2v": wan22_module.Wan22Pipeline, + "i2v": i2v_module.Wan22I2VPipeline, + "vace": vace_module.Wan22VACEPipeline, + }[pipeline_kind] + monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None) + od_config = OmniDiffusionConfig( + model=str(tmp_path), + model_class_name=pipeline_cls.__name__, + dtype=torch.bfloat16, + quantization_config=active, + boundary_ratio=0.875, + flow_shift=12.0, + enable_diffusion_pipeline_profiler=False, + ) + od_config.enrich_config() + pipeline = pipeline_cls(od_config=od_config) + return pipeline + + +@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"]) +@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")]) +@pytest.mark.parametrize("offline", [False, True]) +def test_pipeline_constructor_routes_each_expert_quantization( + pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker +): + """Run the actual pipeline constructors, expert factories and quant-method dispatch.""" + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.quantization import build_quant_config + + active = build_quant_config( + { + "transformer": { + "method": method, + "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"], + "w4a8_fallback_steps": [0, 1], + }, + "transformer_2": { + "method": method, + "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"], + "w4a8_fallback_steps": [], + }, + } + ) + disk_quant_config = None + if offline: + disk_quant_config = { + "quant_method": method, + flag: True, + "ignored_layers": ["blocks.12.ffn.net_2"], + "w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"], + "w4a8_fallback_steps": [9], + } + pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch) + for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)): + received = transformer.received_quant_config + assert received.get_name() == method + assert getattr(received, flag) == offline + assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else []) + selected = f"blocks.{10 + index}.attn1.to_qkv" + other = f"blocks.{11 - index}.attn1.to_qkv" + assert received.w4a8_fallback_layers == [selected] + assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer + assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer + if offline: + assert isinstance( + received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod + ) + assert pipeline.od_config.quantization_config is active + + +@pytest.mark.parametrize("explicit_none", [False, True]) +def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch): + from vllm_omni.quantization import build_quant_config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None} + active = build_quant_config(components) + pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active)) + pipeline._create_transformer(_MIN_CFG, component="transformer_2") + assert captured[-1].get("quant_config") is None + with pytest.raises(ValueError, match="disabled for component 'transformer_2'"): + pipeline._create_transformer( + {**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}}, + component="transformer_2", + ) + + +def test_pipeline_component_default_reaches_expert_factory(monkeypatch): + from vllm_omni.quantization import build_quant_config + + FakeTransformer, captured = _make_fake_transformer() + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}}) + pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active)) + pipeline._create_transformer(_MIN_CFG, component="transformer_2") + assert captured[-1]["quant_config"] is active.default_config + + +@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"]) +@pytest.mark.parametrize( + "configured_expert, with_default", + [ + ("transformer", False), + ("transformer", True), + ("transformer_2", False), + ("transformer_2", True), + ], +) +def test_pipeline_partial_expert_config_does_not_match_sibling( + pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch +): + from vllm_omni.quantization import build_quant_config + + component_spec = { + configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}, + } + if with_default: + component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]} + active = build_quant_config(component_spec) + pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch) + sibling = "transformer_2" if configured_expert == "transformer" else "transformer" + selected = getattr(pipeline, configured_expert).received_quant_config + other = getattr(pipeline, sibling).received_quant_config + assert selected is active.component_configs[configured_expert] + assert other is active.default_config + if with_default: + assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"] + else: + assert other is None + + +@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"]) +@pytest.mark.parametrize("different_ignored_layers", [False, True]) +def test_enriched_auto_quantization_preserves_each_expert_policy( + pipeline_kind, different_ignored_layers, tmp_path, monkeypatch +): + """Startup auto-detection must not turn the first expert into a user override.""" + high = { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "ignored_layers": ["blocks.0.ffn.net_2"], + "w4a8_fallback_steps": [], + "w4a8_fallback_layers": [], + "mxfp4_scale_alg": 0, + "require_smooth_scale": False, + } + low = { + **high, + "ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"], + "w4a8_fallback_steps": [37], + "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"], + "mxfp4_scale_alg": 2, + "require_smooth_scale": True, + } + pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch) + assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert. + for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]: + config = transformer.received_quant_config + for policy in ( + "ignored_layers", + "w4a8_fallback_steps", + "w4a8_fallback_layers", + "mxfp4_scale_alg", + "require_smooth_scale", + ): + assert getattr(config, policy) == expected[policy], policy diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 526eeeb15..58a5e961b 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @pytest.fixture(autouse=True) -def _cpu_platform(monkeypatch): +def _cpu_pipeline_runtime(monkeypatch): + # These lightweight pipelines have no accelerator or distributed groups. module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2") monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False) + monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1)) class _StubTransformer(nn.Module): @@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None: assert torch.equal(result, torch.full_like(latents, 10.0)) +@pytest.mark.parametrize("fail_first_request", [False, True]) +def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, + override_forward_context, + ) + from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step + + contexts: list[ForwardContext] = [] + events = [] + + class InjectedPredictionError(Exception): + pass + + class RecordingTransformer(_StubTransformer): + def __init__(self, expert): + super().__init__() + self.expert = expert + + def forward(self, hidden_states, encoder_hidden_states, **kwargs): + context = get_forward_context() + assert context is contexts[-1] + request_index = len(contexts) - 1 + step = context.denoise_step_idx + positive = bool(encoder_hidden_states[0, 0, 0] > 0) + events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2]))) + if fail_first_request and request_index == 0 and step == 1 and not positive: + raise InjectedPredictionError + return (torch.ones_like(hidden_states) * (1 if positive else -1),) + + def encode_prompt(**kwargs): + context = get_forward_context() + assert context.denoise_step_idx is None + contexts.append(context) + embeds = torch.ones(1, 8, 8) + return embeds, -embeds + + pipeline = _make_pipeline() + pipeline.transformer = RecordingTransformer("high") + pipeline.transformer_2 = RecordingTransformer("low") + pipeline.scheduler = _StubScheduler([900, 500, 100]) + monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt) + monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents) + + # Use the real runner context manager, Wan forward/diffuse, and CFG + # dispatch. Only model computation, scheduler math and unrelated KV I/O + # are replaced; no test code publishes a denoise step. + runner = object.__new__(DiffusionModelRunner) + runner.pipeline = pipeline + runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32) + runner.vllm_config = None + runner.cache_backend = None + monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None) + with override_forward_context(None): + for request_index in range(2): + request = OmniDiffusionRequest( + prompt="a fox walks", + request_id=f"request-{request_index}", + sampling_params=OmniDiffusionSamplingParams( + num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent" + ), + ) + + def execute(): + return runner._execute_request_list( + [request], + od_config=runner.od_config, + allow_single_output=True, + require_request_batch_support=False, + record_name="test_wan_step_context", + record_output_peak_memory=False, + ) + + if fail_first_request and request_index == 0: + with pytest.raises(InjectedPredictionError): + execute() + else: + execute() + assert not is_forward_context_available() + + assert len(contexts) == 2 and contexts[0] is not contexts[1] + expected = [ + (request, step, "high" if step == 0 else "low", positive, step in (0, 2)) + for request in range(2) + for step in range(2 if fail_first_request and request == 0 else 3) + for positive in (True, False) + ] + assert events == expected + + class _StubDMDScheduler: def __init__(self) -> None: self.predict_clean_calls: list[tuple[float, float, float]] = [] diff --git a/tests/diffusion/quantization/test_mxfp4_checkpoint.py b/tests/diffusion/quantization/test_mxfp4_checkpoint.py new file mode 100644 index 000000000..e183df078 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp4_checkpoint.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract.""" + +import sys +from types import ModuleType + +import pytest +import torch +from vllm.config.load import LoadConfig +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, +) +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4Config, + DiffusionMXFP4DualScaleMixedConfig, + NPUMxfp4DualScaleLinearMethod, + NPUMxfp4LinearMethod, + NPUMxfp4OnlineLinearMethod, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _patch_tp_state(monkeypatch): + 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) + + +class _CheckpointLayer(torch.nn.Module): + def __init__(self, method, dtype=torch.bfloat16): + super().__init__() + method.create_weights( + self, + input_size_per_partition=512, + output_partition_sizes=[2], + input_size=512, + output_size=2, + params_dtype=dtype, + weight_loader=default_weight_loader, + ) + + def load_weights(self, weights): + params = dict(self.named_parameters()) + loaded = set() + for name, value in weights: + if name in params: + params[name].weight_loader(params[name], value) + loaded.add(name) + return loaded + + +def _load_checkpoint(layer, config, weights, monkeypatch): + od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.counter_before_loading_weights = 0.0 + monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items())) + loader.load_weights(layer) + + +@pytest.mark.parametrize("serialized", [False, True]) +def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized) + # Exercise the shared implementation directly as well as the online MRO + # below: checkpoint requirements must follow the configured storage mode. + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)} + if serialized: + with pytest.raises(ValueError, match="Required weights.*weight_scale"): + _load_checkpoint(layer, config, weights, monkeypatch) + else: + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch): + config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg) + method = NPUMxfp4OnlineLinearMethod(config) + layer = _CheckpointLayer(method) + assert set(dict(layer.named_parameters())) == {"weight"} + assert layer.weight.is_meta + + weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512) + packed_weight = torch.zeros((2, 256), dtype=torch.uint8) + generated_scale = torch.full((2, 16), 127, dtype=torch.uint8) + quantized_inputs = [] + + def quantize(value, **kwargs): + quantized_inputs.append(value.clone()) + expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg) + if scale_alg == 2: + expected["dst_type_max"] = 7.25 + assert kwargs == expected + return packed_weight, generated_scale + + # Keep real creation, lazy loading, processing, and strict loader checks; + # substitute only the device quantization operator for this CPU test. + npu = ModuleType("torch_npu") + npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize) + monkeypatch.setitem(sys.modules, "torch_npu", npu) + _load_checkpoint(layer, config, {"weight": weight}, monkeypatch) + + assert len(quantized_inputs) == 1 + torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0) + torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0) + torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0) + assert not getattr(layer.weight_scale, "is_checkpoint_required", False) + method.process_weights_after_loading(layer) + assert len(quantized_inputs) == 1 + + +@pytest.mark.parametrize("require_smooth", [False, True]) +@pytest.mark.parametrize("has_smooth", [False, True]) +def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + scale = torch.linspace(0.5, 3.0, 512) + if has_smooth: + weights["mul_scale"] = scale + else: + weights.pop("mul_scale") + + if require_smooth and not has_smooth: + with pytest.raises(ValueError, match="mul_scale"): + _load_checkpoint(layer, config, weights, monkeypatch) + else: + _load_checkpoint(layer, config, weights, monkeypatch) + torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0) + + +@pytest.mark.parametrize("bad_scale", ["shape", "name"]) +def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + if bad_scale == "shape": + weights["mul_scale"] = torch.ones(511) + with pytest.raises(AssertionError): + _load_checkpoint(layer, config, weights, monkeypatch) + else: + weights["multiply_scale"] = weights.pop("mul_scale") + with pytest.raises(ValueError, match="mul_scale"): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("required", [False, True]) +@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0]) +def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + weights["mul_scale"][7] = bad_value + with pytest.raises(ValueError, match="finite, strictly positive"): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)]) +def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True) + layer = _CheckpointLayer(NPUMxfp4LinearMethod(config)) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()} + weights["mul_scale"] = bad_scale + with pytest.raises(ValueError, match="one-dimensional floating-point"): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("bad_value", [1e10, 1e-10]) +def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch): + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True) + method = NPUMxfp4LinearMethod(config) + layer = _CheckpointLayer(method, torch.float16) + layer.mul_scale.data.fill_(bad_value) + # Validation must fail before any device packing or GEMM. + monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu")) + with pytest.raises(ValueError, match="finite, strictly positive"): + method.process_weights_after_loading(layer) + + +@pytest.mark.parametrize( + ("mode", "missing"), + [ + ("single", "weight"), + ("single", "weight_scale"), + ("dual", "weight"), + ("dual", "weight_scale"), + ("dual", "weight_dual_scale"), + ("dual", "mul_scale"), + ], +) +def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch): + if mode == "single": + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True) + method = NPUMxfp4LinearMethod(config) + else: + config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True) + method = NPUMxfp4DualScaleLinearMethod(config) + layer = _CheckpointLayer(method) + weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing} + with pytest.raises(ValueError, match=missing): + _load_checkpoint(layer, config, weights, monkeypatch) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_single_scale_smooth_uses_real_row_parallel_loader(rank): + from vllm.model_executor.layers.linear import RowParallelLinear + + layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True))) + layer.tp_rank = rank + checkpoint_scale = torch.arange(1024, dtype=torch.float32) + RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale) + torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("fallback", [False, True]) +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_single_scale_smooth_precedes_quantization_and_preserves_output( + dtype, fallback, with_bias, scale_alg, monkeypatch +): + config = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg + ) + method = NPUMxfp4LinearMethod(config) + layer = _CheckpointLayer(method, dtype) + weights = { + "weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512), + "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8), + "mul_scale": torch.linspace(0.5, 3.0, 512), + } + _load_checkpoint(layer, config, weights, monkeypatch) + x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512) + bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None + quantized_inputs = [] + + # Only the two device operators are substituted. Exercise real apply(), + # step selection, Smooth multiplication, bias handling, and reshape logic. + npu = ModuleType("torch_npu") + + def quantize(value, **kwargs): + quantized_inputs.append(value.clone()) + expected = dict( + dst_type=torch.float8_e4m3fn if fallback else "fp4", + axis=-1, + block_size=32, + round_mode="rint", + scale_alg=0 if fallback else scale_alg, + ) + if not fallback and scale_alg == 2: + expected["dst_type_max"] = 7.25 + assert kwargs == expected + return value, torch.ones((6, 16), dtype=torch.uint8) + + def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs): + if fallback: + assert output_dtype == torch.bfloat16 + if bias is not None: + assert bias.dtype == output_dtype + assert bias.shape == (1, 2) + result = x_q.float() @ weight.float() + if bias is not None: + result += bias.float() + return result.to(output_dtype) + + npu.__dict__.update( + float8_e4m3fn=torch.float8_e4m3fn, + float8_e8m0fnu=torch.uint8, + float4_e2m1fn_x2="fp4", + npu_dynamic_mx_quant=quantize, + npu_quant_matmul=matmul, + ) + monkeypatch.setitem(sys.modules, "torch_npu", npu) + with override_forward_context(ForwardContext()): + set_forward_context_denoise_step_idx(0 if fallback else 1) + result = method.apply(layer, x, bias) + + expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype) + assert len(quantized_inputs) == 1 + torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0) + expected = expected_input.float() @ weights["weight"].float().T + if bias is not None: + expected += bias.float() + if fallback: + expected = expected.bfloat16() + torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0) + assert result.dtype == dtype diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py index 240ee2509..c51b2013b 100644 --- a/tests/diffusion/quantization/test_mxfp4_config.py +++ b/tests/diffusion/quantization/test_mxfp4_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config.""" import pytest @@ -27,12 +27,42 @@ def test_mxfp4_config_get_name(): assert DiffusionMXFP4Config.get_name() == "mxfp4" +@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3]) +def test_mxfp4_scale_alg_rejects_invalid_values(value): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="mxfp4_scale_alg"): + build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value}) + + +@pytest.mark.parametrize("runtime_alg", [0, 2]) +@pytest.mark.parametrize("serialized", [False, True]) +def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg) + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "mxfp4_scale_alg": 2, + "ignored_layers": ["proj_out"], + }, + ) + assert resolved.is_checkpoint_mxfp4_serialized + assert resolved.mxfp4_scale_alg == runtime_alg + assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == [] + + 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 == [] + assert cfg.w4a8_fallback_layers == [] def test_mxfp4_config_from_config_serialized(): @@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized(): assert cfg.is_checkpoint_mxfp4_serialized is True +def test_mxfp4_config_requires_offline_smooth(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True} + assert DiffusionMXFP4Config.from_config(values).require_smooth_scale + assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale + with pytest.raises(ValueError, match="offline"): + build_quant_config({"method": "mxfp4", "require_smooth_scale": True}) + with pytest.raises(ValueError, match="boolean"): + DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"}) + + +@pytest.mark.parametrize("serialized", [False, True]) +def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1]) + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "require_smooth_scale": True, + }, + ) + assert resolved.require_smooth_scale + assert resolved.is_checkpoint_mxfp4_serialized + assert resolved.w4a8_fallback_steps == [1] + + +@pytest.mark.parametrize("disk_requirement", [None, False, True]) +def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True) + disk = { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "ignored_layers": ["blocks.0.attn2.to_q"], + } + if disk_requirement is not None: + disk["require_smooth_scale"] = disk_requirement + resolved = resolve_quant_config_from_disk(active, disk) + assert resolved.require_smooth_scale + assert resolved.ignored_layers == disk["ignored_layers"] + + +@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}]) +def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + active = DiffusionMXFP4Config( + ignored_layers=["blocks.0.attn2.to_q"], + w4a8_fallback_steps=[1], + w4a8_fallback_layers=["blocks.1.attn1.to_qkv"], + ) + + resolved = resolve_quant_config_from_disk(active, disk_qc) + + assert resolved is active + assert resolved.ignored_layers == ["blocks.0.attn2.to_q"] + assert resolved.w4a8_fallback_steps == [1] + assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"] + + +def test_quantization_disk_string_rejects_mismatched_active_method(): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active = build_quant_config("mxfp4") + + with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"): + resolve_quant_config_from_disk(active, "mxfp4_dualscale") + + +def test_quantization_disk_string_accepts_equivalent_method_alias(): + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + class ActiveAliasConfig: + @staticmethod + def get_name() -> str: + return "inc" + + active = ActiveAliasConfig() + + assert resolve_quant_config_from_disk(active, "auto-round") is active + + def test_mxfp4_config_from_config_ignored_layers(): from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config @@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback(): assert cfg.ignored_layers == ["proj_out"] +@pytest.mark.parametrize("serialized", [False, True]) +@pytest.mark.parametrize( + "policy", + [ + {"w4a8_fallback_steps": [0]}, + {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}, + {"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}, + ], +) +def test_dualscale_rejects_w4a8_configuration(serialized, policy): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig + + config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy} + with pytest.raises(ValueError, match="does not support W4A8"): + DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy) + with pytest.raises(ValueError, match="does not support W4A8"): + DiffusionMXFP4DualScaleMixedConfig.from_config(config) + with pytest.raises(ValueError, match="does not support W4A8"): + build_quant_config({"transformer": config}) + with pytest.raises(ValueError, match="does not support W4A8"): + resolve_quant_config_from_disk(None, config) + + +@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"]) +@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)]) +def test_mxfp4_rejects_invalid_fallback_steps(method, steps): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="non-negative integer"): + build_quant_config({"method": method, "w4a8_fallback_steps": steps}) + + +def test_mxfp4_fallback_steps_follow_request_context(): + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, + ) + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step + + steps = [2, 0, 2] + cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps}) + steps.append(1) + assert cfg.w4a8_fallback_steps == [0, 2] + with override_forward_context(None): + assert not _is_w4a8_fallback_step([]) + with pytest.raises(RuntimeError, match="publishes its denoise step"): + _is_w4a8_fallback_step(cfg.w4a8_fallback_steps) + for _request in range(2): + with override_forward_context(ForwardContext()): + for step, expected in [(0, True), (1, False), (2, True), (3, False)]: + set_forward_context_denoise_step_idx(step) + # Both CFG passes and both experts read the same index. + assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [ + expected, + expected, + ] + with pytest.raises(RuntimeError, match="publishes its denoise step"): + _is_w4a8_fallback_step(cfg.w4a8_fallback_steps) + + +@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")]) +@pytest.mark.parametrize("steps", [[], [0, 2]]) +@pytest.mark.parametrize("already_serialized", [False, True]) +def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + layers = ["blocks.10.attn1.to_qkv"] if steps else [] + active = build_quant_config( + { + "method": method, + flag: already_serialized, + "w4a8_fallback_steps": steps, + "w4a8_fallback_layers": layers, + } + ) + # Separate experts may have distinct high precision layers. Both must + # retain the requested step policy during online/offline reconciliation. + for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]: + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": method, + flag: True, + "ignored_layers": ignored, + "w4a8_fallback_steps": [9], + "w4a8_fallback_layers": ["blocks.9.ffn.net_2"], + }, + ) + assert getattr(resolved, flag) + assert resolved.ignored_layers == ignored + assert resolved.w4a8_fallback_steps == steps + assert resolved.w4a8_fallback_layers == layers + assert active.ignored_layers == [] + + +@pytest.mark.parametrize( + "runtime_policy", + [{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}], +) +def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active = build_quant_config({"method": "mxfp4", **runtime_policy}) + resolved_experts = [] + for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]): + disk = { + "quant_method": "mxfp4", + "is_checkpoint_mxfp4_serialized": True, + "ignored_layers": ignored, + "w4a8_fallback_steps": [9], + "w4a8_fallback_layers": ["blocks.9.ffn.net_2"], + } + resolved = resolve_quant_config_from_disk(active, disk) + assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", []) + assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", []) + assert resolved.ignored_layers == ignored + detected = resolve_quant_config_from_disk(None, disk) + assert detected.w4a8_fallback_steps == [9] + assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"] + resolved_experts.append(resolved) + assert resolved_experts[0] is not resolved_experts[1] + assert active.is_checkpoint_mxfp4_serialized is False + assert active.ignored_layers == [] + + +@pytest.mark.parametrize( + "method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")] +) +def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active_spec = { + "method": method, + flag: True, + "ignored_layers": ["blocks.0.attn2.to_q"], + "w4a8_fallback_steps": [1] if method == "mxfp4" else [], + "w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [], + } + if method == "mxfp4": + active_spec["require_smooth_scale"] = True + active = build_quant_config(active_spec) + + resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True}) + + assert resolved is not active + assert resolved.ignored_layers == [] + assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps + assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers + if method == "mxfp4": + assert resolved.require_smooth_scale + + +@pytest.mark.parametrize( + "method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")] +) +@pytest.mark.parametrize("with_active", [False, True]) +def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + active = None + disk_steps = [9] if method == "mxfp4" else [] + active_steps = [1] if method == "mxfp4" else [] + expected_steps = disk_steps + if with_active: + active = build_quant_config( + { + "method": method, + flag: True, + "ignored_layers": ["blocks.0.attn2.to_q"], + "w4a8_fallback_steps": active_steps, + } + ) + expected_steps = active_steps + + resolved = resolve_quant_config_from_disk( + active, + { + "quant_method": method, + flag: True, + "modules_to_not_convert": ["blocks.2.ffn.net_2"], + "w4a8_fallback_steps": disk_steps, + }, + ) + + assert resolved.ignored_layers == ["blocks.2.ffn.net_2"] + assert resolved.w4a8_fallback_steps == expected_steps + + # --------------------------------------------------------------------------- # build_quant_config integration # --------------------------------------------------------------------------- @@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch): assert isinstance(method, ROCmMxfp4OnlineLinearMethod) -def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch): +@pytest.mark.parametrize("steps", [[], [0, 2]]) +def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps): """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe.""" from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") - cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"]) + cfg = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps + ) assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod) +def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch): + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2]) + + with pytest.raises(NotImplementedError, match="only supported on NPU"): + cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q") + + def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch): """MXFP4 on ROCm requires gfx950; any other arch must raise.""" from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config @@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform): assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) assert layer.input_size_per_partition == _TP2_K // _TP2 assert layer.output_size_per_partition == _TP2_N + + +@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"]) +@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]]) +def test_mxfp4_rejects_invalid_fallback_layers(method, layers): + from vllm_omni.quantization import build_quant_config + + with pytest.raises(ValueError, match="exact runtime Linear paths"): + build_quant_config({"method": method, "w4a8_fallback_layers": layers}) + + +@pytest.mark.parametrize("method", ["mxfp4"]) +def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method): + from vllm_omni.quantization import build_quant_config + from vllm_omni.quantization.component_config import resolve_component_quant_config + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"] + components = build_quant_config( + {"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}} + ) + names.append("blocks.11.attn1.to_qkv") + assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]] + assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == [] + detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]}) + assert detected.w4a8_fallback_layers == [names[0]] + assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]] + + +@pytest.mark.parametrize( + "kind, serialized_flag", + [("mxfp4", "is_checkpoint_mxfp4_serialized")], +) +@pytest.mark.parametrize("offline", [False, True]) +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_mxfp4_layer_and_step_union_routes_actual_methods( + kind, serialized_flag, offline, scale_alg, monkeypatch, mocker +): + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.quantization import build_quant_config, mxfp4_config + + monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True) + selected = "blocks.10.attn1.to_qkv" + other = "blocks.11.attn2.to_q" + preserved = "blocks.12.ffn.net_2" + config = build_quant_config( + { + "method": kind, + "mxfp4_scale_alg": scale_alg, + serialized_flag: offline, + "w4a8_fallback_layers": [selected, preserved], + "ignored_layers": [preserved], + "w4a8_fallback_steps": [1], + } + ) + # Real config dispatch creates separate methods for selected and other layers. + linear = mocker.Mock(spec=LinearBase) + assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod) + x = torch.ones(2, 64, dtype=torch.bfloat16) + layer = torch.nn.Module() + for name in ("weight", "weight_scale", "mul_scale"): + layer.register_buffer(name, torch.ones(64, dtype=x.dtype)) + calls = [] + + def record_a8(*args): + calls.append("A8") + return x + + def record_a4(*args): + calls.append("A4") + return x + + monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8) + for name in (selected, other): + method = config.get_quant_method(linear, name) + assert method.is_w4a8_fallback_layer == (name == selected) + scales = (x, x) + monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales) + monkeypatch.setattr(method, "_quant_matmul", record_a4) + for step in (0, 1, 2): + with override_forward_context(ForwardContext(denoise_step_idx=step)): + method.apply(layer, x) + assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"] + # A permanently selected layer works without a request context, even if steps are also configured. + with override_forward_context(None): + config.get_quant_method(linear, selected).apply(layer, x) + assert calls[-1] == "A8" + + +@pytest.mark.parametrize("online", [False, True]) +@pytest.mark.parametrize("selected", [False, True]) +@pytest.mark.parametrize("scale_alg", [0, 2]) +def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch): + import sys + from types import SimpleNamespace + + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.quantization import mxfp4_config + + monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False) + prepared = [] + packed = [] + + def quantize(x, **kwargs): + prepared.append(kwargs) + return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8) + + def pack(x, dtype): + packed.append(x) + return x.to(dtype) + + fake_npu = SimpleNamespace( + float4_e2m1fn_x2=torch.int8, + float8_e8m0fnu=torch.uint8, + float8_e4m3fn=torch.float8_e4m3fn, + npu_dtype_cast=pack, + npu_dynamic_mx_quant=quantize, + ) + monkeypatch.setitem(sys.modules, "torch_npu", fake_npu) + config = mxfp4_config.DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=not online, + w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [], + w4a8_fallback_steps=[1], + mxfp4_scale_alg=scale_alg, + ) + cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod + method = cls(config, prefix="blocks.0.attn1.to_qkv") + layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16) + layer.orig_dtype = torch.bfloat16 + if not online: + layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False)) + layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False)) + method.process_weights_after_loading(layer) + pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) + method.process_weights_after_loading(layer) + assert len(prepared) == int(online) + expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg) + if scale_alg == 2: + expected["dst_type_max"] = 7.25 + assert prepared == ([expected] if online else []) + assert len(packed) == int(not online) + branches = [] + + def a4(x_q, x_scale, target, bias, dtype): + assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers + branches.append("A4") + return torch.zeros(x_q.shape[0], 2, dtype=dtype) + + def a8(x, weight, scale, bias, dtype): + assert (weight.data_ptr(), scale.data_ptr()) == pointers + branches.append("A8") + return torch.zeros(x.shape[0], 2, dtype=dtype) + + monkeypatch.setattr(method, "_quant_matmul", a4) + monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8) + monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x)) + monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization")) + monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking")) + for step in (0, 1, 2): + with override_forward_context(ForwardContext(denoise_step_idx=step)): + method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16)) + assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"]) + assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers + assert set(layer.state_dict()) == ( + {"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"} + ) + assert not hasattr(layer, "w4a8_weight") + assert not hasattr(layer, "w4a8_weight_scale") + + +@pytest.mark.parametrize("online", [False, True]) +def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch): + import sys + from types import SimpleNamespace + + from vllm_omni.quantization import mxfp4_config + + monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False) + original_ones = torch.ones + monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"})) + calls = [] + + def quantize_dual(x, **kw): + calls.append(kw.get("smooth_scale")) + return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8) + + fake_npu = SimpleNamespace( + float4_e2m1fn_x2=torch.int8, + npu_dtype_cast=lambda x, dtype: x.to(dtype), + npu_format_cast=lambda x, *args, **kw: x, + npu_dynamic_dual_level_mx_quant=quantize_dual, + ) + monkeypatch.setitem(sys.modules, "torch_npu", fake_npu) + config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig( + is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[] + ) + cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod + method = cls(config) + layer = torch.nn.Module() + layer.input_size_per_partition = 512 + for name, value in { + "weight": torch.ones(2, 512, dtype=torch.bfloat16), + "weight_scale": torch.full((2, 16), 127, dtype=torch.uint8), + "weight_dual_scale": torch.ones(2, 1, 1), + "mul_scale": torch.ones(512), + }.items(): + layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False)) + method.process_weights_after_loading(layer) + method.process_weights_after_loading(layer) + assert len(calls) == int(online) + assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"} + assert not hasattr(layer, "w4a8_weight") + monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16)) + method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16)) + assert len(calls) == 1 + int(online) + torch.testing.assert_close(calls[-1], layer.mul_scale) diff --git a/tests/diffusion/quantization/test_wan_mxfp4_loading.py b/tests/diffusion/quantization/test_wan_mxfp4_loading.py new file mode 100644 index 000000000..fc7a43d43 --- /dev/null +++ b/tests/diffusion/quantization/test_wan_mxfp4_loading.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU.""" + +import pytest +import torch +from vllm.config.load import LoadConfig +from vllm.model_executor.layers.linear import QKVParallelLinear + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel +from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _patch_tp_state(monkeypatch): + 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) + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True) + prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer" + monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1) + + +class _WanCheckpoint(torch.nn.Module): + load_weights = WanTransformer3DModel.load_weights + _log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary + + def __init__(self, config, *, num_blocks=1): + super().__init__() + blocks = [] + for block_idx in range(num_blocks): + block = torch.nn.Module() + block.attn1 = torch.nn.Module() + block.attn1.to_qkv = QKVParallelLinear( + hidden_size=512, + head_size=32, + total_num_heads=4, + bias=False, + params_dtype=torch.bfloat16, + quant_config=config, + prefix=f"blocks.{block_idx}.attn1.to_qkv", + disable_tp=True, + ) + blocks.append(block) + self.blocks = torch.nn.ModuleList(blocks) + + +def _checkpoint(mode, fused): + if mode == "single": + config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True) + elif mode == "dual": + config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True) + else: + config = None + model = _WanCheckpoint(config) + weights = {} + expected = {} + for name, param in model.named_parameters(): + suffix = name.rsplit(".", 1)[1] + output_dim = getattr(param, "output_dim", None) + if output_dim is None: + value = torch.full_like(param, 2.5) + shards = [value] * 3 + else: + parts = param.chunk(3, dim=output_dim) + shards = [ + torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts) + ] + value = torch.cat(shards, dim=output_dim) + expected[name] = value + if fused: + weights[name] = value + else: + for shard_id, shard in zip(("q", "k", "v"), shards): + weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard + return model, config, weights, expected + + +def _load(model, config, weights, monkeypatch): + loader = DiffusersPipelineLoader( + LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config) + ) + loader.counter_before_loading_weights = 0.0 + monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items())) + loader.load_weights(model) + + +@pytest.mark.parametrize("mode", ["single", "dual", "bf16"]) +@pytest.mark.parametrize("fused", [False, True]) +def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch): + model, config, weights, expected = _checkpoint(mode, fused) + _load(model, config, weights, monkeypatch) + for name, param in model.named_parameters(): + torch.testing.assert_close(param, expected[name], rtol=0, atol=0) + + +@pytest.mark.parametrize("mode", ["single", "dual", "bf16"]) +def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch): + model, config, weights, _ = _checkpoint(mode, fused=False) + weights.pop("blocks.0.attn1.to_k.weight") + weights.pop("blocks.0.attn1.to_v.weight") + with pytest.raises(ValueError, match="to_qkv.weight"): + _load(model, config, weights, monkeypatch) + + +@pytest.mark.parametrize("mode", ["single", "dual"]) +def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch): + model, config, weights, _ = _checkpoint(mode, fused=False) + weights.pop("blocks.0.attn1.to_k.weight_scale") + weights.pop("blocks.0.attn1.to_v.weight_scale") + with pytest.raises(ValueError, match="to_qkv.weight_scale"): + _load(model, config, weights, monkeypatch) + + +@pytest.mark.parametrize("mode", ["single", "dual"]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch): + model, config, weights, _ = _checkpoint(mode, fused=False) + weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2 + if reverse: + weights = dict(reversed(weights.items())) + with pytest.raises(ValueError, match="must share the same Smooth tensor"): + _load(model, config, weights, monkeypatch) + + +@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"]) +def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name): + model, _, _, _ = _checkpoint("single", fused=True) + loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))]) + assert loaded == set() + + +@pytest.mark.parametrize("fused", [False, True]) +def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch): + model, config, weights, _ = _checkpoint("single", fused=fused) + projection = "to_qkv" if fused else "to_q" + weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1) + with pytest.raises(ValueError, match="cannot load DualScale tensor"): + _load(model, config, weights, monkeypatch) + + +def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override(): + config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"]) + model = _WanCheckpoint(config) + # Uses the same validation called at the end of WanTransformer3DModel construction. + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"): + config.w4a8_fallback_layers = [bad] + with pytest.raises(ValueError, match="exact runtime Linear paths"): + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + + +def test_wan_layer_policy_defers_only_other_pp_stage_paths(): + from vllm.model_executor.models.utils import PPMissingLayer + + config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"]) + model = _WanCheckpoint(config) + model.blocks.append(PPMissingLayer()) + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + config.w4a8_fallback_layers = ["blocks.1"] + with pytest.raises(ValueError, match="exact runtime Linear paths"): + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + + +@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"]) +def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker): + selected = "blocks.0.attn1.to_qkv" + bf16 = "blocks.1.attn1.to_qkv" + transformer_config = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=is_checkpoint_serialized, + w4a8_fallback_layers=[selected], + ) + transformer_2_config = DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=is_checkpoint_serialized, + ignored_layers=[bf16], + w4a8_fallback_layers=[selected, bf16], + ) + pipeline = torch.nn.Module() + pipeline.transformer = _WanCheckpoint(transformer_config) + pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2) + WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config) + WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config) + + loader = DiffusersPipelineLoader( + LoadConfig(), + OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config), + ) + events = [] + summaries = [] + + def process_weights(model, _target_device): + events.append("process") + for module in model.modules(): + method = getattr(module, "quant_method", None) + if getattr(method, "is_w4a8_fallback_layer", False): + module._already_called_process_weights_after_loading = True + + def capture_summary(message, *args): + rendered = message % args + summaries.append(rendered) + events.append(f"summary:{args[0]}") + + loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign] + loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = process_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + mocker.patch( + "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info", + side_effect=capture_summary, + ) + + assert loader.load_model(load_device="cpu") is pipeline + assert events == ["load", "process", "summary:transformer", "summary:transformer_2"] + assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0] + assert "BF16 overrides=0 []" in summaries[0] + assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0] + assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1] + assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1] + assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1] + + +def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker): + selected = "blocks.0.attn1.to_qkv" + config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected]) + model = _WanCheckpoint(config) + WanTransformer3DModel._validate_w4a8_fallback_layers(model, config) + logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info") + + WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer") + + message, *args = logger_info.call_args.args + rendered = message % tuple(args) + assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 23ff796e2..6f37f823d 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.fixture(autouse=True) +def _local_model_paths(monkeypatch): + # These tests exercise config transport, not model repository resolution. + monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model) + + def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: """Simulate the real path: create_default_diffusion → OmniDiffusionConfig. diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py index c588b7a38..e0c381bf3 100644 --- a/tests/examples/offline_inference/test_text_to_video_prompts.py +++ b/tests/examples/offline_inference/test_text_to_video_prompts.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import functools import importlib.util +import sys +from types import SimpleNamespace import pytest import torch @@ -24,6 +26,51 @@ def _load_text_to_video(): return module +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + (["--quantization", "mxfp4"], {"quantization": "mxfp4"}), + ( + ["--quantization-config", '{"method":"mxfp4","w4a8_fallback_steps":[0,2]}'], + {"quantization_config": {"method": "mxfp4", "w4a8_fallback_steps": [0, 2]}}, + ), + ], +) +def test_quantization_cli_reaches_omni(arguments, expected, monkeypatch): + mod = _load_text_to_video() + captured = {} + + class StopBeforeModelLoadError(Exception): + pass + + def capture_omni(**kwargs): + captured.update(kwargs) + raise StopBeforeModelLoadError + + monkeypatch.setattr(sys, "argv", ["text_to_video.py", "--model", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", *arguments]) + monkeypatch.setattr(mod, "current_omni_platform", SimpleNamespace(device_type="cpu")) + monkeypatch.setattr(mod, "Omni", capture_omni) + + with pytest.raises(StopBeforeModelLoadError): + mod.main() + + assert {key: captured[key] for key in ("quantization", "quantization_config") if key in captured} == expected + + +def test_quantization_cli_options_are_mutually_exclusive(monkeypatch): + mod = _load_text_to_video() + monkeypatch.setattr( + sys, + "argv", + ["text_to_video.py", "--quantization", "mxfp4", "--quantization-config", '{"method":"mxfp4"}'], + ) + + with pytest.raises(SystemExit) as exc: + mod.parse_args() + + assert exc.value.code == 2 + + @pytest.mark.parametrize( ("negative_prompt", "expected_negative_prompt"), [(None, None), ("", ""), ("blurry", "blurry")], diff --git a/tests/platforms/npu/quant/test_mxfp4_step_fallback.py b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py new file mode 100644 index 000000000..8431ebf17 --- /dev/null +++ b/tests/platforms/npu/quant/test_mxfp4_step_fallback.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights.""" + +import pytest +import torch + +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu] + +torch_npu = pytest.importorskip("torch_npu") + + +@torch.inference_mode() +def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2): + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, + ) + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4Config, + NPUMxfp4LinearMethod, + NPUMxfp4OnlineLinearMethod, + ) + + prefix = "blocks.10.attn1.to_qkv" + steps = [1] if policy in ("step", "layer_and_step") else [] + layers = [prefix] if policy in ("layer", "layer_and_step") else [] + fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg} + torch.manual_seed(42) + layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype) + if mode == "single": + method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix) + elif mode == "single_offline": + method = NPUMxfp4LinearMethod( + DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix + ) + layer.orig_dtype = dtype + values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0]) + layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight)) + layer.register_parameter( + "weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False) + ) + layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False)) + else: + raise ValueError(f"Unsupported single-scale test mode: {mode}") + method.process_weights_after_loading(layer) + x = torch.randn(input_shape, device="npu", dtype=dtype) + pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) + method.process_weights_after_loading(layer) + quantize = torch_npu.npu_dynamic_mx_quant + activation_types = [] + + def record_quantize(x, **kwargs): + dst = kwargs["dst_type"] + expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg + assert kwargs["scale_alg"] == expected_alg + assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint" + if expected_alg == 2: + assert kwargs["dst_type_max"] == 7.25 + else: + assert "dst_type_max" not in kwargs + activation_types.append(dst) + return quantize(x, **kwargs) + + monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize) + results = [] + for _request in range(2): + with override_forward_context(ForwardContext()): + for step in range(3): + set_forward_context_denoise_step_idx(step) + output = method.apply(layer, x, layer.bias) + torch.npu.synchronize() + assert output.shape == (*input_shape[:-1], 128) + assert output.dtype == dtype + assert torch.isfinite(output).all() + results.append(output.cpu()) + expected_types = [ + torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3) + ] + assert activation_types == expected_types * 2 + assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers + assert not hasattr(layer, "w4a8_weight") + assert not hasattr(layer, "w4a8_weight_scale") + for first, second in zip(results[:3], results[3:]): + torch.testing.assert_close(first, second, rtol=0, atol=0) + + +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)]) +@pytest.mark.parametrize("mode", ["single", "single_offline"]) +def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch): + _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step") + + +@pytest.mark.cards_1 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_single_scale_step_fallback_matches_independent_dense_reference(dtype): + if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name(): + pytest.skip("MXFP4 kernels require an Ascend950 NPU") + + from vllm_omni.diffusion.forward_context import ( + ForwardContext, + override_forward_context, + set_forward_context_denoise_step_idx, + ) + from vllm_omni.quantization.mxfp4_config import ( + DiffusionMXFP4Config, + NPUMxfp4LinearMethod, + ) + + # Synthetic serialized single-scale values and E8M0 exponent bytes. + # Powers of two make both activation precisions exact: this independent + # CPU dense oracle needs no production packing, dequantization or apply. + # Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate. + values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512) + weight = torch.cat((values, values), dim=1) + fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8) + smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0))) + bias = torch.tensor([0.25, -0.25]).repeat(64) + x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64 + expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias) + + layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype) + layer.weight.copy_(weight) + layer.bias.copy_(bias) + layer.orig_dtype = dtype + for name, value in {"weight_scale": fine, "mul_scale": smooth}.items(): + layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False)) + method = NPUMxfp4LinearMethod( + DiffusionMXFP4Config( + is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2 + ) + ) + method.process_weights_after_loading(layer) + prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) + assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"} + assert not hasattr(layer, "w4a8_weight") + assert not hasattr(layer, "w4a8_weight_scale") + x = x.to(device="npu", dtype=dtype) + + with override_forward_context(ForwardContext()): + for step in range(3): + set_forward_context_denoise_step_idx(step) + actual = method.apply(layer, x, layer.bias).cpu() + # A8's documented GEMM output is BF16 before restoring the caller + # dtype; the reference does not reuse quantization or unpacking. + reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype) + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + assert actual.dtype == dtype + assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers + + +@pytest.mark.cards_1 +@pytest.mark.parametrize("mode", ["single", "single_offline"]) +@pytest.mark.parametrize("policy", ["layer", "layer_and_step"]) +def test_mxfp4_layer_fallback(mode, policy, monkeypatch): + """Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix.""" + _exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy) + + +@pytest.mark.cards_1 +@pytest.mark.parametrize("mode", ["single", "single_offline"]) +def test_mxfp4_ocp_compatibility(mode, monkeypatch): + """OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests.""" + _exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 56b40b836..c13935489 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -817,6 +817,8 @@ class _DiffusionConfigProjection: enable_stage_verification: bool = True prompt_file_path: str | None = None quantization_config: _QuantizationConfigType = None + # Internal provenance, retained across config projection and worker transport. + quantization_config_is_auto_detected: bool = False extras: dict[str, Any] = field(default_factory=dict) @field_validator("kv_transfer_config", mode="before") @@ -964,6 +966,8 @@ class _DiffusionConfigProjection: or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config)) ) if should_use_checkpoint_config: + if self.quantization_config is None: + self.quantization_config_is_auto_detected = True self.quantization_config = quant_config if quant_method is not None: self.additional_config.setdefault("auto_detected_quant_method", quant_method) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e1ce550d4..0ffd0c54e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1021,6 +1021,8 @@ class OmniDiffusionConfig: # str is resolved to {"method": } internally. # Per-component: {"transformer": {"method": "fp8"}, "vae": None} quantization_config: str | QuantizationConfig | dict[str, Any] | None = None + # Internal provenance, retained across config projection and worker transport. + quantization_config_is_auto_detected: bool = False # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint # has already resolved to vLLM's ModelOpt FP8 linear method. @@ -1358,6 +1360,8 @@ class OmniDiffusionConfig: or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config)) ) if should_use_checkpoint_config: + if self.quantization_config is None: + self.quantization_config_is_auto_detected = True self.quantization_config = tf_config.quant_config logger.info( "Auto-detected quantization '%s' from model config", diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 78247e7b9..834887f41 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -7,7 +7,7 @@ import json import os import re import time -from collections.abc import Generator, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable, Sequence from pathlib import Path from typing import cast @@ -481,8 +481,14 @@ class DiffusersPipelineLoader(HWRLoaderMixin): ) def _get_expected_parameter_names(self, model: nn.Module) -> set[str]: - """Return parameter names that should be covered by strict load checks.""" - all_parameter_names = {name for name, _ in model.named_parameters()} + """Return parameter names that should be covered by strict load checks. + + A parameter with an initialized checkpoint default can explicitly set + ``is_checkpoint_optional``. It is still loaded when present on disk. + """ + all_parameter_names = { + name for name, param in model.named_parameters() if not getattr(param, "is_checkpoint_optional", False) + } sources = self._get_weight_sources(model) # Keep strict behavior if no source metadata exists. @@ -853,9 +859,24 @@ class DiffusersPipelineLoader(HWRLoaderMixin): del model return self.load_fresh_canonical_model() raise + self._log_w4a8_fallback_load_summaries(model) self._attach_offload_startup_state(model) return model + @staticmethod + def _log_w4a8_fallback_load_summaries(model: nn.Module) -> None: + """Ask discovered DiTs to report W4A8 state at the common load exit. + + Each DiT derives readiness from its processed layers, so a deferred + weight plan cannot be mistaken for the ordinary or HSDP load path. + """ + components = ModuleDiscovery.discover(model) + for component_name, dit in zip(components.dit_names, components.dits): + candidate = getattr(dit, "_log_w4a8_fallback_load_summary", None) + if callable(candidate): + reporter = cast(Callable[[str], None], candidate) + reporter(component_name) + @staticmethod def _request_offload_after_quant(model: nn.Module) -> int: """Ask online-quant layers to return to host memory once quantized. @@ -1065,6 +1086,16 @@ class DiffusersPipelineLoader(HWRLoaderMixin): # that have loaded weights tracking currently. if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights + # Offline formats can require scales that older online loaders + # were allowed to synthesize. Do not apply that legacy tolerance + # to an explicitly required checkpoint tensor. + required_missing = { + name + for name, param in model.named_parameters() + if name in weights_not_loaded and getattr(param, "is_checkpoint_required", False) + } + if required_missing: + raise ValueError(f"Required weights were not initialized from checkpoint: {required_missing}") # NOTE: if the model is quantized, ignore not_loaded check for scale # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a # static activation `input_scale`, which the quant method may 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 15b036864..915fb45f7 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -179,8 +179,36 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc return {} +def resolve_wan_transformer_quant_config( + config: dict, quant_config: QuantizationConfig | None, component: str +) -> QuantizationConfig | None: + """Resolve the expert before applying its checkpoint's storage contract.""" + from vllm_omni.quantization.component_config import ComponentQuantizationConfig + from vllm_omni.quantization.factory import resolve_quant_config_from_disk + + # Wan experts are siblings: "transformer_2" must not inherit "transformer" + # through the generic layer-prefix resolver. Select the whole expert name. + component_quant_config = ( + quant_config.component_configs.get(component, quant_config.default_config) + if isinstance(quant_config, ComponentQuantizationConfig) + else quant_config + ) + quantization_disabled = isinstance(quant_config, ComponentQuantizationConfig) and component_quant_config is None + resolved_quant_config = resolve_quant_config_from_disk(component_quant_config, config.get("quantization_config")) + if quantization_disabled and resolved_quant_config is not None: + raise ValueError( + f"Quantization is disabled for component {component!r}, but its checkpoint declares quantization. " + "Use a BF16 checkpoint for this component or enable its matching quantization method." + ) + return resolved_quant_config + + def create_transformer_from_config( - config: dict, quant_config: QuantizationConfig | None = None, prefix: str = "" + config: dict, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + *, + component: str = "transformer", ) -> WanTransformer3DModel: """Create WanTransformer3DModel from config dict.""" kwargs: dict = {} @@ -216,11 +244,7 @@ def create_transformer_from_config( if "pos_embed_seq_len" in config: kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"] - if "quantization_config" in config: - from vllm_omni.quantization.factory import resolve_quant_config_from_disk - - quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"]) - + quant_config = resolve_wan_transformer_quant_config(config, quant_config, component) if quant_config is not None: kwargs["quant_config"] = quant_config if prefix: @@ -460,13 +484,13 @@ class Wan22Pipeline( # Initialize transformers with correct config (weights loaded via load_weights) if load_transformer: transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = self._create_transformer(transformer_config) + self.transformer = self._create_transformer(transformer_config, component="transformer") else: self.transformer = None if load_transformer_2: transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only) - self.transformer_2 = self._create_transformer(transformer_2_config) + self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2") else: self.transformer_2 = None @@ -504,10 +528,17 @@ class Wan22Pipeline( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) - def _create_transformer(self, config: dict) -> WanTransformer3DModel: + def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel: """Create a transformer from a config dict. Respects od_config.quantization_config.""" quant_config = getattr(self.od_config, "quantization_config", None) - return create_transformer_from_config(config, quant_config=quant_config) + # Startup metadata describes the first expert, not a user policy for both. + if getattr(self.od_config, "quantization_config_is_auto_detected", False): + quant_config = None + return create_transformer_from_config( + config, + quant_config=quant_config, + component=component, + ) @property def guidance_scale(self): diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 01bd53bb8..1a75c40da 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -304,25 +304,10 @@ class Wan22I2VPipeline( # Transformers (weights loaded via load_weights) # Load config from model directory or HF Hub to get correct in_channels for I2V models transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = create_transformer_from_config( - transformer_config, - quant_config=od_config.quantization_config, - ) + self.transformer = self._create_transformer(transformer_config, component="transformer") if self.has_transformer_2: transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only) - t2_quant = transformer_2_config.get("quantization_config") - if isinstance(t2_quant, dict) and "quant_method" in t2_quant: - from vllm_omni.quantization.factory import build_quant_config - - method = t2_quant["quant_method"] - kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"} - t2_quant = build_quant_config(method, **kwargs) - else: - t2_quant = None - self.transformer_2 = create_transformer_from_config( - transformer_2_config, - quant_config=t2_quant, - ) + self.transformer_2 = self._create_transformer(transformer_2_config, component="transformer_2") else: self.transformer_2 = None @@ -462,10 +447,13 @@ class Wan22I2VPipeline( image_embeds = self.image_encoder(pixel_values, output_hidden_states=True) return image_embeds.hidden_states[-2] - def _create_transformer(self, config: dict) -> WanTransformer3DModel: + def _create_transformer(self, config: dict, component: str = "transformer") -> WanTransformer3DModel: """Create a transformer from a config dict. Respects od_config.quantization_config.""" quant_config = getattr(self.od_config, "quantization_config", None) - return create_transformer_from_config(config, quant_config=quant_config) + # Startup metadata describes the first expert, not a user policy for both. + if getattr(self.od_config, "quantization_config_is_auto_detected", False): + quant_config = None + return create_transformer_from_config(config, quant_config=quant_config, component=component) def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: sampling_params_list = req.sampling_params_list diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index aeecff608..0d5f3cd67 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -28,6 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( Wan22Pipeline, load_wan_weights_with_optional_gate, + resolve_wan_transformer_quant_config, retrieve_latents, ) from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( @@ -60,6 +61,8 @@ def create_vace_transformer_from_config( config: dict, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + component: str = "transformer", ) -> WanVACETransformer3DModel: """Create WanVACETransformer3DModel from config dict.""" kwargs = {} @@ -97,6 +100,7 @@ def create_vace_transformer_from_config( kwargs["vace_layers"] = config["vace_layers"] if "vace_in_channels" in config: kwargs["vace_in_channels"] = config["vace_in_channels"] + quant_config = resolve_wan_transformer_quant_config(config, quant_config, component) if quant_config is not None: kwargs["quant_config"] = quant_config if prefix: @@ -205,10 +209,13 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): super().__init__(od_config=od_config, prefix=prefix) - def _create_transformer(self, config: dict) -> WanVACETransformer3DModel: + def _create_transformer(self, config: dict, component: str = "transformer") -> WanVACETransformer3DModel: """Build VACE transformer. Respects od_config.quantization_config.""" quant_config = getattr(self.od_config, "quantization_config", None) - return create_vace_transformer_from_config(config, quant_config=quant_config) + # Startup metadata describes the first expert, not a user policy for both. + if getattr(self.od_config, "quantization_config_is_auto_detected", False): + quant_config = None + return create_vace_transformer_from_config(config, quant_config=quant_config, component=component) def diffuse( self, diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index dbc555604..2808dcee7 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -18,7 +18,7 @@ from vllm.distributed import ( ) from vllm.logger import init_logger from vllm.model_executor.layers.conv import Conv3dLayer -from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase, QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.utils import ( @@ -45,6 +45,7 @@ from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.mxfp4_config import NPUMxfp4LinearMethod logger = init_logger(__name__) @@ -985,6 +986,68 @@ class WanTransformer3DModel(nn.Module): # ROPE helper self._cached_rope_emb = None self._cached_rope_resolution = None + self._validate_w4a8_fallback_layers(quant_config) + + def _validate_w4a8_fallback_layers(self, quant_config: QuantizationConfig | None) -> None: + """Reject checkpoint aliases/typos before loading weights or running inference.""" + requested = tuple(getattr(quant_config, "w4a8_fallback_layers", [])) + self._w4a8_fallback_layers = requested + self._local_w4a8_fallback_layers: tuple[str, ...] = () + if not requested: + return + modules = dict(self.named_modules()) + missing_pp_prefixes = [name for name, module in modules.items() if isinstance(module, PPMissingLayer)] + local_layers = [ + name for name in requested if not any(name.startswith(prefix + ".") for prefix in missing_pp_prefixes) + ] + self._local_w4a8_fallback_layers = tuple(local_layers) + invalid = [name for name in local_layers if not isinstance(modules.get(name), LinearBase)] + if invalid: + raise ValueError( + "w4a8_fallback_layers requires exact runtime Linear paths relative to each Wan transformer; " + f"unknown/non-Linear paths: {invalid}. Use attn1.to_qkv for fused self-attention, " + "not checkpoint attn1.to_q/to_k/to_v names." + ) + + def _log_w4a8_fallback_load_summary(self, component_name: str) -> None: + """Log post-load transforms/cache readiness, not numerical acceptance.""" + requested = self._w4a8_fallback_layers + if not requested: + return + modules = dict(self.named_modules()) + local_layers = self._local_w4a8_fallback_layers + selected_layers = [ + name for name in local_layers if getattr(modules[name].quant_method, "is_w4a8_fallback_layer", False) + ] + bf16_overrides = [name for name in local_layers if name not in selected_layers] + processed = [ + name + for name in selected_layers + if getattr(modules[name], "_already_called_process_weights_after_loading", False) + ] + if not selected_layers: + processing_state = "not-required" + elif len(processed) == len(selected_layers): + processing_state = "ready" + else: + processing_state = "not-ready" + + def qualified(names: Iterable[str]) -> list[str]: + return [f"{component_name}.{name}" for name in names] + + logger.info( + "Wan W4A8 post-load summary: component=%s; selected=%d %s; BF16 overrides=%d %s; " + "weight/A8 processing state=%s (%d/%d selected layers processed); other PP ranks=%d", + component_name, + len(selected_layers), + qualified(selected_layers), + len(bf16_overrides), + qualified(bf16_overrides), + processing_state, + len(processed), + len(selected_layers), + len(requested) - len(local_layers), + ) @property def dtype(self) -> torch.dtype: @@ -1155,7 +1218,22 @@ class WanTransformer3DModel(nn.Module): } params_dict = dict(self.named_parameters()) + single_scale_prefixes = { + name + for name, module in self.named_modules() + if isinstance(getattr(module, "quant_method", None), NPUMxfp4LinearMethod) + } + + def reject_dualscale_tensor(name: str) -> None: + if name.endswith(".weight_dual_scale") and name.rsplit(".", 1)[0] in single_scale_prefixes: + raise ValueError( + f"Single-scale mxfp4 cannot load DualScale tensor {name}; " + "changing quant_method does not convert a checkpoint." + ) + loaded_params: set[str] = set() + loaded_qkv_shards: dict[str, set[str]] = {} + qkv_smooth_scales: dict[str, torch.Tensor] = {} for name, loaded_weight in weights: name = weight_name_remapping.get(name, name) @@ -1166,15 +1244,32 @@ class WanTransformer3DModel(nn.Module): # Pre-fused to_qkv tensors (from offline MXFP8 merged checkpoint) fall # through to the else branch and are loaded directly. for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in original_name: + if f"{weight_name}." not in original_name: continue - lookup_name = original_name.replace(weight_name, param_name) + lookup_name = original_name.replace(f"{weight_name}.", f"{param_name}.", 1) + reject_dualscale_tensor(lookup_name) # Skip weights that belong to PP stages other than this one if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict: break param = params_dict[lookup_name] + if lookup_name.endswith(".mul_scale"): + previous_scale = qkv_smooth_scales.get(lookup_name) + if previous_scale is not None and not torch.equal(previous_scale, loaded_weight): + raise ValueError(f"Fused Q/K/V must share the same Smooth tensor: {lookup_name}") + if previous_scale is None: + qkv_smooth_scales[lookup_name] = loaded_weight.clone() weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) + loaded_params.add(original_name) + if getattr(param, "output_dim", None) is None and not getattr(param, "needs_scalar_to_array", False): + # Shared input-channel Smooth is a complete tensor even + # when supplied under a single Q/K/V source name. + loaded_params.add(lookup_name) + else: + shards = loaded_qkv_shards.setdefault(lookup_name, set()) + shards.add(shard_id) + if shards == {"q", "k", "v"}: + loaded_params.add(lookup_name) break else: # diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight @@ -1198,6 +1293,7 @@ class WanTransformer3DModel(nn.Module): if is_pp_missing_parameter(lookup_name, self): continue + reject_dualscale_tensor(lookup_name) if lookup_name not in params_dict: logger.warning(f"Skipping weight {original_name} -> {lookup_name}") continue @@ -1222,8 +1318,6 @@ class WanTransformer3DModel(nn.Module): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) - - loaded_params.add(original_name) - loaded_params.add(lookup_name) + loaded_params.update((original_name, lookup_name)) return loaded_params diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 8f15aaf50..c0975f651 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -446,6 +446,33 @@ def _disk_marks_serialized(qc_kwargs: dict[str, Any], quant_config: object) -> b return False +_MXFP4_SERIALIZED_FLAGS = { + "mxfp4": "is_checkpoint_mxfp4_serialized", + "mxfp4_dualscale": "is_checkpoint_serialized", +} + + +def _normalize_serialized_mxfp4_layer_policy(qc_method: str, qc_kwargs: dict[str, Any]) -> None: + """Canonicalize the BF16 layer policy owned by serialized MXFP4 metadata. + + A serialized checkpoint config is a complete storage contract: an omitted + layer list means no explicit BF16 layers. Method-only online metadata does + not carry that contract and must leave caller-provided layer policy intact. + """ + method = _normalize_quant_method_alias(qc_method) + if method is None: + return + serialized_flag = _MXFP4_SERIALIZED_FLAGS.get(method) + if serialized_flag is None or not qc_kwargs.get(serialized_flag): + return + + ignored_layers = qc_kwargs.get("ignored_layers") + if not ignored_layers: + ignored_layers = qc_kwargs.get("modules_to_not_convert") + qc_kwargs["ignored_layers"] = ignored_layers or [] + qc_kwargs.pop("modules_to_not_convert", None) + + def resolve_quant_config_from_disk( quant_config: QuantizationConfig | None, disk_qc: dict[str, Any] | str | None, @@ -460,22 +487,24 @@ def resolve_quant_config_from_disk( - 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). + - Serialized MXFP4 metadata owns the complete BF16 layer policy; omitted + ignored_layers means the checkpoint default (no explicit BF16 layers). + - Method-only metadata carries no storage policy and preserves caller + ignored_layers and runtime fallback policies. """ 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 + qc_method = disk_qc + qc_kwargs: dict[str, Any] = {} + else: + if not isinstance(disk_qc, Mapping) or "quant_method" not in disk_qc: + return quant_config + qc_method = disk_qc["quant_method"] + qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"} - 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"} + _normalize_serialized_mxfp4_layer_policy(qc_method, qc_kwargs) if quant_config is None: logger.info( @@ -494,6 +523,26 @@ def resolve_quant_config_from_disk( "Pass a matching --quantization flag or omit it for auto-detection." ) + if isinstance(disk_qc, str): + return quant_config + + # Serialized checkpoint metadata owns storage and BF16 layer routing. The + # active config owns runtime step/layer policies, including explicit empty + # lists. Preserve them when each cascade transformer rebuilds from disk. + for policy in ("w4a8_fallback_steps", "w4a8_fallback_layers"): + if hasattr(quant_config, policy): + qc_kwargs[policy] = list(getattr(quant_config, policy)) + if hasattr(quant_config, "mxfp4_scale_alg"): + # Runtime activation policy is independent of how offline W4 was made. + qc_kwargs["mxfp4_scale_alg"] = quant_config.mxfp4_scale_alg + if hasattr(quant_config, "require_smooth_scale"): + disk_requires_smooth = qc_kwargs.get("require_smooth_scale", False) + if not isinstance(disk_requires_smooth, bool): + raise ValueError("require_smooth_scale must be a boolean") + # Either a calibrated checkpoint or an explicit caller can require + # Smooth. Per-expert rebuilds must never weaken that requirement. + qc_kwargs["require_smooth_scale"] = quant_config.require_smooth_scale or disk_requires_smooth + if _disk_marks_serialized(qc_kwargs, quant_config): logger.info( "config.json marks checkpoint as serialized; switching to offline %s mode.", @@ -501,6 +550,14 @@ def resolve_quant_config_from_disk( ) return build_quant_config(qc_method, **qc_kwargs) + if ( + "require_smooth_scale" in qc_kwargs + and hasattr(quant_config, "require_smooth_scale") + and qc_kwargs["require_smooth_scale"] != quant_config.require_smooth_scale + ): + logger.info("config.json Smooth requirement differs from active config; rebuilding quant_config.") + return build_quant_config(qc_method, **qc_kwargs) + # AutoRound MXFP checkpoints use data_type="mx_fp" instead of # is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is # selected according to the checkpoint's bit width. diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py index ac93eaa00..43f3efa9d 100644 --- a/vllm_omni/quantization/mxfp4_config.py +++ b/vllm_omni/quantization/mxfp4_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers. Architecture mirrors mxfp8_config.py: @@ -41,7 +41,7 @@ Dual-scale (W4A4_MXFP4_DUALSCALE): 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 : (N, K) float8_e4m3fn – numeric FP4 values (packed after loading) 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 @@ -51,9 +51,9 @@ Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/q from __future__ import annotations -import re -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast +import regex as re import torch from torch.nn import Module from vllm.logger import init_logger @@ -83,9 +83,54 @@ from vllm_omni.quantization.mxfp8_config import ( if TYPE_CHECKING: from vllm.model_executor.models.utils import WeightsMapper + logger = init_logger(__name__) +def _validate_w4a8_fallback_steps(steps: list[int] | None) -> list[int]: + if steps is None: + return [] + if not isinstance(steps, list) or any(type(step) is not int or step < 0 for step in steps): + raise ValueError("w4a8_fallback_steps must be a list of non-negative integer denoise step indices.") + return sorted(set(steps)) + + +def _validate_w4a8_fallback_layers(layers: list[str] | None) -> list[str]: + """Names are exact runtime Linear paths, not checkpoint aliases or patterns.""" + if layers is None: + return [] + if not isinstance(layers, list) or any( + not isinstance(name, str) or re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*(?:\.[A-Za-z_0-9]+)*", name) is None + for name in layers + ): + raise ValueError("w4a8_fallback_layers must be a list of non-empty exact runtime Linear paths.") + return sorted(set(layers)) + + +def _is_w4a8_fallback_step(steps: list[int]) -> bool: + """Read the request's zero-based denoise index; never keep state on a Linear.""" + if not steps: + return False + + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + step = get_forward_context().denoise_step_idx if is_forward_context_available() else None + if step is None: + raise RuntimeError("W4A8 step fallback requires a pipeline that publishes its denoise step in ForwardContext.") + return step in steps + + +def _validate_smooth_scale(scale: torch.Tensor) -> None: + """Reject malformed calibration before loading or activation dtype conversion.""" + if scale.ndim != 1 or not scale.is_floating_point(): + raise ValueError("MXFP4 mul_scale must be a one-dimensional floating-point Smooth tensor.") + if not bool(torch.isfinite(scale).all()) or not bool((scale > 0).all()): + raise ValueError("MXFP4 mul_scale must contain finite, strictly positive Smooth values.") + + # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @@ -105,10 +150,24 @@ class DiffusionMXFP4Config(QuantizationConfig): self, is_checkpoint_mxfp4_serialized: bool = False, ignored_layers: list[str] | None = None, + w4a8_fallback_steps: list[int] | None = None, + require_smooth_scale: bool = False, + w4a8_fallback_layers: list[str] | None = None, + mxfp4_scale_alg: int = 0, ) -> None: super().__init__() self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized self.ignored_layers = ignored_layers or [] + self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps) + self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers) + if not isinstance(require_smooth_scale, bool): + raise ValueError("require_smooth_scale must be a boolean") + if require_smooth_scale and not self.is_checkpoint_mxfp4_serialized: + raise ValueError("require_smooth_scale requires an offline MXFP4 checkpoint") + self.require_smooth_scale = require_smooth_scale + if type(mxfp4_scale_alg) is not int or mxfp4_scale_alg not in (0, 2): + raise ValueError("mxfp4_scale_alg must be 0 (OCP MX) or 2 (UOS with dst_type_max=7.25).") + self.mxfp4_scale_alg = mxfp4_scale_alg @classmethod def get_name(cls) -> QuantizationMethods: @@ -139,6 +198,10 @@ class DiffusionMXFP4Config(QuantizationConfig): return cls( is_checkpoint_mxfp4_serialized=is_serialized, ignored_layers=ignored_layers, + w4a8_fallback_steps=config.get("w4a8_fallback_steps"), + w4a8_fallback_layers=config.get("w4a8_fallback_layers"), + require_smooth_scale=config.get("require_smooth_scale", False), + mxfp4_scale_alg=config.get("mxfp4_scale_alg", 0), ) def get_quant_method( @@ -155,9 +218,15 @@ class DiffusionMXFP4Config(QuantizationConfig): return UnquantizedLinearMethod() if current_omni_platform.is_npu(): if self.is_checkpoint_mxfp4_serialized: - return NPUMxfp4LinearMethod(self) - return NPUMxfp4OnlineLinearMethod(self) + return NPUMxfp4LinearMethod(self, prefix=prefix) + return NPUMxfp4OnlineLinearMethod(self, prefix=prefix) if current_omni_platform.is_rocm(): + if self.mxfp4_scale_alg != 0: + raise NotImplementedError("MXFP4 UOS quantization is currently only supported on NPU (Ascend).") + if self.w4a8_fallback_steps or self.w4a8_fallback_layers: + raise NotImplementedError( + "MXFP4 W4A8 step/layer fallback is currently only supported on NPU (Ascend)." + ) gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName if "gfx950" not in gcn_arch: raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}") @@ -176,6 +245,56 @@ class DiffusionMXFP4Config(QuantizationConfig): # --------------------------------------------------------------------------- +def _npu_quantize_mxfp4(x: torch.Tensor, scale_alg: int) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight or activation using the same explicit FP4 algorithm.""" + import torch_npu + + # Universal Optimal Scaling (Qmax=7.25), MXAttention Section 4.1: + # https://arxiv.org/abs/2607.24377 + scale_kwargs = {"dst_type_max": 7.25} if scale_alg == 2 else {} + return torch_npu.npu_dynamic_mx_quant( + x, + dst_type=torch_npu.float4_e2m1fn_x2, + axis=-1, + block_size=32, + round_mode="rint", + scale_alg=scale_alg, + **scale_kwargs, + ) + + +def _npu_w4a8_matmul( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + output_dtype: torch.dtype, +) -> torch.Tensor: + """MXFP8 activations with prepared single-scale MXFP4 weights (N, K).""" + import torch_npu + + x_q, x_scale = torch_npu.npu_dynamic_mx_quant( + x, dst_type=torch_npu.float8_e4m3fn, axis=-1, block_size=32, round_mode="rint", scale_alg=0 + ) + # Use the BF16 output/bias pair supported by the mixed FP8/FP4 kernel + # in the target A5 runtime. Restore the caller's dtype afterwards. + if bias is not None: + bias = bias.to(torch.bfloat16).reshape(1, -1) + output = torch_npu.npu_quant_matmul( + x_q, + weight.transpose(0, 1), + weight_scale.transpose(0, 1), + scale_dtype=torch_npu.float8_e8m0fnu, + x2_dtype=torch_npu.float4_e2m1fn_x2, + pertoken_scale=x_scale, + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, + bias=bias, + output_dtype=torch.bfloat16, + group_sizes=[0, 0, 32], + ) + return output.to(output_dtype) + + class NPUMxfp4LinearMethod(MXFPLinearMethodBase): """NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints. @@ -187,7 +306,8 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared. """ - def __init__(self, quant_config: DiffusionMXFP4Config) -> None: + def __init__(self, quant_config: DiffusionMXFP4Config, *, prefix: str = "") -> None: + self.is_w4a8_fallback_layer = prefix in quant_config.w4a8_fallback_layers self.quant_config = quant_config self.out_dtype = torch.get_default_dtype() @@ -210,11 +330,14 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): 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), + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=params_dtype, + ), input_dim=1, output_dim=0, weight_loader=weight_loader, @@ -234,12 +357,42 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): ), ) + # Older single-scale checkpoints have no Smooth tensor. Keep their + # identity transform while allowing calibrated checkpoints to load it. + def smooth_weight_loader( + param: ModelWeightParameter, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any + ) -> None: + _validate_smooth_scale(loaded_weight) + if not callable(weight_loader): + raise ValueError("Loading MXFP4 Smooth requires a callable weight_loader") + weight_loader(param, loaded_weight, *args, **kwargs) + + layer.register_parameter( + "mul_scale", + ModelWeightParameter( + data=torch.ones(input_size_per_partition, dtype=torch.float32), + input_dim=0, + output_dim=None, + weight_loader=smooth_weight_loader, + ), + ) + setattr(layer.mul_scale, "ignore_warning", True) + setattr(layer.mul_scale, "is_checkpoint_optional", not self.quant_config.require_smooth_scale) + if self.quant_config.is_checkpoint_mxfp4_serialized: + setattr(layer.weight_scale, "is_checkpoint_required", True) + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return import torch_npu + smooth_param = cast(torch.Tensor, layer.mul_scale) + _validate_smooth_scale(smooth_param) + smooth = smooth_param.data.to(dtype=cast(torch.dtype, layer.orig_dtype)).contiguous() + # Finite FP32 values can overflow or underflow in the activation dtype. + _validate_smooth_scale(smooth) + # NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose. w = layer.weight if w.dtype != torch_npu.float4_e2m1fn_x2: @@ -258,15 +411,31 @@ class NPUMxfp4LinearMethod(MXFPLinearMethodBase): replace_parameter(layer, "weight", w) replace_parameter(layer, "weight_scale", s) + # Keep loader metadata while converting the calibration tensor once. + layer.mul_scale.data = smooth.to(device=cast(torch.Tensor, w).device) 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 + def _apply_inner( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + # Online quantization has no calibrated Smooth tensor. For offline + # weights the same pre-scale must precede both A4 and A8 quantization. + if (mul_scale := getattr(layer, "mul_scale", None)) is not None: + x = x * mul_scale.to(x.dtype) + if self.is_w4a8_fallback_layer or _is_w4a8_fallback_step(self.quant_config.w4a8_fallback_steps): + return _npu_w4a8_matmul( + x, cast(torch.Tensor, layer.weight), cast(torch.Tensor, layer.weight_scale), bias, ori_dtype + ) + return super()._apply_inner(layer, x, bias, ori_dtype) - # No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2. - return torch_npu.npu_dynamic_mx_quant(x) + def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return _npu_quantize_mxfp4(x, self.quant_config.mxfp4_scale_alg) def _quant_matmul( self, @@ -318,8 +487,6 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): 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), @@ -331,8 +498,10 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod): 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) + # Prepare once; offline checkpoints keep their existing weight algorithm. + weight_fp4, weight_scale_raw = _npu_quantize_mxfp4( + cast(torch.Tensor, layer.weight), self.quant_config.mxfp4_scale_alg + ) # 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. @@ -497,7 +666,7 @@ 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 : (N, K) float8_e4m3fn – numeric FP4 values; packed into NZ after loading 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) @@ -509,7 +678,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): Reference: MindIE-SD W4A4MXFP4DualQuantLinear. """ - def __init__(self, quant_config: Any) -> None: + def __init__(self, quant_config: DiffusionMXFP4DualScaleMixedConfig) -> None: self.quant_config = quant_config self.out_dtype = torch.get_default_dtype() @@ -532,7 +701,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): 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 + # Checkpoint FP4 values are stored numerically in FP8; BF16 preserves them exactly. layer.register_parameter( "weight", ModelWeightParameter( @@ -587,6 +756,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): ), ) setattr(layer.mul_scale, "ignore_warning", True) + setattr(layer.weight_scale, "is_checkpoint_required", True) def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): @@ -594,7 +764,7 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): import torch_npu - # float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29). + # Numeric FP4 values → 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) @@ -631,7 +801,10 @@ class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase): """ 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) + # The dual-level activation kernel requires x and Smooth to have the + # same dtype. Prepared calibration scales are BF16, including online. + smooth_scale = cast(torch.Tensor, layer.mul_scale).to(x.dtype) + x_q, l0_scale, l1_scale = self._quantize_activation(x, smooth_scale) return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype) def _quantize_activation( # type: ignore[override] @@ -785,11 +958,20 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig): is_checkpoint_serialized: bool = False, ignored_layers: list[str] | None = None, num_bf16_fallback_layers: int = 5, + w4a8_fallback_steps: list[int] | None = None, + w4a8_fallback_layers: list[str] | None = None, ) -> 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 + self.w4a8_fallback_steps = _validate_w4a8_fallback_steps(w4a8_fallback_steps) + self.w4a8_fallback_layers = _validate_w4a8_fallback_layers(w4a8_fallback_layers) + if self.w4a8_fallback_steps or self.w4a8_fallback_layers: + raise ValueError( + "mxfp4_dualscale does not support W4A8 fallback. Use mxfp4 with single-scale weights; " + "changing a DualScale checkpoint's method name does not convert its storage format." + ) @classmethod def get_name(cls) -> QuantizationMethods: @@ -822,6 +1004,8 @@ class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig): is_checkpoint_serialized=is_serialized, ignored_layers=ignored_layers, num_bf16_fallback_layers=num_bf16_fallback_layers, + w4a8_fallback_steps=config.get("w4a8_fallback_steps"), + w4a8_fallback_layers=config.get("w4a8_fallback_layers"), ) def get_quant_method(