From 50e8627af21ce9314592fd5595be783b42501c32 Mon Sep 17 00:00:00 2001
From: Chakshu Dhannawat <chakshu.dhannawat1@gmail.com>
Date: Fri, 28 Aug 2026 10:13:41 +0900
Subject: [PATCH] fix(modelopt): expand is_layer_excluded for fused and
 model.-prefixed names

ModelOpt FP4/FP8 mixed-precision checkpoints list excluded modules by unfused projection names (e.g. q_proj) or model.visual.* prefixes, but is_layer_excluded only matched the literal prefix. Expand the candidate set using packed_modules_mapping and also test model.-prefixed variants so fused modules and vision towers are correctly excluded.

Fixes sgl-project/sglang#36596
---
 .../srt/layers/quantization/modelopt_quant.py | 21 ++++++++
 .../quantization/test_modelopt_nvfp4.py       | 49 +++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py
index 38508638d88a..b6aa293e406a 100755
--- a/python/sglang/srt/layers/quantization/modelopt_quant.py
+++ b/python/sglang/srt/layers/quantization/modelopt_quant.py
@@ -349,6 +349,27 @@ def is_layer_excluded(self, prefix: str) -> bool:
         if prefix.startswith("language_model."):
             prefixes_to_check.append(prefix.removeprefix("language_model."))
 
+        # Expand fused module names to their checkpoint shard names. The exclude
+        # list may reference per-projection names (e.g. "q_proj") while the model
+        # builds a fused module (e.g. "qkv_proj"), so matching must consider the
+        # unfused constituents.
+        if self.packed_modules_mapping:
+            expanded = []
+            for p in prefixes_to_check:
+                head, _, tail = p.rpartition(".")
+                for shard_name in self.packed_modules_mapping.get(tail, []):
+                    expanded_prefix = f"{head}.{shard_name}" if head else shard_name
+                    expanded.append(expanded_prefix)
+                    if expanded_prefix.startswith("language_model."):
+                        expanded.append(expanded_prefix.removeprefix("language_model."))
+            prefixes_to_check.extend(expanded)
+
+        # Vision-language checkpoints sometimes rename "model.visual.*" to
+        # "visual.*" during load; also test the "model."-prefixed variant.
+        prefixes_to_check.extend(
+            "model." + p for p in prefixes_to_check if not p.startswith("model.")
+        )
+
         # Fused module patterns: the exclude list may reference a sub-component
         # (e.g., "q_a_proj") that is fused into a combined parameter name
         # (e.g., "fused_qkv_a_proj_with_mqa"). We check if the last segment of
diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py
index cb2f4bdd0a84..08060519a0e5 100644
--- a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py
+++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py
@@ -17,6 +17,55 @@
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
+class TestModelOptIsLayerExcluded(unittest.TestCase):
+    def test_exact_match(self):
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            exclude_modules=["lm_head"],
+        )
+        self.assertTrue(config.is_layer_excluded("lm_head"))
+        self.assertFalse(config.is_layer_excluded("embed_tokens"))
+
+    def test_language_model_prefix_stripping(self):
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            exclude_modules=["model.layers.0.self_attn.q_proj"],
+        )
+        self.assertTrue(
+            config.is_layer_excluded("language_model.model.layers.0.self_attn.q_proj")
+        )
+
+    def test_packed_modules_mapping_expansion(self):
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            exclude_modules=["model.layers.0.self_attn.q_proj"],
+            packed_modules_mapping={"qkv_proj": ["q_proj", "k_proj", "v_proj"]},
+        )
+        self.assertTrue(config.is_layer_excluded("model.layers.0.self_attn.qkv_proj"))
+
+    def test_model_prefix_variant(self):
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            exclude_modules=["model.visual.encoder.layers.0.attn.qkv_proj"],
+            packed_modules_mapping={"qkv_proj": ["q_proj", "k_proj", "v_proj"]},
+        )
+        self.assertTrue(
+            config.is_layer_excluded("visual.encoder.layers.0.attn.qkv_proj")
+        )
+
+    def test_wildcard_pattern(self):
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            exclude_modules=["model.visual.*"],
+        )
+        self.assertTrue(config.is_layer_excluded("visual.encoder.layers.0.attn.q_proj"))
+
+
 class TestModelOptNvfp4(CustomTestCase):
     def _make_layer(self):
         return MergedColumnParallelLinear(