Fix MXFP4 scale placeholder initialization (#33500)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-08-06 01:46:54 -07:00
committed by GitHub
co-authored by weireweire
parent efc99a86ff
commit fe55d78b7d
+18 -8
View File
@@ -69,6 +69,10 @@ from sglang.srt.utils.custom_op import register_custom_op
has_triton_kernels = is_triton_kernels_available()
# Serialized MXFP4 scales use raw UE8M0 bytes. Keep fresh parameters valid for
# post-load transforms and dummy initialization; 127 is the neutral scale (1.0).
_UE8M0_ONE = 127
if is_flashinfer_available():
from flashinfer import (
@@ -472,10 +476,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
set_weight_attrs(w13_weight, extra_weight_attrs)
w13_weight_scale = torch.nn.Parameter(
torch.zeros(
layer.num_local_experts,
2 * intermediate_size_per_partition_after_pad,
hidden_size // mxfp4_block,
torch.full(
(
layer.num_local_experts,
2 * intermediate_size_per_partition_after_pad,
hidden_size // mxfp4_block,
),
fill_value=_UE8M0_ONE,
dtype=scale_dtype,
),
requires_grad=False,
@@ -511,10 +518,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
set_weight_attrs(w2_weight, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.zeros(
layer.num_local_experts,
hidden_size,
intermediate_size_per_partition_after_pad // mxfp4_block,
torch.full(
(
layer.num_local_experts,
hidden_size,
intermediate_size_per_partition_after_pad // mxfp4_block,
),
fill_value=_UE8M0_ONE,
dtype=scale_dtype,
),
requires_grad=False,