From fa5698d7916497288af8fe5a5b57bc4ee7e6fb37 Mon Sep 17 00:00:00 2001 From: Zheng Li Date: Wed, 18 Feb 2026 11:44:25 +0800 Subject: [PATCH] feat: [Qwen3.5] Support block-wise FP8 quantization and model adaptation (#18926) --- python/sglang/srt/layers/linear.py | 48 ++++++++++++++++++++ python/sglang/srt/layers/quantization/fp8.py | 7 ++- python/sglang/srt/models/qwen3_5_mtp.py | 7 +-- python/sglang/srt/models/qwen3_vl.py | 7 ++- 4 files changed, 57 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 3b40a6067..1b4382a9b 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -728,6 +728,51 @@ class MergedColumnParallelLinear(ColumnParallelLinear): ) self.weight_loader_v2(param, loaded_weight_shard, shard_id) + def _load_merged_block_scale( + self, param: BasevLLMParameter, loaded_weight: torch.Tensor + ): + """ + Handle block-wise scale loading for MergedColumnParallelLinear. + Similar to QKVParallelLinear._load_qkv_block_scale, but for merged column layers. + """ + weight_block_size = self.quant_method.quant_config.weight_block_size + block_n, _ = weight_block_size[0], weight_block_size[1] + block_n = 1 if getattr(param, "format_ue8m0", False) else block_n + + # Calculate block sizes for each shard + shard_block_sizes = [] + shard_block_offsets = [] + current_block_offset = 0 + for output_size in self.output_sizes: + shard_block_size = (output_size + block_n - 1) // block_n + shard_block_sizes.append(shard_block_size) + shard_block_offsets.append(current_block_offset) + current_block_offset += shard_block_size + + # Load each shard + for shard_id, (shard_block_offset, shard_block_size) in enumerate( + zip(shard_block_offsets, shard_block_sizes) + ): + # Extract the shard from loaded_weight + loaded_weight_shard = loaded_weight.narrow( + param.output_dim, shard_block_offset, shard_block_size + ) + + # Calculate per-rank offset and size (considering TP) + rank_shard_offset = shard_block_offset // self.tp_size + rank_shard_size = shard_block_size // self.tp_size + + # Load into the parameter + param.load_merged_column_weight( + loaded_weight=loaded_weight_shard, + shard_id=shard_id, + shard_offset=rank_shard_offset, + shard_size=rank_shard_size, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + use_presharded_weights=self.use_presharded_weights, + ) + def weight_loader_v2( self, param: BasevLLMParameter, @@ -743,6 +788,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear): tp_size=self.tp_size, ) return + elif isinstance(param, BlockQuantScaleParameter): + self._load_merged_block_scale(param, loaded_weight) + return elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_merged_column_weight( loaded_weight=loaded_weight, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index cb3ca7e0d..4cbc3a7e3 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -171,8 +171,11 @@ class Fp8Config(QuantizationConfig): config, ["ignored_layers", "modules_to_not_convert"], None ) if ignored_layers: - # hack for ministral - ignored_layers = [layer.replace("model.", "") for layer in ignored_layers] + if "mistral3" in config.get("model_type", ""): + # hack for ministral + ignored_layers = [ + layer.replace("model.", "") for layer in ignored_layers + ] weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) if use_mxfp8 and weight_block_size is not None: logger.warning( diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 069ba761e..b65e1d6f5 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -64,7 +64,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): self.model = Qwen3_5ForCausalLM( config, quant_config, - prefix=add_prefix("model", prefix), + prefix=add_prefix("mtp", prefix), ) if get_pp_group().is_last_rank: @@ -214,16 +214,11 @@ class Qwen3_5ForCausalLMMTP(nn.Module): if "mtp" not in name: continue - # Some checkpoints use model.language_model.mtp.* prefix - if "language_model" in name: - name = name.replace(r"model.language_model.", r"model.") - if name.startswith("mtp."): # Remove the mtp. prefix for processing name = name.replace("mtp.", "model.") name = name.replace("model.fc", "fc") - name = name.replace("model.norm", "norm") name = name.replace("model.pre_fc", "pre_fc") if ".self_attn." in name: diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 2a8628dbc..d641826e3 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -786,9 +786,9 @@ class Qwen3VLForConditionalGeneration(nn.Module): config.vision_config, # NOTE: Qwen3-VL vision encoder currently supports BitsAndBytes 4-bit quantization. # Other quantization methods (e.g., GPTQ, AWQ) are untested and may not be supported. - quant_config=quant_config, + quant_config=None, norm_eps=getattr(config, "rms_norm_eps", 1e-6), - prefix=add_prefix("visual", prefix), + prefix=add_prefix("model.visual", prefix), use_data_parallel=self.use_data_parallel, ) @@ -804,7 +804,7 @@ class Qwen3VLForConditionalGeneration(nn.Module): self.model = language_model_cls( config=self.config, quant_config=quant_config, - prefix=add_prefix("model", prefix), + prefix=add_prefix("model.language_model", prefix), ) if self.pp_group.is_last_rank: if self.pp_group.world_size == 1 and self.config.tie_word_embeddings: @@ -1110,7 +1110,6 @@ class Qwen3VLForConditionalGeneration(nn.Module): if "visual" in name: # adapt to VisionAttention name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") - name = name.replace(r"model.visual.", r"visual.") try: # Skip loading extra bias for GPTQ models.