Add Mistral Large 3 Eagle Support (#14466)

Co-authored-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com>

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]

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()

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_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
-                    weight=weight, weight_scale=weight_scale_inv
+                    weight=weight, weight_scale=weight_scale
                 )
             layer.weight = Parameter(weight.data, requires_grad=False)
-            layer.weight_scale_inv = Parameter(
+            layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)
-                weight_scale_inv.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,
             )

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!
+        if self.block_quant:
-        a_sf_t = a_sf.t().contiguous()
+            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.
+            if self.block_quant:
-            # It ignored the `output`` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325
+                # FIXME: there is a bug in the trtllm_fp8_block_scale_moe.
-            # so we put the whole function under the ``use_symmetric_memory`` context manager.
+                # It ignored the `output`` argument. https://github.com/flashinfer-ai/flashinfer/blob/da01b1bd8f9f22aec8c0eea189ad54860b034947/flashinfer/fused_moe/core.py#L1323-L1325
-            # If the bug is fixed, we can only put the output tensor allocation under the context manager.
+                # so we put the whole function under the ``use_symmetric_memory`` context manager.
-            return trtllm_fp8_block_scale_moe(
+                # If the bug is fixed, we can only put the output tensor allocation under the context manager.
-                routing_logits=(
+                return trtllm_fp8_block_scale_moe(
-                    router_logits.to(torch.float32)
+                    routing_logits=(
-                    if routing_method_type == RoutingMethodType.DeepSeekV3
+                        router_logits.to(torch.float32)
-                    else router_logits
+                        if routing_method_type == RoutingMethodType.DeepSeekV3
-                ),
+                        else router_logits
-                routing_bias=correction_bias,
+                    ),
-                hidden_states=a_q,
+                    routing_bias=correction_bias,
-                hidden_states_scale=a_sf_t,
+                    hidden_states=a_q,
-                gemm1_weights=layer.w13_weight,
+                    hidden_states_scale=a_sf_t,
-                gemm1_weights_scale=layer.w13_weight_scale_inv,
+                    gemm1_weights=layer.w13_weight,
-                gemm2_weights=layer.w2_weight,
+                    gemm1_weights_scale=layer.w13_weight_scale_inv,
-                gemm2_weights_scale=layer.w2_weight_scale_inv,
+                    gemm2_weights=layer.w2_weight,
-                num_experts=layer.num_experts,
+                    gemm2_weights_scale=layer.w2_weight_scale_inv,
-                top_k=topk_config.top_k,
+                    num_experts=layer.num_experts,
-                n_group=topk_config.num_expert_group,
+                    top_k=topk_config.top_k,
-                topk_group=topk_config.topk_group,
+                    n_group=topk_config.num_expert_group,
-                intermediate_size=layer.w2_weight.shape[2],
+                    topk_group=topk_config.topk_group,
-                local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
+                    intermediate_size=layer.w2_weight.shape[2],
-                local_num_experts=layer.num_local_experts,
+                    local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
-                routed_scaling_factor=(
+                    local_num_experts=layer.num_local_experts,
-                    routed_scaling_factor if routed_scaling_factor is not None else 1.0
+                    routed_scaling_factor=(
-                ),
+                        routed_scaling_factor
-                tile_tokens_dim=None,
+                        if routed_scaling_factor is not None
-                routing_method_type=routing_method_type,
+                        else 1.0
-                use_shuffled_weight=False,
+                    ),
-            )
+                    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,

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 (

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
 
 

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]

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(
+                    if model_arch not in [
-                        "DeepSeek MTP does not require setting speculative_draft_model_path."
+                        "MistralLarge3ForCausalLM",
-                    )
+                        "PixtralForConditionalGeneration",
+                    ]:
+                        logger.warning(
+                            "DeepSeek MTP does not require setting speculative_draft_model_path."
+                        )
 
             if self.speculative_num_steps is None:
                 assert (

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