Skip to content

Commit ef5b9aa

Browse files
committed
[MISC] Fix misleading batch_size_capture_list when cuda_graph_sizes < 4
Signed-off-by: billishyahao <[email protected]>
1 parent da63274 commit ef5b9aa

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

vllm/config/__init__.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -617,9 +617,10 @@ def _set_cudagraph_sizes(self):
617617
not self.model_config.enforce_eager:
618618
cuda_graph_sizes = self.scheduler_config.cuda_graph_sizes
619619
if len(cuda_graph_sizes) == 1:
620-
batch_size_capture_list = [1, 2, 4] + [
621-
i for i in range(8, cuda_graph_sizes[0] + 1, 8)
622-
]
620+
max_graph_size = cuda_graph_sizes[0]
621+
batch_size_capture_list = [
622+
i for i in [1, 2, 4] if i <= max_graph_size
623+
] + [i for i in range(8, max_graph_size + 1, 8)]
623624
elif len(cuda_graph_sizes) > 1:
624625
batch_size_capture_list = sorted(cuda_graph_sizes)
625626
else:

0 commit comments

Comments
 (0)