From 4a102a2b02f120fc84c481019e35b2b041ce06f7 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 10 Jun 2025 06:32:41 -0700 Subject: [PATCH] Minor style fix in cuda_graph_runner.py (#7053) --- python/sglang/srt/model_executor/cuda_graph_runner.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index d42fd8466..be8a1ad14 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -152,10 +152,11 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner): ) gpu_mem = get_device_memory_capacity() - if gpu_mem is not None and gpu_mem > 96 * 1024: - capture_bs += list(range(160, 257, 8)) - if gpu_mem is not None and gpu_mem > 180 * 1000: - capture_bs += list(range(256, 513, 16)) + if gpu_mem is not None: + if gpu_mem > 90 * 1024: # H200 + capture_bs += list(range(160, 257, 8)) + if gpu_mem > 160 * 1000: # B200, MI300 + capture_bs += list(range(256, 513, 16)) 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