chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add Mistral Large 3 Eagle Support (#14466)

Browse Source 
Co-authored-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com>
This commit is contained in:
elvischenv
2025-12-05 23:11:41 +08:00
committed by GitHub
parent 7235a7fbe9
commit 205f041e96
9 changed files with 313 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/configs/model_config.py
View File

@@ -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]

3
python/sglang/srt/layers/attention/trtllm_mla_backend.py
View File

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

16
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
View File

@@ -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,
)

197
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -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,

20
python/sglang/srt/models/deepseek_v2.py
View File

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

3
python/sglang/srt/models/mistral_large_3.py
View File

@@ -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

105
python/sglang/srt/models/mistral_large_3_eagle.py Normal file
View File

@@ -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]

10
python/sglang/srt/server_args.py
View File

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

9
python/sglang/srt/utils/mistral_utils.py
View File

@@ -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