chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[hotfix] missing w13_weight_fp8 and w2_weight_fp8 in UE8M0 requantization (#12259)

Browse Source
This commit is contained in:
Cheng Wan
2025-10-28 19:10:38 -07:00
committed by GitHub
parent 334543ff3b
commit 83087247d1
4 changed files with 30 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -131,23 +131,6 @@ class DeepEPMoE(FusedMoE):
)
# the last one is invalid rank_id
self.expert_mask[:-1] = 1
elif not _is_npu:
self.w13_weight_fp8 = (
self.w13_weight,
(
self.w13_weight_scale_inv
if self.use_block_quant or self.use_w4afp8
else self.w13_weight_scale
),
)
self.w2_weight_fp8 = (
self.w2_weight,
(
self.w2_weight_scale_inv
if self.use_block_quant or self.use_w4afp8
else self.w2_weight_scale
),
)
def forward(
self,

15
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
View File

@@ -227,15 +227,16 @@ class DeepGemmRunnerCore(MoeRunnerCore):
hidden_states_device = running_state["hidden_states_device"]
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
b, s_mn, s_k = hidden_states_scale.shape
assert (
s_mn % 4 == 0 and s_k % 4 == 0
), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})"
# GroupGemm-0
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
hidden_states_scale = _cast_to_e8m0_with_rounding_up(hidden_states_scale)
if hidden_states_scale.dtype != torch.int:
b, s_mn, s_k = hidden_states_scale.shape
assert (
s_mn % 4 == 0 and s_k % 4 == 0
), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})"
hidden_states_scale = _cast_to_e8m0_with_rounding_up(
hidden_states_scale
)
else:
hidden_states_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
hidden_states_scale

24
python/sglang/srt/models/deepseek_v2.py
View File

@@ -3289,8 +3289,8 @@ class DeepseekV2ForCausalLM(nn.Module):
experts = layer.mlp.experts
if isinstance(experts, DeepEPMoE):
for w in [
experts.w13_weight_fp8,
experts.w2_weight_fp8,
(experts.w13_weight, experts.w13_weight_scale_inv),
(experts.w2_weight, experts.w2_weight_scale_inv),
]:
requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size)
else:
@@ -3338,10 +3338,26 @@ class DeepseekV2ForCausalLM(nn.Module):
)
experts = layer.mlp.experts
w13_weight_fp8 = (
experts.w13_weight,
(
experts.w13_weight_scale_inv
if hasattr(experts, "w13_weight_scale_inv")
else experts.w13_weight_scale
),
)
w2_weight_fp8 = (
experts.w2_weight,
(
experts.w2_weight_scale_inv
if hasattr(experts, "w2_weight_scale_inv")
else experts.w2_weight_scale
),
)
if isinstance(experts, DeepEPMoE):
for w in [
experts.w13_weight_fp8,
experts.w2_weight_fp8,
w13_weight_fp8,
w2_weight_fp8,
]:
transform_scale_ue8m0_inplace(w[1], mn=w[0].shape[-2])

4
python/sglang/srt/models/longcat_flash.py
View File

@@ -821,8 +821,8 @@ class LongcatFlashForCausalLM(nn.Module):
experts = layer.mlp.experts
if isinstance(experts, DeepEPMoE):
for w in [
experts.w13_weight_fp8,
experts.w2_weight_fp8,
(experts.w13_weight, experts.w13_weight_scale_inv),
(experts.w2_weight, experts.w2_weight_scale_inv),
]:
requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size)