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 14, 05:30 UTC
[V1] [Hybrid] Enable Full CUDA graph by default for hybrid models in V1 (#22594)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Aug 27, 2025 at 01:28 UTC
5f1af97f86021cf2819e5ab2d84722dac53c2257
1 file changed
+42
vllm/model_executor/models/config.py
+42
@@ -4,6 +4,7 @@ from copy import deepcopy
4
from typing import TYPE_CHECKING
5
6
import vllm.envs as envs
7
+from vllm.config.compilation import CUDAGraphMode
8
from vllm.logger import init_logger
9
from vllm.model_executor.models import ModelRegistry
10
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
@@ -275,6 +276,42 @@ class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
276
"%d for performance.", 1024)
277
278
279
+class MambaModelConfig(VerifyAndUpdateConfig):
280
+
281
+ @classmethod
282
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
283
+ """
284
+ Enable FULL_AND_PIECEWISE cuda graph mode by default (required
285
+ to get good performance for mamba layers in V1).
286
+
287
+ Args:
288
+ vllm_config: vLLM Config
289
+ """
290
+
291
+ if not envs.VLLM_USE_V1:
292
+ return
293
+
294
+ model_config = vllm_config.model_config
295
+ compilation_config = vllm_config.compilation_config
296
+
297
+ model_cls, _ = ModelRegistry.resolve_model_cls(
298
+ model_config.architecture,
299
+ model_config=model_config,
300
+ )
301
+
302
+ # TODO(tdoublep): remove as full cuda graph support is added
303
+ FCG_NOT_SUPPORTED_MODELS = [
304
+ "Lfm2ForCausalLM", "MiniMaxText01ForCausalLM"
305
+ ]
306
+
307
+ if (model_config.architecture not in FCG_NOT_SUPPORTED_MODELS
308
+ and compilation_config.cudagraph_mode is None):
309
+ logger.info(
310
+ "Hybrid or mamba-based model detected: setting cudagraph mode "
311
+ "to FULL_AND_PIECEWISE in order to optimize performance.")
312
+ compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
313
+
314
+
315
class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
316
317
@classmethod
@@ -293,6 +330,9 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
330
if not envs.VLLM_USE_V1:
331
return
332
333
+ # Enable FULL_AND_PIECEWISE by default
334
+ MambaModelConfig.verify_and_update_config(vllm_config)
335
+
336
cache_config = vllm_config.cache_config
337
model_config = vllm_config.model_config
338
parallel_config = vllm_config.parallel_config
@@ -374,4 +414,6 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
414
"JambaForSequenceClassification": JambaForSequenceClassificationConfig,
415
"GraniteMoeHybridForCausalLM": GraniteMoeHybridModelConfig,
416
"GptOssForCausalLM": GptOssForCausalLMConfig,
417
+ "MambaForCausalLM": MambaModelConfig,
418
+ "Mamba2ForCausalLM": MambaModelConfig,
419
}