From 8fa56a0ab145a9e3626cd90bbcb554dc03a6cdb0 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Wed, 20 May 2026 23:01:28 -0700 Subject: [PATCH] Fix FlashInfer A2A token cap sizing (#25907) Co-authored-by: Jaewon <52840625+jaewonlee-fb@users.noreply.github.com> --- .../sglang/srt/layers/moe/token_dispatcher/flashinfer.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 5ec7e4d52..dc1f2ccbc 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -101,9 +101,12 @@ class FlashinferDispatcher(BaseDispatcher): self.payload_in_workspace = get_moe_runner_backend().is_flashinfer_cutlass() # TODO: Can this be a server arg and shared with deepep/mooncakeep? - self.max_num_tokens = ( - get_int_env_var("SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 1024) - * self.ep_size + # 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. + self.max_num_tokens = get_int_env_var( + "SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 4096 ) # Calculate workspace size. For eagle mode, use the larger workspace size since nextn layer will be unquantized.