feat: support Ascend W8A8-MXFP8 rollout#355
Conversation
Documentation build overview
45 files changed ·
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements online W8A8-MXFP8 rollout weight updates for Ascend 950 in both colocated IPC and decoupled NCCL deployments. It introduces a lazy-imported Ascend utility module and a vLLM worker extension (AscendMXFP8WorkerExtension) to handle on-demand quantization of BF16 weights using torch_npu.npu_dynamic_mx_quant and coordinate layout preparation and finalization. The Vime engine orchestration is updated to inject this extension and manage the update lifecycle. Comprehensive unit tests and bilingual documentation are also added. Feedback is provided regarding a potential shape mismatch bug in quantize_mxfp8_weights where flattening and squeezing the scale tensor can over-squeeze it when the number of groups is exactly 1; simplifying this to scale.squeeze(-1) is recommended.
| axis=-1, | ||
| dst_type=torch_npu.float8_e4m3fn, | ||
| ) | ||
| scale = scale.flatten(-2, -1).squeeze(-1) |
There was a problem hiding this comment.
Using scale.flatten(-2, -1).squeeze(-1) can lead to a shape mismatch bug when the number of groups along the quantized dimension is exactly 1 (i.e., N // 32 == 1 or M // 32 == 1).
In that case:
scalestarts with shape(..., 1, 1).scale.flatten(-2, -1)flattens the last two dimensions to(..., 1).scale.squeeze(-1)removes the remaining dimension of size 1, resulting in(...)(e.g., a 1D tensor instead of a 2D tensor).
This unexpected reduction in dimensionality will cause shape mismatches or errors during layout restoration or inference in vllm-ascend. Since npu_dynamic_mx_quant always appends a singleton dimension of size 1 at the end of the scale tensor, simply calling scale.squeeze(-1) is both safer and cleaner.
| scale = scale.flatten(-2, -1).squeeze(-1) | |
| scale = scale.squeeze(-1) |
1 participant
Summary
MXFP4 and external vLLM servers are intentionally outside this change. The implementation targets local vllm-ascend/main at commit
8bedc666.Testing
101 passed, 1 warningacross MXFP8, vLLM engine, IPC, and NCCL focused testsgit diff --checkpassedReal Ascend 950 hardware validation remains required; the documented acceptance procedure covers both IPC and NCCL modes.