From cfa4aa988fe38370eeb3b4b2cd7e4196f15931fb Mon Sep 17 00:00:00 2001
From: cctry <csycfl@gmail.com>
Date: Wed, 17 Jun 2026 17:11:05 -0700
Subject: [PATCH] Revert "revert the head_dim assignment from PR 23862"
 (#28583)

---
 python/sglang/srt/configs/model_config.py     | 41 +++++------
 .../unit/configs/test_model_config_shapes.py  | 71 +++++++++++++++++++
 2 files changed, 92 insertions(+), 20 deletions(-)
 create mode 100644 test/registered/unit/configs/test_model_config_shapes.py

diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py
index a354fa14c..abc43d354 100644
--- a/python/sglang/srt/configs/model_config.py
+++ b/python/sglang/srt/configs/model_config.py
@@ -595,27 +595,28 @@ class ModelConfig:
 
     def _derive_model_shapes(self):
         # Unify the config keys for hf_text_config
-        self.head_dim = getattr(
-            self.hf_text_config,
-            "head_dim",
-            self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads,
-        )
-        self.v_head_dim = getattr(
-            self.hf_text_config,
-            "v_head_dim",
-            self.head_dim,
-        )
+        self.head_dim = getattr(self.hf_text_config, "head_dim", None)
+        if self.head_dim is None:
+            self.head_dim = (
+                self.hf_text_config.hidden_size
+                // self.hf_text_config.num_attention_heads
+            )
+            setattr(self.hf_text_config, "head_dim", self.head_dim)
 
-        self.swa_head_dim = getattr(
-            self.hf_text_config,
-            "swa_head_dim",
-            self.head_dim,
-        )
-        self.swa_v_head_dim = getattr(
-            self.hf_text_config,
-            "swa_v_head_dim",
-            self.swa_head_dim,
-        )
+        self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None)
+        if self.v_head_dim is None:
+            self.v_head_dim = self.head_dim
+            setattr(self.hf_text_config, "v_head_dim", self.v_head_dim)
+
+        self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None)
+        if self.swa_head_dim is None:
+            self.swa_head_dim = self.head_dim
+            setattr(self.hf_text_config, "swa_head_dim", self.swa_head_dim)
+
+        self.swa_v_head_dim = getattr(self.hf_text_config, "swa_v_head_dim", None)
+        if self.swa_v_head_dim is None:
+            self.swa_v_head_dim = self.swa_head_dim
+            setattr(self.hf_text_config, "swa_v_head_dim", self.swa_v_head_dim)
         # FIXME: temporary special judge for MLA architecture
         if (
             "DeepseekV2ForCausalLM" in self.hf_config.architectures
diff --git a/test/registered/unit/configs/test_model_config_shapes.py b/test/registered/unit/configs/test_model_config_shapes.py
new file mode 100644
index 000000000..ec505f378
--- /dev/null
+++ b/test/registered/unit/configs/test_model_config_shapes.py
@@ -0,0 +1,71 @@
+"""Unit tests for ModelConfig shape normalization."""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.configs.model_config import ModelConfig
+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")
+
+
+def _make_text_config(**overrides):
+    defaults = dict(
+        architectures=["MixtralForCausalLM"],
+        model_type="mixtral",
+        hidden_size=4096,
+        num_attention_heads=32,
+        num_hidden_layers=2,
+        vocab_size=32000,
+        num_key_value_heads=8,
+    )
+    defaults.update(overrides)
+    return SimpleNamespace(**defaults)
+
+
+class TestModelConfigShapes(CustomTestCase):
+    def _derive_shapes(self, text_config):
+        model_config = ModelConfig.__new__(ModelConfig)
+        model_config.hf_config = text_config
+        model_config.hf_text_config = text_config
+        model_config._derive_model_shapes()
+        return model_config
+
+    def test_optional_head_dims_default_when_none(self):
+        text_config = _make_text_config(
+            head_dim=None,
+            v_head_dim=None,
+            swa_head_dim=None,
+            swa_v_head_dim=None,
+        )
+
+        model_config = self._derive_shapes(text_config)
+
+        self.assertEqual(model_config.head_dim, 128)
+        self.assertEqual(model_config.v_head_dim, 128)
+        self.assertEqual(model_config.swa_head_dim, 128)
+        self.assertEqual(model_config.swa_v_head_dim, 128)
+        self.assertEqual(text_config.head_dim, 128)
+        self.assertEqual(text_config.v_head_dim, 128)
+        self.assertEqual(text_config.swa_head_dim, 128)
+        self.assertEqual(text_config.swa_v_head_dim, 128)
+
+    def test_explicit_head_dims_are_preserved(self):
+        text_config = _make_text_config(
+            head_dim=128,
+            v_head_dim=96,
+            swa_head_dim=64,
+            swa_v_head_dim=48,
+        )
+
+        model_config = self._derive_shapes(text_config)
+
+        self.assertEqual(model_config.head_dim, 128)
+        self.assertEqual(model_config.v_head_dim, 96)
+        self.assertEqual(model_config.swa_head_dim, 64)
+        self.assertEqual(model_config.swa_v_head_dim, 48)
+
+
+if __name__ == "__main__":
+    unittest.main()