[Bugfix] Max-load throughput cliff when `max_num_seqs` is not a multiple of 8 (#57355)
Signed-off-by: Andy Lo <andy@mistral.ai> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andy Lo committed
Sep 17, 2026 at 19:09 UTC
09c379ab6183085fd90e8195792cc63f5247cdf3
2 files changed
+24
-3
tests/compile/test_config.py
+19
@@ -541,6 +541,25 @@ def _mock_config_for_cudagraph_sizes(
541
return config
542
543
544
+@pytest.mark.parametrize("max_num_seqs", [100, 101])
545
+def test_default_cudagraph_capture_sizes_cover_off_stride_max_num_seqs(
546
+ max_num_seqs: int,
547
+) -> None:
548
+ compilation_config = CompilationConfig(
549
+ cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
550
+ )
551
+ config = _mock_config_for_cudagraph_sizes(
552
+ max_num_seqs=max_num_seqs,
553
+ num_speculative_tokens=0,
554
+ max_num_batched_tokens=32768,
555
+ compilation_config=compilation_config,
556
+ )
557
+
558
+ VllmConfig._set_cudagraph_sizes(config)
559
+
560
+ assert max_num_seqs in compilation_config.cudagraph_capture_sizes
561
+
562
+
563
@pytest.mark.parametrize(
564
("max_num_seqs", "num_speculative_tokens", "widest_is_captured"),
565
[
vllm/config/vllm.py
+5
-3
@@ -2223,9 +2223,9 @@ class VllmConfig:
2223
max_cudagraph_capture_size = (
2224
self.compilation_config.max_cudagraph_capture_size
2225
)
2226
- # Decode sizes to cover, in tokens. Populated only when a request
2227
- # is more than one token wide and only when the default is computed
2228
- # here, so an explicit capture range is left exactly as configured.
2226
+ # Decode sizes to cover, in tokens. Populated only when the default
2227
+ # is computed here, so an explicit capture range is left exactly as
2228
+ # configured.
2229
uniform_decode_sizes: list[int] = []
2230
if max_cudagraph_capture_size is None:
2231
from vllm.platforms import current_platform
@@ -2316,6 +2316,8 @@ class VllmConfig:
2316
if n * query_len <= max_cudagraph_capture_size
2317
}
2318
)
2319
+ elif max_num_seqs <= max_cudagraph_capture_size:
2320
+ uniform_decode_sizes = [max_num_seqs]
2321
max_num_tokens = self.scheduler_config.max_num_batched_tokens
2322
max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)
2323