diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index bbb9cb1fcb..d822b68c50 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -405,6 +405,9 @@ def test_should_split(): (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0), # truncated to nearest multiple of 8 or 16 (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + # max_num_batched_tokens <= max_cudagraph_capture_size should always be + # captured even if not landing on a 16-stride step + (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), # max from list ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), # SP forces full-graph compilation, sizes are filtered by TP diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f591605d08..56123542ce 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1432,6 +1432,10 @@ class VllmConfig: cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list( range(256, max_graph_size + 1, 16)) + `max_num_batched_tokens` is also appended to the list if it fits + within `max_cudagraph_capture_size`, so the max batch size is captured + even when off-stride. + In the end, `vllm_config.compilation_config.cudagraph_capture_sizes` will be the final sizes to capture cudagraph (in ascending order). @@ -1520,6 +1524,12 @@ class VllmConfig: cudagraph_capture_sizes += list( range(256, max_cudagraph_capture_size + 1, 16) ) + # ensure max_num_tokens is captured if within max capture size + if ( + max_num_tokens <= max_cudagraph_capture_size + and max_num_tokens not in cudagraph_capture_sizes + ): + cudagraph_capture_sizes.append(max_num_tokens) # de-duplicate and sort the sizes cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))