Skip to content

Commit 60511fe

Browse files
committed
update
1 parent 7c32086 commit 60511fe

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

python/sglang/srt/model_executor/cuda_graph_runner.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
139139
gpu_mem = get_whatever_gpu_memory_capacity() / 1024
140140

141141
if gpu_mem > 120:
142-
capture_bs += list(range(160, 320, 8))
142+
capture_bs += list(range(160, 256, 8))
143143

144144
if max(capture_bs) > model_runner.req_to_token_pool.size:
145145
# In some case (e.g., with a small GPU or --max-running-requests), the #max-running-requests
@@ -196,6 +196,9 @@ def __init__(self, model_runner: ModelRunner):
196196

197197
# Batch sizes to capture
198198
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
199+
200+
print(f"\x1b[32mx)bs={self.capture_bs}\x1b[0m")
201+
199202
self.capture_forward_mode = ForwardMode.DECODE
200203
self.capture_hidden_mode = CaptureHiddenMode.NULL
201204
self.num_tokens_per_bs = 1

0 commit comments

Comments
 (0)