[Blackwell] Make mxint4 flashinfer_trtllm moe gemm set by default on blackwell (#18136)

Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
This commit is contained in:
Yi Zhong
2026-03-01 00:21:01 -05:00
committed by GitHub
parent 38dc372dae
commit 894e887e4a

View File

@@ -1287,6 +1287,20 @@ class ServerArgs:
# Set moe backend for DeepSeek
if is_sm100_supported():
quant_method = get_quantization_config(hf_config)
quant_cfg = getattr(hf_config, "quantization_config", None) or {}
config_groups = quant_cfg.get("config_groups", {})
group0 = config_groups.get("group_0", {})
weights_cfg = group0.get("weights", {})
# this also apply to kimi k2.5
# since it follow the compressed tensor int4 recipe
# but not kimi k2 instruct or 0905 instruct.
is_kimi_k2_k25_thinking_int4 = (
quant_method == "compressed-tensors"
and weights_cfg.get("num_bits") == 4
and weights_cfg.get("group_size") == 32
and weights_cfg.get("strategy") == "group"
and weights_cfg.get("type") == "int"
)
if self.quantization is None:
# Default DeepSeek V3/R1 native FP8 when not explicitly set,
# Because we need this condition for an assertion in
@@ -1301,12 +1315,20 @@ class ServerArgs:
if (
self.moe_a2a_backend == "none"
and self.moe_runner_backend == "auto"
and self.quantization in ["fp8", "modelopt_fp8", "modelopt_fp4"]
and (
self.quantization in ["fp8", "modelopt_fp8", "modelopt_fp4"]
or is_kimi_k2_k25_thinking_int4
)
):
self.moe_runner_backend = "flashinfer_trtllm"
logger.info(
"Use flashinfer_trtllm as MoE runner backend on sm100 for DeepseekV3ForCausalLM"
)
if is_kimi_k2_k25_thinking_int4:
logger.info(
"Use flashinfer_trtllm as MoE runner backend on Blackwell for Kimi K2 / K2.5 thinking int4"
)
else:
logger.info(
"Use flashinfer_trtllm as MoE runner backend on sm100 for DeepseekV3ForCausalLM"
)
elif is_hip():
if not self.enable_dp_attention and self.nnodes == 1:
# TODO (Hubert): Put this back later