Fix MXFP4 scale placeholder initialization (#33500)
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
@@ -69,6 +69,10 @@ from sglang.srt.utils.custom_op import register_custom_op
|
|||||||
|
|
||||||
has_triton_kernels = is_triton_kernels_available()
|
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():
|
if is_flashinfer_available():
|
||||||
from flashinfer import (
|
from flashinfer import (
|
||||||
@@ -472,10 +476,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||||
|
|
||||||
w13_weight_scale = torch.nn.Parameter(
|
w13_weight_scale = torch.nn.Parameter(
|
||||||
torch.zeros(
|
torch.full(
|
||||||
|
(
|
||||||
layer.num_local_experts,
|
layer.num_local_experts,
|
||||||
2 * intermediate_size_per_partition_after_pad,
|
2 * intermediate_size_per_partition_after_pad,
|
||||||
hidden_size // mxfp4_block,
|
hidden_size // mxfp4_block,
|
||||||
|
),
|
||||||
|
fill_value=_UE8M0_ONE,
|
||||||
dtype=scale_dtype,
|
dtype=scale_dtype,
|
||||||
),
|
),
|
||||||
requires_grad=False,
|
requires_grad=False,
|
||||||
@@ -511,10 +518,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||||
|
|
||||||
w2_weight_scale = torch.nn.Parameter(
|
w2_weight_scale = torch.nn.Parameter(
|
||||||
torch.zeros(
|
torch.full(
|
||||||
|
(
|
||||||
layer.num_local_experts,
|
layer.num_local_experts,
|
||||||
hidden_size,
|
hidden_size,
|
||||||
intermediate_size_per_partition_after_pad // mxfp4_block,
|
intermediate_size_per_partition_after_pad // mxfp4_block,
|
||||||
|
),
|
||||||
|
fill_value=_UE8M0_ONE,
|
||||||
dtype=scale_dtype,
|
dtype=scale_dtype,
|
||||||
),
|
),
|
||||||
requires_grad=False,
|
requires_grad=False,
|
||||||
|
|||||||
Reference in New Issue
Block a user