From ab627e5d756094311973f7e00adda6075503115d Mon Sep 17 00:00:00 2001
From: Shaun Kotek <93727115+shaunkotek@users.noreply.github.com>
Date: Wed, 15 Jul 2026 22:12:04 +0300
Subject: [PATCH] fix: load the right mtp lm head quantization (#30976)

---
 python/sglang/srt/models/nemotron_h_mtp.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py
index 61b1e3bba..d38de9b9d 100644
--- a/python/sglang/srt/models/nemotron_h_mtp.py
+++ b/python/sglang/srt/models/nemotron_h_mtp.py
@@ -369,5 +369,10 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
     ):
         super().load_weights(weights, is_mtp=True)
 
+    def set_lm_head_from_target(self, target_lm_head: nn.Module) -> None:
+        if self.config.tie_word_embeddings:
+            return
+        self.lm_head = target_lm_head
+
 
 EntryClass = [NemotronHForCausalLMMTP]