From 9fec359a60aa5ece10a97f09c0891a328e2f109b Mon Sep 17 00:00:00 2001
From: Auroter <7332587+auroter@users.noreply.github.com>
Date: Mon, 13 Jul 2026 05:37:36 -0700
Subject: [PATCH] [Fix] Load HunyuanV3 NextN final_layernorm into the draft
 head's output norm (#30331)

Co-authored-by: Auroter <auroter@users.noreply.github.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
---
 python/sglang/srt/models/hunyuan_v3_nextn.py  |  4 +
 .../test_hunyuan_v3_nextn_weight_loading.py   | 93 +++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 test/registered/unit/models/test_hunyuan_v3_nextn_weight_loading.py

diff --git a/python/sglang/srt/models/hunyuan_v3_nextn.py b/python/sglang/srt/models/hunyuan_v3_nextn.py
index c1b159982..1b2388391 100644
--- a/python/sglang/srt/models/hunyuan_v3_nextn.py
+++ b/python/sglang/srt/models/hunyuan_v3_nextn.py
@@ -186,6 +186,10 @@ class HYV3ForCausalLMNextN(nn.Module):
                 subname = name[len(nextn_prefix) :]
                 if any(subname.startswith(s) for s in spec_weight_names):
                     name = f"model.{subname}"
+                elif subname.startswith("final_layernorm"):
+                    # Released checkpoints store the draft head's output norm
+                    # as model.layers.<N>.final_layernorm.weight.
+                    name = "model.shared_head.norm.weight"
                 else:
                     name = f"model.decoder.{subname}"
             elif name == "model.shared_head.norm.weight":
diff --git a/test/registered/unit/models/test_hunyuan_v3_nextn_weight_loading.py b/test/registered/unit/models/test_hunyuan_v3_nextn_weight_loading.py
new file mode 100644
index 000000000..f417f6de8
--- /dev/null
+++ b/test/registered/unit/models/test_hunyuan_v3_nextn_weight_loading.py
@@ -0,0 +1,93 @@
+"""
+Unit tests for HYV3ForCausalLMNextN.load_weights.
+
+Regression test for the released Hy3 MTP checkpoint key
+``model.layers.<num_hidden_layers>.final_layernorm.weight``, which must load
+into the draft head's output norm (``model.shared_head.norm``) instead of
+being remapped to a nonexistent decoder parameter and silently dropped.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=4, suite="base-a-test-cpu")
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.models.hunyuan_v3_nextn import HYV3ForCausalLMNextN
+
+
+class _FakeParam:
+    def __init__(self):
+        self.loaded = None
+
+    def weight_loader(self, param, loaded_weight, *args, **kwargs):
+        self.loaded = (param, loaded_weight, args, kwargs)
+
+
+class TestHunyuanV3NextNWeightLoading(unittest.TestCase):
+    def _make_minimal_model(self, named_parameters=()):
+        model = object.__new__(HYV3ForCausalLMNextN)
+        model.config = SimpleNamespace(num_hidden_layers=80, num_experts=2)
+        model.named_parameters = lambda: iter(named_parameters)
+        return model
+
+    def test_final_layernorm_loads_into_shared_head_norm(self):
+        param = _FakeParam()
+        model = self._make_minimal_model([("model.shared_head.norm.weight", param)])
+        loaded_weight = torch.ones(1)
+
+        model.load_weights([("model.layers.80.final_layernorm.weight", loaded_weight)])
+
+        self.assertEqual(param.loaded, (param, loaded_weight, (), {}))
+
+    def test_spec_weights_map_to_model_prefix(self):
+        params = {
+            "model.enorm.weight": _FakeParam(),
+            "model.hnorm.weight": _FakeParam(),
+            "model.eh_proj.weight": _FakeParam(),
+        }
+        model = self._make_minimal_model(list(params.items()))
+        weights = [
+            ("model.layers.80.enorm.weight", torch.ones(1)),
+            ("model.layers.80.hnorm.weight", torch.ones(1)),
+            ("model.layers.80.eh_proj.weight", torch.ones(1)),
+        ]
+
+        model.load_weights(weights)
+
+        for name, param in params.items():
+            self.assertIsNotNone(param.loaded, f"{name} was not loaded")
+
+    def test_decoder_layer_weight_maps_to_decoder_prefix(self):
+        param = _FakeParam()
+        model = self._make_minimal_model(
+            [("model.decoder.input_layernorm.weight", param)]
+        )
+        loaded_weight = torch.ones(1)
+
+        model.load_weights([("model.layers.80.input_layernorm.weight", loaded_weight)])
+
+        self.assertEqual(param.loaded, (param, loaded_weight, (), {}))
+
+    def test_embed_tokens_and_lm_head_are_skipped(self):
+        params = {
+            "model.embed_tokens.weight": _FakeParam(),
+            "lm_head.weight": _FakeParam(),
+        }
+        model = self._make_minimal_model(list(params.items()))
+        weights = [
+            ("model.embed_tokens.weight", torch.ones(1)),
+            ("lm_head.weight", torch.ones(1)),
+        ]
+
+        model.load_weights(weights)
+
+        for name, param in params.items():
+            self.assertIsNone(param.loaded, f"{name} should have been skipped")
+
+
+if __name__ == "__main__":
+    unittest.main()