diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 0a3eccf57..e924d646c 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1552,6 +1552,13 @@ def requant_weight_ue8m0( ): assert weight_block_size == [128, 128] + # 3D+ weights stack multiple experts (e.g. MoE); requant each group separately. + # 2D weights are a single matrix and fall through to the direct path below. + if weight.dim() > 2: + return _requant_weight_ue8m0_grouped( + weight, weight_scale_inv, weight_block_size + ) + *_, n, k = weight.shape weight_dequant = block_quant_dequant( @@ -1571,6 +1578,47 @@ def requant_weight_ue8m0( return out_w, out_s +def _requant_weight_ue8m0_grouped( + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + weight_block_size: List[int], +): + *group_dims, n, k = weight.shape + w_groups = weight.reshape(-1, n, k) + s_groups = weight_scale_inv.reshape(-1, *weight_scale_inv.shape[-2:]) + num_groups = w_groups.shape[0] + + out_w = None + out_s = None + for g in range(num_groups): + weight_dequant = block_quant_dequant( + w_groups[g], + s_groups[g], + weight_block_size, + torch.bfloat16, + ) + w_g, s_g = quant_weight_ue8m0( + weight_dequant=weight_dequant, + weight_block_size=weight_block_size, + ) + if out_w is None: + out_w = torch.empty( + (num_groups, *w_g.shape), dtype=w_g.dtype, device=w_g.device + ) + out_s = torch.empty( + (num_groups, *s_g.shape), dtype=s_g.dtype, device=s_g.device + ) + out_w[g] = w_g + out_s[g] = s_g + + out_w = out_w.view(*group_dims, n, k) + out_s = out_s.view(*group_dims, *out_s.shape[-2:]) + + out_s = transform_scale_ue8m0(out_s, mn=n) + + return out_w, out_s + + def quant_weight_ue8m0( weight_dequant: torch.Tensor, weight_block_size: List[int],