diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 87fc66b63..6483487fd 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -338,6 +338,7 @@ class ModelConfig: or "DotsVLMForCausalLM" in self.hf_config.architectures or "MistralLarge3ForCausalLM" in self.hf_config.architectures or "PixtralForConditionalGeneration" in self.hf_config.architectures + or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures ): self.head_dim = 256 self.attention_arch = AttentionArch.MLA @@ -702,7 +703,16 @@ class ModelConfig: if self.quantization is None: self.quantization = quant_method elif self.quantization != quant_method: - if ( + # Allow auto-detection of quantization from checkpoint for draft model + # even if it differs from main model's quantization + if self.is_draft_model: + logger.info( + f"Draft model quantization ({quant_method}) differs from " + f"main model quantization ({self.quantization}). " + f"Using draft model's detected quantization: {quant_method}" + ) + self.quantization = quant_method + elif ( self.quantization not in compatible_quantization_methods or quant_method not in compatible_quantization_methods[self.quantization] diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 7afa0a09e..b409cfdbf 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -964,7 +964,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # Apply llama 4 scaling if provided if llama_4_scaling is not None: - q *= llama_4_scaling + q = q.to(self.q_data_type) * llama_4_scaling + q = q.to(self.data_type) if ( forward_batch.forward_mode.is_target_verify() diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index 2effad036..423e74132 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -106,14 +106,12 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): weight_loader=weight_loader, ) weight_scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("weight_scale", weight_scale) elif self.strategy == QuantizationStrategy.TENSOR: weight_scale = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) weight_scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("weight_scale", weight_scale) elif self.strategy == QuantizationStrategy.BLOCK: assert layer.weight_block_size is not None block_n, block_k = layer.weight_block_size[0], layer.weight_block_size[1] @@ -130,8 +128,8 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): ) weight_scale.format_ue8m0 = False weight_scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("weight_scale_inv", weight_scale) + layer.register_parameter("weight_scale", weight_scale) # INPUT SCALE if self.is_static_input_scheme: input_scale = PerTensorScaleParameter( @@ -190,16 +188,14 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): elif self.strategy == QuantizationStrategy.BLOCK: assert self.is_static_input_scheme is False weight = layer.weight - weight_scale_inv = layer.weight_scale_inv + weight_scale = layer.weight_scale if is_fp8_fnuz(): - weight, weight_scale_inv, _ = normalize_e4m3fn_to_e4m3fnuz( - weight=weight, weight_scale=weight_scale_inv + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=weight, weight_scale=weight_scale ) layer.weight = Parameter(weight.data, requires_grad=False) - layer.weight_scale_inv = Parameter( - weight_scale_inv.data, requires_grad=False - ) + layer.weight_scale = Parameter(weight_scale.data, requires_grad=False) else: raise ValueError(f"Unknown quantization strategy {self.strategy}") @@ -221,7 +217,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): input=x, weight=layer.weight, block_size=self.weight_block_size, - weight_scale=layer.weight_scale_inv, + weight_scale=layer.weight_scale, input_scale=layer.input_scale, bias=bias, ) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 3338d5a41..ad6d81cb8 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -927,6 +927,86 @@ class Fp8MoEMethod(FusedMoEMethodBase): if _is_hip: self.process_weights_hip_scale_padding(layer) + + # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled + if get_moe_runner_backend().is_flashinfer_trtllm(): + from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a + + # Note: No need to swap W13 halves, they are already in the correct order: [Gate, Up] + num_experts, two_n, hidden = layer.w13_weight.shape + + # 2) Reorder rows for fused gated activation (W13) + w13_interleaved = [ + reorder_rows_for_gated_act_gemm(layer.w13_weight[i]) + for i in range(num_experts) + ] + w13_interleaved = torch.stack(w13_interleaved).reshape( + num_experts, two_n, hidden + ) + + # 3) Shuffle weights for transposed MMA output (both W13, W2) + epilogue_tile_m = 128 + w13_shuffled = [ + shuffle_matrix_a( + w13_interleaved[i].view(torch.uint8), epilogue_tile_m + ) + for i in range(num_experts) + ] + w2_shuffled = [ + shuffle_matrix_a( + layer.w2_weight[i].view(torch.uint8), epilogue_tile_m + ) + for i in range(num_experts) + ] + + layer.w13_weight = Parameter( + torch.stack(w13_shuffled).view(torch.float8_e4m3fn), + requires_grad=False, + ) + layer.w2_weight = Parameter( + torch.stack(w2_shuffled).view(torch.float8_e4m3fn), + requires_grad=False, + ) + + # Precompute and register per-expert output scaling factors for FI MoE + # Note: w13_input_scale and w2_input_scale are scalar Parameters post-reduction + assert ( + hasattr(layer, "w13_input_scale") + and layer.w13_input_scale is not None + ) + assert ( + hasattr(layer, "w2_input_scale") + and layer.w2_input_scale is not None + ) + assert ( + hasattr(layer, "w13_weight_scale") + and layer.w13_weight_scale is not None + ) + assert ( + hasattr(layer, "w2_weight_scale") + and layer.w2_weight_scale is not None + ) + + input_scale = layer.w13_input_scale.to(torch.float32) + activation_scale = layer.w2_input_scale.to(torch.float32) + w13_weight_scale = layer.w13_weight_scale.to(torch.float32) + w2_weight_scale = layer.w2_weight_scale.to(torch.float32) + + output1_scales_scalar = ( + w13_weight_scale * input_scale * (1.0 / activation_scale) + ) + output1_scales_gate_scalar = w13_weight_scale * input_scale + output2_scales_scalar = activation_scale * w2_weight_scale + + layer.output1_scales_scalar = Parameter( + output1_scales_scalar, requires_grad=False + ) + layer.output1_scales_gate_scalar = Parameter( + output1_scales_gate_scalar, requires_grad=False + ) + layer.output2_scales_scalar = Parameter( + output2_scales_scalar, requires_grad=False + ) return def process_weights_hip_int4(self, layer: Module): @@ -1230,7 +1310,10 @@ class Fp8MoEMethod(FusedMoEMethodBase): activation = self.moe_runner_config.activation routed_scaling_factor = self.moe_runner_config.routed_scaling_factor - from flashinfer.fused_moe import trtllm_fp8_block_scale_moe + from flashinfer.fused_moe import ( + trtllm_fp8_block_scale_moe, + trtllm_fp8_per_tensor_scale_moe, + ) from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType @@ -1241,9 +1324,15 @@ class Fp8MoEMethod(FusedMoEMethodBase): assert ( activation == "silu" ), "Only silu is supported for flashinfer blockscale fp8 moe" - a_q, a_sf = per_token_group_quant_fp8(x, self.quant_config.weight_block_size[1]) - # NOTE: scales of hidden states have to be transposed! - a_sf_t = a_sf.t().contiguous() + + if self.block_quant: + a_q, a_sf = per_token_group_quant_fp8( + x, self.quant_config.weight_block_size[1] + ) + # NOTE: scales of hidden states have to be transposed! + a_sf_t = a_sf.t().contiguous() + else: + a_q, _ = scaled_fp8_quant(x, layer.w13_input_scale) correction_bias = ( None @@ -1251,43 +1340,79 @@ class Fp8MoEMethod(FusedMoEMethodBase): else topk_config.correction_bias.to(x.dtype) ) - routing_method_type = getattr(layer, "routing_method_type") + routing_method_type = getattr( + layer, "routing_method_type", RoutingMethodType.DeepSeekV3 + ) with use_symmetric_memory( get_tp_group(), disabled=not is_allocation_symmetric() ): - # FIXME: there is a bug in the trtllm_fp8_block_scale_moe. - # It ignored the `output`` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325 - # so we put the whole function under the ``use_symmetric_memory`` context manager. - # If the bug is fixed, we can only put the output tensor allocation under the context manager. - return trtllm_fp8_block_scale_moe( - routing_logits=( - router_logits.to(torch.float32) - if routing_method_type == RoutingMethodType.DeepSeekV3 - else router_logits - ), - routing_bias=correction_bias, - hidden_states=a_q, - hidden_states_scale=a_sf_t, - gemm1_weights=layer.w13_weight, - gemm1_weights_scale=layer.w13_weight_scale_inv, - gemm2_weights=layer.w2_weight, - gemm2_weights_scale=layer.w2_weight_scale_inv, - num_experts=layer.num_experts, - top_k=topk_config.top_k, - n_group=topk_config.num_expert_group, - topk_group=topk_config.topk_group, - intermediate_size=layer.w2_weight.shape[2], - local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, - local_num_experts=layer.num_local_experts, - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - tile_tokens_dim=None, - routing_method_type=routing_method_type, - use_shuffled_weight=False, - ) + if self.block_quant: + # FIXME: there is a bug in the trtllm_fp8_block_scale_moe. + # It ignored the `output`` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325 + # so we put the whole function under the ``use_symmetric_memory`` context manager. + # If the bug is fixed, we can only put the output tensor allocation under the context manager. + return trtllm_fp8_block_scale_moe( + routing_logits=( + router_logits.to(torch.float32) + if routing_method_type == RoutingMethodType.DeepSeekV3 + else router_logits + ), + routing_bias=correction_bias, + hidden_states=a_q, + hidden_states_scale=a_sf_t, + gemm1_weights=layer.w13_weight, + gemm1_weights_scale=layer.w13_weight_scale_inv, + gemm2_weights=layer.w2_weight, + gemm2_weights_scale=layer.w2_weight_scale_inv, + num_experts=layer.num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=layer.w2_weight.shape[2], + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + local_num_experts=layer.num_local_experts, + routed_scaling_factor=( + routed_scaling_factor + if routed_scaling_factor is not None + else 1.0 + ), + tile_tokens_dim=None, + routing_method_type=routing_method_type, + use_shuffled_weight=False, + ) + else: + routing_bias_cast = ( + None + if correction_bias is None + else correction_bias.to(torch.bfloat16) + ) + + return trtllm_fp8_per_tensor_scale_moe( + routing_logits=router_logits.to(torch.bfloat16), + routing_bias=routing_bias_cast, + hidden_states=a_q, + gemm1_weights=layer.w13_weight, + output1_scales_scalar=layer.output1_scales_scalar, + output1_scales_gate_scalar=layer.output1_scales_gate_scalar, + gemm2_weights=layer.w2_weight, + output2_scales_scalar=layer.output2_scales_scalar, + num_experts=layer.num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=layer.w2_weight.shape[2], + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + local_num_experts=layer.num_local_experts, + routed_scaling_factor=( + routed_scaling_factor + if routed_scaling_factor is not None + else 1.0 + ), + use_routing_scales_on_input=False, + routing_method_type=routing_method_type, + ) def maybe_apply_hip_fused_experts( self, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c2213cf95..737dd5cf1 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -658,7 +658,9 @@ class DeepseekV2MoE(nn.Module): layer_id=self.layer_id, quant_config=quant_config, routed_scaling_factor=self.routed_scaling_factor, - routing_method_type=RoutingMethodType.DeepSeekV3, + routing_method_type=getattr( + config, "routing_method_type", RoutingMethodType.DeepSeekV3 + ), prefix=add_prefix("experts", prefix), ) @@ -3180,6 +3182,7 @@ class DeepseekV2Model(nn.Module): class DeepseekV2ForCausalLM(nn.Module): # for quark model load packed_modules_mapping = {} + model_cls = DeepseekV2Model def __init__( self, @@ -3188,7 +3191,6 @@ class DeepseekV2ForCausalLM(nn.Module): prefix: str = "", ) -> None: super().__init__() - # for quark model load # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None self.fuse_qkv_a_proj = ( @@ -3206,7 +3208,7 @@ class DeepseekV2ForCausalLM(nn.Module): self.quant_config = quant_config self.determine_num_fused_shared_experts() self.use_nsa = is_deepseek_nsa(config) - self.model = DeepseekV2Model( + self.model = self.model_cls( config, quant_config, prefix=add_prefix("model", prefix) ) if self.pp_group.is_last_rank: @@ -3400,16 +3402,22 @@ class DeepseekV2ForCausalLM(nn.Module): selected_quant_config, "weight_block_size", None ) if weight_block_size is not None: - assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") or hasattr( + self_attn.kv_b_proj, "weight_scale" + ) + weight_scale = ( + self_attn.kv_b_proj.weight_scale + if hasattr(self_attn.kv_b_proj, "weight_scale") + else self_attn.kv_b_proj.weight_scale_inv + ) if _is_fp8_fnuz: weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( weight=w, - weight_scale=self_attn.kv_b_proj.weight_scale_inv, + weight_scale=weight_scale, input_scale=None, ) else: weight = w - weight_scale = self_attn.kv_b_proj.weight_scale_inv # In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading. if ( diff --git a/python/sglang/srt/models/mistral_large_3.py b/python/sglang/srt/models/mistral_large_3.py index 719d47915..fd60ef61f 100644 --- a/python/sglang/srt/models/mistral_large_3.py +++ b/python/sglang/srt/models/mistral_large_3.py @@ -72,9 +72,6 @@ class MistralLarge3ForCausalLM(DeepseekV3ForCausalLM): elif name.endswith(".qscale_weight"): name = re.sub(r"\.qscale_weight$", ".weight_scale", name) - if name.endswith(".weight_scale") and ".experts." not in name: - name = re.sub(r"\.weight_scale$", ".weight_scale_inv", name) - yield name, loaded_weight diff --git a/python/sglang/srt/models/mistral_large_3_eagle.py b/python/sglang/srt/models/mistral_large_3_eagle.py new file mode 100644 index 000000000..f136640fd --- /dev/null +++ b/python/sglang/srt/models/mistral_large_3_eagle.py @@ -0,0 +1,105 @@ +from typing import Optional + +import torch +from torch import nn +from transformers import PretrainedConfig + +from python.sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import RowParallelLinear +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.models.deepseek_v2 import DeepseekV2DecoderLayer, DeepseekV2Model +from sglang.srt.models.mistral_large_3 import MistralLarge3ForCausalLM +from sglang.srt.utils import add_prefix + + +class MistralLarge3Model(DeepseekV2Model): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + nn.Module.__init__(self) + + self.config = config + self.vocab_size = config.vocab_size + assert get_pp_group().world_size == 1 + self.pp_group = get_pp_group() + self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + + self.layers = nn.ModuleList( + [ + DeepseekV2DecoderLayer( + config=config, + prefix=add_prefix(prefix, f"layers.{i}"), + quant_config=quant_config, + layer_id=i, + ) + for i in range(self.config.num_hidden_layers) + ] + ) + self.start_layer = 0 + self.end_layer = self.config.num_hidden_layers + + self.fc = RowParallelLinear( + self.config.hidden_size * 2, + self.config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix(prefix, "fc"), + input_is_parallel=False, + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layers_to_capture = [] + self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + if input_embeds is None: + input_embeds = self.embed_tokens(input_ids) + input_embeds, _ = self.fc( + torch.cat((input_embeds, forward_batch.spec_info.hidden_states), dim=-1) + ) + output = super().forward( + input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors + ) + assert isinstance(output, torch.Tensor) + return output + + +class MistralLarge3ForCausalLMEagle(MistralLarge3ForCausalLM): + remapping = MistralLarge3ForCausalLM.remapping | { + r"eagle_linear\.weight": r"model.fc.weight", + r"eagle_linear\.qscale_act": r"model.fc.input_scale", + r"eagle_linear\.qscale_weight": r"model.fc.weight_scale", + } + + def __init__( + self, + *, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + config.quant_config = quant_config + self.model_cls = MistralLarge3Model + super().__init__(config=config, quant_config=quant_config, prefix=prefix) + + +EntryClass = [MistralLarge3ForCausalLMEagle] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f3ca2833e..90af3ad1c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1722,9 +1722,13 @@ class ServerArgs: self.speculative_draft_model_path = self.model_path self.speculative_draft_model_revision = self.revision else: - logger.warning( - "DeepSeek MTP does not require setting speculative_draft_model_path." - ) + if model_arch not in [ + "MistralLarge3ForCausalLM", + "PixtralForConditionalGeneration", + ]: + logger.warning( + "DeepSeek MTP does not require setting speculative_draft_model_path." + ) if self.speculative_num_steps is None: assert ( diff --git a/python/sglang/srt/utils/mistral_utils.py b/python/sglang/srt/utils/mistral_utils.py index e23abb53c..ecce3042d 100644 --- a/python/sglang/srt/utils/mistral_utils.py +++ b/python/sglang/srt/utils/mistral_utils.py @@ -22,11 +22,15 @@ def adapt_config_dict( is_mistral_large_3 = ( is_moe and (config_dict["moe"].get("num_shared_experts") or 0) > 0 ) + is_eagle = "eagle" in model.lower() if is_moe: if is_mistral_large_3: config_dict = _remap_moe_args(config_dict) config_dict["model_type"] = "deepseek_v3" - config_dict["architectures"] = ["MistralLarge3ForCausalLM"] + if is_eagle: + config_dict["architectures"] = ["MistralLarge3ForCausalLMEagle"] + else: + config_dict["architectures"] = ["MistralLarge3ForCausalLM"] assert ( "llama_4_scaling" in config_dict @@ -77,6 +81,8 @@ def adapt_config_dict( config_dict = _remap_mistral_vision_args(config_dict) if is_audio: config_dict = _remap_mistral_audio_args(config_dict) + if is_eagle: + config_dict["routing_method_type"] = 1 # RoutingMethodType.Renormalize config = PretrainedConfig.from_dict(config_dict) @@ -227,7 +233,6 @@ def _remap_moe_args(config: dict) -> dict: config[new_name] = value config["topk_method"] = None - config["routing_method_type"] = 1 # RoutingMethodType.Renormalize config["scoring_func"] = "softmax" return config