From 5d6b35eabb1b72342281ed9485678ace3b58d8fc Mon Sep 17 00:00:00 2001
From: cctry <csycfl@gmail.com>
Date: Wed, 17 Jun 2026 15:46:28 -0700
Subject: [PATCH] revert the head_dim assignment from PR 23862 (#28571)

---
 python/sglang/srt/configs/model_config.py     | 41 ++++++-----
 .../unit/configs/test_model_config_shapes.py  | 71 -------------------
 2 files changed, 20 insertions(+), 92 deletions(-)
 delete 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 abc43d354..a354fa14c 100644
--- a/python/sglang/srt/configs/model_config.py
+++ b/python/sglang/srt/configs/model_config.py
@@ -595,28 +595,27 @@ 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", 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.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.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)
+        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,
+        )
         # 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
deleted file mode 100644
index ec505f378..000000000
--- a/test/registered/unit/configs/test_model_config_shapes.py
+++ /dev/null
@@ -1,71 +0,0 @@
-"""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()