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
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
+13 -29
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:
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
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