diff --git a/python/sglang/srt/state_capturer/routed_experts.py b/python/sglang/srt/state_capturer/routed_experts.py index 553bb30ff..44f339e00 100644 --- a/python/sglang/srt/state_capturer/routed_experts.py +++ b/python/sglang/srt/state_capturer/routed_experts.py @@ -59,12 +59,14 @@ class RoutedExpertsCapturer(BaseTopkCapturer): num_layers = model_config.hf_text_config.num_hidden_layers server_args = get_global_server_args() - # FIXME: spec decoding is not accounted for here. The device buffer can - # overflow when max_running_requests * num_verify_tokens exceeds - # chunked_prefill_size * dp_size. + # Scale by dp_size so the buffer covers the full DP-concatenated batch. + # _get_local_slice indexes into [attention_dp_rank * cuda_graph_batch, ...) + # and otherwise overflows on dp_rank > 0 when max_running_requests > + # chunked_prefill_size. + # FIXME: spec decoding's num_verify_tokens is still not accounted for. max_batch_size = max( server_args.chunked_prefill_size * server_args.dp_size, - max_running_requests, + max_running_requests * server_args.dp_size, ) super().__init__(