sgl-project--sglang/python/sglang/srt/models/mistral_eagle.py

# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""EAGLE draft model for GQA Mistral targets (e.g. Mistral Medium 3.5).

Reuses ``LlamaForCausalLMEagle`` for the EAGLE machinery (lm_head/embed_tokens
construction, optional tied embeddings, capture-aux-hidden-states plumbing) but
swaps in a Mistral-specific draft model body that:

- runs through the standard :class:`LlamaDecoderLayer` (GQA), not the layernorm
  -less variant ``llama_eagle.LlamaDecoderLayer`` — Mistral's EAGLE checkpoint
  ships ``layers.0.attention_norm.weight``, so layer 0 expects the input
  layernorm to be present.
- uses ``RowParallelLinear`` for the EAGLE fc fusion layer with a
  ``quant_config``, so the FP8-quantized ``eagle_linear`` weights from the
  Mistral native checkpoint load via the standard quant pipeline (``LlamaModel``
  in ``llama_eagle.py`` uses a plain :class:`torch.nn.Linear` which cannot
  consume FP8 e4m3 tensors).

The weight name remapping mirrors :class:`MistralForCausalLMMistralFormat` and
adds the eagle-specific entries for ``eagle_linear`` → ``model.fc``.
"""

import logging
from collections.abc import Iterable
from typing import Optional, Tuple

import regex as re
import torch
from torch import nn
from transformers import PretrainedConfig

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.llama import LlamaDecoderLayer, LlamaForCausalLM
from sglang.srt.models.llama_eagle import LlamaForCausalLMEagle
from sglang.srt.utils import add_prefix

logger = logging.getLogger(__name__)


class MistralEagleModel(nn.Module):
    """GQA EAGLE draft body with the input-embed ⊕ target-hidden-state fusion."""

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.vocab_size = config.vocab_size
        assert (
            get_pp_group().world_size == 1
        ), "MistralForCausalLMEagle currently does not support pipeline parallelism"
        self.pp_group = get_pp_group()
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            prefix=add_prefix("embed_tokens", prefix),
        )
        self.layers = nn.ModuleList(
            [
                LlamaDecoderLayer(
                    config=config,
                    layer_id=i,
                    prefix=add_prefix(f"layers.{i}", prefix),
                    quant_config=quant_config,
                )
                for i in range(config.num_hidden_layers)
            ]
        )
        self.start_layer = 0
        self.end_layer = config.num_hidden_layers
        self.fc = RowParallelLinear(
            config.hidden_size * 2,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("fc", prefix),
            input_is_parallel=False,
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: Optional[torch.Tensor] = None,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> torch.Tensor:
        if input_embeds is None:
            hidden_states = self.embed_tokens(input_ids)
        else:
            hidden_states = input_embeds

        # EAGLE fusion: concat input embedding with target's previous hidden
        # state, project back to hidden_size before going through the draft's
        # transformer layers.
        hidden_states, _ = self.fc(
            torch.cat(
                (hidden_states, forward_batch.spec_info.hidden_states),
                dim=-1,
            )
        )

        residual = None
        for layer in self.layers:
            hidden_states, residual = layer(
                positions, hidden_states, forward_batch, residual
            )
        return hidden_states + residual


class MistralForCausalLMEagle(LlamaForCausalLMEagle):
    """EAGLE draft for GQA Mistral targets.

    Inherits LlamaForCausalLMEagle for the lm_head/embed_tokens setup and the
    capture-aux-hidden-state hooks, then overrides ``self.model`` with the
    quant-aware :class:`MistralEagleModel` and applies Mistral native-format
    weight remapping during ``load_weights``.
    """

    # fmt: off
    remapping = {
        r"layers\.(\d+)\.attention_norm\.weight": r"model.layers.\1.input_layernorm.weight",
        r"layers\.(\d+)\.attention\.wq\.(\w+)": r"model.layers.\1.self_attn.q_proj.\2",
        r"layers\.(\d+)\.attention\.wk\.(\w+)": r"model.layers.\1.self_attn.k_proj.\2",
        r"layers\.(\d+)\.attention\.wv\.(\w+)": r"model.layers.\1.self_attn.v_proj.\2",
        r"layers\.(\d+)\.attention\.wo\.(\w+)": r"model.layers.\1.self_attn.o_proj.\2",
        r"layers\.(\d+)\.ffn_norm\.weight": r"model.layers.\1.post_attention_layernorm.weight",
        r"layers\.(\d+)\.feed_forward\.w1\.(\w+)": r"model.layers.\1.mlp.gate_proj.\2",
        r"layers\.(\d+)\.feed_forward\.w2\.(\w+)": r"model.layers.\1.mlp.down_proj.\2",
        r"layers\.(\d+)\.feed_forward\.w3\.(\w+)": r"model.layers.\1.mlp.up_proj.\2",
        r"norm\.weight": "model.norm.weight",
        # Eagle-specific: the fc layer that fuses input embeds and target
        # hidden states is named `eagle_linear` in the Mistral checkpoint.
        # Its FP8 weights live alongside per-tensor activation/weight scales.
        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",
        # tok_embeddings and output are intentionally absent — EAGLE shares
        # both with the target model and the framework ties them at runtime.
    }
    # fmt: on

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        # Run LlamaForCausalLMEagle.__init__ to set up lm_head/embed_tokens/etc.
        # then replace self.model (which uses a plain torch.nn.Linear for fc and
        # cannot consume FP8 weights) with our quant-aware draft body.
        super().__init__(config=config, quant_config=quant_config, prefix=prefix)
        self.model = MistralEagleModel(
            config,
            quant_config=quant_config,
            prefix=add_prefix("model", prefix),
        )

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        # Bypass LlamaForCausalLMEagle.load_weights' "prepend model." behaviour
        # because our remap already emits fully-qualified target names.
        return LlamaForCausalLM.load_weights(
            self, self._remap_mistral_to_llama(weights)
        )

    def _remap_mistral_to_llama(
        self, weights: Iterable[Tuple[str, torch.Tensor]]
    ) -> Iterable[Tuple[str, torch.Tensor]]:
        for name, loaded_weight in weights:
            if name.startswith("model.") or name.startswith("lm_head."):
                yield name, loaded_weight
                continue
            for k, v in self.remapping.items():
                match = re.fullmatch(k, name)
                if match:
                    name = match.expand(v)
                    break
            else:
                logger.warning(f"Unrecognized weight: {name}. Skipping.")
                continue
            if name.endswith(".qscale_act"):
                name = re.sub(r"\.qscale_act$", ".input_scale", name)
            elif name.endswith(".qscale_weight"):
                name = re.sub(r"\.qscale_weight$", ".weight_scale", name)
            yield name, loaded_weight


EntryClass = [MistralForCausalLMEagle]