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

"""EAGLE3 / EAGLE3.1 draft model with MLA attention for Kimi-K2.x.

The ``kimi-k2.5-eagle3-mla`` checkpoint pairs an EAGLE3 layout
(concatenated [embed_norm, hidden_norm] pre-attention input, fc projection
over the concatenated multi-layer aux hidden states, single decoder layer,
dense MLP) with DeepSeek-V2 multi-latent attention. Sharing the MLA layout
with the Kimi-K2.x target keeps the draft KV cache small.

The eagle3.1 variant (e.g. ``kimi-k2.6-eagle3.1-mla``) adds two optional
config flags on top of the same layout:

* ``fc_norm``: per-chunk RMSNorm applied to each auxiliary hidden state
  before the fc projection.
* ``norm_output``: emit post-norm (rather than pre-norm) hidden states as
  the auxiliary output consumed by the next draft step.
"""

import copy
import logging
import re
from typing import Iterable, List, Optional, Tuple

import torch
from torch import nn
from transformers import PretrainedConfig

from sglang.srt.distributed import get_pp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ColumnParallelLinear, ReplicatedLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import (
    ParallelLMHead,
    VocabParallelEmbedding,
    get_embedding_tp_kwargs,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP
from sglang.srt.utils import BumpAllocator, add_prefix

logger = logging.getLogger(__name__)


def _get_eagle_aux_layer_count(config: PretrainedConfig) -> int:
    """Number of target layers whose hidden states get concatenated into fc."""
    eagle_config = getattr(config, "eagle_config", None)
    if isinstance(eagle_config, dict):
        layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids")
    else:
        layer_ids = getattr(eagle_config, "eagle_aux_hidden_state_layer_ids", None)
    if layer_ids is None:
        layer_ids = getattr(config, "eagle_aux_hidden_state_layer_ids", None)
    if layer_ids is None:
        return 3
    return len(layer_ids)


class Eagle3MLADecoderLayer(nn.Module):
    """One EAGLE3 draft layer that uses DeepSeek-V2 multi-latent attention.

    Pre-attention concatenates the input embedding and the target hidden
    state along the channel dim, doubling the input width to MLA's fused
    QKV-down projection.
    """

    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int = 0,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size

        if hasattr(config, "rope_parameters") and config.rope_parameters is not None:
            rope_params = config.rope_parameters
            rope_theta = rope_params.get("rope_theta", 10000)
            rope_scaling = (
                rope_params if rope_params.get("rope_type") != "default" else None
            )
        else:
            rope_theta = config.rope_theta
            rope_scaling = config.rope_scaling
        max_position_embeddings = config.max_position_embeddings

        self.self_attn = DeepseekV2AttentionMLA(
            config=config,
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            qk_nope_head_dim=config.qk_nope_head_dim,
            qk_rope_head_dim=config.qk_rope_head_dim,
            v_head_dim=config.v_head_dim,
            q_lora_rank=config.q_lora_rank,
            kv_lora_rank=config.kv_lora_rank,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            layer_id=layer_id,
            reduce_results=True,
            prefix=add_prefix("self_attn", prefix),
        )

        # EAGLE3 doubles MLA's QKV-down input by concatenating
        # input_layernorm(embed) and hidden_norm(target_hidden) along the
        # feature dim. Replace the projection that DeepseekV2AttentionMLA
        # built for a single-hidden input.
        attn = self.self_attn
        if attn.q_lora_rank is None:
            raise ValueError(
                "Eagle3 MLA layer requires q_lora_rank in the draft config"
            )
        attn.fused_qkv_a_proj_with_mqa = ReplicatedLinear(
            2 * config.hidden_size,
            attn.q_lora_rank + attn.kv_lora_rank + attn.qk_rope_head_dim,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("self_attn.fused_qkv_a_proj_with_mqa", prefix),
        )
        # Recompute fused-proj-dependent flags so they reflect the new input dim.
        attn.has_fused_proj = True
        attn.use_min_latency_fused_a_gemm = False
        quant_method = getattr(attn.fused_qkv_a_proj_with_mqa, "quant_method", None)
        attn.is_packed_weight = (
            quant_method is not None
            and hasattr(quant_method, "quant_config")
            and quant_method.quant_config is not None
            and quant_method.quant_config.get_name()
            in {"awq", "awq_marlin", "moe_wna16"}
        )

        self.mlp = DeepseekV2MLP(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
            quant_config=quant_config,
            prefix=add_prefix("mlp", prefix),
        )

        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        positions: torch.Tensor,
        embeds: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        zero_allocator: BumpAllocator,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        residual = hidden_states
        embeds = self.input_layernorm(embeds)
        hidden_states = self.hidden_norm(hidden_states)
        attn_input = torch.cat([embeds, hidden_states], dim=-1)

        # MLA's forward_absorb_prepare reads the qkv-down projection result
        # from attn_tp_context. We bypass LayerCommunicator here (the eagle3
        # draft layer is one isolated layer with custom pre-attention norms),
        # so publish the attention input ourselves.
        get_attn_tp_context().set_attn_inputs(
            AttentionInputs(
                attn_input,
                forward_batch,
                self.self_attn.prepare_qkv_latent,
            )
        )

        attn_out = self.self_attn(
            positions=positions,
            hidden_states=attn_input,
            forward_batch=forward_batch,
            zero_allocator=zero_allocator,
        )
        if isinstance(attn_out, tuple):
            attn_out = attn_out[0]

        hidden_states, residual = self.post_attention_layernorm(attn_out, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


class Eagle3MLAModel(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            prefix=add_prefix("embed_tokens", prefix),
            **get_embedding_tp_kwargs(),
        )

        target_hidden_size = (
            getattr(config, "target_hidden_size", None) or config.hidden_size
        )
        self.num_aux_hidden_states = _get_eagle_aux_layer_count(config)
        self.fc = ColumnParallelLinear(
            target_hidden_size * self.num_aux_hidden_states,
            config.hidden_size,
            bias=getattr(config, "bias", False),
            gather_output=False,
            quant_config=quant_config,
            prefix=add_prefix("fc", prefix),
        )
        # Guarded multimem all-gather for the fc output; buffer covers prefill.
        self._fc_gatherer = triton_symm_mem_ag.MultimemAllGatherer(
            max_tokens=triton_symm_mem_ag.recommended_max_tokens(
                include_prefill=True, floor=512
            ),
            skip_entry_sync=False,
        )

        # Per-aux RMSNorm before fc; enabled via `fc_norm` or legacy
        # `use_aux_norm` flag. Matches the eagle3.1 layout.
        use_fc_norm = getattr(config, "fc_norm", None) or getattr(
            config, "use_aux_norm", False
        )
        if use_fc_norm:
            self.fc_norm = nn.ModuleList(
                [
                    RMSNorm(target_hidden_size, eps=config.rms_norm_eps)
                    for _ in range(self.num_aux_hidden_states)
                ]
            )
        else:
            self.fc_norm = None

        if config.num_hidden_layers != 1:
            raise ValueError("EAGLE3 currently only supports 1 layer")
        self.midlayer = Eagle3MLADecoderLayer(
            config,
            layer_id=0,
            quant_config=quant_config,
            prefix=prefix,
        )

        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # Draft decode captures pre-norm hidden by default; eagle3.1 opts for
        # post-norm via `norm_output: true`.
        self.norm_output = getattr(config, "norm_output", False)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: torch.Tensor = None,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        if input_embeds is None:
            # MM positions in input_ids hold MM_PAD_SHIFT_VALUE+hash sentinels (far above
            # vocab_size). Use target-produced mm_input_embeds for these positions and
            # only call embed_tokens on the appended next-token to avoid embed OOB.
            embeds = forward_batch.mm_input_embeds
            if (
                forward_batch.forward_mode.is_extend()
                and forward_batch.contains_mm_inputs()
                and not forward_batch.forward_mode.is_draft_extend_v2()
            ):
                assert embeds is not None
                last_indices = (
                    forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1
                ).long()
                embeds[last_indices] = self.embed_tokens(input_ids[last_indices])
            if embeds is None:
                embeds = self.embed_tokens(input_ids)
        else:
            embeds = input_embeds

        hidden_states = forward_batch.spec_info.hidden_states
        if hidden_states.shape[-1] != embeds.shape[-1]:
            if self.fc_norm is not None:
                chunks = hidden_states.chunk(self.num_aux_hidden_states, dim=-1)
                hidden_states = torch.cat(
                    [norm(chunk) for norm, chunk in zip(self.fc_norm, chunks)],
                    dim=-1,
                )
            hidden_states, _ = self.fc(hidden_states)
            hidden_states = self._fc_gatherer(hidden_states)

        if hidden_states.shape[0] == 0:
            return hidden_states, [hidden_states]

        zero_allocator = BumpAllocator(
            buffer_size=2,
            dtype=torch.float32,
            device=embeds.device,
        )
        hidden_states, residual = self.midlayer(
            positions, embeds, hidden_states, forward_batch, zero_allocator
        )

        hidden_states_to_logits, hidden_states_to_aux = self.norm(
            hidden_states, residual
        )
        aux = hidden_states_to_logits if self.norm_output else hidden_states_to_aux
        return hidden_states_to_logits, [aux]


class Eagle3DeepseekV2ForCausalLM(nn.Module):
    """EAGLE3 draft model architecture with DeepSeek-V2 MLA attention.

    Used by checkpoints like ``kimi-k2.5-eagle3-mla`` that pair
    an EAGLE3 layout with multi-latent attention so the draft KV cache shape
    matches the target's MLA cache.
    """

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        # Match deepseek_nextn behavior: modelopt_fp4 is target-only and the
        # bf16 draft must not inherit the FP4 quant method.
        if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
            logger.warning(
                "Overriding Eagle3DeepseekV2ForCausalLM quant config for "
                "modelopt_fp4 target; draft weights are bf16."
            )
            quant_config = None
        self.quant_config = quant_config
        self.pp_group = get_pp_group()

        self.model = Eagle3MLAModel(
            config, quant_config=quant_config, prefix=add_prefix("model", prefix)
        )

        # llama_eagle3 sets a load-from-target flag when draft_vocab_size is
        # missing. This checkpoint declares its own draft head, so keep ours.
        self.load_lm_head_from_target = False
        draft_vocab_size = getattr(config, "draft_vocab_size", None)
        if config.tie_word_embeddings:
            self.lm_head = self.model.embed_tokens
        else:
            if draft_vocab_size is None:
                self.load_lm_head_from_target = True
                draft_vocab_size = config.vocab_size
                config.draft_vocab_size = draft_vocab_size
            self.lm_head = ParallelLMHead(
                draft_vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=add_prefix("lm_head", prefix),
            )

        # Logits processor sees the draft vocab.
        config_for_logits = copy.deepcopy(config)
        config_for_logits.vocab_size = draft_vocab_size or config.vocab_size
        self.logits_processor = LogitsProcessor(config_for_logits)

        self.capture_aux_hidden_states = True
        self.hot_token_id = None

    @torch.no_grad()
    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:
        with get_attn_tp_context().maybe_input_scattered(forward_batch):
            hidden_states = self.model(
                input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
            )
        aux_hidden_states = None
        if isinstance(hidden_states, tuple):
            hidden_states, aux_hidden_states = hidden_states
        return self.logits_processor(
            input_ids,
            hidden_states,
            self.lm_head,
            forward_batch,
            aux_hidden_states,
        )

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

    def get_embed_and_head(self):
        return self.model.embed_tokens.weight, self.lm_head.weight

    def set_embed(self, embed: torch.Tensor) -> None:
        # If draft hidden size != target hidden size, embeddings can't be shared.
        if (
            hasattr(self.config, "target_hidden_size")
            and self.config.target_hidden_size is not None
            and self.config.target_hidden_size != self.config.hidden_size
        ):
            return
        del self.model.embed_tokens.weight
        self.model.embed_tokens.weight = embed
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None:
        del self.model.embed_tokens.weight
        del self.lm_head.weight
        self.model.embed_tokens.weight = embed
        self.lm_head.weight = head
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    def get_hot_token_id(self):
        return self.hot_token_id

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
        params_dict = dict(self.named_parameters())
        stacked_params_mapping = [
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]
        cached_a_proj: dict[str, torch.Tensor] = {}

        for name, loaded_weight in weights:
            if name == "d2t" or name.endswith(".d2t"):
                # d2t stores diffs between draft id and target id; absent in
                # checkpoints whose draft_vocab_size equals vocab_size.
                self.hot_token_id = loaded_weight + torch.arange(loaded_weight.shape[0])
                continue
            if name == "t2d" or name.endswith(".t2d"):
                continue

            # Map checkpoint layout (layers.0.X, embed_tokens.X, fc, norm) to the
            # internal layout (model.midlayer.X, model.embed_tokens.X, model.fc,
            # model.norm). lm_head stays at the top level.
            mapped_name = re.sub(r"^layers\.0\.", "midlayer.", name)
            if not mapped_name.startswith("lm_head."):
                mapped_name = f"model.{mapped_name}"

            handled = False
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in mapped_name:
                    continue
                target_name = mapped_name.replace(weight_name, param_name)
                if target_name not in params_dict:
                    continue
                param = params_dict[target_name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight, shard_id)
                handled = True
                break
            if handled:
                continue

            if "q_a_proj" in mapped_name or "kv_a_proj_with_mqa" in mapped_name:
                cached_a_proj[mapped_name] = loaded_weight
                q_name = (
                    mapped_name
                    if "q_a_proj" in mapped_name
                    else mapped_name.replace("kv_a_proj_with_mqa", "q_a_proj")
                )
                kv_name = (
                    mapped_name
                    if "kv_a_proj_with_mqa" in mapped_name
                    else mapped_name.replace("q_a_proj", "kv_a_proj_with_mqa")
                )
                if q_name in cached_a_proj and kv_name in cached_a_proj:
                    fused_weight = torch.cat(
                        [cached_a_proj[q_name], cached_a_proj[kv_name]], dim=0
                    )
                    fused_name = q_name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa")
                    if fused_name in params_dict:
                        param = params_dict[fused_name]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, fused_weight)
                    cached_a_proj.pop(q_name)
                    cached_a_proj.pop(kv_name)
                continue

            if mapped_name not in params_dict:
                logger.warning("Eagle3 MLA: skipping unexpected weight %s", name)
                continue
            param = params_dict[mapped_name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)

        self.post_load_weights()

    def post_load_weights(self) -> None:
        """Split kv_b_proj into w_kc / w_vc tensors used by MLA absorb_core.

        DeepseekV2 normally does this in DeepseekV2WeightLoaderMixin.post_load_weights;
        we re-implement the bf16 fast-path directly here to keep the eagle3 draft
        path independent of the full DeepseekV2 weight loader.
        """
        self_attn = self.model.midlayer.self_attn
        w = self_attn.kv_b_proj.weight
        if w.dtype not in (torch.bfloat16, torch.float16, torch.float32):
            raise NotImplementedError(
                f"Eagle3 MLA draft post_load_weights only supports float dtypes, got {w.dtype}"
            )
        w_kc, w_vc = w.unflatten(
            0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
        ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
        self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2)
        self_attn.w_vc = w_vc.contiguous().transpose(1, 2)


EntryClass = [Eagle3DeepseekV2ForCausalLM]