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()