""" Copyright 2023-2024 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. """ from sglang.srt.utils import add_prefix # Adapted from # https://github.com/SafeAILab/EAGLE/blob/main/eagle/model/cnets.py """Inference-only LLaMA-EAGLE model compatible with HuggingFace weights.""" import copy from typing import Iterable, Optional, Tuple import torch from torch import nn from transformers import LlamaConfig from sglang.srt.distributed import get_pp_group from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import QKVParallelLinear 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.llama import LlamaDecoderLayer, LlamaForCausalLM, LlamaMLP from sglang.srt.runtime_context import get_server_args class LlamaDecoderLayer(LlamaDecoderLayer): def __init__( self, config: LlamaConfig, layer_id: int = 0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__(config, layer_id, quant_config=quant_config, prefix=prefix) # Input layer concats embeds + target_hidden before qkv (input dim 2x). self.is_input_layer = layer_id == 0 hidden_size = 2 * self.hidden_size if self.is_input_layer else self.hidden_size # override qkv self.self_attn.qkv_proj = QKVParallelLinear( hidden_size, self.self_attn.head_dim, self.self_attn.total_num_heads, self.self_attn.total_num_kv_heads, bias=False, quant_config=quant_config, prefix=add_prefix("qkv_proj", prefix), ) if config.model_type == "llama4_text": inter_size = config.intermediate_size_mlp else: inter_size = config.intermediate_size self.mlp = LlamaMLP( config.hidden_size, inter_size, config.hidden_act, quant_config, prefix ) self.hidden_norm = 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, residual: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: if self.is_input_layer: # Input layer consumes target hidden states; no carried residual to fuse. residual = hidden_states hidden_states = self.hidden_norm(hidden_states) embeds = self.input_layernorm(embeds) hidden_states = torch.cat([embeds, hidden_states], dim=-1) else: # Fuse the previous layer's MLP residual add into hidden_norm. hidden_states, residual = self.hidden_norm(hidden_states, residual) # Self Attention hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) # Fully Connected hidden_states = self.mlp(hidden_states) return hidden_states, residual class LlamaModel(nn.Module): def __init__( self, config: LlamaConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config rope_parameters = getattr(config, "rope_parameters", None) if rope_parameters is not None: rope_scaling = rope_parameters else: rope_scaling = getattr(config, "rope_scaling", None) self.is_mrope_enabled = ( rope_scaling is not None and "mrope_section" in rope_scaling ) # fix rope_scaling for qwen2.5-vl if self.is_mrope_enabled: rope_scaling["rope_type"] = "default" self.vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, prefix=add_prefix("embed_tokens", prefix), ) if hasattr(config, "target_hidden_size"): self.hidden_size_in = config.target_hidden_size else: self.hidden_size_in = config.hidden_size # num_aux resolution: explicit attr > eagle_config layer_ids > default 3. self.num_aux_hidden_states = getattr(config, "num_aux_hidden_states", None) if self.num_aux_hidden_states is None: eagle_config = getattr(config, "eagle_config", None) or {} layer_ids = eagle_config.get("eagle_aux_hidden_state_layer_ids") self.num_aux_hidden_states = len(layer_ids) if layer_ids else 3 self.fc = torch.nn.Linear( self.hidden_size_in * self.num_aux_hidden_states, config.hidden_size, bias=getattr(config, "bias", False), ) # Per-aux RMSNorm before fc; enabled via `fc_norm` or legacy `use_aux_norm` flag. 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(self.hidden_size_in, eps=config.rms_norm_eps) for _ in range(self.num_aux_hidden_states) ] ) else: self.fc_norm = None self.layers = nn.ModuleList( [ LlamaDecoderLayer(config, i, quant_config, prefix) for i in range(config.num_hidden_layers) ] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) 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, ) -> torch.Tensor: if input_embeds is None: 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 if self.is_mrope_enabled: positions = forward_batch.mrope_positions 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) # idle batch if hidden_states.shape[0] == 0: return hidden_states, [hidden_states] residual = None for layer in self.layers: hidden_states, residual = layer( positions, embeds, hidden_states, forward_batch, residual, ) hidden_states_to_logits, hidden_states_to_aux = self.norm( hidden_states, residual ) # Draft decode captures pre-norm hidden by default; `norm_output` opts for normed. aux = hidden_states_to_logits if self.norm_output else hidden_states_to_aux return hidden_states_to_logits, [aux] class LlamaForCausalLMEagle3(LlamaForCausalLM): def __init__( self, config: LlamaConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: nn.Module.__init__(self) self.config = config self.quant_config = quant_config self.pp_group = get_pp_group() # Cache draft SWA size from server args once; consumed both by the post-init # attention patch below and by `get_attention_sliding_window_size` later. self._draft_window_size: Optional[int] = ( get_server_args().speculative_draft_window_size ) self.model = LlamaModel( config, quant_config=quant_config, prefix=add_prefix("model", prefix), ) if self._draft_window_size is not None: for layer in self.model.layers: layer.self_attn.attn.sliding_window_size = self._draft_window_size # Llama 3.2 1B Instruct set tie_word_embeddings to True # Llama 3.1 8B Instruct set tie_word_embeddings to False self.load_lm_head_from_target = False if self.config.tie_word_embeddings: self.lm_head = self.model.embed_tokens else: if config.draft_vocab_size is None: self.load_lm_head_from_target = True config.draft_vocab_size = config.vocab_size self.lm_head = ParallelLMHead( config.draft_vocab_size, config.hidden_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), ) config_ = copy.deepcopy(config) config_.vocab_size = ( config_.draft_vocab_size ) # draft logits processor has it's own vocab size self.logits_processor = LogitsProcessor(config_) self.capture_aux_hidden_states = True self.hot_token_id = None def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None: params_dict = dict(self.named_parameters()) # Define the parameter mapping for stacked parameters stacked_params_mapping = [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), ] # Legacy weight names -> new module attribute names (backwards compat). legacy_name_map = { "midlayer": "layers.0", "aux_norm_low": "fc_norm.0", "aux_norm_mid": "fc_norm.1", "aux_norm_high": "fc_norm.2", } for name, loaded_weight in weights: for legacy, new in legacy_name_map.items(): if legacy in name: name = name.replace(legacy, new) if "d2t" in name: # d2t stores diffs between draft id and target id self.hot_token_id = loaded_weight + torch.arange(loaded_weight.shape[0]) continue if "t2d" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) param_name = f"model.{name}" if name not in params_dict else name if param_name in params_dict: param = params_dict[param_name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight, shard_id) break else: # Handle regular parameters param_name = name if name in params_dict else f"model.{name}" if param_name in params_dict: param = params_dict[param_name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) def get_hot_token_id(self): return self.hot_token_id def get_attention_sliding_window_size(self) -> Optional[int]: return self._draft_window_size EntryClass = [LlamaForCausalLMEagle3]