minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Henning Thieß
2026-07-30 14:30:26 -07:00
committed by GitHub
parent 5339450ed4
commit c4af6cf263
2 changed files with 177 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/qwen3_5.py
+20 -14
View File
@@ -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
+157
View File
@@ -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()