Feat/support nemotron h mtp (#17013)

Signed-off-by: Shahar Mor <smor@nvidia.com>
This commit is contained in:
shaharmor98
2026-01-14 10:30:35 +02:00
committed by GitHub
parent b025cff441
commit ba625c2d90
4 changed files with 378 additions and 1 deletions

View File

@@ -1485,6 +1485,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
@property
def mamba2_config(self):
config = self.model_config.hf_config
if isinstance(config, NemotronHConfig) and self.is_draft_worker:
# NemotronH MTP draft models have no Mamba layers (pattern like "*E")
# so they shouldn't use HybridLinearAttnBackend
pattern = getattr(config, "mtp_hybrid_override_pattern", None)
if pattern is not None and "M" not in pattern:
return None
if isinstance(config, FalconH1Config | NemotronHConfig):
return config
if isinstance(config, NemotronH_Nano_VL_V2_Config):