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

Fix FlashInfer A2A token cap sizing (#25907)

Browse Source 
Co-authored-by: Jaewon <52840625+jaewonlee-fb@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2026-05-20 23:01:28 -07:00
committed by GitHub
co-authored by Jaewon
parent e8608bdcb5
commit 8fa56a0ab1
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
+6 -3
View File
@@ -101,9 +101,12 @@ class FlashinferDispatcher(BaseDispatcher):
self.payload_in_workspace = get_moe_runner_backend().is_flashinfer_cutlass() self.payload_in_workspace = get_moe_runner_backend().is_flashinfer_cutlass()
# TODO: Can this be a server arg and shared with deepep/mooncakeep? # TODO: Can this be a server arg and shared with deepep/mooncakeep?
self.max_num_tokens = ( # FlashInfer sizes the workspace from the maximum dispatched tokens per
get_int_env_var("SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 1024) # EP rank. See FlashInfer's moe_a2a_get_workspace_size_per_rank(),
* self.ep_size # 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. # Calculate workspace size. For eagle mode, use the larger workspace size since nextn layer will be unquantized.