diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 2529012f7..da18583e2 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -608,6 +608,7 @@ class ModelConfig: ) or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures or "KimiK25ForConditionalGeneration" in self.hf_config.architectures + or "Eagle3DeepseekV2ForCausalLM" in self.hf_config.architectures ): self.head_dim = 256 self.attention_arch = AttentionArch.MLA diff --git a/python/sglang/srt/models/kimi_k25_eagle3.py b/python/sglang/srt/models/kimi_k25_eagle3.py new file mode 100644 index 000000000..93c77d4cc --- /dev/null +++ b/python/sglang/srt/models/kimi_k25_eagle3.py @@ -0,0 +1,458 @@ +"""EAGLE3 draft model with MLA attention for Kimi-K2.5. + +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.5 target keeps the draft KV cache small. +""" + +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.layers.communicator import AttentionInputs, get_attn_tp_context +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import 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, +) +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), + ) + + target_hidden_size = ( + getattr(config, "target_hidden_size", None) or config.hidden_size + ) + num_fc_input = _get_eagle_aux_layer_count(config) + self.fc = nn.Linear( + target_hidden_size * num_fc_input, + config.hidden_size, + bias=getattr(config, "bias", False), + ) + + 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) + + 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: + 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]: + hidden_states = self.fc(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 + ) + return hidden_states_to_logits, [hidden_states_to_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] diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index cd8729798..7f1589756 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -119,6 +119,12 @@ try: _CONFIG_REGISTRY["deepseek_v32"] = _DeepseekV32ConfigAlias _CONFIG_REGISTRY["deepseek_v4"] = _DeepseekV4ConfigAlias + + # For kimi_k25_eagle3 + class _KimiK2ConfigAlias(_HFDeepseekV3Config): + model_type = "kimi_k2" + + _CONFIG_REGISTRY["kimi_k2"] = _KimiK2ConfigAlias except ImportError: pass