A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 13:45 UTC
[Bugfix] Fix max_num_batched_token not captured in cuda graph (#40734)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Wei Zhao (Engrg-Hardware 1) <weizha@login-bia02.bia.clusters.nvidia.com>
Wei Zhao committed
Apr 29, 2026 at 00:33 UTC
8b49cf3a37eb1a267a06b0df23328909330af1e6
2 files changed
+13
tests/compile/test_config.py
+3
@@ -405,6 +405,9 @@ def test_should_split():
405
(None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0),
406
# truncated to nearest multiple of 8 or 16
407
(None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
408
+ # max_num_batched_tokens <= max_cudagraph_capture_size should always be
409
+ # captured even if not landing on a 16-stride step
410
+ (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257),
411
# max from list
412
([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
413
# SP forces full-graph compilation, sizes are filtered by TP
vllm/config/vllm.py
+10
@@ -1432,6 +1432,10 @@ class VllmConfig:
1432
cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
1433
range(256, max_graph_size + 1, 16))
1434
1435
+ `max_num_batched_tokens` is also appended to the list if it fits
1436
+ within `max_cudagraph_capture_size`, so the max batch size is captured
1437
+ even when off-stride.
1438
+
1439
In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
1440
will be the final sizes to capture cudagraph (in ascending order).
1441
@@ -1520,6 +1524,12 @@ class VllmConfig:
1524
cudagraph_capture_sizes += list(
1525
range(256, max_cudagraph_capture_size + 1, 16)
1526
)
1527
+ # ensure max_num_tokens is captured if within max capture size
1528
+ if (
1529
+ max_num_tokens <= max_cudagraph_capture_size
1530
+ and max_num_tokens not in cudagraph_capture_sizes
1531
+ ):
1532
+ cudagraph_capture_sizes.append(max_num_tokens)
1533
# de-duplicate and sort the sizes
1534
cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))
1535