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, 12:15 UTC
NemotronH default mamba_ssm_cache_dtype=float32; enable auto-hook for NemotronHNanoVLV2Config (#39032)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Netanel Haber committed
Apr 6, 2026 at 22:47 UTC
dfa5062a8f372ea78e48197939289c15247c1840
1 file changed
+26
-7
vllm/model_executor/models/config.py
+26
-7
@@ -7,7 +7,9 @@ from vllm.logger import init_logger
7
from vllm.utils.math_utils import round_up
8
9
if TYPE_CHECKING:
10
- from vllm.config import ModelConfig, VllmConfig
10
+ from transformers import PretrainedConfig
11
+
12
+ from vllm.config import CacheConfig, ModelConfig, VllmConfig
13
14
15
logger = init_logger(__name__)
@@ -346,17 +348,20 @@ class MambaModelConfig(VerifyAndUpdateConfig):
348
349
350
class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
349
- @staticmethod
350
- def verify_and_update_config(vllm_config: "VllmConfig") -> None:
351
+ DEFAULT_MAMBA_SSM_CACHE_DTYPE = "float32"
352
+ """Only `float32` is known to have no accuracy issues by default."""
353
+
354
+ @classmethod
355
+ def update_mamba_ssm_cache_dtype(
356
+ cls, *, cache_config: "CacheConfig", hf_config: "PretrainedConfig"
357
+ ) -> None:
358
"""Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto'
359
(or not explicitly set), to the value specified in the HF config, or to
353
- float16 if not specified.
360
+ `float32` if not specified.
361
"""
355
- cache_config = vllm_config.cache_config
362
if cache_config.mamba_ssm_cache_dtype == "auto":
357
- hf_config = vllm_config.model_config.hf_config
363
mamba_ssm_cache_dtype = getattr(
359
- hf_config, "mamba_ssm_cache_dtype", "float16"
364
+ hf_config, "mamba_ssm_cache_dtype", cls.DEFAULT_MAMBA_SSM_CACHE_DTYPE
365
)
366
logger.info(
367
"Updating mamba_ssm_cache_dtype to '%s' for NemotronH model",
@@ -364,8 +369,22 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
369
)
370
cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype
371
372
+ @classmethod
373
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
374
+ cls.update_mamba_ssm_cache_dtype(
375
+ cache_config=vllm_config.cache_config,
376
+ hf_config=vllm_config.model_config.hf_config,
377
+ )
378
+
379
380
class NemotronHNanoVLV2Config(VerifyAndUpdateConfig):
381
+ @classmethod
382
+ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
383
+ NemotronHForCausalLMConfig.update_mamba_ssm_cache_dtype(
384
+ cache_config=vllm_config.cache_config,
385
+ hf_config=vllm_config.model_config.hf_config.text_config,
386
+ )
387
+
388
@staticmethod
389
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
390
mm_config = model_config.multimodal_config