From 343998865ae224922a80aefeeb0ff1c97b5d593b Mon Sep 17 00:00:00 2001 From: Yilong Zhao <74357408+happierpig@users.noreply.github.com> Date: Sun, 22 Mar 2026 17:06:16 -0700 Subject: [PATCH] perf: pad max-num-requests in decode cuda graph for higher coverage (#20978) --- .../srt/model_executor/cuda_graph_runner.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index ebe4c41d7..05e2d5887 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -454,14 +454,9 @@ def set_torch_compile_config(): def get_batch_sizes_to_capture(model_runner: ModelRunner, num_tokens_per_bs=1): server_args = model_runner.server_args capture_bs = server_args.cuda_graph_bs - - if max(capture_bs) > model_runner.req_to_token_pool.size: - # In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests - # is very small. We add more values here to make sure we capture the maximum bs. - capture_bs += [model_runner.req_to_token_pool.size] + num_max_requests = model_runner.req_to_token_pool.size mul_base = 1 - if server_args.enable_two_batch_overlap: mul_base *= 2 num_tokens_per_bs = 1 # tbo not test, set num_tokens_per_bs to 1 @@ -472,11 +467,18 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner, num_tokens_per_bs=1): if mul_base % get_attention_cp_size() != 0: mul_base *= get_attention_cp_size() + # pad `num_max_requests` to avoid being filtered out + num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base + if max(capture_bs) > num_max_requests: + # In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests + # is very small. We add more values here to make sure we capture the maximum bs. + capture_bs += [num_max_requests] + # Model input token count = bs * num_tokens_per_bs; must be a multiple of attn_tp_size. capture_bs = [bs for bs in capture_bs if bs * num_tokens_per_bs % mul_base == 0] - - capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size] + capture_bs = [bs for bs in capture_bs if bs <= num_max_requests] capture_bs = list(sorted(set(capture_bs))) + assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}" compile_bs = ( [bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]