From 0ff24159a5bcd0872a65135b0b1eeb3350273949 Mon Sep 17 00:00:00 2001
From: danielafrimi <45691845+danielafrimi@users.noreply.github.com>
Date: Mon, 16 Feb 2026 18:59:50 +0200
Subject: [PATCH] Fix modelopt FP8 create weights (#18447)

Signed-off-by: root <dafrimi@nvidia.com>
---
 .../srt/layers/quantization/modelopt_quant.py     |  2 ++
 python/sglang/srt/models/nemotron_h.py            | 15 +++++++++++----
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py
index d6b3dd68d..0ab272209 100755
--- a/python/sglang/srt/layers/quantization/modelopt_quant.py
+++ b/python/sglang/srt/layers/quantization/modelopt_quant.py
@@ -461,6 +461,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
         layer: torch.nn.Module,
         input_size_per_partition: int,
         output_partition_sizes: List[int],
+        input_size: Optional[int],
+        output_size: Optional[int],
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ) -> None:
diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py
index 0ff9be822..625722b5c 100644
--- a/python/sglang/srt/models/nemotron_h.py
+++ b/python/sglang/srt/models/nemotron_h.py
@@ -633,7 +633,12 @@ class NemotronHForCausalLM(nn.Module):
     }
 
     remap_prefix = {"backbone": "model"}
-    remap_substr = {"A_log": "A", "embeddings": "embed_tokens"}
+    remap_substr = {
+        "A_log": "A",
+        "embeddings": "embed_tokens",
+        "k_proj.k_scale": "attn.k_scale",
+        "v_proj.v_scale": "attn.v_scale",
+    }
 
     def __init__(
         self,
@@ -645,6 +650,7 @@ class NemotronHForCausalLM(nn.Module):
         super().__init__()
         lora_config = None
         self.config = config
+        self.quant_config = quant_config
         self.model = self._init_model(
             config=config, quant_config=quant_config, prefix=prefix
         )
@@ -776,9 +782,10 @@ class NemotronHForCausalLM(nn.Module):
                 continue
 
             if "scale" in name:
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
-                    continue
+                if name not in params_dict:
+                    name = maybe_remap_kv_scale_name(name, params_dict)
+                    if name is None:
+                        continue
 
             layer_id = get_layer_id(name)
             if (