Revert "revert the head_dim assignment from PR 23862" (#28583)

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()