From 2634149388dde1dec9e4fe341527f312803ce80c Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Tue, 27 Oct 2026 13:44:05 -0700 Subject: [PATCH 1/9] Switch NVLinkOneSidedAllToAll to use moe-ep flashinfer API --- .../layers/moe/token_dispatcher/flashinfer.py | 209 ++++++------------ .../layers/moe/test_flashinfer_a2a_wide_ep.py | 8 +- .../layers/moe/test_flashinfer_dispatcher.py | 56 +++-- .../test_muse_spark_runner_dispatch.py | 26 ++- 4 files changed, 120 insertions(+), 179 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 261e9fc2bd61..4fad510c6a79 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -41,16 +41,18 @@ get_flags, get_forward, get_parallel, - get_spec, max_prefill_buffer_tokens, ) -from sglang.srt.speculative.spec_info import SpeculativeAlgorithm try: from flashinfer import nvfp4_block_scale_interleave - from flashinfer.comm import MoeAlltoAll, moe_a2a_get_workspace_size_per_rank - from flashinfer.comm.mapping import Mapping - from flashinfer.comm.mnnvl import MnnvlConfig + from flashinfer.moe_ep import ( + BootstrapConfig, + MoEEpCommParams, + NVLinkOneSidedAlltoAll, + NVLinkOneSidedConfig, + create_communication, + ) from sglang.srt.layers.quantization.fp4_utils import fp4_quantize @@ -60,11 +62,11 @@ logger = logging.getLogger(__name__) -# FlashInfer keys MNNVL allocations by workspace size; aligned tail padding gives -# concurrently live paths distinct persistent workspaces without extra token work. -_WORKSPACE_NAMESPACE_ALIGNMENT = 128 - +# FlashInfer shares one workspace among communication instances with the same +# geometry, and those instances must not interleave dispatch/combine rounds. +# Padding the token capacity by the namespace gives concurrently live paths +# distinct persistent workspaces without extra token work. class _WorkspaceNamespace(enum.IntEnum): DECODE = 0 SPECULATIVE = 1 @@ -92,8 +94,8 @@ def _scattered_source_token_counts( return counts -def _workspace_size_for_namespace( - workspace_size: int, +def _max_tokens_for_namespace( + max_num_tokens: int, *, namespace: Optional[_WorkspaceNamespace] = None, speculative: Optional[bool] = None, @@ -106,7 +108,7 @@ def _workspace_size_for_namespace( if speculative else _WorkspaceNamespace.DECODE ) - return workspace_size + int(namespace) * _WORKSPACE_NAMESPACE_ALIGNMENT + return max_num_tokens + int(namespace) class FlashinferDispatchOutput(NamedTuple): @@ -195,90 +197,45 @@ def __init__( moe_runner_config = MoeRunnerConfig( num_experts=num_experts, num_local_experts=num_local_experts, + num_fused_shared_experts=0, hidden_size=hidden_size, top_k=router_topk, ) self.prefill_dispatcher = StandardDispatcher(moe_runner_config) # FlashInfer sizes the workspace from the maximum dispatched tokens per - # EP rank. See FlashInfer's moe_a2a_get_workspace_size_per_rank(), - # which reserves ep_size * max_num_tokens * payload bytes, and the C++ - # dispatch op's epSize * runtimeMaxTokensPerRank payload buffer. + # EP rank: its dispatch receive buffer is ep_size * max_tokens_per_rank rows. # # Cover both the largest eager prefill chunk (including PP dynamic # chunking) and the largest decode batch using upstream's shared bound. self.max_num_tokens = flashinfer_a2a_max_dispatch_tokens_per_rank( max_prefill_buffer_tokens() ) - speculative_algo = SpeculativeAlgorithm.from_string( - get_spec().speculative_algorithm + bootstrap = BootstrapConfig( + world_size=self.ep_size, rank=self.ep_rank, process_group=group ) - can_use_quantized_dispatch = not speculative_algo.is_eagle() - topk_id_and_weight_bytes = self.router_topk * 4 + self.router_topk * 4 - bf16_dispatch_payload_size_per_token = ( - hidden_size * 2 + topk_id_and_weight_bytes # bf16 hidden states + a2a_config = NVLinkOneSidedConfig( + comm_backend=TorchDistributedCommBackend(group) ) - if ( - self.dispatch_type == FlashinferA2ADispatchType.NVFP4 - and can_use_quantized_dispatch - ): - quantized_dispatch_payload_size_per_token = ( - hidden_size // 2 # nvfp4 hidden states - + hidden_size // 16 # uint8 scaling factors - + topk_id_and_weight_bytes - ) - total_dispatch_payload_size_per_token = max( - quantized_dispatch_payload_size_per_token, - bf16_dispatch_payload_size_per_token, - ) - elif ( - self.dispatch_type == FlashinferA2ADispatchType.MXFP8 - and can_use_quantized_dispatch - ): - quantized_dispatch_payload_size_per_token = ( - hidden_size # fp8 hidden states - + hidden_size // 32 # ue8m0 scaling factors - + topk_id_and_weight_bytes - ) - total_dispatch_payload_size_per_token = max( - quantized_dispatch_payload_size_per_token, - bf16_dispatch_payload_size_per_token, - ) - else: - total_dispatch_payload_size_per_token = bf16_dispatch_payload_size_per_token - combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states - self.workspace_size = moe_a2a_get_workspace_size_per_rank( - ep_size=self.ep_size, - max_num_tokens=self.max_num_tokens, - total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token, - combine_payload_size_per_token=combine_payload_size_per_token, - ) - - self.mapping = Mapping( - rank=self.ep_rank, - tp_size=self.ep_size, - moe_ep_size=self.ep_size, - world_size=self.ep_size, - gpus_per_node=torch.cuda.device_count(), - pp_size=1, - cp_size=1, - ) - mnnvl_config = MnnvlConfig(comm_backend=TorchDistributedCommBackend(group)) is_speculative_model = get_flags().moe.speculative_context - def make_moe_a2a(namespace: _WorkspaceNamespace) -> MoeAlltoAll: - workspace_size = _workspace_size_for_namespace( - self.workspace_size, - namespace=namespace, - ) - return MoeAlltoAll( - mapping=self.mapping, - max_num_tokens=self.max_num_tokens, - top_k=self.router_topk, + def make_moe_a2a(namespace: _WorkspaceNamespace) -> NVLinkOneSidedAlltoAll: + params = MoEEpCommParams( num_experts=self.num_experts, - workspace_size_per_rank=workspace_size, - mnnvl_config=mnnvl_config, + top_k=self.router_topk, + max_tokens_per_rank=_max_tokens_for_namespace( + self.max_num_tokens, namespace=namespace + ), + hidden_size=self.hidden_size, + # No dispatch_format: BF16 rows bound every quantized format and + # cover layers whose dispatch falls back to BF16 (draft, EAGLE). + dtype=torch.bfloat16, + invalid_expert_id=self.invalid_token_expert_id, ) + # FlashInfer's host-side workspace metadata must stay on the CPU even + # when layers are built under a CUDA default-device context. + with torch.device("cpu"): + return create_communication(bootstrap, params, backend=a2a_config) decode_namespace = ( _WorkspaceNamespace.SPECULATIVE @@ -309,11 +266,7 @@ def _start_a2a(self, *, sp_prefill: bool) -> None: ) def _clear_active_a2a(self) -> None: - for attr in ( - "runtime_max_tokens_per_rank", - "active_moe_a2a", - "active_a2a_namespace", - ): + for attr in ("active_moe_a2a", "active_a2a_namespace"): if hasattr(self, attr): delattr(self, attr) @@ -424,8 +377,8 @@ def dispatch( # FlashInfer dispatch requires materialized top-k IDs and weights. if TopKOutputChecker.format_is_bypassed(topk_output): topk_output = topk_output.to_standard() - # FlashInfer MoeAlltoAll's expert-ID ABI is int32. This dispatcher is - # only selected for moe_a2a_backend="flashinfer". + # FlashInfer NVLinkOneSidedAlltoAll's expert-ID ABI is int32. This + # dispatcher is only selected for moe_a2a_backend="flashinfer". topk_ids = topk_output.topk_ids.to(torch.int32) topk_weights = topk_output.topk_weights @@ -460,19 +413,9 @@ def dispatch( 0, self.hidden_size // 32, dtype=torch.uint8, device=x.device ) - payloads = [] - payloads.append(x) - if x_sf is not None: - payloads.append(x_sf) - expert_id_payload_index = 2 - else: - expert_id_payload_index = 1 - payloads.append(topk_ids) - payloads.append(topk_weights) - # runtime_max_tokens_per_rank selection # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # MoeAlltoAll uses fixed-geometry buffers shaped + # NVLinkOneSidedAlltoAll uses fixed-geometry buffers shaped # [ep_size, runtime_max_tokens_per_rank, ...], so every EP rank must pass # the SAME value. This code (Python) runs during eager forwards and during # CUDA-graph *capture*; on *replay* dispatch() is not re-executed and the @@ -507,8 +450,8 @@ def dispatch( ) runtime_max_tokens_per_rank = x.shape[0] - # MoeAlltoAll does not resize its max_num_tokens workspace; reject larger - # runtime geometry here before it becomes an illegal memory access. + # The FlashInfer workspace is sized once for max_num_tokens; reject larger + # runtime geometry here with a sizing hint rather than FlashInfer's error. assert runtime_max_tokens_per_rank <= self.max_num_tokens, ( "FlashInfer A2A runtime token geometry exceeds its fixed workspace: " f"runtime_max_tokens_per_rank={runtime_max_tokens_per_rank} > " @@ -525,54 +468,42 @@ def dispatch( # count) rather than a benign case. assert runtime_max_tokens_per_rank >= x.shape[0], ( f"runtime_max_tokens_per_rank={runtime_max_tokens_per_rank} < " - f"x.shape[0]={x.shape[0]}: MoeAlltoAll recv buffer would overflow." + f"x.shape[0]={x.shape[0]}: A2A recv buffer would overflow." ) - # Passing topk_ids + invalid_token_expert_id triggers the sanitize step - # inside moe_a2a. The recv buffer has shape - # [ep_size, max_tokens_per_rank, ...], so any rank below max leaves - # padding slots whose expert_id would otherwise route to a real expert - # and waste downstream MoE compute. Sanitizing the padding to a - # sentinel id is structural, not optional. - self.runtime_max_tokens_per_rank = runtime_max_tokens_per_rank + # The recv buffer has shape [ep_size * runtime_max_tokens_per_rank, ...], + # so any rank below max leaves padding slots; FlashInfer sanitizes their + # expert ids to MoEEpCommParams.invalid_expert_id so that downstream MoE + # compute skips them. self._start_a2a(sp_prefill=is_sp_prefill) try: - recv_tensors = self.active_moe_a2a.dispatch( + received = self.active_moe_a2a.dispatch( + x, topk_ids, - payloads, - runtime_max_tokens_per_rank, - invalid_token_expert_id=self.invalid_token_expert_id, - expert_id_payload_index=expert_id_payload_index, + topk_weights, + hidden_states_scale=x_sf, + max_tokens_per_rank=runtime_max_tokens_per_rank, ) - if x_sf is not None: - x_recv, x_sf_recv, topk_ids_recv, topk_weights_recv = recv_tensors - x_sf = x_sf_recv.view(-1, x_sf_recv.shape[-1]) - # TODO: fuse interleave into cutlass moe - if dispatch_type == FlashinferA2ADispatchType.NVFP4 and ( - self.runner_backend is not None - and self.runner_backend.is_flashinfer_cutlass() - ): - x_sf = nvfp4_block_scale_interleave(x_sf) - else: - x_recv, topk_ids_recv, topk_weights_recv = recv_tensors - x = x_recv.view(-1, x_recv.shape[-1]) - if dispatch_type == FlashinferA2ADispatchType.MXFP8: - x = x.view(torch.float8_e4m3fn) - topk_ids = topk_ids_recv.view(-1, topk_ids_recv.shape[-1]) - topk_weights = topk_weights_recv.view(-1, topk_weights_recv.shape[-1]) + x = received.hidden_states + x_sf = received.hidden_states_scale + # TODO: fuse interleave into cutlass moe + if ( + dispatch_type == FlashinferA2ADispatchType.NVFP4 + and self.runner_backend is not None + and self.runner_backend.is_flashinfer_cutlass() + ): + x_sf = nvfp4_block_scale_interleave(x_sf) # Provide an output tensor to fused_moe so it writes directly to our buffer moe_output = None if self.payload_in_workspace: - moe_output = ( - self.active_moe_a2a.get_combine_payload_tensor_in_workspace( - runtime_max_tokens_per_rank, self.hidden_size, output_dtype - ).view(-1, self.hidden_size) - ) + moe_output = self.active_moe_a2a.get_combine_input_buffer(output_dtype) return FlashinferDispatchOutput( x, x_sf, - StandardTopKOutput(topk_weights, topk_ids, topk_output.router_logits), + StandardTopKOutput( + received.topk_weights, received.topk_ids, topk_output.router_logits + ), moe_output, output_dtype, ) @@ -621,16 +552,8 @@ def combine( "FlashInfer A2A FP8 MoE combine payload must be BF16, but " f"received {hidden_states.dtype}." ) - output_hidden_size = hidden_states.shape[-1] try: - return self.active_moe_a2a.combine( - hidden_states.view( - self.ep_size, - self.runtime_max_tokens_per_rank, - output_hidden_size, - ), - self.runtime_max_tokens_per_rank, - payload_in_workspace=self.payload_in_workspace, - ) + # FlashInfer detects a payload already in its combine input buffer. + return self.active_moe_a2a.combine(hidden_states) finally: self._clear_active_a2a() diff --git a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py index 82f9ac39f249..2f3e68aefc06 100644 --- a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py +++ b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py @@ -9,7 +9,7 @@ ) from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( _max_tokens_per_scattered_source, - _workspace_size_for_namespace, + _max_tokens_for_namespace, ) from sglang.srt.layers.quantization import fp8 # noqa: F401 from sglang.test.ci.ci_register import register_cuda_ci @@ -38,11 +38,11 @@ def test_dp4_tp4_uses_physical_source_rank_geometry(self): self.assertEqual(_max_tokens_per_scattered_source([512] * 16, 1), 512) def test_target_and_draft_decode_use_distinct_workspaces(self): - sizes = { - _workspace_size_for_namespace(4096, speculative=speculative) + capacities = { + _max_tokens_for_namespace(4096, speculative=speculative) for speculative in (False, True) } - self.assertEqual(sizes, {4096, 4224}) + self.assertEqual(capacities, {4096, 4097}) def test_deepgemm_dispatch_marks_empty_expert_lanes_invalid(self): topk_ids = torch.tensor([-1, 0, -1, 1], dtype=torch.int32, device="cuda") diff --git a/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py b/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py index fee0392a99be..e2d885be10d8 100644 --- a/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py +++ b/test/registered/unit/layers/moe/test_flashinfer_dispatcher.py @@ -21,12 +21,26 @@ register_cpu_ci(est_time=11, suite="base-a-test-cpu") -class _RecordingMoeAlltoAll: - def dispatch(self, token_selected_experts, payloads, *_args, **kwargs): - self.token_selected_experts = token_selected_experts - self.payloads = payloads - self.expert_id_payload_index = kwargs["expert_id_payload_index"] - return payloads +class _RecordingCommunication: + def dispatch( + self, + hidden_states, + topk_ids, + topk_weights=None, + *, + hidden_states_scale=None, + **_kwargs, + ): + self.hidden_states = hidden_states + self.hidden_states_scale = hidden_states_scale + self.topk_ids = topk_ids + self.topk_weights = topk_weights + return SimpleNamespace( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) def _make_dispatcher(*, dispatch_type, max_num_tokens, quant_config): @@ -39,7 +53,7 @@ def _make_dispatcher(*, dispatch_type, max_num_tokens, quant_config): dispatcher.invalid_token_expert_id = 8 dispatcher.payload_in_workspace = False dispatcher.quant_config = quant_config - dispatcher.moe_a2a = _RecordingMoeAlltoAll() + dispatcher.moe_a2a = _RecordingCommunication() dispatcher.decode_a2a_namespace = _WorkspaceNamespace.DECODE dispatcher.sp_prefill_moe_a2a = None return dispatcher @@ -78,16 +92,22 @@ def test_empty_mxfp8_dispatch_uses_same_payload_dtype_as_nonempty_rank(self): topk_output=topk_output, ) + moe_a2a = dispatcher.moe_a2a self.assertEqual( - [payload.dtype for payload in dispatcher.moe_a2a.payloads], + [ + moe_a2a.hidden_states.dtype, + moe_a2a.hidden_states_scale.dtype, + moe_a2a.topk_ids.dtype, + moe_a2a.topk_weights.dtype, + ], [torch.float8_e4m3fn, torch.uint8, torch.int32, torch.float32], ) self.assertEqual(output.hidden_states.dtype, torch.float8_e4m3fn) self.assertEqual(output.hidden_states_scale.dtype, torch.uint8) def test_int64_expert_ids_reach_moe_a2a_as_int32(self): - """Expert IDs reach MoeAlltoAll as int32 even when routing emits int64, - as static EP dispatch does; MoeAlltoAll rejects any other dtype.""" + """Expert IDs reach the A2A as int32 even when routing emits int64, + as static EP dispatch does; FlashInfer's kernel rejects any other dtype.""" dispatcher = _make_dispatcher( dispatch_type=FlashinferA2ADispatchType.BF16, max_num_tokens=4, @@ -107,16 +127,11 @@ def test_int64_expert_ids_reach_moe_a2a_as_int32(self): ) moe_a2a = dispatcher.moe_a2a - self.assertEqual(moe_a2a.token_selected_experts.dtype, torch.int32) - self.assertEqual( - moe_a2a.payloads[moe_a2a.expert_id_payload_index].dtype, torch.int32 - ) - self.assertTrue( - torch.equal(moe_a2a.token_selected_experts, topk_ids.to(torch.int32)) - ) + self.assertEqual(moe_a2a.topk_ids.dtype, torch.int32) + self.assertTrue(torch.equal(moe_a2a.topk_ids, topk_ids.to(torch.int32))) def test_bypassed_topk_is_routed_before_moe_a2a(self): - """Runners that fuse routing hand over bypassed top-k; MoeAlltoAll still + """Runners that fuse routing hand over bypassed top-k; the A2A still needs explicit expert IDs and weights, so routing must run first.""" dispatcher = _make_dispatcher( dispatch_type=FlashinferA2ADispatchType.BF16, @@ -150,13 +165,12 @@ def test_bypassed_topk_is_routed_before_moe_a2a(self): ) moe_a2a = dispatcher.moe_a2a - expert_ids = moe_a2a.token_selected_experts + expert_ids = moe_a2a.topk_ids self.assertEqual(expert_ids.dtype, torch.int32) self.assertEqual( expert_ids.sort(dim=-1).values.tolist(), [[0, 1], [2, 3], [1, 3]] ) - topk_weights = moe_a2a.payloads[moe_a2a.expert_id_payload_index + 1] - self.assertEqual(topk_weights.shape, (3, 2)) + self.assertEqual(moe_a2a.topk_weights.shape, (3, 2)) class TestFlashinferRunnerDispatchFormats(CustomTestCase): diff --git a/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py b/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py index c76859f05386..3cc59ecd3027 100644 --- a/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py +++ b/test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py @@ -22,13 +22,17 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu") -class _CpuMoeAlltoAll: - def __init__(self, **_kwargs): - self.invalid_token_expert_id = None +class _CpuCommunication: + def __init__(self, _bootstrap, params, **_kwargs): + self.params = params - def dispatch(self, _topk_ids, payloads, _max_tokens, **kwargs): - self.invalid_token_expert_id = kwargs["invalid_token_expert_id"] - return payloads + def dispatch(self, hidden_states, topk_ids, topk_weights, **_kwargs): + return SimpleNamespace( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) class TestMuseSparkRunnerDispatch(CustomTestCase): @@ -83,11 +87,11 @@ def test_custom_expert_config_selects_trtllm_receive_sentinel(self): with ExitStack() as stack: for name, value in { "use_flashinfer": True, - "moe_a2a_get_workspace_size_per_rank": lambda **_: 128, - "Mapping": lambda **_: None, - "MnnvlConfig": lambda **_: None, + "BootstrapConfig": lambda **_: None, + "MoEEpCommParams": SimpleNamespace, + "NVLinkOneSidedConfig": lambda **_: None, "TorchDistributedCommBackend": lambda _: None, - "MoeAlltoAll": _CpuMoeAlltoAll, + "create_communication": _CpuCommunication, "get_flashinfer_a2a_dispatch_type": lambda: ( FlashinferA2ADispatchType.BF16 ), @@ -118,7 +122,7 @@ def test_custom_expert_config_selects_trtllm_receive_sentinel(self): torch.zeros((2, 128), dtype=torch.bfloat16), topk ) self.assertEqual( - dispatcher.sp_prefill_moe_a2a.invalid_token_expert_id, -1 + dispatcher.sp_prefill_moe_a2a.params.invalid_expert_id, -1 ) torch.testing.assert_close( output.topk_output.topk_ids, topk.topk_ids From 4d0a1702ba350e38e27179ad1b7a383d67cbc273 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Wed, 28 Oct 2026 12:32:49 -0700 Subject: [PATCH 2/9] fix lint --- test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py index 2f3e68aefc06..7c1468f253ac 100644 --- a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py +++ b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py @@ -8,8 +8,8 @@ PermuteMethodPool, ) from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( - _max_tokens_per_scattered_source, _max_tokens_for_namespace, + _max_tokens_per_scattered_source, ) from sglang.srt.layers.quantization import fp8 # noqa: F401 from sglang.test.ci.ci_register import register_cuda_ci From b559e7e6321df18f269319345f765919fd49df7d Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Wed, 4 Nov 2026 13:45:47 -0800 Subject: [PATCH 3/9] use runtime_max_tokens_per_rank instead of max_tokens_per_rank in dispatch --- python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 4fad510c6a79..174ba5b2e8a2 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -482,7 +482,7 @@ def dispatch( topk_ids, topk_weights, hidden_states_scale=x_sf, - max_tokens_per_rank=runtime_max_tokens_per_rank, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, ) x = received.hidden_states x_sf = received.hidden_states_scale From fce266d92f8e665f5946c4ca329f5d06889461b8 Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Wed, 4 Nov 2026 15:05:18 -0800 Subject: [PATCH 4/9] Bump FlashInfer to 0.7.3 Co-Authored-By: Claude Opus 5.5 --- docker/Dockerfile | 2 +- docker/Dockerfile.cu134 | 2 +- python/pyproject.toml | 2 +- python/sglang/srt/entrypoints/engine.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index b71488c0d7a9..c7af0c48b8a1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -18,7 +18,7 @@ ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=0 -ARG FLASHINFER_VERSION=0.7.2 +ARG FLASHINFER_VERSION=0.7.3 ARG MOONCAKE_VERSION=0.3.14 ARG MSCCLPP_VERSION=sglang-v0.10.1 diff --git a/docker/Dockerfile.cu134 b/docker/Dockerfile.cu134 index b5eaff036a66..cb803c5e1c78 100644 --- a/docker/Dockerfile.cu134 +++ b/docker/Dockerfile.cu134 @@ -34,7 +34,7 @@ ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=1 -ARG FLASHINFER_VERSION=0.7.2 +ARG FLASHINFER_VERSION=0.7.3 ARG MOONCAKE_VERSION=0.3.13.post1 ARG MSCCLPP_VERSION=sglang-v0.10.1 diff --git a/python/pyproject.toml b/python/pyproject.toml index 90cf70c505f1..c4090adf5411 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "einops", "fastapi", "flash-attn-4>=4.0.0b18", - "flashinfer_python[cu13]==0.7.2", # keep it aligned with jit-cache version in Dockerfile + "flashinfer_python[cu13]==0.7.3", # keep it aligned with jit-cache version in Dockerfile "gguf", "humming-kernels[cu13]==0.1.18", "interegular", diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 296658170f89..4321e7f20e37 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1992,7 +1992,7 @@ def _set_envs_and_config(server_args: ServerArgs): ): assert_pkg_version( "flashinfer_python", - "0.7.2", + "0.7.3", "Please uninstall the old version and " "reinstall the latest version by following the instructions " "at https://docs.flashinfer.ai/installation.html.", From 542af95ad9425da7b667bc83243eb012ca167cf0 Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Wed, 4 Nov 2026 17:12:12 -0800 Subject: [PATCH 5/9] speculative: Pass last_valid_index to FlashInfer 0.7.3 DeviceSamplingFromProb FlashInfer 0.7.3 (flashinfer-ai/flashinfer#5352) added an int& last_valid_index parameter to flashinfer::sampling::DeviceSamplingFromProb, so the JIT tree speculative sampling kernel no longer compiled and EAGLE verify crashed the scheduler. Keep the per-thread candidate across the scan, as FlashInfer's own chain speculative sampling kernel now does. Co-Authored-By: Claude Opus 5.5 --- .../jit/include/sgl_kernel/speculative/sampling.cuh | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh b/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh index ac46f636bcb5..a5999485550a 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh @@ -134,6 +134,8 @@ __global__ void TreeSpeculativeSamplingTargetOnly( DType u = coin * sum_relu_q_minus_p; DType aggregate_relu_q_minus_p(0); + // Each thread's fallback candidate; FlashInfer reduces it only if the scan selects nothing. + int last_valid_index = -1; for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { q_vec.fill(DType(0)); p_vec.fill(DType(0)); @@ -152,7 +154,14 @@ __global__ void TreeSpeculativeSamplingTargetOnly( } DeviceSamplingFromProb( - i, d, [&](DType x) { return x > 0; }, u, relu_q_minus_p_vec, aggregate_relu_q_minus_p, &temp_storage); + i, + d, + [&](DType x) { return x > 0; }, + u, + relu_q_minus_p_vec, + aggregate_relu_q_minus_p, + last_valid_index, + &temp_storage); if (aggregate_relu_q_minus_p > u) { break; } From c8d1c296dab99312ecc64e7df379e768e5f3330b Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Thu, 5 Nov 2026 05:56:06 -0800 Subject: [PATCH 6/9] flashinfer: Pass fixed_uniform_q_len in the fast prefill plan for FlashInfer 0.7.3 FlashInfer 0.7.3 appends a fixed_uniform_q_len argument to the FA2 paged prefill planner, so fast_prefill_plan's 20-argument call fails with "Expected 21 but got 20 arguments" on the first EAGLE verify or draft-extend CUDA graph replay. Also clear _eagle_verify_fp8kv_sm90_state, as 0.7.3's plan() does on every call, so run() never launches the SM90 fp8-KV verify kernel with state from an earlier plan. Co-Authored-By: Claude Opus 5.5 --- python/sglang/srt/layers/attention/flashinfer_backend.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index f221b1a59439..c0a42f420857 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -304,6 +304,10 @@ def fast_prefill_plan( ) self._cached_o_data_type = o_data_type self._block_tables = None + # FlashInfer's plan() clears this on every call and may set it again for its + # SM90 fp8-KV verify kernel. This plan never takes that route, so clear it + # here too rather than let run() launch it with a previous plan's state. + self._eagle_verify_fp8kv_sm90_state = None args = [ self._float_workspace_buffer, @@ -326,6 +330,7 @@ def fast_prefill_plan( False, # disable_split_kv 0, # num_colocated_ctas 0, # uniform_q_len + False, # fixed_uniform_q_len ] self._plan_info = self._cached_module.plan(*args) From f7e46a525893a0cc7a2e14643972fbbea9b59b29 Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Thu, 5 Nov 2026 05:56:15 -0800 Subject: [PATCH 7/9] minimax: Zero the NVFP4 payload when the block scale underflows, as FlashInfer 0.7.3 does flashinfer-ai/flashinfer#6996 (in 0.7.3) changed nvfp4_kv_quantize to guard the rounded E4M3 scale instead of the block amax, so a block whose scale underflows to zero now stores an all-zero payload instead of a saturated one. The fused MiniMax-M3.1 KV store reproduces that quantizer bit for bit, so make the same change; test_fused_store_matches_separate_stores failed on its _tiny case. Dequantized values are unchanged: a zero scale reads back as zero either way. Co-Authored-By: Claude Opus 5.5 --- .../sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh | 5 +++-- .../registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh b/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh index 21c35647290c..1bb006e3f086 100644 --- a/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh +++ b/python/sglang/kernels/jit/csrc/minimax/nvfp4_store_kv_index.cuh @@ -63,7 +63,8 @@ SGL_DEVICE uint32_t f32x8_to_e2m1x8(const float (&x)[8]) { /// /// Reproduces the reference operation order exactly: /// `sf = rcp(gs) * (amax * rcp(6))`, rounded to E4M3 and read back, -/// `osc = amax != 0 ? rcp(sf * gs) : 0`, +/// `osc = sf != 0 ? rcp(sf * gs) : 0`, so a block whose scale underflows +/// to zero stores an all-zero payload, /// each value scaled by `osc` then `cvt.rn.satfinite.e2m1x2.f32`. /// `amax` is exact regardless of reduction order because `fabsf` of a BF16 /// value is representable in BF16. @@ -86,7 +87,7 @@ SGL_DEVICE uint64_t quantize_block(const bf16_t* __restrict__ src, float global_ const __nv_fp8_e4m3 sf_e4m3 = __nv_fp8_e4m3(sf); *scale_out = sf_e4m3.__x; sf = static_cast(sf_e4m3); - const float osc = amax != 0.0f ? rcp_approx(sf * global_scale) : 0.0f; + const float osc = sf != 0.0f ? rcp_approx(sf * global_scale) : 0.0f; float lo[8], hi[8]; #pragma unroll diff --git a/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py b/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py index 4b646f6eb8a4..839ac9d2ed5b 100644 --- a/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py +++ b/test/registered/kernels/ops/kvcache/test_minimax_m31_nvfp4.py @@ -48,7 +48,7 @@ def _zeros(shape): def _tiny(shape): # Below the E4M3 scale underflow point, where the reference writes a zero - # scale and a saturated payload. + # scale and a zero payload. return torch.full(shape, 1e-4, dtype=torch.bfloat16, device="cuda") From ef199b4775347549e61df9fc4bcd34cb7882e2bf Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Thu, 5 Nov 2026 11:17:49 -0800 Subject: [PATCH 8/9] moe: Sum the NVFP4 trtllm MoE outputs in a fixed order in deterministic mode FlashInfer 0.7.2.post1 and 0.7.3 added a small-batch MoE finalize kernel that adds the top-k expert outputs in a different order from the large-batch kernel (flashinfer-ai/flashinfer#8831). A token's output then depends on its batch size, so prefill and decode logprobs no longer match. Under --enable-deterministic-inference, the pre-routed NVFP4 path now asks FlashInfer for the unsummed outputs and sums them with moe_finalize_fuse_shared, which adds them in the same order at every batch size. It writes into the usual output buffer through a new out= argument. Co-Authored-By: Claude Opus 5.5 --- .../ops/moe/moe_finalize_fuse_shared.py | 13 ++++-- .../moe/moe_runner/flashinfer_trtllm.py | 40 ++++++++++++++++++- 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/python/sglang/kernels/ops/moe/moe_finalize_fuse_shared.py b/python/sglang/kernels/ops/moe/moe_finalize_fuse_shared.py index 9e862f2c7500..b21c20ce65ab 100644 --- a/python/sglang/kernels/ops/moe/moe_finalize_fuse_shared.py +++ b/python/sglang/kernels/ops/moe/moe_finalize_fuse_shared.py @@ -24,6 +24,7 @@ def moe_finalize_fuse_shared( shared_output: Optional[torch.Tensor], top_k: int, enable_pdl: bool = False, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: assert gemm2_out.dtype == torch.bfloat16 assert expert_weights.dtype in (torch.float32, torch.bfloat16) @@ -42,9 +43,15 @@ def moe_finalize_fuse_shared( hidden_dim = shared_output.shape[1] assert hidden_dim <= gemm2_out.shape[1] - out = torch.empty( - num_tokens, hidden_dim, dtype=torch.bfloat16, device=gemm2_out.device - ) + if out is None: + out = torch.empty( + num_tokens, hidden_dim, dtype=torch.bfloat16, device=gemm2_out.device + ) + else: + assert out.dtype == torch.bfloat16 and out.is_contiguous() + assert out.dim() == 2 and out.shape[0] == num_tokens + assert out.shape[1] <= gemm2_out.shape[1] + assert shared_output is None or out.shape[1] == hidden_dim if shared_output is None: shared_output = gemm2_out.new_empty((0, 0), dtype=torch.bfloat16) 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 d3a85b7d2d9c..bf161b13f56e 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -34,7 +34,7 @@ register_fused_func, ) from sglang.srt.layers.utils import copy_or_rebind_param -from sglang.srt.runtime_context import get_forward, get_parallel +from sglang.srt.runtime_context import get_exec, get_forward, get_parallel from sglang.srt.utils.common import ( is_flashinfer_available, next_power_of_2, @@ -132,6 +132,24 @@ def _make_deferred_finalize_output( ) +def _finalize_in_order(result, *, top_k: int, out: torch.Tensor) -> torch.Tensor: + """Sum FlashInfer's ``do_finalize=False`` outputs into ``out``, adding the + top-k terms in the same order whatever the batch size.""" + from sglang.kernels.ops.moe.moe_finalize_fuse_shared import moe_finalize_fuse_shared + + deferred = _make_deferred_finalize_output(result, top_k=top_k) + num_tokens = out.shape[0] + return moe_finalize_fuse_shared( + deferred.gemm2_out.view(-1, deferred.gemm2_out.shape[-1]), + deferred.expanded_idx_to_permuted_idx.view(-1)[: num_tokens * top_k], + deferred.expert_weights.view(-1, top_k)[:num_tokens], + None, + top_k, + enable_pdl=trtllm_moe_enable_pdl(num_tokens), + out=out, + ) + + def round_up_to_multiple(x: int, m: int) -> int: """Round up *x* to the nearest multiple of *m*.""" return (x + m - 1) // m * m @@ -1561,6 +1579,20 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( symmetric=hidden_pad == 0, ) + # FlashInfer's small-batch finalize sums the top-k outputs in a different order + # from its large-batch finalize (flashinfer-ai/flashinfer#8831), so a token's + # output would depend on its batch size. In deterministic mode, take the + # unsummed outputs and sum them in one fixed order at every batch size. Only + # the pre-routed call: its routing weights already carry any routed scaling. + finalize_in_order = ( + use_routed_topk + and not defer_finalize + and hidden_pad == 0 + and symm_output.dtype == torch.bfloat16 + and get_exec().deterministic.enable_deterministic_inference + ) + flashinfer_finalize = not (defer_finalize or finalize_in_order) + if use_routed_topk: routing = _get_routing_for_flashinfer_routed(topk_output) moe_kwargs = dict( @@ -1590,7 +1622,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( local_num_experts=quant_info.local_num_experts, routed_scaling_factor=None, routing_method_type=1, # Unused, but must be 1 to pass validation. - do_finalize=not defer_finalize, + do_finalize=flashinfer_finalize, activation_type=activation_type, tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), enable_pdl=trtllm_moe_enable_pdl(hs_fp4.shape[0]), @@ -1603,6 +1635,10 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( result, top_k=_routing_top_k(routing), ) + elif finalize_in_order: + result = _finalize_in_order( + result, top_k=_routing_top_k(routing), out=symm_output + ) else: result = result[0] else: From cab9499ffe65b1f7f078f900bc01d721d336d594 Mon Sep 17 00:00:00 2001 From: Lee Nau Date: Thu, 5 Nov 2026 11:56:22 -0800 Subject: [PATCH 9/9] moe: Skip the deterministic-mode check when no server config is published Kernel tests such as test_minimax_m31_nvfp4_trtllm_moe.py call the trtllm NVFP4 MoE runner without publishing the server's runtime config, so get_exec() raised. Check that the exec namespace exists first, as vision.py does, and treat a missing config as deterministic mode off. Co-Authored-By: Claude Opus 5.5 --- .../srt/layers/moe/moe_runner/flashinfer_trtllm.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 bf161b13f56e..8a6adba4d876 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -34,7 +34,12 @@ register_fused_func, ) from sglang.srt.layers.utils import copy_or_rebind_param -from sglang.srt.runtime_context import get_exec, get_forward, get_parallel +from sglang.srt.runtime_context import ( + get_context, + get_exec, + get_forward, + get_parallel, +) from sglang.srt.utils.common import ( is_flashinfer_available, next_power_of_2, @@ -1589,6 +1594,8 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( and not defer_finalize and hidden_pad == 0 and symm_output.dtype == torch.bfloat16 + # Kernel tests call this runner without a published server config. + and get_context().is_config_namespace_published("exec") and get_exec().deterministic.enable_deterministic_inference ) flashinfer_finalize = not (defer_finalize or finalize_in_order)