diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index d869d61bc..fab3d1bab 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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"), diff --git a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py new file mode 100644 index 000000000..59cc8e3c5 --- /dev/null +++ b/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()