iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][Spec Decode] Keep default CUDA graph sizes memory-safe (#54418)

Browse Source 
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com>
This commit is contained in:
Kevin H. Luu
2026-08-30 18:02:31 +00:00
committed by GitHub
parent 9d0fe9bac8
commit b2dc864bb6
3 changed files with 73 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/compile/test_config.py
+64 -56
View File
@@ -542,34 +542,36 @@ def _mock_config_for_cudagraph_sizes(
@pytest.mark.parametrize(
("max_num_seqs", "num_speculative_tokens"),
("max_num_seqs", "num_speculative_tokens", "widest_is_captured"),
[
# No speculation: the 2x headroom under the platform ceiling, unchanged.
(8, 0),
(32, 0),
(8, 0, True),
(32, 0, True),
# Speculating, but the widest decode batch still fits under the ceiling.
(64, 7),
(256, 1),
# Widest decode batch above the ceiling, which must not cut below it.
(32, 16),
(64, 16),
(64, 7, True),
(256, 1, True),
# Wider decode batches must not raise the memory-safety ceiling.
(32, 16, False),
(64, 16, False),
# Exact H200 regressions: DFlash (2176) and suffix decoding (6400).
(128, 16, False),
(256, 24, False),
# Widest decode batch off the capture stride, above the ceiling.
(33, 16),
(33, 16, False),
# ... and off the stride while below it, where the ceiling stands but
# the generated sizes would otherwise stop at 400.
(24, 16),
(24, 16, True),
],
)
def test_default_cudagraph_capture_size_covers_widest_uniform_decode(
max_num_seqs, num_speculative_tokens
def test_default_cudagraph_capture_size_respects_platform_ceiling(
max_num_seqs, num_speculative_tokens, widest_is_captured
):
"""The widest uniform decode batch has to be inside the captured range.
"""Uniform decode coverage must not raise the platform default ceiling.
A decode step presents up to `max_num_seqs * (1 + num_speculative_tokens)`
tokens. If that size is not captured the decode path silently falls back to
eager while the resolved mode still reports full decode graphs, so neither
the platform's token-denominated ceiling nor the capture stride may stop
short of it.
tokens. Off-stride sizes within the platform ceiling are captured exactly.
Larger batches fall back to eager execution instead of expanding the
default capture range and risking OOM during engine initialization.
"""
compilation_config = CompilationConfig(
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
@@ -592,20 +594,21 @@ def test_default_cudagraph_capture_size_covers_widest_uniform_decode(
default_max_graph_size,
)
widest_uniform_decode = max_num_seqs * decode_query_len
expected_max_size = max(token_grid_max, widest_uniform_decode)
assert compilation_config.max_cudagraph_capture_size == expected_max_size
assert compilation_config.max_cudagraph_capture_size == token_grid_max
assert (
compilation_config.max_cudagraph_capture_size
== compilation_config.cudagraph_capture_sizes[-1]
)
assert widest_uniform_decode in compilation_config.cudagraph_capture_sizes
assert (
widest_uniform_decode in compilation_config.cudagraph_capture_sizes
) is widest_is_captured
def test_default_cudagraph_capture_sizes_keep_token_grid_bounded():
"""Keep the token grid bounded; only request counts may exceed the default.
def test_default_cudagraph_capture_sizes_keep_all_sizes_bounded():
"""Keep both token and uniform-decode grids under the platform default.
This guards the measured 581 versus 100 capture-size regression at
max_num_seqs=512 with 16 draft tokens.
This guards both the measured 581 versus 100 capture-count regression and
the graph-memory regression from capturing shapes up to 8704 tokens.
"""
max_num_seqs = 512
decode_query_len = 17
@@ -640,12 +643,19 @@ def test_default_cudagraph_capture_sizes_keep_token_grid_bounded():
request_counts += list(range(256, max_request_count + 1, 16))
request_counts.append(max_request_count)
expected_sizes = sorted(
set(token_grid + [count * decode_query_len for count in request_counts])
set(
token_grid
+ [
count * decode_query_len
for count in request_counts
if count * decode_query_len <= default_max_graph_size
]
)
)
assert compilation_config.cudagraph_capture_sizes == expected_sizes
assert all(
size <= default_max_graph_size or size % decode_query_len == 0
size <= default_max_graph_size
for size in compilation_config.cudagraph_capture_sizes
)
@@ -675,12 +685,13 @@ def test_cudagraph_capture_sizes_respect_sequence_parallelism():
VllmConfig._set_cudagraph_sizes(config)
assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes)
assert 544 in compilation_config.cudagraph_capture_sizes
assert 544 not in compilation_config.cudagraph_capture_sizes
assert compilation_config.max_cudagraph_capture_size == 512
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
def test_real_vllm_config_captures_widest_ngram_decode_batch():
"""Real config post-init preserves the widest static ngram decode batch."""
def test_real_vllm_config_caps_widest_ngram_decode_batch():
"""Real config post-init preserves the default capture ceiling."""
model_config = ModelConfig(model="facebook/opt-125m")
speculative_config = SpeculativeConfig(
prompt_lookup_min=1,
@@ -708,11 +719,11 @@ def test_real_vllm_config_captures_widest_ngram_decode_batch():
compilation_config=compilation_config,
)
assert 544 in config.compilation_config.cudagraph_capture_sizes
assert 544 not in config.compilation_config.cudagraph_capture_sizes
assert (
config.compilation_config.max_cudagraph_capture_size
== config.compilation_config.cudagraph_capture_sizes[-1]
== 544
== 512
)
@@ -750,13 +761,12 @@ def test_default_cudagraph_capture_size_unchanged_without_speculation(max_num_se
)
def test_default_cudagraph_capture_size_covers_a_single_speculative_token():
"""One speculative token is already enough to lose the widest batch.
def test_single_speculative_token_does_not_raise_default_capture_size():
"""One speculative token must not raise the default capture ceiling.
MTP at depth 1 gives a query length of 2, so 300 requests is a 600-token
decode batch against the platform's token ceiling. Gating the unit conversion on a
deeper speculative width would leave this common configuration dispatching
its largest decode steps eager.
decode batch against the platform's token ceiling. The widest batch falls
back to eager execution rather than expanding the capture range.
"""
compilation_config = CompilationConfig(
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
@@ -774,21 +784,19 @@ def test_default_cudagraph_capture_size_covers_a_single_speculative_token():
1024 if current_platform.is_device_capability_family(100) else 512
)
token_grid_max = min(300 * 2 * 2, default_max_graph_size)
widest_uniform_decode = 300 * 2
expected_max_size = max(token_grid_max, widest_uniform_decode)
assert compilation_config.max_cudagraph_capture_size == expected_max_size
assert compilation_config.max_cudagraph_capture_size == token_grid_max
assert (
compilation_config.max_cudagraph_capture_size
== compilation_config.cudagraph_capture_sizes[-1]
)
assert 600 in compilation_config.cudagraph_capture_sizes
assert 600 not in compilation_config.cudagraph_capture_sizes
def test_default_cudagraph_capture_size_caps_requests_not_tokens():
"""The ceiling admits at most the platform default number of requests.
def test_default_cudagraph_capture_size_caps_tokens():
"""The ceiling remains a platform-bounded token count.
`max_num_seqs` far above the platform default must not push the capture
range up without bound just because each request is now several tokens wide.
`max_num_seqs` and speculative width must not push the default capture range
above the platform's memory-safety guard.
"""
compilation_config = CompilationConfig(
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
@@ -805,16 +813,15 @@ def test_default_cudagraph_capture_size_caps_requests_not_tokens():
default_max_graph_size = (
1024 if current_platform.is_device_capability_family(100) else 512
)
decode_query_len = 17
token_grid_max = min(1024 * decode_query_len * 2, default_max_graph_size)
widest_covered_decode = min(1024, default_max_graph_size) * decode_query_len
expected_max_size = max(token_grid_max, widest_covered_decode)
assert compilation_config.max_cudagraph_capture_size == expected_max_size
assert compilation_config.max_cudagraph_capture_size == default_max_graph_size
assert (
compilation_config.max_cudagraph_capture_size
== compilation_config.cudagraph_capture_sizes[-1]
)
assert expected_max_size == default_max_graph_size * 17
assert all(
size <= default_max_graph_size
for size in compilation_config.cudagraph_capture_sizes
)
def _widest_covered_request_count(capture_sizes, query_len, max_num_seqs):
@@ -844,20 +851,21 @@ def test_default_cudagraph_capture_sizes_cover_every_dynamic_decode_width():
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
)
config = _mock_config_for_cudagraph_sizes(
max_num_seqs=256,
max_num_seqs=128,
num_speculative_tokens=16,
max_num_batched_tokens=32768,
compilation_config=compilation_config,
# 16 draft tokens up to a batch of 16, then 2 out to 256.
num_speculative_tokens_per_batch_size=[(1, 16, 16), (17, 256, 2)],
# 16 draft tokens up to a batch of 16, then 2 out to 128. Both
# tier maxima fit under the 512-token default capture ceiling.
num_speculative_tokens_per_batch_size=[(1, 16, 16), (17, 128, 2)],
)
VllmConfig._set_cudagraph_sizes(config)
sizes = compilation_config.cudagraph_capture_sizes
# The wide tier only ever runs to a batch of 16, the narrow one to 256.
assert _widest_covered_request_count(sizes, 17, 256) >= 16
assert _widest_covered_request_count(sizes, 3, 256) == 256
# The wide tier only ever runs to a batch of 16, the narrow one to 128.
assert _widest_covered_request_count(sizes, 17, 128) >= 16
assert _widest_covered_request_count(sizes, 3, 128) == 128
def test_dynamic_decode_capture_clamps_configured_width():
vllm/config/compilation.py
+2 -3
View File
@@ -703,9 +703,8 @@ class CompilationConfig:
If not specified, max_cudagraph_capture_size is capped at 512 by default,
or 1024 on data center Blackwell GPUs. This avoids OOM in tight memory
scenarios with small max_num_seqs, and limits capture of large graphs that
increase startup time and memory usage. When the uniform decode query length
exceeds one, uniform decode sizes are appended and can raise the final value
above this default.
increase startup time and memory usage. Uniform decode sizes are appended
only within this default ceiling.
"""
dynamic_shapes_config: DynamicShapesConfig = field(
vllm/config/vllm.py
+7 -8
View File
@@ -1953,10 +1953,9 @@ class VllmConfig:
`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. Likewise, when the uniform decode query length
exceeds one, the widest uniform decode batch (`max_num_seqs *
decode_query_len`) is appended when it fits, since it need not land on
an 8- or 16-token stride.
even when off-stride. Uniform decode sizes are appended when they fit
within the platform's default capture ceiling, since they need not land
on an 8- or 16-token stride.
In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
will be the final sizes to capture cudagraph (in ascending order).
@@ -2085,6 +2084,7 @@ class VllmConfig:
n * query_len
for query_len, tier_max_reqs in decode_tiers
for n in request_counts(tier_max_reqs)
if n * query_len <= max_cudagraph_capture_size
}
)
max_num_tokens = self.scheduler_config.max_num_batched_tokens
@@ -2134,10 +2134,9 @@ class VllmConfig:
and max_num_tokens not in cudagraph_capture_sizes
):
cudagraph_capture_sizes.append(max_num_tokens)
# These extend past the token-strided ceiling, which counts one
# token per request. valid_max_size below raises the final
# max_cudagraph_capture_size to the widest of them. They remain
# filtered by max_num_tokens.
# Preserve the platform's default capture ceiling. Larger
# uniform decode batches fall back to eager execution unless
# users explicitly configure wider capture sizes.
cudagraph_capture_sizes += [
size for size in uniform_decode_sizes if size <= max_num_tokens
]