[Fix] Transformers-fallback (GPT-NeoX) + KV pool config (DeepSeek-VL2) (#35244)

Co-authored-by: YanbingJiang <yanbing.jiang@intel.com>

python/sglang/srt/configs/model_config.py

@@ -907,7 +907,7 @@ class ModelConfig:
             setattr(self.hf_text_config, "head_dim", self.head_dim)
 
         self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
-        if self.v_head_dim is None:
+        if self.v_head_dim is None or self.v_head_dim == 0:
             self.v_head_dim = self.head_dim
             setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
 

python/sglang/srt/models/transformers.py

@@ -544,6 +544,8 @@ class TransformersBase(nn.Module):
             "model.score.": "classifier.",
             "model.classifier.": "classifier.",
             "transformer.": "model.",
+            "gpt_neox.": "model.",
+            "embed_out.": "lm_head.",
             "model.": "model.",
             "lm_head.": "lm_head.",
             "score.": "classifier.",
@@ -581,7 +583,9 @@ class TransformersBase(nn.Module):
         self.skip_substrs: list[str] = []
         self.ignore_unexpected_prefixes: list[str] = []
         self.ignore_unexpected_suffixes: list[str] = []
-        self.skip_substrs.extend([".attn.bias", ".attn.masked_bias", ".masked_bias"])
+        self.skip_substrs.extend(
+            [".attn.bias", ".attn.masked_bias", ".attention.bias", ".masked_bias"]
+        )
         self.ignore_unexpected_prefixes.extend(["classifier.", "score."])
 
         if self.quant_config is not None:

test/registered/unit/configs/test_model_config_shapes.py

@@ -69,6 +69,20 @@ class TestModelConfigShapes(CustomTestCase):
         self.assertEqual(model_config.swa_head_dim, 64)
         self.assertEqual(model_config.swa_v_head_dim, 48)
 
+    def test_v_head_dim_zero_falls_back_to_head_dim(self):
+        # deepseek-vl2-tiny's language_config sets use_mla=False + v_head_dim=0
+        # as an MLA-disabled sentinel. The non-MLA KV pool sizes its V buffer
+        # from v_head_dim directly, so a literal 0 collapses V-cache to a
+        # zero-width tensor and `v_cache[indices] = v` fails with a shape
+        # mismatch inside `_set_kv_buffer_impl`. v_head_dim=0 must be treated
+        # identically to `None` and fall back to head_dim.
+        text_config = _make_text_config(head_dim=128, v_head_dim=0)
+
+        model_config = self._derive_shapes(text_config)
+
+        self.assertEqual(model_config.v_head_dim, 128)
+        self.assertEqual(text_config.v_head_dim, 128)
+
     def test_ling_mla_nope_shapes(self):
         text_config = _make_text_config(
             architectures=["BailingMoeV3ForCausalLM"],

test/registered/unit/model_loader/test_transformers_fallback.py

@@ -0,0 +1,77 @@
+"""Unit tests for the shared Transformers-fallback loader path in SGLang."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.models.transformers import TransformersBase
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestTransformersFallbackWeightMapper(CustomTestCase):
+    def test_gpt_neox_prefix_rewrites(self):
+        # GPT-NeoX checkpoints (Pythia, gpt-neox-20b) wrap the backbone under
+        # `gpt_neox.` and expose the LM head as `embed_out.`. `AutoModel`
+        # returns the backbone as `self.model` and `CausalMixin` installs a
+        # `ParallelLMHead` at `self.lm_head`, so the fallback mapper must
+        # rewrite both prefixes. Without these rules, `load_weights` raises
+        # unexpected-key errors on every `gpt_neox.*` / `embed_out.*` entry.
+        mapped = TransformersBase.hf_to_sglang_mapper.apply_list(
+            [
+                "gpt_neox.embed_in.weight",
+                "gpt_neox.layers.0.mlp.dense_h_to_4h.weight",
+                "gpt_neox.final_layer_norm.weight",
+                "embed_out.weight",
+            ]
+        )
+
+        self.assertEqual(
+            mapped,
+            [
+                "model.embed_in.weight",
+                "model.layers.0.mlp.dense_h_to_4h.weight",
+                "model.final_layer_norm.weight",
+                "lm_head.weight",
+            ],
+        )
+
+
+class TestTransformersFallbackSkipSubstrs(CustomTestCase):
+    def test_init_registers_attention_bias_skip(self):
+        # Older GPT-NeoX checkpoints (pythia-1.4b, gpt-neox-20b) ship a
+        # persistent `attention.bias` causal-mask buffer. Newer transformers
+        # builds register it as `persistent=False`, so it is absent from the
+        # constructed module tree and `AutoWeightsLoader` raises on the
+        # unexpected key unless the fallback tells it to skip. The pre-existing
+        # `.attn.bias` covers GPT-2, not NeoX, so `.attention.bias` must be its
+        # own entry.
+        stub = TransformersBase.__new__(TransformersBase)
+
+        class _Stop(RuntimeError):
+            pass
+
+        with (
+            patch(
+                "sglang.srt.models.transformers.get_pp_group",
+                return_value=SimpleNamespace(),
+            ),
+            patch(
+                "sglang.srt.models.transformers.get_hf_text_config",
+                return_value=SimpleNamespace(),
+            ),
+            patch(
+                "sglang.srt.models.transformers._resolve_attention_backend_model_cls",
+                side_effect=_Stop,
+            ),
+            self.assertRaises(_Stop),
+        ):
+            TransformersBase.__init__(stub, config=SimpleNamespace())
+
+        self.assertIn(".attention.bias", stub.skip_substrs)
+
+
+if __name__ == "__main__":
+    unittest.main()