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

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 7, 07:45 UTC

[Bugfix] Cap FULL decode cudagraph sizes for Mamba/hybrid models (#34094) (#34571)

Signed-off-by: haosdent <haosdent@gmail.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>

haosdent committed Mar 4, 2026 at 18:56 UTC d6e04f4c43612eb9fbbafc3723da8144d54dfde9

4 files changed +230

tests/compile/test_config.py

+42

@@ -570,3 +570,45 @@ def test_compile_sizes_padding_validation():

570 assert sorted(config.compile_sizes) == [3, 5, 7]

571 dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config))

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

573 +

574 +

575 +@pytest.mark.parametrize(

576 + "capture_sizes, max_size, num_blocks, expected_sizes, expected_max",

577 + [

578 + # Normal capping: sizes filtered to <= num_blocks

579 + (

580 + [1, 2, 4, 8, 16, 32, 64, 128, 256, 512],

581 + 512,

582 + 200,

583 + [1, 2, 4, 8, 16, 32, 64, 128],

584 + 128,

585 + ),

586 + # No capping needed: num_blocks >= max

587 + ([1, 2, 4, 8, 16], 16, 1000, [1, 2, 4, 8, 16], 16),

588 + # Exact boundary: num_blocks == max (no capping)

589 + ([1, 2, 4, 8, 16, 32], 32, 32, [1, 2, 4, 8, 16, 32], 32),

590 + # All sizes capped: num_blocks < smallest size

591 + ([8, 16, 32], 32, 4, [], 0),

592 + # num_blocks <= 0: early return, no change

593 + ([1, 2, 4], 4, 0, [1, 2, 4], 4),

594 + ],

595 +)

596 +def test_adjust_cudagraph_sizes_for_mamba_cache(

597 + capture_sizes, max_size, num_blocks, expected_sizes, expected_max

598 +):

599 + """Test that cudagraph capture sizes are correctly capped to fit

600 + available Mamba cache blocks.

601 +

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

603 + """

604 + config = CompilationConfig(

605 + cudagraph_capture_sizes=capture_sizes,

606 + max_cudagraph_capture_size=max_size,

607 + cudagraph_mode=CUDAGraphMode.NONE,

608 + )

609 + config.adjust_cudagraph_sizes_for_mamba_cache(num_blocks)

610 + assert config.cudagraph_capture_sizes == expected_sizes

611 + assert config.max_cudagraph_capture_size == expected_max

612 + # Invariant: last element == max_cudagraph_capture_size

613 + if expected_sizes:

614 + assert config.cudagraph_capture_sizes[-1] == config.max_cudagraph_capture_size

tests/v1/worker/test_gpu_model_runner.py

+120

@@ -1199,3 +1199,123 @@ def test_is_uniform_decode() -> None:

1199 num_reqs=15,

1200 force_uniform_decode=False,

1201 )

1202 +

1203 +

1204 +@pytest.mark.skipif(

1205 + current_platform.is_rocm(),

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

1207 +)

1208 +def test_cudagraph_sizes_capped_for_mamba_cache():

1209 + """Test that cudagraph capture sizes are capped to num_blocks for

1210 + hybrid models with Mamba layers.

1211 +

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

1213 + """

1214 + set_random_seed(42)

1215 +

1216 + update_environment_variables(

1217 + {

1218 + "RANK": "0",

1219 + "LOCAL_RANK": "0",

1220 + "WORLD_SIZE": "1",

1221 + "MASTER_ADDR": "localhost",

1222 + "MASTER_PORT": "12345",

1223 + }

1224 + )

1225 + from tests.utils import ensure_current_vllm_config

1226 +

1227 + with ensure_current_vllm_config():

1228 + init_distributed_environment()

1229 + initialize_model_parallel(tensor_model_parallel_size=1)

1230 + torch.set_default_dtype(torch.float16)

1231 +

1232 + model_config = ModelConfig(

1233 + model="ibm-granite/granite-4.0-tiny-preview",

1234 + dtype="float16",

1235 + )

1236 + scheduler_config = SchedulerConfig(

1237 + max_num_seqs=10,

1238 + max_num_batched_tokens=512,

1239 + max_model_len=512,

1240 + is_encoder_decoder=model_config.is_encoder_decoder,

1241 + )

1242 + cache_config = CacheConfig(

1243 + block_size=BLOCK_SIZE,

1244 + gpu_memory_utilization=0.9,

1245 + swap_space=0,

1246 + cache_dtype="auto",

1247 + )

1248 + parallel_config = ParallelConfig()

1249 + attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)

1250 + vllm_config = VllmConfig(

1251 + model_config=model_config,

1252 + cache_config=cache_config,

1253 + scheduler_config=scheduler_config,

1254 + parallel_config=parallel_config,

1255 + attention_config=attention_config,

1256 + )

1257 +

1258 + with set_current_vllm_config(vllm_config):

1259 + hf_config = vllm_config.model_config.hf_config

1260 + fwd_context = {}

1261 + for key in ["model.layers.0.self_attn.attn", "model.layers.1.self_attn.attn"]:

1262 + fwd_context[key] = Attention(

1263 + num_heads=model_config.get_num_attention_heads(parallel_config),

1264 + num_kv_heads=model_config.get_num_kv_heads(parallel_config),

1265 + head_size=model_config.get_head_size(),

1266 + scale=1.0,

1267 + prefix=key,

1268 + )

1269 + for key in [

1270 + "model.layers.2.mixer",

1271 + "model.layers.3.mixer",

1272 + "model.layers.4.mixer",

1273 + "model.layers.5.mixer",

1274 + ]:

1275 + fwd_context[key] = MambaMixer2(

1276 + hidden_size=hf_config.hidden_size,

1277 + ssm_state_size=hf_config.mamba_d_state,

1278 + conv_kernel_size=hf_config.mamba_d_conv,

1279 + intermediate_size=hf_config.mamba_expand * hf_config.hidden_size,

1280 + use_conv_bias=hf_config.mamba_conv_bias,

1281 + use_bias=hf_config.mamba_proj_bias,

1282 + n_groups=hf_config.mamba_n_groups,

1283 + num_heads=hf_config.mamba_n_heads,

1284 + head_dim=hf_config.mamba_d_head,

1285 + rms_norm_eps=hf_config.rms_norm_eps,

1286 + activation=hf_config.hidden_act,

1287 + cache_config=cache_config,

1288 + model_config=model_config,

1289 + prefix=key,

1290 + )

1291 + assert fwd_context is not None

1292 +

1293 + runner = GPUModelRunner(vllm_config, DEVICE)

1294 + kv_cache_spec = runner.get_kv_cache_spec()

1295 +

1296 + available_memory = 5 * GiB_bytes

1297 + kv_cache_config = get_kv_cache_configs(

1298 + vllm_config, [kv_cache_spec], [available_memory]

1299 + )[0]

1300 + num_blocks = kv_cache_config.num_blocks

1301 +

1302 + # Set max_cudagraph_capture_size to a value larger than num_blocks

1303 + # to trigger the Mamba capping logic.

1304 + large_max = num_blocks + 100

1305 + compilation_config = vllm_config.compilation_config

1306 + compilation_config.max_cudagraph_capture_size = large_max

1307 + compilation_config.cudagraph_capture_sizes = [

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

1309 + ]

1310 +

1311 + runner.initialize_kv_cache(kv_cache_config)

1312 +

1313 + # After initialization, cudagraph sizes should be capped

1314 + assert compilation_config.max_cudagraph_capture_size <= num_blocks

1315 + assert all(s <= num_blocks for s in compilation_config.cudagraph_capture_sizes)

1316 + # Invariant: last element == max

1317 + if compilation_config.cudagraph_capture_sizes:

1318 + assert (

1319 + compilation_config.cudagraph_capture_sizes[-1]

1320 + == compilation_config.max_cudagraph_capture_size

1321 + )

vllm/config/compilation.py

+52

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

1190 self.max_cudagraph_capture_size = rounded_sizes[-1]

1191 self.cudagraph_capture_sizes = rounded_sizes

1192

1193 + def adjust_cudagraph_sizes_for_mamba_cache(

1194 + self, num_mamba_cache_blocks: int

1195 + ) -> None:

1196 + """Cap cudagraph capture sizes to available Mamba cache blocks.

1197 +

1198 + For hybrid Mamba/attention models, the Mamba conv_state and

1199 + ssm_state tensors have their first dimension equal to num_blocks

1200 + (from KVCacheConfig). During CUDA graph capture the decode batch

1201 + size equals num_tokens, so capture sizes exceeding num_blocks

1202 + would cause out-of-bounds access in Mamba kernels.

1203 +

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

1205 + """

1206 + if not self.cudagraph_capture_sizes or num_mamba_cache_blocks <= 0:

1207 + return

1208 +

1209 + assert self.max_cudagraph_capture_size is not None

1210 +

1211 + if num_mamba_cache_blocks >= self.max_cudagraph_capture_size:

1212 + return

1213 +

1214 + capped_sizes = [

1215 + s for s in self.cudagraph_capture_sizes if s <= num_mamba_cache_blocks

1216 + ]

1217 +

1218 + if len(capped_sizes) == 0:

1219 + logger.warning(

1220 + "No valid cudagraph capture sizes remain after capping "

1221 + "to Mamba cache blocks (%d). The smallest capture size "

1222 + "was %d. Disabling cudagraph capture. Consider reducing "

1223 + "max_num_seqs or increasing available GPU memory.",

1224 + num_mamba_cache_blocks,

1225 + self.cudagraph_capture_sizes[0],

1226 + )

1227 + self.cudagraph_capture_sizes = []

1228 + self.max_cudagraph_capture_size = 0

1229 + return

1230 +

1231 + logger.warning(

1232 + "Capping cudagraph capture sizes from max %d to %d to fit "

1233 + "Mamba cache blocks (%d blocks available). This limits the "

1234 + "maximum batch size that can use CUDA graphs. To increase "

1235 + "this limit, reduce max_num_seqs or increase available GPU "

1236 + "memory.",

1237 + self.max_cudagraph_capture_size,

1238 + capped_sizes[-1],

1239 + num_mamba_cache_blocks,

1240 + )

1241 +

1242 + self.max_cudagraph_capture_size = capped_sizes[-1]

1243 + self.cudagraph_capture_sizes = capped_sizes

1244 +

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

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

1247 if self.compile_ranges_split_points is None:

vllm/v1/worker/gpu_model_runner.py

+16

@@ -5700,6 +5700,22 @@ class GPUModelRunner(

5700 self.uniform_decode_query_len, self.parallel_config.tensor_parallel_size

5701 )

5702

5703 + # If the model has Mamba layers and cudagraph mode includes FULL

5704 + # decode, cap cudagraph capture sizes to the number of available

5705 + # Mamba cache blocks. Each decode request needs one conv_state

5706 + # cache line, so capture batch sizes cannot exceed num_blocks.

5707 + # Only FULL decode graphs are affected because PIECEWISE captures

5708 + # run GDN/Mamba ops eagerly (prefill path, no causal_conv1d_update).

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

5710 + if cudagraph_mode.has_full_cudagraphs():

5711 + has_mamba = any(

5712 + isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_groups

5713 + )

5714 + if has_mamba and self.kv_cache_config is not None:

5715 + self.compilation_config.adjust_cudagraph_sizes_for_mamba_cache(

5716 + self.kv_cache_config.num_blocks

5717 + )

5718 +

5719 # Trigger cudagraph dispatching keys initialization after

5720 # resolved cudagraph mode.

5721 self.compilation_config.cudagraph_mode = cudagraph_mode