From b8c25bfaa7922343fe03dba12569f0d740dfe977 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Mon, 29 Jun 2026 16:23:58 -0700 Subject: [PATCH] [NVIDIA] Support flashinfer a2a with flashinfer_trtllm_routed moe (#22394) --- .../advanced_features/expert_parallelism.mdx | 6 +- .../moe/moe_runner/flashinfer_trtllm.py | 71 +++++++++- .../layers/moe/token_dispatcher/flashinfer.py | 42 ++---- python/sglang/srt/server_args.py | 15 ++- test/registered/ep/test_flashinfer_a2a.py | 125 ++++++++++++++++++ 5 files changed, 218 insertions(+), 41 deletions(-) create mode 100644 test/registered/ep/test_flashinfer_a2a.py diff --git a/docs_new/docs/advanced_features/expert_parallelism.mdx b/docs_new/docs/advanced_features/expert_parallelism.mdx index 0f012640d..15d03f412 100644 --- a/docs_new/docs/advanced_features/expert_parallelism.mdx +++ b/docs_new/docs/advanced_features/expert_parallelism.mdx @@ -112,12 +112,12 @@ Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support case flashinfer_trtllm_routed - FlashInfer integrated with TensorRT-LLM for accelerated routed MoE computations, consuming SGLang-computed top-k expert assignments and weights. + FlashInfer integrated with TensorRT-LLM for accelerated routed MoE computations, consuming SGLang-computed top-k expert assignments and weights. Compatible with flashinfer all-to-all. Blackwell with TRT-LLM. `flashinfer_cutlass` - FlashInfer combined with CUTLASS for high-performance grouped GEMMs in MoE layers, handling FP4/FP8 quantization efficiently. + FlashInfer combined with CUTLASS for high-performance grouped GEMMs in MoE layers, handling FP4/FP8 quantization efficiently. Compatible with flashinfer all-to-all. Blackwell with FP4/FP8 models. @@ -127,7 +127,7 @@ Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support case `flashinfer_cutedsl` - FlashInfer with a custom DSL for flexible and efficient MoE kernel generation, integrated with ModelOpt FP4 quantization. + FlashInfer with a custom DSL for flexible and efficient MoE kernel generation, integrated with ModelOpt FP4 quantization. Compatible with flashinfer all-to-all. Low-precision models with NVFP4. 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 4d44dfe6d..5c80dc22d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -96,6 +96,10 @@ if TYPE_CHECKING: StandardCombineInput, StandardDispatchOutput, ) + from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + FlashinferCombineInput, + FlashinferDispatchOutput, + ) if is_flashinfer_available(): from sglang.srt.layers.quantization.fp4_utils import fp4_quantize @@ -926,7 +930,18 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( topk_output = dispatch_output.topk_output # Quantize hidden states to FP4 - if quant_info.use_per_token_activation: + hidden_states_scale = ( + dispatch_output.hidden_states_scale + if hasattr(dispatch_output, "hidden_states_scale") + else None + ) + per_token_scale = None + if hidden_states_scale is not None: + # NVFP4 dispatch, inputs are already quantized. + hs_fp4 = hidden_states + hs_scale_linear = hidden_states_scale + elif quant_info.use_per_token_activation: + # Enable FlashInfer TRTLLM per-token NVFP4 activation scaling; ignores checkpoint activation FP32 scale by treating it as from flashinfer import SfLayout, nvfp4_quantize e4m3_max = 448.0 @@ -949,7 +964,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( seq_len, hidden_size // 16 ) else: - per_token_scale = None hs_fp4, hs_scale_linear = quantize_hidden_states_fp4( hidden_states, quant_info.w13_input_scale_quant ) @@ -976,12 +990,15 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( hidden_size = ( hs_fp4.shape[-1] * 2 if hs_fp4.dtype == torch.uint8 else hs_fp4.shape[-1] ) + output_dtype = ( + hidden_states.dtype if hidden_states_scale is None else torch.bfloat16 + ) _provided = _moe_output_buf.get() _symm_required = is_allocation_symmetric() if ( _provided is not None and _provided.shape == (num_tokens, hidden_size) - and _provided.dtype == hidden_states.dtype + and _provided.dtype == output_dtype and _provided.device == hs_fp4.device and ( not _symm_required @@ -995,7 +1012,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( symm_output = torch.empty( hs_fp4.shape[0], hidden_size, - dtype=hidden_states.dtype, + dtype=output_dtype, device=hs_fp4.device, ) @@ -1285,6 +1302,52 @@ def fused_experts_none_to_flashinfer_trtllm_routed( ) +@register_fused_func("flashinfer", "flashinfer_trtllm_routed") +def fused_experts_flashinfer_to_flashinfer_trtllm_routed( + dispatch_output: FlashinferDispatchOutput, + quant_info: MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> FlashinferCombineInput: + """Fused function for flashinfer A2A + flashinfer_trtllm_routed runner. + + FlashinferDispatchOutput and StandardDispatchOutput share the same field + layout (hidden_states, hidden_states_scale, topk_output), so the existing + FP8/FP4/BF16 implementations work unchanged. We wrap the returned + StandardCombineInput into a FlashinferCombineInput for the FlashinferDispatcher + combine path. + """ + from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + FlashinferCombineInput, + ) + + if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo): + result = fused_experts_none_to_flashinfer_trtllm_fp4( + dispatch_output, + quant_info, + runner_config, + use_routed_topk=True, + ) + elif isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo): + result = fused_experts_none_to_flashinfer_trtllm_fp8( + dispatch_output, + quant_info, + runner_config, + use_routed_topk=True, + ) + elif isinstance(quant_info, FlashInferTrtllmBf16MoeQuantInfo): + result = fused_experts_none_to_flashinfer_trtllm_bf16( + dispatch_output, + quant_info, + runner_config, + use_routed_topk=True, + ) + else: + raise TypeError( + f"Unexpected quant_info type for flashinfer a2a + flashinfer_trtllm_routed: {type(quant_info)}" + ) + return FlashinferCombineInput(hidden_states=result.hidden_states) + + # Register the experimental experimental_sgl_trtllm MoE fused-func (MoeRunner needs it at # build time even for LoRA); gated by the master switch so the upstream path is untouched. if _SGLANG_EXPERIMENTAL_LORA_OPTI: diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 702485ce4..d6e6a7e40 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -98,7 +98,11 @@ class FlashinferDispatcher(BaseDispatcher): self.hidden_size = hidden_size self.num_experts = num_experts self.num_local_experts = num_local_experts - + self.invalid_token_expert_id = ( + -1 + if get_moe_runner_backend().is_flashinfer_trtllm_routed() + else self.num_experts + ) # TODO: Can other moe runners use payload_in_workspace too? self.payload_in_workspace = get_moe_runner_backend().is_flashinfer_cutlass() @@ -164,19 +168,6 @@ class FlashinferDispatcher(BaseDispatcher): mnnvl_config=MnnvlConfig(comm_backend=TorchDistributedCommBackend(group)), ) - self.dummy_topk_ids = torch.full( - (1, self.router_topk), self.num_experts, dtype=torch.int32, device="cuda" - ) - self.dummy_topk_ids_current_rank = torch.full( - (1, self.router_topk), - self.ep_rank * self.num_local_experts, - dtype=torch.int32, - device="cuda", - ) - self.dummy_topk_weights = torch.zeros( - (1, self.router_topk), dtype=torch.float32, device="cuda" - ) - @debug_kernel_api def dispatch( self, hidden_states: torch.Tensor, topk_output: TopKOutput @@ -187,15 +178,14 @@ class FlashinferDispatcher(BaseDispatcher): topk_ids = topk_output.topk_ids topk_weights = topk_output.topk_weights - self.has_dummy_token = x.shape[0] == 0 - if self.has_dummy_token: - x = hidden_states.new_zeros((1, self.hidden_size)) - topk_ids = self.dummy_topk_ids - topk_weights = self.dummy_topk_weights - global_scale = self.quant_config.get("input_global_scale", None) if global_scale is not None: - x, x_sf = fp4_quantize(x, global_scale, is_sf_swizzled_layout=False) + if x.shape[0] > 0: + x, x_sf = fp4_quantize(x, global_scale, is_sf_swizzled_layout=False) + else: + x_col = x.shape[1] + x = torch.zeros(0, x_col // 2, dtype=torch.uint8, device=x.device) + x_sf = torch.zeros(0, x_col // 16, dtype=torch.uint8, device=x.device) payloads = [] payloads.append(x) @@ -248,8 +238,6 @@ class FlashinferDispatcher(BaseDispatcher): else: # Case 3 self.runtime_max_tokens_per_rank = x.shape[0] - if self.has_dummy_token: - self.runtime_max_tokens_per_rank = max(self.runtime_max_tokens_per_rank, 1) # Passing topk_ids + invalid_token_expert_id triggers the sanitize step # inside moe_a2a. The recv buffer has shape @@ -258,10 +246,10 @@ class FlashinferDispatcher(BaseDispatcher): # and waste downstream MoE compute. Sanitizing the padding to a # sentinel id is structural, not optional. recv_tensors = self.moe_a2a.dispatch( - self.dummy_topk_ids_current_rank if self.has_dummy_token else topk_ids, + topk_ids, payloads, self.runtime_max_tokens_per_rank, - invalid_token_expert_id=self.num_experts, + invalid_token_expert_id=self.invalid_token_expert_id, expert_id_payload_index=expert_id_payload_index, ) if x_sf is not None: @@ -301,9 +289,5 @@ class FlashinferDispatcher(BaseDispatcher): payload_in_workspace=self.payload_in_workspace, ) - if self.has_dummy_token: - hidden_states = hidden_states[1:, :] - del self.runtime_max_tokens_per_rank - del self.has_dummy_token return hidden_states diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1faa7d23f..ec99f9713 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5351,9 +5351,10 @@ class ServerArgs: ], "The expert parallel size must be 1 or the same as the tensor parallel size" if self.moe_runner_backend == "flashinfer_cutedsl": - assert self.quantization in [ - "modelopt_fp4" - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4'." + assert ( + self.quantization in ["modelopt_fp4"] + or self.get_model_config().nvfp4_moe_meta is not None + ), f"Invalid quantization '{self.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4' or hybrid NVFP4 models." assert self.ep_size in [ 1, self.tp_size, @@ -5574,7 +5575,10 @@ class ServerArgs: ) if self.deepep_mode != "auto": logger.warning("--deepep-mode is ignored for Flashinfer MoE A2A") - if not envs.SGLANG_MOE_NVFP4_DISPATCH.is_set(): + if not envs.SGLANG_MOE_NVFP4_DISPATCH.is_set() and ( + self.quantization == "modelopt_fp4" + or self.get_model_config().nvfp4_moe_meta is not None + ): envs.SGLANG_MOE_NVFP4_DISPATCH.set(True) logger.warning( "SGLANG_MOE_NVFP4_DISPATCH is set to True for Flashinfer MoE A2A" @@ -5582,7 +5586,8 @@ class ServerArgs: assert self.moe_runner_backend in [ "flashinfer_cutlass", "flashinfer_cutedsl", - ], "Flashinfer MoE A2A is only supported with flashinfer_cutlass or flashinfer_cutedsl moe runner backend" + "flashinfer_trtllm_routed", + ], "Flashinfer MoE A2A is only supported with flashinfer_cutlass, flashinfer_cutedsl or flashinfer_trtllm_routed moe runner backend" if self.moe_a2a_backend == "mori": self.ep_size = self.tp_size diff --git a/test/registered/ep/test_flashinfer_a2a.py b/test/registered/ep/test_flashinfer_a2a.py new file mode 100644 index 000000000..7a2a078aa --- /dev/null +++ b/test/registered/ep/test_flashinfer_a2a.py @@ -0,0 +1,125 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-gb300") + +DEEPSEEK_V3_FP4_MODEL = "nvidia/DeepSeek-V3-0324-FP4" +QWEN3_FP8_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" +SERVER_LAUNCH_TIMEOUT = 1000 + + +class TestFlashinferA2ATrtllmRoutedFP4(CustomTestCase): + """flashinfer A2A + flashinfer_trtllm_routed with modelopt_fp4 (DeepSeek V3).""" + + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V3_FP4_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--tp", + "4", + "--ep", + "4", + "--dp", + "4", + "--enable-dp-attention", + "--moe-a2a-backend", + "flashinfer", + "--moe-runner-backend", + "flashinfer_trtllm_routed", + "--quantization", + "modelopt_fp4", + "--disable-flashinfer-autotune", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + ], + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["score"], 0.90) + + +class TestFlashinferA2ATrtllmRoutedFP8(CustomTestCase): + """flashinfer A2A + flashinfer_trtllm_routed with fp8 (Qwen3-Next).""" + + @classmethod + def setUpClass(cls): + cls.model = QWEN3_FP8_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--tp", + "4", + "--ep", + "4", + "--dp", + "4", + "--enable-dp-attention", + "--moe-a2a-backend", + "flashinfer", + "--moe-runner-backend", + "flashinfer_trtllm_routed", + "--attention-backend", + "triton", + "--mem-fraction-static", + "0.7", + "--mamba-ssm-dtype", + "bfloat16", + "--disable-flashinfer-autotune", + ], + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["score"], 0.93) + + +if __name__ == "__main__": + unittest.main()