Add Super VL and Mamba DisAgg compatibility to vLLM 0.28

tests/models/registry.py

@@ -1244,6 +1244,10 @@ def check_available_online(
     "NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo(
         "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False
     ),
+    # TODO: Change repo id once pertinent archs are public.
+    "NemotronH_Omni_Reasoning_V3": _HfExamplesInfo(
+        "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False
+    ),
     "OpenCUAForConditionalGeneration": _HfExamplesInfo(
         "xlangai/OpenCUA-7B",
         trust_remote_code=True,

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1054,7 +1054,7 @@ def test_mamba_n1_d_side_builds_decode_metadata():
 
 @pytest.mark.cpu_test
 def test_mamba_n1_p_side_truncation():
-    """P-side: Mamba truncates prompt to N-1, sets max_tokens=1.
+    """P-side truncates to N-1 before the scheduler's prefix-cache lookup.
 
     Also verifies idempotency (calling again is a no-op) which is
     needed for preemption safety via the _p_side_truncated guard,
@@ -1065,25 +1065,28 @@ def test_mamba_n1_p_side_truncation():
     req.max_tokens = 128
     original_len = len(req.prompt_token_ids)
 
+    sched.on_new_request(req)
+    assert len(req.prompt_token_ids) == original_len - 1
+    assert req.num_prompt_tokens == original_len - 1
+    assert req.max_tokens == 1
+    assert req.kv_transfer_params["_p_side_truncated"] is True
+
     count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
 
     assert count == 0
     assert is_async is False
     assert len(req.prompt_token_ids) == original_len - 1
-    assert req.num_prompt_tokens == original_len - 1
-    assert req.max_tokens == 1
-    assert req.kv_transfer_params["_p_side_truncated"] is True
 
-    # Idempotency: second call must not truncate further
-    sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
+    # Idempotency: re-adding a preempted request must not truncate further.
+    sched.on_new_request(req)
     assert len(req.prompt_token_ids) == original_len - 1
 
     # Non-Mamba: truncation is skipped
     fa_sched = make_nixl_scheduler(has_mamba=False, is_hma_required=False)
     fa_req = create_request(num_tokens=10, do_remote_decode=True)
     fa_original = len(fa_req.prompt_token_ids)
 
-    fa_sched.get_num_new_matched_tokens(fa_req, num_computed_tokens=0)
+    fa_sched.on_new_request(fa_req)
     assert len(fa_req.prompt_token_ids) == fa_original
 
 

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -46,7 +46,7 @@
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 from vllm.v1.outputs import KVConnectorOutput
 
-from .utils import make_nixl_push_scheduler
+from .utils import create_request, make_nixl_push_scheduler
 
 # ----------------------------------------------------------------- #
 #  Helpers / fakes                                                   #
@@ -117,6 +117,24 @@ def _stub_sw_clipping(scheduler) -> None:
 
 
 class TestPushScheduler:
+    def test_p_side_mamba_truncates_before_cache_lookup(self):
+        """Push mode normalizes P-side Mamba requests before cache lookup."""
+        sched = make_nixl_push_scheduler(has_mamba=True)
+        request = create_request(num_tokens=10, do_remote_decode=True)
+        original_len = len(request.prompt_token_ids)
+
+        sched.on_new_request(request)
+
+        assert len(request.prompt_token_ids) == original_len - 1
+        assert request.kv_transfer_params["_p_side_truncated"] is True
+
+        with patch.object(
+            sched,
+            "_truncate_mamba_request_for_prefill",
+            side_effect=AssertionError("must not truncate after cache lookup"),
+        ):
+            assert sched.get_num_new_matched_tokens(request, 0) == (0, False)
+
     def test_d_side_update_state_after_alloc_stages_registration(self):
         """D scheduler stashes registration data + arms watchdog deadline."""
         sched = make_nixl_push_scheduler()

vllm/config/speculative.py

@@ -484,7 +484,10 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
                 {"n_predict": n_predict, "architectures": ["ErnieMTPModel"]}
             )
 
-        if hf_config.architectures[0] == "NemotronH_Super_Omni_Reasoning_V3":
+        if hf_config.architectures[0] in (
+            "NemotronH_Super_Omni_Reasoning_V3",
+            "NemotronH_Omni_Reasoning_V3",
+        ):
             # Promote VLM's text_config so MTP detection below fires correctly
             hf_config = hf_config.text_config
 

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

@@ -194,6 +194,9 @@ def shutdown(self):
     def on_new_request(self, request: "Request") -> None:
         """Track a request that may need heartbeats."""
         params = request.kv_transfer_params
+        if params is not None and params.get("do_remote_decode") and self._has_mamba:
+            self._truncate_mamba_request_for_prefill(request)
+
         # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
         # effectively disabled for Bidirectional KV transfer.
         if params is None or not params.get("do_remote_prefill"):

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -65,9 +65,6 @@ def get_num_new_matched_tokens(
             if count > 0:
                 return count, True
 
-        if params is not None and params.get("do_remote_decode") and self._has_mamba:
-            self._truncate_mamba_request_for_prefill(request)
-
         if (
             params is not None
             and params.get("do_remote_decode")

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py

@@ -123,9 +123,6 @@ def get_num_new_matched_tokens(
             if count > 0:
                 return count, True
 
-        if params is not None and params.get("do_remote_decode") and self._has_mamba:
-            self._truncate_mamba_request_for_prefill(request)
-
         return 0, False
 
     def update_state_after_alloc(

vllm/model_executor/models/config.py

@@ -677,6 +677,27 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
 
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
+        config = model_config.hf_config
+        vision_config = config.vision_config
+        if not hasattr(vision_config, "args"):
+            image_processor_config = model_config.hf_image_processor_config
+            vision_config.args = {
+                "model": "vit_huge_patch16_224",
+                "qkv_bias": vision_config.qkv_bias,
+                "layer_norm_eps": vision_config.layer_norm_eps,
+                "initializer_factor": vision_config.layerscale_value,
+                "hidden_act": vision_config.hidden_act,
+                "cpe_max_size": vision_config.max_img_size,
+                "num_cls_tokens": vision_config.num_cls_tokens,
+                "num_registers": vision_config.num_registers,
+                "summary_idxs": vision_config.summary_idxs,
+                "min_num_patches": image_processor_config["min_num_patches"],
+                "max_num_patches": image_processor_config["max_num_patches"],
+            }
+            config.norm_mean = vision_config.norm_mean
+            config.norm_std = vision_config.norm_std
+            config.use_thumbnail = False
+
         mm_config = model_config.multimodal_config
         if mm_config is not None:
             video_kwargs = mm_config.media_io_kwargs.setdefault("video", {})
@@ -933,6 +954,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
     "NemotronHForCausalLM": NemotronHForCausalLMConfig,
     "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
     "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config,
+    "NemotronH_Omni_Reasoning_V3": NemotronHNanoVLV2Config,
     "NomicBertModel": NomicBertModelConfig,
     "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
     "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,

vllm/model_executor/models/nano_nemotron_vl.py

@@ -1515,11 +1515,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         def is_llm(name: str) -> bool:
             return name.startswith("language_model")
 
+        adapter_mapping = {
+            "vision_projector.mlp1.norm.weight": "0.weight",
+            "vision_projector.mlp1.linear1.weight": "1.weight",
+            "vision_projector.mlp1.linear2.weight": "3.weight",
+        }
+
         def is_adapter_weights(weight: tuple[str, torch.Tensor]):
             return weight[0].startswith("mlp1")
 
         def is_vision_weights(name: str) -> bool:
-            return name.startswith("vision_model.radio_model.")
+            return name.startswith("vision_model.")
 
         def is_sound_weights(name: str) -> bool:
             return name.startswith("sound")
@@ -1544,10 +1550,14 @@ def llm_weights_gen():
                         continue
                     trimmed_name = ".".join(name.split(".")[1:])
                     adapter_weights.append((trimmed_name, w.detach().clone()))
+                elif name in adapter_mapping:
+                    if not load_multimodal_weights:
+                        continue
+                    adapter_weights.append((adapter_mapping[name], w.detach().clone()))
                 elif is_vision_weights(name):
                     if not load_multimodal_weights:
                         continue
-                    # Convert: vision_model.radio_model.* → radio_model.*
+                    # Strip the multimodal wrapper prefix.
                     hf_key = name[len("vision_model.") :]
                     vision_weights.append((hf_key, w.detach().clone()))
                 elif is_sound_weights(name):

vllm/model_executor/models/nemotron_h_mtp.py

@@ -216,7 +216,7 @@ class NemotronHMultiTokenPredictor(nn.Module):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
-        config = vllm_config.model_config.hf_config
+        config = vllm_config.model_config.hf_config.get_text_config()
 
         self.config = config
         self.vocab_size = config.vocab_size
@@ -322,7 +322,7 @@ class NemotronHMTP(nn.Module, SupportsPP):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        config = vllm_config.model_config.hf_config
+        config = vllm_config.model_config.hf_config.get_text_config()
         self.vllm_config = vllm_config
         self.config = config
         self.quant_config = vllm_config.quant_config
@@ -414,6 +414,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loaded_params: set[str] = set()
 
         for name, loaded_weight in weights:
+            # MTP weights are nested in "language_model."
+            # in Multimodal Nemotron-H checkpoints.
+            name = name.removeprefix("language_model.")
+
             # Only process MTP weights - skip all non-MTP weights
             if not name.startswith("mtp.") and "embeddings" not in name:
                 continue

vllm/model_executor/models/radio.py

@@ -545,14 +545,18 @@ def __init__(
         )
         self.temporal_patch_size = config.video_temporal_patch_size
         unique_teachers = set(t["name"] for t in config.teachers)
+        num_cls_tokens = getattr(config, "num_cls_tokens", None)
+        if num_cls_tokens is None:
+            num_cls_tokens = len(unique_teachers) if config.cls_token_per_teacher else 1
         self.patch_generator = ViTPatchGenerator(
             config.patch_size,
             config.hidden_size,
             input_dims=self.img_size,
             max_input_dims=max_img_size,
             cls_token=True,
-            num_cls_tokens=len(unique_teachers) if config.cls_token_per_teacher else 1,
+            num_cls_tokens=num_cls_tokens,
             register_multiple=config.register_multiple,
+            num_registers=getattr(config, "num_registers", None),
             temporal_patch_size=self.temporal_patch_size,
             separate_video_embedder=config.separate_video_embedder,
         )
@@ -669,13 +673,15 @@ def __init__(
             prefix=prefix,
         )
 
-        summary_idxs = None
-        if config.teachers:
+        summary_idxs = getattr(config, "summary_idxs", None)
+        if summary_idxs is not None:
+            summary_idxs = torch.tensor(summary_idxs)
+        elif config.teachers:
             summary_idxs = torch.tensor(
                 [i for i, t in enumerate(config.teachers) if t.get("use_summary", True)]
             )
-            if summary_idxs.numel() > 0:
-                self.register_buffer("summary_idxs", summary_idxs)
+        if summary_idxs is not None and summary_idxs.numel() > 0:
+            self.register_buffer("summary_idxs", summary_idxs)
         self.summary_idxs = summary_idxs
 
     def forward(
@@ -697,32 +703,47 @@ def load_weights(self, weights) -> set[str]:
         loaded_params: set[str] = set()
         params_dict = dict(self.named_parameters())
 
+        native_embedding_mapping = {
+            "embeddings.patch_projection.": "model.patch_generator.embedder.",
+            "embeddings.video_patch_projection.": (
+                "model.patch_generator.video_embedder."
+            ),
+            "embeddings.position_embedding": "model.patch_generator.pos_embed",
+            "embeddings.cls_register_token": "model.patch_generator.cls_token.token",
+        }
+        native_layer_mapping = {
+            "attention.output.dense.": ("attn.proj.", None),
+            "attention.attention.query.": ("attn.qkv.", "q"),
+            "attention.attention.key.": ("attn.qkv.", "k"),
+            "attention.attention.value.": ("attn.qkv.", "v"),
+            "layer_scale1.lambda1": ("ls1", None),
+            "layer_scale2.lambda1": ("ls2", None),
+        }
+
         if isinstance(weights, dict):
             weights_list = list(weights.items())
         else:
             weights_list = list(weights)
 
         for name, weight in weights_list:
-            if not name.startswith("radio_model."):
-                # Skip non-radio weights
-                continue
-
-            sub = name[len("radio_model.") :]  # drop "radio_model." prefix
+            is_legacy = name.startswith("radio_model.")
+            sub = name.removeprefix("radio_model.")
 
             # Skip buffers not used in vLLM
             if sub in {"summary_idxs"}:
                 continue
-            if sub.startswith("input_conditioner."):
+            if is_legacy and sub.startswith("input_conditioner."):
                 # we normalize in the input processor,
                 # based on norm and std values from the config
                 continue
 
             vllm_key = None
-            if sub.startswith("model.patch_generator."):
+            shard_id = None
+            if is_legacy and sub.startswith("model.patch_generator."):
                 vllm_key = f"model.patch_generator.{sub.split('.', 2)[-1]}"
-            elif sub.startswith("input_conditioner."):
+            elif is_legacy and sub.startswith("input_conditioner."):
                 vllm_key = f"input_conditioner.{sub.split('.', 1)[-1]}"
-            elif sub.startswith("model.blocks."):
+            elif is_legacy and sub.startswith("model.blocks."):
                 # Encoder blocks: HF 'model.blocks.{i}.' ->
                 # vLLM 'model.encoder.layers.{i}.'
                 parts = sub.split(".")
@@ -733,11 +754,32 @@ def load_weights(self, weights) -> set[str]:
                     if suffix in {"ls1", "ls2"} or suffix.startswith(("ls1.", "ls2.")):
                         continue
                     vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}"
+            elif not is_legacy:
+                for source_prefix, target_prefix in native_embedding_mapping.items():
+                    if sub.startswith(source_prefix):
+                        vllm_key = sub.replace(source_prefix, target_prefix, 1)
+                        break
+
+                if sub.startswith("encoder.layer."):
+                    parts = sub.split(".")
+                    layer_idx = parts[2]
+                    suffix = ".".join(parts[3:])
+                    for source_prefix, (
+                        target_prefix,
+                        shard_id,
+                    ) in native_layer_mapping.items():
+                        if suffix.startswith(source_prefix):
+                            suffix = suffix.replace(source_prefix, target_prefix, 1)
+                            break
+                    vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}"
 
             if vllm_key and vllm_key in params_dict:
                 param = params_dict[vllm_key]
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, weight)
+                if shard_id is None:
+                    weight_loader(param, weight)
+                else:
+                    weight_loader(param, weight, shard_id)
                 loaded_params.add(vllm_key)
 
         return loaded_params

vllm/model_executor/models/registry.py

@@ -530,6 +530,7 @@
     "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
     "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
     "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
+    "NemotronH_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
     "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
     "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"),
     "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"),

vllm/transformers_utils/config.py

@@ -23,6 +23,7 @@
 )
 from transformers.models.auto.tokenization_auto import get_tokenizer_config
 from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME
+from transformers.utils import IMAGE_PROCESSOR_NAME
 
 from vllm import envs
 from vllm.logger import init_logger
@@ -1050,9 +1051,12 @@ def get_hf_image_processor_config(
     # ModelScope does not provide an interface for image_processor
     if envs.VLLM_USE_MODELSCOPE:
         return dict()
-    return get_image_processor_config(
+    image_processor_config = get_image_processor_config(
         model, token=hf_token, revision=revision, **kwargs
     )
+    if image_processor_config:
+        return image_processor_config
+    return get_hf_file_to_dict(IMAGE_PROCESSOR_NAME, model, revision) or {}
 
 
 def get_hf_text_config(config: PretrainedConfig):

vllm/v1/attention/backends/mamba2_attn.py

@@ -141,6 +141,7 @@ def build(
             common_attn_metadata,
             num_accepted_tokens=kwargs.get("num_accepted_tokens"),
             prev_last_scheduled_idx=kwargs.get("prev_last_scheduled_idx"),
+            num_decode_draft_tokens_cpu=kwargs.get("num_decode_draft_tokens_cpu"),
         )
 
         seq_idx_p = None