From 935cbf24ec55e2d473610be7b43de6d02a6f0b70 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:23:15 +0800 Subject: [PATCH] Accept MXFP8 dispatch in FlashInfer A2A TRT-LLM MoE (#39613) Co-authored-by: Mohammad Angkad --- .../moe/moe_runner/flashinfer_trtllm.py | 31 +++++++++++++------ 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 66de3a131..4d307df24 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -1707,16 +1707,27 @@ def fused_experts_flashinfer_to_flashinfer_trtllm( use_routed_topk=True, ) elif isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo): - if dispatch_output.hidden_states.dtype != torch.bfloat16: - raise TypeError( - "FlashInfer A2A + TRT-LLM Gen FP8 MoE requires a BF16 " - f"dispatch payload, got {dispatch_output.hidden_states.dtype}." - ) - if dispatch_output.hidden_states_scale is not None: - raise ValueError( - "FlashInfer A2A + TRT-LLM Gen FP8 MoE quantizes locally; " - "the BF16 dispatch payload must not carry activation scales." - ) + mxfp8_dispatch = ( + quant_info.use_mxfp8 + and dispatch_output.hidden_states.dtype == torch.float8_e4m3fn + ) + if mxfp8_dispatch: + if dispatch_output.hidden_states_scale is None: + raise ValueError( + "FlashInfer A2A + TRT-LLM Gen MXFP8 MoE requires activation " + "scales alongside the FP8 dispatch payload." + ) + else: + if dispatch_output.hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer A2A + TRT-LLM Gen FP8 MoE requires a BF16 " + f"dispatch payload, got {dispatch_output.hidden_states.dtype}." + ) + if dispatch_output.hidden_states_scale is not None: + raise ValueError( + "FlashInfer A2A + TRT-LLM Gen FP8 MoE quantizes locally; " + "the BF16 dispatch payload must not carry activation scales." + ) result = fused_experts_none_to_flashinfer_trtllm_fp8( dispatch_output, quant_info,