EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Minor style fix in cuda_graph_runner.py (#7053)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-06-10 06:32:41 -07:00
committed by GitHub
parent 6406408a70
commit 4a102a2b02
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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