@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cc06b4e86b

Fork 0

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 5, 21:00 UTC

[Mamba][Bugfix] Raise on insufficient cache blocks instead of silently capping cudagraph sizes (#38270)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Mar 30, 2026 at 11:41 UTC cc06b4e86b2beb04fbee3e6d9167cc97f1491b1f

4 files changed +41 -131

tests/compile/test_config.py

-42

@@ -577,48 +577,6 @@ def test_compile_sizes_padding_validation():

577 dispatcher.initialize_cudagraph_keys(CUDAGraphMode.NONE) # Should not raise

578

579

580 -@pytest.mark.parametrize(

581 - "capture_sizes, max_size, num_blocks, expected_sizes, expected_max",

582 - [

583 - # Normal capping: sizes filtered to <= num_blocks

584 - (

585 - [1, 2, 4, 8, 16, 32, 64, 128, 256, 512],

586 - 512,

587 - 200,

588 - [1, 2, 4, 8, 16, 32, 64, 128],

589 - 128,

590 - ),

591 - # No capping needed: num_blocks >= max

592 - ([1, 2, 4, 8, 16], 16, 1000, [1, 2, 4, 8, 16], 16),

593 - # Exact boundary: num_blocks == max (no capping)

594 - ([1, 2, 4, 8, 16, 32], 32, 32, [1, 2, 4, 8, 16, 32], 32),

595 - # All sizes capped: num_blocks < smallest size

596 - ([8, 16, 32], 32, 4, [], 0),

597 - # num_blocks <= 0: early return, no change

598 - ([1, 2, 4], 4, 0, [1, 2, 4], 4),

599 - ],

600 -)

601 -def test_adjust_cudagraph_sizes_for_mamba_cache(

602 - capture_sizes, max_size, num_blocks, expected_sizes, expected_max

603 -):

604 - """Test that cudagraph capture sizes are correctly capped to fit

605 - available Mamba cache blocks.

606 -

607 - See: https://github.com/vllm-project/vllm/issues/34094

608 - """

609 - config = CompilationConfig(

610 - cudagraph_capture_sizes=capture_sizes,

611 - max_cudagraph_capture_size=max_size,

612 - cudagraph_mode=CUDAGraphMode.NONE,

613 - )

614 - config.adjust_cudagraph_sizes_for_mamba_cache(num_blocks)

615 - assert config.cudagraph_capture_sizes == expected_sizes

616 - assert config.max_cudagraph_capture_size == expected_max

617 - # Invariant: last element == max_cudagraph_capture_size

618 - if expected_sizes:

619 - assert config.cudagraph_capture_sizes[-1] == config.max_cudagraph_capture_size

620 -

621 -

580 def test_inductor_asserts_default_disabled(monkeypatch):

581 """Test that inductor runtime asserts are disabled by default

582 (INFO logging level) on torch < 2.12."""

tests/v1/worker/test_gpu_model_runner.py

+7 -22

@@ -1191,9 +1191,9 @@ def test_is_uniform_decode() -> None:

1191 current_platform.is_rocm(),

1192 reason="Attention backend FLASHINFER is not supported on ROCm.",

1193 )

1194 -def test_cudagraph_sizes_capped_for_mamba_cache():

1195 - """Test that cudagraph capture sizes are capped to num_blocks for

1196 - hybrid models with Mamba layers.

1194 +def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():

1195 + """Test that a ValueError is raised when max_num_seqs exceeds the

1196 + available Mamba cache blocks for hybrid models with FULL cudagraphs.

1197

1198 See: https://github.com/vllm-project/vllm/issues/34094

1199 """

@@ -1284,23 +1284,8 @@ def test_cudagraph_sizes_capped_for_mamba_cache():

1284 )[0]

1285 num_blocks = kv_cache_config.num_blocks

1286

1287 - # Set max_cudagraph_capture_size to a value larger than num_blocks

1288 - # to trigger the Mamba capping logic.

1289 - large_max = num_blocks + 100

1290 - compilation_config = vllm_config.compilation_config

1291 - compilation_config.max_cudagraph_capture_size = large_max

1292 - compilation_config.cudagraph_capture_sizes = [

1293 - s for s in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] if s <= large_max

1294 - ]

1287 + # Force max_num_seqs to exceed num_blocks so the check triggers.

1288 + runner.max_num_reqs = num_blocks + 100

1289

1296 - runner.initialize_kv_cache(kv_cache_config)

1297 -

1298 - # After initialization, cudagraph sizes should be capped

1299 - assert compilation_config.max_cudagraph_capture_size <= num_blocks

1300 - assert all(s <= num_blocks for s in compilation_config.cudagraph_capture_sizes)

1301 - # Invariant: last element == max

1302 - if compilation_config.cudagraph_capture_sizes:

1303 - assert (

1304 - compilation_config.cudagraph_capture_sizes[-1]

1305 - == compilation_config.max_cudagraph_capture_size

1306 - )

1290 + with pytest.raises(ValueError, match="max_num_seqs"):

1291 + runner.initialize_kv_cache(kv_cache_config)

vllm/config/compilation.py

-52

@@ -1279,58 +1279,6 @@ class CompilationConfig:

1279 self.max_cudagraph_capture_size = rounded_sizes[-1]

1280 self.cudagraph_capture_sizes = rounded_sizes

1281

1282 - def adjust_cudagraph_sizes_for_mamba_cache(

1283 - self, num_mamba_cache_blocks: int

1284 - ) -> None:

1285 - """Cap cudagraph capture sizes to available Mamba cache blocks.

1286 -

1287 - For hybrid Mamba/attention models, the Mamba conv_state and

1288 - ssm_state tensors have their first dimension equal to num_blocks

1289 - (from KVCacheConfig). During CUDA graph capture the decode batch

1290 - size equals num_tokens, so capture sizes exceeding num_blocks

1291 - would cause out-of-bounds access in Mamba kernels.

1292 -

1293 - See: https://github.com/vllm-project/vllm/issues/34094

1294 - """

1295 - if not self.cudagraph_capture_sizes or num_mamba_cache_blocks <= 0:

1296 - return

1297 -

1298 - assert self.max_cudagraph_capture_size is not None

1299 -

1300 - if num_mamba_cache_blocks >= self.max_cudagraph_capture_size:

1301 - return

1302 -

1303 - capped_sizes = [

1304 - s for s in self.cudagraph_capture_sizes if s <= num_mamba_cache_blocks

1305 - ]

1306 -

1307 - if len(capped_sizes) == 0:

1308 - logger.warning(

1309 - "No valid cudagraph capture sizes remain after capping "

1310 - "to Mamba cache blocks (%d). The smallest capture size "

1311 - "was %d. Disabling cudagraph capture. Consider reducing "

1312 - "max_num_seqs or increasing available GPU memory.",

1313 - num_mamba_cache_blocks,

1314 - self.cudagraph_capture_sizes[0],

1315 - )

1316 - self.cudagraph_capture_sizes = []

1317 - self.max_cudagraph_capture_size = 0

1318 - return

1319 -

1320 - logger.warning(

1321 - "Capping cudagraph capture sizes from max %d to %d to fit "

1322 - "Mamba cache blocks (%d blocks available). This limits the "

1323 - "maximum batch size that can use CUDA graphs. To increase "

1324 - "this limit, reduce max_num_seqs or increase available GPU "

1325 - "memory.",

1326 - self.max_cudagraph_capture_size,

1327 - capped_sizes[-1],

1328 - num_mamba_cache_blocks,

1329 - )

1330 -

1331 - self.max_cudagraph_capture_size = capped_sizes[-1]

1332 - self.cudagraph_capture_sizes = capped_sizes

1333 -

1282 def get_compile_ranges(self) -> list[Range]:

1283 """Get the compile ranges for the compilation config."""

1284 if self.compile_ranges_endpoints is None:

vllm/v1/worker/gpu_model_runner.py

+34 -15

@@ -5800,7 +5800,7 @@ class GPUModelRunner(

5800 )

5801 self.cache_config.num_gpu_blocks_override = saved_override

5802

5803 - self.initialize_kv_cache(minimal_config)

5803 + self.initialize_kv_cache(minimal_config, is_profiling=True)

5804 self.cache_config.num_gpu_blocks = minimal_config.num_blocks

5805

5806 logger.debug("Initialized minimal KV cache for CUDA graph profiling")

@@ -6121,7 +6121,11 @@ class GPUModelRunner(

6121 torch.accelerator.synchronize()

6122 self.maybe_remove_all_loras(self.lora_config)

6123

6124 - def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:

6124 + def initialize_attn_backend(

6125 + self,

6126 + kv_cache_config: KVCacheConfig,

6127 + is_profiling: bool = False,

6128 + ) -> None:

6129 """

6130 Initialize the attention backends and attention metadata builders.

6131 """

@@ -6193,7 +6197,9 @@ class GPUModelRunner(

6197

6198 # Resolve cudagraph_mode before actually initialize metadata_builders

6199 self._check_and_update_cudagraph_mode(

6196 - attention_backend_list, kv_cache_config.kv_cache_groups

6200 + attention_backend_list,

6201 + kv_cache_config.kv_cache_groups,

6202 + is_profiling=is_profiling,

6203 )

6204

6205 # Check if attention backend supports PCP&DCP and related features.

@@ -6237,6 +6243,7 @@ class GPUModelRunner(

6243 self,

6244 attention_backends: list[set[type[AttentionBackend]]],

6245 kv_cache_groups: list[KVCacheGroupSpec],

6246 + is_profiling: bool = False,

6247 ) -> None:

6248 """

6249 Resolve the cudagraph_mode when there are multiple attention

@@ -6377,21 +6384,29 @@ class GPUModelRunner(

6384 self.uniform_decode_query_len, self.parallel_config.tensor_parallel_size

6385 )

6386

6380 - # If the model has Mamba layers and cudagraph mode includes FULL

6381 - # decode, cap cudagraph capture sizes to the number of available

6382 - # Mamba cache blocks. Each decode request needs one conv_state

6383 - # cache line, so capture batch sizes cannot exceed num_blocks.

6384 - # Only FULL decode graphs are affected because PIECEWISE captures

6385 - # run GDN/Mamba ops eagerly (prefill path, no causal_conv1d_update).

6387 + # For Mamba models with FULL decode cudagraphs, each decode

6388 + # sequence needs one Mamba cache block. The decode cudagraph

6389 + # dispatcher already caps batch sizes at max_num_seqs, so we just

6390 + # need to verify that enough blocks exist. Raising here instead

6391 + # of silently capping cudagraph_capture_sizes avoids unintended

6392 + # restrictions on PIECEWISE (prefill) cudagraphs.

6393 # See: https://github.com/vllm-project/vllm/issues/34094

6387 - if cudagraph_mode.has_full_cudagraphs():

6394 + if cudagraph_mode.has_full_cudagraphs() and not is_profiling:

6395 has_mamba = any(

6396 isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_groups

6397 )

6398 if has_mamba and self.kv_cache_config is not None:

6392 - self.compilation_config.adjust_cudagraph_sizes_for_mamba_cache(

6393 - self.kv_cache_config.num_blocks

6394 - )

6399 + num_blocks = self.kv_cache_config.num_blocks

6400 + if self.max_num_reqs > num_blocks:

6401 + raise ValueError(

6402 + f"max_num_seqs ({self.max_num_reqs}) exceeds "

6403 + f"available Mamba cache blocks ({num_blocks}). "

6404 + f"Each decode sequence requires one Mamba cache "

6405 + f"block, so CUDA graph capture cannot proceed. "

6406 + f"Please lower max_num_seqs to at most "

6407 + f"{num_blocks} or increase "

6408 + f"gpu_memory_utilization."

6409 + )

6410

6411 # Trigger cudagraph dispatching keys initialization after

6412 # resolved cudagraph mode.

@@ -6752,7 +6767,11 @@ class GPUModelRunner(

6767 else:

6768 break

6769

6755 - def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:

6770 + def initialize_kv_cache(

6771 + self,

6772 + kv_cache_config: KVCacheConfig,

6773 + is_profiling: bool = False,

6774 + ) -> None:

6775 """

6776 Initialize KV cache based on `kv_cache_config`.

6777 Args:

@@ -6764,7 +6783,7 @@ class GPUModelRunner(

6783 self._mamba_copy_bufs = None

6784 self.may_add_encoder_only_layers_to_kv_cache_config()

6785 self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)

6767 - self.initialize_attn_backend(kv_cache_config)

6786 + self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)

6787 # The kernel block size for all KV cache groups. For example, if

6788 # kv_cache_manager uses block_size 256 for a given group, but the attention

6789 # backends for that group only supports block_size 64, we will return