Add MXFP4 (including Quark W4A4) quantization support for DeepSeek-architecture on ROCm - #21529
JohnQinAMD wants to merge 1 commit into
Conversation
… models Enable MXFP4 quantized models (e.g. GLM-5-MXFP4-Q8, Kimi-K2.5-MXFP4) to load and run on AMD MI355X with ROCm. These models use Quark's W4A4 MXFP4 scheme with block-scaled FP4 for MoE expert weights and BF16 for attention. Key changes: mxfp4.py: - Return None instead of raising NotImplementedError for attention layers (attention uses BF16 weights, not MXFP4) fused_moe_triton/layer.py: - Skip weight transpose for packed uint8 (MXFP4) tensors during loading - Pre-compute expert_mask_gpu in __init__ for CUDA graph compatibility - Fix expert_id None check (was `not expert_id` which fails for id=0) quark/: - Fix Quark MXFP4 MoE scheme weight creation and loading for block-scaled format with proper shard handling - Add per-expert weight processing support deepseek_v2.py: - Set packed_modules_mapping for Quark weight loader compatibility - Disable shared expert fusion for MXFP4 serialized models (routed experts are quantized but shared experts remain BF16) - Disable shared expert fusion for Quark models (mixed precision) deepseek_weight_loader.py: - Handle fused_qkv_a_proj_with_mqa weight mapping for Quark checkpoints Made-with: Cursor
|
Warning Gemini is experiencing higher than usual traffic and was unable to create the review. Please try again in a few hours by commenting |
BowenBao
left a comment
BowenBao
left a comment
There was a problem hiding this comment.
The deepseek mxfp4 models, such as https://huggingface.co/amd/DeepSeek-R1-MXFP4 and https://huggingface.co/amd/DeepSeek-R1-0528-MXFP4, are well supported by SGLang. I'm not sure what additional functionality this PR is intended to provide.
| else ActivationType.Gelu | ||
| ), | ||
| doweight_stage1=False, | ||
| activation=ActivationType.Swiglu, |
There was a problem hiding this comment.
is this correct to fix as swiglu?
| h, b, d = x.shape | ||
| x, x_scales = dynamic_mxfp4_quant(x.reshape(-1, d)) | ||
| return x.view(h, b, d // 2), x_scales.view(h, b, d // 32) | ||
| x_out = x.view(h, b, d // 2).contiguous() |
There was a problem hiding this comment.
are these necessary changes?
| layer.register_parameter("w2_weight_scale", w2_weight_scale) | ||
|
|
||
| # WEIGHT_BIAS (zeros, matching standard mxfp4 path for CUDA graph compat) | ||
| w13_weight_bias = torch.nn.Parameter( |
There was a problem hiding this comment.
not all models have biases, could we make this conditional?
| self.intermediate_pad = 0 | ||
|
|
||
| def process_weights_after_loading(self, layer: torch.nn.Module) -> None: | ||
| float_dtype = torch.get_default_dtype() |
There was a problem hiding this comment.
I'm not sure about the shuffling as well as activation changes.
|
Thanks @JohnQinAMD. Closing this because it has had no updates in 151 days. Reopen it if the work is still relevant. Some directories moved recently, so an older branch may need retargeting: |
4 participants
Enable MXFP4 (including Quark W4A4) quantized models (e.g. GLM-5-MXFP4-Q8, Kimi-K2.5-MXFP4) to load and run on AMD MI355X with ROCm. These models use Quark's W4A4 MXFP4 scheme with block-scaled FP4 for MoE expert weights and BF16 for attention.
Key changes:
mxfp4.py:
fused_moe_triton/layer.py:
not expert_idwhich fails for id=0)quark/:
deepseek_v2.py:
deepseek_weight_loader.py:
Motivation
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci