From 0d186f49be34c12e133212023f4087f6b66635f1 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Sat, 1 Aug 2026 13:09:37 +0800 Subject: [PATCH] [AMD][Fix] Restore aiter-padded MoE weight dims for serialized checkpoints (#33090) AMD only change --- .../quark/schemes/quark_w4a4_mxfp4_moe.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index df97eeec3..5612bda1a 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -196,9 +196,18 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): extra_weight_attrs["weight_loader"] = weight_loader # WEIGHTS + # Serialized MXFP4 must keep the aiter-aligned dims from + # get_moe_weight_sizes(): w13_weight_scale / w2_weight_scale below are + # still sized off w13_up_dim / w2_down_dim, and `weight_padded` is + # advertised to the loader, so allocating the raw (unpadded) dims here + # desyncs the weights from their scales. w13_shape = ( num_experts, - 2 * intermediate_size_per_partition, + ( + w13_up_dim + if self.is_checkpoint_mxfp4_serialized + else 2 * intermediate_size_per_partition + ), hidden_size // 2 if self.is_checkpoint_mxfp4_serialized else hidden_size, ) w13_weight = torch.nn.Parameter( @@ -217,7 +226,7 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): num_experts, hidden_size, ( - intermediate_size_per_partition // 2 + w2_down_dim if self.is_checkpoint_mxfp4_serialized else intermediate_size_per_partition ),