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

[2/2] Fuse routed scaling factor into select_experts (#8690)

Browse Source
This commit is contained in:
Trevor Morris
2025-08-20 15:10:16 -07:00
committed by GitHub
parent f96413c444
commit a91e90d9a3
6 changed files with 55 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -514,6 +514,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
self.quant_config = quant_config
self.block_quant = self.quant_config.weight_block_size is not None
self.cutlass_fp8_supported = cutlass_fp8_supported()
self.use_cutlass_fused_experts_fp8 = (
get_bool_env_var("SGLANG_CUTLASS_MOE")
and self.cutlass_fp8_supported
and self.block_quant
and (is_sm100_supported() or is_sm90_supported())
)
def create_weights(
self,
@@ -1021,12 +1027,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
if ret is not None:
return ret
if (
get_bool_env_var("SGLANG_CUTLASS_MOE")
and self.cutlass_fp8_supported
and self.block_quant
and (is_sm100_supported() or is_sm90_supported())
):
if self.use_cutlass_fused_experts_fp8:
from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8
topk_weights, topk_ids, _ = topk_output
@@ -1053,9 +1054,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
self.problem_sizes2,
use_fp8_blockscale=True,
)
# TODO: Fuse into select_experts
if moe_runner_config.routed_scaling_factor is not None:
output *= moe_runner_config.routed_scaling_factor
# Scale by routed_scaling_factor is fused into select_experts.
return output
# Expert fusion with FP8 quantization
return fused_experts(

6
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -1305,8 +1305,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
tp_rank=layer.moe_tp_rank,
tune_max_num_tokens=next_power_of_2(x.shape[0]),
)[0]
if moe_runner_config.routed_scaling_factor is not None:
output *= moe_runner_config.routed_scaling_factor
# Scale by routed_scaling_factor is fused into select_experts.
if should_use_flashinfer_cutlass_moe_fp4_allgather():
output, global_output = get_local_dp_buffer(), output
get_tp_group().reduce_scatterv(
@@ -1332,6 +1331,5 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
params=layer.cutlass_moe_params,
apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input,
).to(x.dtype)
if moe_runner_config.routed_scaling_factor is not None:
output *= moe_runner_config.routed_scaling_factor
# Scale by routed_scaling_factor is fused into select_experts.
return output