chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

perf: pad max-num-requests in decode cuda graph for higher coverage (#20978)

Browse Source
This commit is contained in:
Yilong Zhao
2026-03-22 17:06:16 -07:00
committed by GitHub
parent ce0541404f
commit 343998865a
1 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -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]