From d36f6f043ce881d7db7b08856e217cd1250b6c2f Mon Sep 17 00:00:00 2001 From: b8zhong Date: Fri, 16 Jan 2026 17:24:05 -0800 Subject: [PATCH] [Fix] `flashinfer_trtllm` `intermediate_size` assertion with Qwen3 + TP=8 (#16824) --- .../srt/layers/moe/fused_moe_triton/layer.py | 24 +++++++++++- python/sglang/srt/server_args.py | 39 +++++++------------ python/sglang/srt/utils/common.py | 8 ++++ 3 files changed, 44 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 838c4c5af..b5f90f255 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 78ec3fc71..75b79a15b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -42,6 +42,7 @@ from sglang.srt.utils.common import ( get_device_memory_capacity, get_device_name, get_device_sm, + get_quantization_config, is_blackwell_supported, is_cuda, is_fa3_default_architecture, @@ -1196,12 +1197,7 @@ class ServerArgs: # Set moe backend for DeepSeek if is_sm100_supported(): - quantization_config = getattr(hf_config, "quantization_config", None) - quant_method = ( - quantization_config.get("quant_method") - if quantization_config is not None - else None - ) + quant_method = get_quantization_config(hf_config) if self.quantization is None: # Default DeepSeek V3/R1 native FP8 when not explicitly set, # Because we need this condition for an assertion in @@ -1265,11 +1261,8 @@ class ServerArgs: f"- Decode: {decode_attn_backend}\n" ) - quantization_config = getattr(hf_config, "quantization_config", None) - is_mxfp4_quant_format = ( - quantization_config is not None - and quantization_config.get("quant_method") == "mxfp4" - ) + quant_method = get_quantization_config(hf_config) + is_mxfp4_quant_format = quant_method == "mxfp4" if is_mxfp4_quant_format: # use bf16 for mxfp4 triton kernels self.dtype = "bfloat16" @@ -1445,16 +1438,14 @@ class ServerArgs: "Qwen3VLMoeForConditionalGeneration", ]: if is_sm100_supported(): - quantization_config = getattr(hf_config, "quantization_config", None) - quant_method = ( - quantization_config.get("quant_method") - if quantization_config is not None - else None - ) + quant_method = get_quantization_config(hf_config) if self.quantization is None and quant_method is not None: self.quantization = quant_method if ( - self.quantization in ("fp8", "modelopt_fp4") + ( + self.quantization in ("fp8", "modelopt_fp4") + or self.quantization is None + ) and self.moe_a2a_backend == "none" and self.moe_runner_backend == "auto" ): @@ -1465,16 +1456,14 @@ class ServerArgs: ) elif model_arch in ["Qwen3NextForCausalLM"]: if is_sm100_supported(): - quantization_config = getattr(hf_config, "quantization_config", None) - quant_method = ( - quantization_config.get("quant_method") - if quantization_config is not None - else None - ) + quant_method = get_quantization_config(hf_config) if self.quantization is None and quant_method is not None: self.quantization = quant_method if ( - (self.quantization == "fp8" or self.quantization == "modelopt_fp4") + ( + self.quantization in ("fp8", "modelopt_fp4") + or self.quantization is None + ) and self.moe_a2a_backend == "none" and self.moe_runner_backend == "auto" ): diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index b7cc81f75..63e429bca 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -2708,6 +2708,14 @@ def has_hf_quant_config(model_path: str) -> bool: return False +def get_quantization_config(hf_config) -> str | None: + """Extract quantization method from HuggingFace config.""" + quantization_config = getattr(hf_config, "quantization_config", None) + if quantization_config is not None: + return quantization_config.get("quant_method") + return None + + def flatten_nested_list(nested_list): if isinstance(nested_list, list): return [