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 11, 23:45 UTC
[Frontend][Model] Add 'float16' to possible mamba cache dtype values, override mamba SSM cache dtype value for NemotronH (#29978)
Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com>
amitz-nv committed
Dec 5, 2025 at 10:34 UTC
6038b1b04b7ee1b716a75591040fd3ecaa596a3d
3 files changed
+23
-1
vllm/config/cache.py
+1
-1
@@ -29,7 +29,7 @@ CacheDType = Literal[
29
"fp8_inc",
30
"fp8_ds_mla",
31
]
32
-MambaDType = Literal["auto", "float32"]
32
+MambaDType = Literal["auto", "float32", "float16"]
33
PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"]
34
KVOffloadingBackend = Literal["native", "lmcache"]
35
vllm/model_executor/models/config.py
+21
@@ -485,6 +485,26 @@ class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
485
logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
486
487
488
+class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
489
+ @staticmethod
490
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
491
+ """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto'
492
+ (or not explicitly set), to the value specified in the HF config, or to
493
+ float16 if not specified.
494
+ """
495
+ cache_config = vllm_config.cache_config
496
+ if cache_config.mamba_ssm_cache_dtype == "auto":
497
+ hf_config = vllm_config.model_config.hf_config
498
+ mamba_ssm_cache_dtype = getattr(
499
+ hf_config, "mamba_ssm_cache_dtype", "float16"
500
+ )
501
+ logger.info(
502
+ "Updating mamba_ssm_cache_dtype to '%s' for NemotronH model",
503
+ mamba_ssm_cache_dtype,
504
+ )
505
+ cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype
506
+
507
+
508
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
509
"GteModel": SnowflakeGteNewModelConfig,
510
"GteNewModel": GteNewModelConfig,
@@ -502,4 +522,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
522
"Mamba2ForCausalLM": MambaModelConfig,
523
"FalconMambaForCausalLM": MambaModelConfig,
524
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
525
+ "NemotronHForCausalLM": NemotronHForCausalLMConfig,
526
}
vllm/utils/torch_utils.py
+1
@@ -28,6 +28,7 @@ else:
28
STR_DTYPE_TO_TORCH_DTYPE = {
29
"float32": torch.float32,
30
"half": torch.half,
31
+ "float16": torch.float16,
32
"bfloat16": torch.bfloat16,
33
"float": torch.float,
34
"fp8": torch.uint8,