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

[Fix] flashinfer_trtllm intermediate_size assertion with Qwen3 + TP=8 (#16824)

Browse Source
This commit is contained in:
b8zhong
2026-01-16 17:24:05 -08:00
committed by GitHub
parent b0701f02b3
commit d36f6f043c
3 changed files with 44 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -206,6 +206,16 @@ class FusedMoE(torch.nn.Module):
self.use_triton_kernels = get_moe_runner_backend().is_triton_kernels()
self.use_flashinfer_trtllm_moe = get_moe_runner_backend().is_flashinfer_trtllm()
# flashinfer_trtllm kernel requires intermediate_size to be a multiple of 128
# Pad the intermediate_size_per_partition if necessary
if (
self.use_flashinfer_trtllm_moe
and self.intermediate_size_per_partition % 128 != 0
):
self.intermediate_size_per_partition = round_up(
self.intermediate_size_per_partition, 128
)
self.quant_config = quant_config
self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4()
# TODO maybe we should remove this `if`, since `Mxfp4MoEMethod` does another round-up logic
@@ -395,7 +405,12 @@ class FusedMoE(torch.nn.Module):
else:
start = 0
if _is_cpu:
# Use narrow_padded_param_and_loaded_weight for:
# 1. CPU (always)
# 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128)
# This handles the case where the loaded weights are smaller than the padded expert_data
use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe
if use_padded_loading:
expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
expert_data,
loaded_weight,
@@ -464,7 +479,12 @@ class FusedMoE(torch.nn.Module):
# for w2 in TP, it shards the input_features, i.e., shard_dim=2
shard_size = expert_data.shape[shard_dim]
if _is_cpu:
# Use narrow_padded_param_and_loaded_weight for:
# 1. CPU (always)
# 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128)
# This handles the case where the loaded weights are smaller than the padded expert_data
use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe
if use_padded_loading:
expert_data, loaded_weight = narrow_padded_param_and_loaded_weight(
expert_data,
loaded_weight,