diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index e9f736a18..81b91e50d 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -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, diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index ed1569cba..9533eda78 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index fb83f30c0..827e2e96b 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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]) diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 88a8cad3a..23bfed025 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -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)