Fix server args for gpt oss so users can override the moe runner backend (#12696)

This commit is contained in:
Lianmin Zheng
2025-11-05 11:36:59 -08:00
committed by GitHub
parent 7c45b8b4bb
commit fb2e816e83
+15 -18
View File
@@ -959,30 +959,27 @@ class ServerArgs:
quantization_config is not None quantization_config is not None
and quantization_config.get("quant_method") == "mxfp4" and quantization_config.get("quant_method") == "mxfp4"
) )
if is_mxfp4_quant_format:
# use bf16 for mxfp4 triton kernels
self.dtype = "bfloat16"
if is_blackwell_supported() and is_mxfp4_quant_format: if self.moe_runner_backend == "auto":
self.moe_runner_backend = "flashinfer_mxfp4" if is_blackwell_supported() and is_mxfp4_quant_format:
logger.warning( self.moe_runner_backend = "flashinfer_mxfp4"
"Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel." logger.warning(
) "Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."
else: )
if self.moe_runner_backend == "triton_kernel": elif self.ep_size == 1 and is_triton_kernels_available():
assert (
self.ep_size == 1
), "Triton kernel MoE is only supported when ep_size == 1"
if (
self.moe_runner_backend == "auto"
and self.ep_size == 1
and is_triton_kernels_available()
):
self.moe_runner_backend = "triton_kernel" self.moe_runner_backend = "triton_kernel"
logger.warning( logger.warning(
"Detected GPT-OSS model, enabling triton_kernels MOE kernel." "Detected GPT-OSS model, enabling triton_kernels MOE kernel."
) )
if self.moe_runner_backend == "triton_kernel":
assert (
self.ep_size == 1
), "Triton kernel MoE is only supported when ep_size == 1"
self.disable_hybrid_swa_memory = True self.disable_hybrid_swa_memory = True
if is_mxfp4_quant_format:
# use bf16 for mxfp4 triton kernels
self.dtype = "bfloat16"
elif "Llama4" in model_arch and self.device != "cpu": elif "Llama4" in model_arch and self.device != "cpu":
assert self.attention_backend in { assert self.attention_backend in {