Qwen3.5-MoE: support modelopt_fp4 checkpoints that quantize attention (+ load baked FP8 KV scales) (#31220)

python/sglang/srt/models/qwen3_5.py

@@ -87,6 +87,7 @@ from sglang.srt.models.qwen2_moe import (
 # Models
 from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
 from sglang.srt.models.utils import (
+    WeightsMapper,
     fused_qk_gemma_rmsnorm,
     fused_qk_gemma_rmsnorm_with_gate,
 )
@@ -700,13 +701,8 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
         self.config = config
         self.layer_id = layer_id
 
-        linear_attn_quant_config = (
-            None
-            if quant_config and quant_config.get_name() == "modelopt_fp4"
-            else quant_config
-        )
         self.linear_attn = Qwen3_5GatedDeltaNet(
-            config, layer_id, linear_attn_quant_config, alt_stream, prefix
+            config, layer_id, quant_config, alt_stream, prefix
         )
 
         # NOTE: Determine the MLP type based on the model type
@@ -886,19 +882,13 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
             dtype=torch.get_default_dtype(),
         )
 
-        attn_quant_config = (
-            None
-            if quant_config and quant_config.get_name() == "modelopt_fp4"
-            else quant_config
-        )
-
         self.qkv_proj = QKVParallelLinear(
             config.hidden_size,
             self.head_dim,
             self.total_num_heads * (1 + self.attn_output_gate),
             self.total_num_kv_heads,
             bias=False,
-            quant_config=attn_quant_config,
+            quant_config=quant_config,
             tp_rank=self.attn_tp_rank,
             tp_size=self.attn_tp_size,
             prefix=add_prefix("qkv_proj", prefix),
@@ -908,7 +898,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
             self.total_num_heads * self.head_dim,
             config.hidden_size,
             bias=False,
-            quant_config=attn_quant_config,
+            quant_config=quant_config,
             reduce_results=False,
             tp_rank=self.attn_tp_rank,
             tp_size=self.attn_tp_size,
@@ -922,6 +912,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
             prefix=f"{prefix}.attn",
+            quant_config=quant_config,
         )
 
         # Dense MLP for non-MoE variant
@@ -1236,6 +1227,17 @@ ALL_DECODER_LAYER_TYPES = {
     "linear_attention": Qwen3_5LinearDecoderLayer,
 }
 
+# ModelOpt FP4 checkpoints bake the per-layer KV-cache scales under the HF
+# attention projections; in sglang they live on RadixAttention. Apply this to the
+# weight stream at the top of load_weights(), before ".self_attn" is stripped and
+# before the stacked qkv_proj matching would consume the name.
+QWEN3_5_KV_SCALE_MAPPER = WeightsMapper(
+    orig_to_new_substr={
+        ".self_attn.k_proj.k_scale": ".attn.k_scale",
+        ".self_attn.v_proj.v_scale": ".attn.v_scale",
+    },
+)
+
 
 class Qwen3_5ForCausalLM(nn.Module):
     """Qwen3.5 Model with support for dense variant."""
@@ -1476,6 +1478,7 @@ class Qwen3_5ForCausalLM(nn.Module):
         return hidden_states, aux_hidden_states
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights)
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -1564,6 +1567,7 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
         super().__init__(config=config, quant_config=quant_config, prefix=prefix)
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights)
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -1825,6 +1829,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration):
             torch.cuda.synchronize()
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights)
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -1984,6 +1989,7 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
             torch.cuda.synchronize()
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights)
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),

test/registered/unit/models/test_qwen3_5_modelopt_fp4.py

@@ -0,0 +1,157 @@
+"""Unit tests for modelopt_fp4 checkpoints that quantize Qwen3.5 attention.
+
+Covers three things:
+  1. ModelOptFp4Config.is_layer_excluded() decides per prefix whether attention is
+     quantized or kept in BF16.
+  2. RadixAttention registers k_scale/v_scale when built with a quant_config that
+     declares kv_cache_quant_algo; without them, baked KV scales have nowhere to
+     load into and silently fall back to 1.0.
+  3. QWEN3_5_KV_SCALE_MAPPER remaps the checkpoint's baked KV-scale names onto the
+     RadixAttention parameter names.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4Config
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.qwen3_5 import QWEN3_5_KV_SCALE_MAPPER
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestModelOptFp4AttentionExclusion(CustomTestCase):
+    def test_moe_only_checkpoint_excludes_attention(self):
+        # NVIDIA's Qwen3.5 NVFP4 checkpoints: attention and lm_head are excluded,
+        # MoE experts are not.
+        cfg = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            kv_cache_quant_algo="FP8",
+            group_size=16,
+            exclude_modules=["*self_attn*", "lm_head"],
+        )
+
+        self.assertTrue(cfg.is_layer_excluded("model.layers.0.self_attn.qkv_proj"))
+        self.assertTrue(cfg.is_layer_excluded("lm_head"))
+        self.assertFalse(
+            cfg.is_layer_excluded("model.layers.0.mlp.experts.3.gate_up_proj")
+        )
+
+    def test_uniform_w4a4_checkpoint_quantizes_attention(self):
+        # Uniform W4A4 checkpoint: only lm_head is excluded, attention is quantized.
+        cfg = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            kv_cache_quant_algo="FP8",
+            group_size=16,
+            exclude_modules=["lm_head"],
+        )
+
+        self.assertFalse(cfg.is_layer_excluded("model.layers.0.self_attn.qkv_proj"))
+        self.assertFalse(
+            cfg.is_layer_excluded("model.layers.0.linear_attn.in_proj_qkvz")
+        )
+        self.assertTrue(cfg.is_layer_excluded("lm_head"))
+
+
+class TestRadixAttentionKvScaleRegistration(CustomTestCase):
+    def _make_attn(self, quant_config):
+        return RadixAttention(
+            num_heads=2,
+            head_dim=8,
+            scaling=1.0,
+            num_kv_heads=2,
+            layer_id=0,
+            quant_config=quant_config,
+            prefix="model.layers.0.attn",
+        )
+
+    def test_with_fp8_kv_quant_config_registers_scale_params(self):
+        cfg = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            kv_cache_quant_algo="FP8",
+            group_size=16,
+            exclude_modules=[],
+        )
+        attn = self._make_attn(cfg)
+
+        self.assertIsInstance(attn.k_scale, torch.nn.Parameter)
+        self.assertIsInstance(attn.v_scale, torch.nn.Parameter)
+        # create_weights seeds -1.0, the sentinel for "checkpoint had no scale".
+        self.assertEqual(attn.k_scale.item(), -1.0)
+        self.assertEqual(attn.v_scale.item(), -1.0)
+
+    def test_without_quant_config_has_no_scale_params(self):
+        attn = self._make_attn(None)
+
+        self.assertIsNone(attn.k_scale)
+        self.assertIsNone(attn.v_scale)
+
+    def test_quant_config_without_kv_cache_algo_has_no_scale_params(self):
+        # Registration is gated on kv_cache_quant_algo, not on quant_config alone.
+        cfg = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            kv_cache_quant_algo=None,
+            group_size=16,
+            exclude_modules=[],
+        )
+        attn = self._make_attn(cfg)
+
+        self.assertIsNone(attn.k_scale)
+        self.assertIsNone(attn.v_scale)
+
+
+class TestQwen3_5KvScaleMapper(CustomTestCase):
+    def test_maps_baked_kv_scale_names_onto_radix_attention(self):
+        # Source names come from ModelOpt's export format, target names from the
+        # sglang module tree; a typo on either side silently zeroes the scales.
+        weights = [
+            ("model.layers.3.self_attn.k_proj.k_scale", torch.tensor(0.0347)),
+            ("model.layers.3.self_attn.v_proj.v_scale", torch.tensor(0.0128)),
+        ]
+
+        mapped = list(QWEN3_5_KV_SCALE_MAPPER.apply(weights))
+
+        self.assertEqual(
+            [name for name, _ in mapped],
+            ["model.layers.3.attn.k_scale", "model.layers.3.attn.v_scale"],
+        )
+        torch.testing.assert_close(mapped[0][1], torch.tensor(0.0347))
+        torch.testing.assert_close(mapped[1][1], torch.tensor(0.0128))
+
+    def test_all_other_names_pass_through_unchanged(self):
+        # A mapping key that is too broad would corrupt regular weight loading.
+        names = [
+            "model.layers.3.self_attn.k_proj.weight",
+            "model.layers.3.self_attn.k_proj.input_scale",
+            "model.layers.3.self_attn.k_proj.weight_scale",
+            "model.layers.2.linear_attn.in_proj_qkvz.weight",
+            "model.layers.0.mlp.experts.5.down_proj.weight",
+            "lm_head.weight",
+        ]
+        weights = [(name, torch.zeros(1)) for name in names]
+
+        mapped = list(QWEN3_5_KV_SCALE_MAPPER.apply(weights))
+
+        self.assertEqual([name for name, _ in mapped], names)
+
+    def test_mapped_scale_loads_via_default_weight_loader(self):
+        # The scale params carry no weight_loader, so load_weights' fallback uses
+        # default_weight_loader; its scalar path is what tolerates the 0-dim param
+        # vs the shape-[1] checkpoint tensor.
+        scale_param = torch.nn.Parameter(
+            torch.tensor(-1.0, dtype=torch.float32), requires_grad=False
+        )
+        loaded_weight = torch.tensor([0.0347], dtype=torch.float32)
+
+        weight_loader = getattr(scale_param, "weight_loader", default_weight_loader)
+        weight_loader(scale_param, loaded_weight)
+
+        self.assertAlmostEqual(scale_param.item(), 0.0347, places=6)
+
+
+if __name__ == "__main__":
+    unittest.main()