From 878e6b8886fff5e23ee89126add7cefb549b74d0 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 18 May 2026 15:59:33 -0700 Subject: [PATCH] [SP] Fix runtime_max_tokens_per_rank for sequence parallelism (#25685) Co-authored-by: Ming Yang Co-authored-by: Yinghai Lu --- .../srt/layers/moe/token_dispatcher/flashinfer.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index a49481ca8..f05201143 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -212,11 +212,16 @@ class FlashinferDispatcher(BaseDispatcher): payloads.append(topk_ids) payloads.append(topk_weights) - self.runtime_max_tokens_per_rank = ( - max(get_dp_global_num_tokens()) - if get_dp_global_num_tokens() is not None - else x.shape[0] - ) + dp_global = get_dp_global_num_tokens() + if dp_global is not None and len(dp_global) > 1: + # DP attention: multiple DP ranks with different token counts. + # Use the max across ranks so the A2A workspace fits the fattest. + self.runtime_max_tokens_per_rank = max(dp_global) + else: + # dp_size=1 or SP: use the actual input tensor size (post-scatter + # in SP mode, full batch otherwise). Avoids the pre-scatter + # scheduler count which can exceed the workspace cap. + self.runtime_max_tokens_per_rank = x.shape[0] recv_tensors = self.moe_a2a.dispatch( self.dummy_topk_ids_current_rank if self.has_dummy_token else topk_ids, payloads,