minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NVIDIA] Support flashinfer a2a with flashinfer_trtllm_routed moe (#22394)

Browse Source
This commit is contained in:
Trevor Morris
2026-06-29 16:23:58 -07:00
committed by GitHub
parent 6bdecb8206
commit b8c25bfaa7
5 changed files with 218 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/advanced_features/expert_parallelism.mdx
+3 -3
View File
@@ -112,12 +112,12 @@ Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support case
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashinfer_trtllm_routed</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>FlashInfer integrated with TensorRT-LLM for accelerated routed MoE computations, consuming SGLang-computed top-k expert assignments and weights.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Blackwell with TRT-LLM.</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`flashinfer_cutlass`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>FlashInfer combined with CUTLASS for high-performance grouped GEMMs in MoE layers, handling FP4/FP8 quantization efficiently.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>FlashInfer combined with CUTLASS for high-performance grouped GEMMs in MoE layers, handling FP4/FP8 quantization efficiently. Compatible with flashinfer all-to-all.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Blackwell with FP4/FP8 models.</td>
</tr>
<tr>
@@ -127,7 +127,7 @@ Currently, DeepEP, Mooncake, NIXL-EP, `ascend_fuseep` and MORI only support case
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`flashinfer_cutedsl`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>FlashInfer with a custom DSL for flexible and efficient MoE kernel generation, integrated with ModelOpt FP4 quantization.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>FlashInfer with a custom DSL for flexible and efficient MoE kernel generation, integrated with ModelOpt FP4 quantization. Compatible with flashinfer all-to-all.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Low-precision models with NVFP4.</td>
</tr>
</tbody>
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+67 -4
View File
@@ -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:
python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
+12 -28
View File
@@ -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:
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
python/sglang/srt/server_args.py
+10 -5
View File
@@ -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
test/registered/ep/test_flashinfer_a2a.py
+125
View File
@@ -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()