✨ [diffusion][npu][quant] Add MXFP8 quantization support for Wan2.2 Diffusion on Ascend NPU (#20922)
Co-authored-by: ronnie_zheng <zl19940307@163.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Junlin Wu
co-authored by
ronnie_zheng
github-actions[bot]
ronnie_zheng
github-actions[bot]
parent
7d397ad23d
commit
80a6014243
@@ -43,6 +43,7 @@ The following table summarizes quantization method support across NVIDIA and AMD
|
||||
| `bitsandbytes` | Yes | Experimental | No | Depends on bitsandbytes ROCm support |
|
||||
| `torchao` (`int4wo`, etc.) | Yes | Partial | No | `int4wo` not supported on AMD; other methods may work |
|
||||
| `modelslim` | No | No | Yes | Ascend quantization; Uses CANN kernels |
|
||||
| `mxfp8` (diffusion) | No | No | Yes (A2/A3) | Ascend NPU only; online MXFP8 quantization for diffusion models (e.g., Wan2.2); requires CANN ≥ 8.0.RC3 |
|
||||
|
||||
On AMD, several of these methods use [Aiter](https://github.com/ROCm/aiter) for acceleration -- set `SGLANG_USE_AITER=1` where noted. See [AMD GPU setup](../platforms/amd_gpu.md) for installation and configuration details.
|
||||
|
||||
@@ -590,6 +591,36 @@ SGLang running on AMD GPUs (CDNA3 or CDNA4 architecture) supports the quantizati
|
||||
|
||||
Other layers (e.g. projections in the attention layers) have their weights quantized online to float8 directly.
|
||||
|
||||
## Diffusion Model Quantization on Ascend NPU
|
||||
|
||||
SGLang-Diffusion supports MXFP8 quantization for diffusion models (such as Wan2.2) on Ascend A5 NPUs, in both online and offline (ModelSlim) modes. This is separate from the LLM serving path and uses the `sglang serve` / `sglang generate` CLI.
|
||||
|
||||
**Requirements:** Ascend A5, CANN ≥ 8.0.RC3
|
||||
|
||||
### Online MXFP8
|
||||
|
||||
Pass `--quantization mxfp8` to dynamically quantize FP16/BF16 transformer weights to MXFP8 at load time:
|
||||
|
||||
```bash
|
||||
sglang serve \
|
||||
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--quantization mxfp8 \
|
||||
--num-gpus 4
|
||||
```
|
||||
|
||||
### Offline MXFP8 (ModelSlim)
|
||||
|
||||
Pre-quantize with [msModelSlim](https://gitcode.com/Ascend/msmodelslim) and load the checkpoint directly — the quantization scheme is auto-detected from `quant_model_description.json`:
|
||||
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path /path/to/wan2_2_mxfp8_diffusers \
|
||||
--prompt "a beautiful sunset" \
|
||||
--save-output
|
||||
```
|
||||
|
||||
For the full quantization + format conversion workflow and a complete list of supported schemes, see [Diffusion Quantization on Ascend NPU](../platforms/ascend/ascend_npu_quantization.md#diffusion-model-quantization-on-ascend-npu) and [SGLang-Diffusion Quantization](../diffusion/quantization.md#modelslim).
|
||||
|
||||
## Reference
|
||||
|
||||
- [GPTQModel](https://github.com/ModelCloud/GPTQModel)
|
||||
|
||||
@@ -395,4 +395,4 @@ MindStudio-ModelSlim (msModelSlim) is a model offline quantization compression t
|
||||
- [x] ```W4A4_DYNAMIC``` linear with online quantization of activations
|
||||
- [x] ```W8A8``` linear with offline quantization of activations
|
||||
- [x] ```W8A8_DYNAMIC``` linear with online quantization of activations
|
||||
- [ ] ```mxfp8``` linear in progress
|
||||
- [x] ```mxfp8``` linear with online/offline MXFP8 quantization (Ascend A5, CANN ≥ 8.0.RC3; see [Ascend NPU quantization](../platforms/ascend/ascend_npu_quantization.md#diffusion-model-quantization-on-ascend-npu))
|
||||
|
||||
@@ -5,16 +5,16 @@ To load already quantized models, simply load the model weights and config. Agai
|
||||
SGLang support **mix-bits** quantization (independently defines and loads each layer depending on the type of quantification specified in the `quant_model_description'.json`). [Advanced mix-bits for MoE](https://github.com/sgl-project/sglang/pull/17361) in progress, will add independent quantization determination for the w13 (up-gate) and w2 (down) layers.
|
||||
|
||||
[ModelSlim on Ascend support](https://github.com/sgl-project/sglang/pull/14504)
|
||||
| Quantization scheme | Layer type | A2 Supported | A3 Supported | A5 Supported | Diffusion models |
|
||||
|-----------------------------------------------------------|--------------------------|:----------------------------------------:|:----------------------------------------:|:------------------------------------------:|:------------------------------------------:|
|
||||
| W4A4 dynamic | Linear | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: green;">√</span>** |
|
||||
| W8A8 static | Linear | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: green;">√</span>** |
|
||||
| W8A8 dynamic | Linear | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: green;">√</span>** |
|
||||
| [MXFP8](https://github.com/sgl-project/sglang/pull/20922) | Linear | **<span style="color: red;">x</span>** | **<span style="color: red;">x</span>** | **<span style="color: blue;">WIP</span>** | **<span style="color: blue;">WIP</span>** |
|
||||
| W4A4 dynamic | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: red;">x</span>** |
|
||||
| W4A8 dynamic | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: red;">x</span>** |
|
||||
| W8A8 dynamic | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: red;">x</span>** |
|
||||
| [MXFP8](https://github.com/sgl-project/sglang/pull/20922) | MoE | **<span style="color: red;">x</span>** | **<span style="color: red;">x</span>** | **<span style="color: blue;">WIP</span>** | **<span style="color: red;">x</span>** |
|
||||
| Quantization scheme | `quant_type` in JSON | Scheme class | Layer type | A2 Supported | A3 Supported | A5 Supported | Diffusion models |
|
||||
|-----------------------------------------------------------|----------------------|--------------------------|--------------------------|:----------------------------------------:|:----------------------------------------:|:------------------------------------------:|:------------------------------------------:|
|
||||
| W4A4 dynamic | `W4A4_DYNAMIC` | `ModelSlimW4A4Int4` | Linear | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: green;">√</span>** |
|
||||
| W8A8 static | `W8A8` | `ModelSlimW8A8Int8` | Linear | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: green;">√</span>** |
|
||||
| W8A8 dynamic | `W8A8_DYNAMIC` | `ModelSlimW8A8Int8` | Linear | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: green;">√</span>** |
|
||||
| [MXFP8](https://github.com/sgl-project/sglang/pull/20922) | `W8A8_MXFP8` | `ModelSlimMXFP8Scheme` | Linear | **<span style="color: red;">x</span>** | **<span style="color: red;">x</span>** | **<span style="color: blue;">WIP</span>** | **<span style="color: green;">√</span>** (A5) |
|
||||
| W4A4 dynamic | `W4A4_DYNAMIC` | `ModelSlimW4A4Int4` | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: red;">x</span>** |
|
||||
| W4A8 dynamic | `W4A8_DYNAMIC` | `ModelSlimW4A8Int8MoE` | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: red;">x</span>** |
|
||||
| W8A8 dynamic | `W8A8_DYNAMIC` | `ModelSlimW8A8Int8` | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | **<span style="color: red;">x</span>** |
|
||||
| [MXFP8](https://github.com/sgl-project/sglang/pull/20922) | `W8A8_MXFP8` | `ModelSlimMXFP8Scheme` | MoE | **<span style="color: red;">x</span>** | **<span style="color: red;">x</span>** | **<span style="color: blue;">WIP</span>** | **<span style="color: red;">x</span>** |
|
||||
|
||||
[AWQ on Ascend support](https://github.com/sgl-project/sglang/pull/10158):
|
||||
| Quantization scheme | Layer type | A2 Supported | A3 Supported | A5 Supported |
|
||||
@@ -54,3 +54,81 @@ Compressed-tensors (LLM Compressor) on Ascend support:
|
||||
| [GGUF (all types)](https://github.com/sgl-project/sglang/pull/17883) | MoE | **<span style="color: green;">√</span>** | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** |
|
||||
|
||||
> Note: On Ascend, GGUF weights are pre-dequantized to FP16/BF16 during model loading to ensure optimal inference performance. This enables support for all GGUF quantization types (Q2_K, Q4_K_M, IQ4_XS, etc.) while maintaining high inference speed.
|
||||
|
||||
in progress
|
||||
|
||||
## Diffusion Model Quantization on Ascend NPU
|
||||
|
||||
SGLang-Diffusion supports MXFP8 online and offline quantization for diffusion models (such as Wan2.2) on Ascend NPUs. MXFP8 requires A5; the ModelSlim W8A8/W4A4 schemes work on A2/A3.
|
||||
|
||||
**Requirements for MXFP8:** CANN ≥ 8.0.RC3, Ascend A5
|
||||
|
||||
| Quantization method | `quant_type` in JSON | Scheme class | Mode | A2/A3 Supported | A5 Supported | Trigger |
|
||||
|---------------------|-----------------------|-------------------------------|---------|:--------------------------------------------:|:----------------------------------------:|---------------------------------------------------|
|
||||
| MXFP8 (W8A8) | — | `MXFP8Config` | Online | **<span style="color: red;">x</span>** | **<span style="color: green;">√</span>** | `--quantization mxfp8` |
|
||||
| MXFP8 (W8A8) | `W8A8_MXFP8` | `ModelSlimMXFP8Scheme` | Offline | **<span style="color: red;">x</span>** | **<span style="color: green;">√</span>** | auto-detected from `quant_model_description.json` |
|
||||
| W8A8 static | `W8A8` | `ModelSlimW8A8Int8` | Offline | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | auto-detected from `quant_model_description.json` |
|
||||
| W8A8 dynamic | `W8A8_DYNAMIC` | `ModelSlimW8A8Int8` | Offline | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | auto-detected from `quant_model_description.json` |
|
||||
| W4A4 dynamic | `W4A4_DYNAMIC` | `ModelSlimW4A4Int4` | Offline | **<span style="color: green;">√</span>** | **<span style="color: yellow;">TBD</span>** | auto-detected from `quant_model_description.json` |
|
||||
|
||||
### Online MXFP8 Quantization
|
||||
|
||||
Online quantization dynamically quantizes FP16/BF16 weights to MXFP8 at load time using `npu_dynamic_mx_quant` + `npu_quant_matmul` CANN kernels. Pass `--quantization mxfp8` to override auto-detection.
|
||||
|
||||
```bash
|
||||
# Start the diffusion server with online MXFP8 quantization
|
||||
sglang serve \
|
||||
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--quantization mxfp8 \
|
||||
--num-gpus 4
|
||||
```
|
||||
|
||||
```bash
|
||||
# One-shot generation
|
||||
sglang generate \
|
||||
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--quantization mxfp8 \
|
||||
--prompt "a beautiful sunset over the mountains" \
|
||||
--save-output
|
||||
```
|
||||
|
||||
### Offline MXFP8 Quantization (ModelSlim)
|
||||
|
||||
For offline quantization, pre-quantize the model with msModelSlim and load the resulting checkpoint. The quantization scheme is auto-detected from `quant_model_description.json`, so no extra `--quantization` flag is needed.
|
||||
|
||||
**Step 1: Quantize with msModelSlim**
|
||||
|
||||
```bash
|
||||
msmodelslim quant \
|
||||
--model_path /path/to/wan2_2_float_weights \
|
||||
--save_path /path/to/wan2_2_mxfp8_weights \
|
||||
--device npu \
|
||||
--model_type Wan2_2 \
|
||||
--quant_type mxfp8 \
|
||||
--trust_remote_code True
|
||||
```
|
||||
|
||||
> Note: SGLang does not support quantized embeddings; disable embedding quantization when using msmodelslim.
|
||||
|
||||
**Step 2: Convert to Diffusers format**
|
||||
|
||||
msModelSlim saves quantized Wan2.2 weights in the original Wan format. Convert to Diffusers format using the provided repack script:
|
||||
|
||||
```bash
|
||||
python python/sglang/multimodal_gen/tools/wan_repack.py \
|
||||
--input-path /path/to/wan2_2_mxfp8_weights \
|
||||
--output-path /path/to/wan2_2_mxfp8_diffusers
|
||||
```
|
||||
|
||||
Then copy all files from the original Diffusers checkpoint (except the `transformer`/`transformer_2` folders) into the output directory.
|
||||
|
||||
**Step 3: Run inference**
|
||||
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path /path/to/wan2_2_mxfp8_diffusers \
|
||||
--prompt "a beautiful sunset over the mountains" \
|
||||
--save-output
|
||||
```
|
||||
|
||||
For pre-quantized checkpoints available on ModelScope, see [modelscope/Eco-Tech](https://modelscope.cn/models/Eco-Tech).
|
||||
|
||||
Reference in New Issue
Block a user