Fix rope config compatibility and VL/transformers-fallback weight loading (#31575)

Co-authored-by: Claude Sonnet 4.5 (1M context) <noreply@anthropic.com>

python/sglang/srt/managers/mm_schedule.py

@@ -9,12 +9,13 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem
 from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
 from sglang.srt.multimodal.evs import EVSEmbeddingResult
 from sglang.srt.runtime_context import get_parallel, get_schedule
-from sglang.srt.utils import is_hip, is_npu
+from sglang.srt.utils import is_hip, is_npu, is_xpu
 from sglang.srt.utils.async_probe import maybe_assert_sum
 from sglang.utils import logger
 
 _is_hip = is_hip()
 _is_npu = is_npu()
+_is_xpu = is_xpu()
 
 embedding_cache: Optional[MultiModalStaticCache] = None
 
@@ -510,9 +511,9 @@ def _get_chunked_prefill_embedding(
 
         is_per_image = all(len(item.offsets) == 1 for item in embedding_items_per_req)
         if is_per_image:
-            if _is_hip or _is_npu:
+            if _is_hip or _is_npu or _is_xpu:
                 # ROCm CI regressed with one large cross-request ViT batch; keep
-                # the previous per-request path on HIP while CUDA uses batching.
+                # the previous per-request path on HIP/NPU/XPU while CUDA uses batching.
                 chunk = _get_chunked_embedding_by_item(
                     data_embedding_func,
                     embedding_items_per_req,

python/sglang/srt/models/ernie45_moe_vl.py

@@ -44,6 +44,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
 from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, make_layers
+from sglang.srt.utils.hf_transformers.common import get_rope_config
 
 logger = logging.getLogger(__name__)
 
@@ -368,8 +369,7 @@ class Ernie4_5_VLMoeDecoderLayer(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        rope_theta = config.rope_parameters["rope_theta"]
-        rope_scaling = config.rope_parameters
+        rope_theta, rope_scaling = get_rope_config(config)
         rope_is_neox_style = getattr(config, "rope_is_neox_style", False)
         freq_allocation = getattr(config, "freq_allocation", 20)
         max_position_embeddings = getattr(config, "max_position_embeddings", 131072)

python/sglang/srt/models/olmo2.py

@@ -49,6 +49,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.runtime_context import get_parallel, get_stream
 from sglang.srt.utils import add_prefix, is_cuda, make_layers
+from sglang.srt.utils.hf_transformers.common import get_rope_config
 
 _is_cuda = is_cuda()
 
@@ -100,7 +101,7 @@ class Olmo2Attention(nn.Module):
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_parameters["rope_theta"]
+        self.rope_theta, _ = get_rope_config(config)
 
         # Attention input projection. Projects x -> (q, k, v)
         self.qkv_proj = QKVParallelLinear(

python/sglang/srt/models/qwen.py

@@ -337,7 +337,10 @@ class QWenLMHeadModel(nn.Module):
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
-                name = name.replace(weight_name, param_name)
+                temp_name = name.replace(weight_name, param_name)
+                if temp_name not in params_dict:
+                    continue
+                name = temp_name
                 # Skip loading extra bias for GPTQ models.
                 if name.endswith(".bias") and name not in params_dict:
                     continue
@@ -349,6 +352,9 @@ class QWenLMHeadModel(nn.Module):
                 # Skip loading extra bias for GPTQ models.
                 if name.endswith(".bias") and name not in params_dict:
                     continue
+                # Skip visual encoder weights (e.g. Qwen-VL-Chat transformer.visual.*)
+                if name not in params_dict:
+                    continue
                 param = params_dict[name]
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)

python/sglang/srt/models/transformers.py

@@ -1345,6 +1345,22 @@ class MultiModalMixin:
         super().__init__(*args, **kwargs)
         self._mm_padding_pattern = MultiModalityDataPaddingPatternMultimodalTokens()
 
+        # transformers v5 flattened SigLIP/CLIP (dropped the "vision_model"
+        # wrapper); older checkpoints still ship "vision_tower.vision_model.*"
+        # keys, so remap them when the live model lacks that sub-module.
+        vt = getattr(self.model, "vision_tower", None)
+        if vt is not None and not any(
+            name == "vision_model" for name, _ in vt.named_children()
+        ):
+            self.weight_mapper = (
+                WeightsMapper(
+                    orig_to_new_prefix={
+                        "vision_tower.vision_model.": "model.vision_tower.",
+                    }
+                )
+                | self.weight_mapper
+            )
+
     def _uses_mrope_positions(self) -> bool:
         rope_scaling = getattr(self.text_config, "rope_scaling", None)
         if isinstance(rope_scaling, Mapping) and "mrope_section" in rope_scaling:
@@ -1497,6 +1513,14 @@ class MultiModalMixin:
         ):
             mm_inputs = forward_batch.mm_inputs
             target_device = next(self.model.parameters()).device
+            # 5D features (num_images, num_patches, C, H, W) can't be flattened
+            # here: anyres models pad num_patches per HF processor call, so a
+            # flattened concat would leave stray padding rows once items with
+            # different tile counts are combined. Defer them and pad to the
+            # batch-wide max instead -- the model's own get_image_features
+            # re-derives each image's real patch count from image_sizes and
+            # slices the padding back out.
+            pending_5d_features: dict = {}
 
             for batch_idx in range(len(mm_inputs or [])):
                 mm_input = mm_inputs[batch_idx]
@@ -1519,6 +1543,11 @@ class MultiModalMixin:
                         feature = item.feature
                         if isinstance(feature, torch.Tensor):
                             feature = feature.to(device=target_device)
+                            if feature.dim() == 5:
+                                pending_5d_features.setdefault(feature_key, []).append(
+                                    feature
+                                )
+                                continue
                         if feature_key not in kwargs:
                             kwargs[feature_key] = feature
                         elif isinstance(feature, torch.Tensor) and isinstance(
@@ -1528,6 +1557,24 @@ class MultiModalMixin:
                                 [kwargs[feature_key], feature], dim=0
                             )
 
+            for feature_key, tensors in pending_5d_features.items():
+                max_patches = max(t.shape[1] for t in tensors)
+                padded = []
+                for t in tensors:
+                    if t.shape[1] < max_patches:
+                        pad = t.new_zeros(
+                            (t.shape[0], max_patches - t.shape[1], *t.shape[2:])
+                        )
+                        t = torch.cat([t, pad], dim=1)
+                    padded.append(t)
+                combined = torch.cat(padded, dim=0)
+                if feature_key in kwargs:
+                    kwargs[feature_key] = torch.cat(
+                        [kwargs[feature_key], combined], dim=0
+                    )
+                else:
+                    kwargs[feature_key] = combined
+
         return kwargs
 
     def _forward_hidden_states(

python/sglang/srt/utils/hf_transformers/common.py

@@ -365,7 +365,8 @@ def get_rope_config(config):
     """
     rope_params = getattr(config, "rope_parameters", None)
     if rope_params is not None:
-        return rope_params["rope_theta"], rope_params
+        rope_theta = rope_params.get("rope_theta", getattr(config, "rope_theta", 10000))
+        return rope_theta, rope_params
     return getattr(config, "rope_theta", 10000), getattr(config, "rope_scaling", None)
 
 

test/manual/models/test_transformers_collect_mm_kwargs.py

@@ -0,0 +1,126 @@
+"""
+Unit tests for MultiModalMixin._collect_mm_kwargs' handling of 5D
+pixel_values features in the generic Transformers fallback backend
+(sglang.srt.models.transformers).
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.models.transformers import MultiModalMixin
+
+
+def _make_item(modality_name, feature, model_specific_data=None):
+    return SimpleNamespace(
+        modality=SimpleNamespace(name=modality_name),
+        feature=feature,
+        model_specific_data=model_specific_data or {},
+    )
+
+
+def _make_mm_input(items):
+    return SimpleNamespace(mm_items=items)
+
+
+def _make_forward_batch(mm_inputs, is_decode=False, contains_mm_inputs=True):
+    return SimpleNamespace(
+        token_type_ids=None,
+        forward_mode=SimpleNamespace(is_decode=lambda: is_decode),
+        mm_inputs=mm_inputs,
+        contains_mm_inputs=lambda: contains_mm_inputs,
+    )
+
+
+def _make_self():
+    """Lightweight stand-in for a TransformersForCausalLM instance -- only
+    `.model` (for the device lookup) and the mixin's own feature-key map
+    are actually used by `_collect_mm_kwargs`."""
+    return SimpleNamespace(
+        model=torch.nn.Linear(1, 1),
+        _mm_feature_kwarg=MultiModalMixin._mm_feature_kwarg,
+    )
+
+
+class TestCollectMmKwargs5DPadding(unittest.TestCase):
+    def test_equal_patch_counts_no_padding(self):
+        """Sanity check: same num_patches across items concatenates cleanly."""
+        item1 = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 1.0))
+        item2 = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 2.0))
+        forward_batch = _make_forward_batch(
+            [_make_mm_input([item1]), _make_mm_input([item2])]
+        )
+
+        kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch)
+
+        pixel_values = kwargs["pixel_values"]
+        self.assertEqual(pixel_values.shape, (2, 3, 3, 4, 4))
+        self.assertTrue(torch.all(pixel_values[0] == 1.0))
+        self.assertTrue(torch.all(pixel_values[1] == 2.0))
+
+    def test_different_patch_counts_padded_to_batch_max(self):
+        """Test: items with a different tile count must be
+        zero-padded to the batch-wide max num_patches, not just concatenated
+        as-is (which would crash on mismatched shapes or misalign data)."""
+        item_small = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 1.0))  # 3 patches
+        item_large = _make_item("IMAGE", torch.full((1, 5, 3, 4, 4), 2.0))  # 5 patches
+        forward_batch = _make_forward_batch(
+            [_make_mm_input([item_small]), _make_mm_input([item_large])]
+        )
+
+        kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch)
+
+        pixel_values = kwargs["pixel_values"]
+        self.assertEqual(pixel_values.shape, (2, 5, 3, 4, 4))
+        # item_small's real 3 patches are preserved...
+        self.assertTrue(torch.all(pixel_values[0, :3] == 1.0))
+        # ...and its padding (patches 3-4) is zeroed, not garbage/leftover data.
+        self.assertTrue(torch.all(pixel_values[0, 3:] == 0.0))
+        # item_large needed no padding at all.
+        self.assertTrue(torch.all(pixel_values[1] == 2.0))
+
+    def test_multi_image_item_with_different_patch_counts_within_one_item(self):
+        """A single multi-image item/request can itself already contain
+        per-image padding applied by the HF processor; the batch-level
+        padding must still pad up to the overall max without disturbing it."""
+        # 2 images already padded to 4 patches by the HF processor, batched
+        # against another item that only needed 2 patches.
+        item_multi_image = _make_item("IMAGE", torch.full((2, 4, 3, 4, 4), 1.0))
+        item_single = _make_item("IMAGE", torch.full((1, 2, 3, 4, 4), 2.0))
+        forward_batch = _make_forward_batch(
+            [_make_mm_input([item_multi_image]), _make_mm_input([item_single])]
+        )
+
+        kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch)
+
+        pixel_values = kwargs["pixel_values"]
+        self.assertEqual(pixel_values.shape, (3, 4, 3, 4, 4))
+        self.assertTrue(torch.all(pixel_values[:2] == 1.0))
+        self.assertTrue(torch.all(pixel_values[2, :2] == 2.0))
+        self.assertTrue(torch.all(pixel_values[2, 2:] == 0.0))
+
+    def test_decode_mode_skips_collection(self):
+        """During decode (no new mm inputs to process this step), no
+        multimodal kwargs should be produced even if mm_inputs is present."""
+        item = _make_item("IMAGE", torch.full((1, 3, 3, 4, 4), 1.0))
+        forward_batch = _make_forward_batch([_make_mm_input([item])], is_decode=True)
+
+        kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch)
+
+        self.assertNotIn("pixel_values", kwargs)
+
+    def test_non_image_modality_uses_correct_feature_key(self):
+        """Video features (also potentially 5D) must land under their own
+        kwarg key, not be mixed in with image pixel_values."""
+        item = _make_item("VIDEO", torch.full((1, 3, 3, 4, 4), 1.0))
+        forward_batch = _make_forward_batch([_make_mm_input([item])])
+
+        kwargs = MultiModalMixin._collect_mm_kwargs(_make_self(), forward_batch)
+
+        self.assertIn("pixel_values_videos", kwargs)
+        self.assertNotIn("pixel_values", kwargs)
+
+
+if __name__ == "__main__":
+    unittest.main()