""" Speculators implementation of EAGLE-3: - https://arxiv.org/abs/2503.01840 Classes: Eagle3SpeculatorConfig: Configuration class for EAGLE-3 speculator model EagleSpeculator3: Main model implementation for EAGLE-3 speculators Eagle3Attention: Custom attention layer for EAGLE-3, processes concatenated embeddings and hidden states Eagle3DecoderLayer: Custom decoder layer for EAGLE-3, processes concatenated embeddings and hidden states with Eagle3Attention and support for moving hidden layernorm before residual """ import os from typing import Any, ClassVar, Literal, Optional, Union import torch from pydantic import Field, field_serializer, field_validator from torch import nn from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import ( LlamaMLP, LlamaRMSNorm, apply_rotary_pos_emb, repeat_kv, ) from speculators import SpeculatorModel, SpeculatorModelConfig __all__ = [ "Eagle3Attention", "Eagle3DecoderLayer", "Eagle3Speculator", "Eagle3SpeculatorConfig", ] @SpeculatorModelConfig.register("eagle3") class Eagle3SpeculatorConfig(SpeculatorModelConfig): """ Configuration for EAGLE-3 speculator with vocabulary mapping. EAGLE-3 features vocabulary mapping between draft (32K) and target (128K) vocabularies, enabling cross-tokenizer speculation. :param transformer_layer_config: Configuration for the transformer decoder layer :param draft_vocab_size: Size of draft model vocabulary for speculation :param norm_before_residual: Apply hidden_norm before storing residual """ speculators_model_type: Literal["eagle3"] = "eagle3" architectures: list[str] = Field( default_factory=lambda: ["Eagle3Speculator"], description="Model architectures that can load these weights", ) transformer_layer_config: PretrainedConfig = Field( default_factory=LlamaConfig, description="Configuration for the transformer decoder layer", ) draft_vocab_size: int = Field( default=32000, description="Size of draft model vocabulary for speculation", ) norm_before_residual: bool = Field( default=False, description="Apply hidden_norm before storing residual", ) target_hidden_size: Optional[int] = Field( default=None, description="Hidden size of the target model (if different from draft model)", ) @property def target_vocab_size(self) -> int: """Get target vocabulary size from transformer config.""" return self.transformer_layer_config.vocab_size @field_serializer("transformer_layer_config") def serialize_transformer_config(self, value: PretrainedConfig) -> dict: """Serialize transformer config to dict.""" return value.to_diff_dict() @field_validator("transformer_layer_config", mode="before") @classmethod def validate_transformer_config(cls, value: Any) -> PretrainedConfig: """Validate and convert transformer config.""" if isinstance(value, dict): config_class: type[PretrainedConfig] = LlamaConfig if "model_type" in value: from transformers import AutoConfig config_class = AutoConfig.for_model( model_type=value["model_type"] ).__class__ return config_class(**value) return value class Eagle3Attention(nn.Module): """ Eagle-3 attention module that processes concatenated embeddings and hidden states. Modified from standard Llama attention to accept 2x hidden_size input for Q/K/V projections while maintaining standard output size. """ def __init__(self, config: PretrainedConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.hidden_size = config.hidden_size self.head_dim = self.hidden_size // self.num_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads input_size = 2 * self.hidden_size self.q_proj = nn.Linear( input_size, self.num_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( input_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias, ) self.v_proj = nn.Linear( input_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias, ) self.o_proj = nn.Linear( self.hidden_size, self.hidden_size, bias=config.attention_bias ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]] = None, output_attentions: bool = False, use_cache: bool = False, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, # noqa: ARG002 ) -> tuple: """ Forward pass for Eagle-3 attention. Taken from Llama Attention but modified to accept 2x hidden_size input. :param hidden_states: Input tensor of shape [batch, seq_len, 2*hidden_size] :param attention_mask: Optional attention mask :param position_ids: Optional position IDs for rotary embeddings :param past_key_value: Optional cached key-value pairs :param output_attentions: Whether to return attention weights :param use_cache: Whether to cache key-value pairs :param position_embeddings: Optional precomputed rotary embeddings :return: Tuple of (hidden_states, [attention_weights], [past_key_value]) """ bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view( bsz, q_len, self.num_heads, self.head_dim ).transpose(1, 2) key_states = key_states.view( bsz, q_len, self.num_key_value_heads, self.head_dim ).transpose(1, 2) value_states = value_states.view( bsz, q_len, self.num_key_value_heads, self.head_dim ).transpose(1, 2) if position_embeddings is not None: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin, position_ids ) past_key_value_out = None if past_key_value is not None: past_key = past_key_value[0] past_value = past_key_value[1] key_states = torch.cat([past_key, key_states], dim=2) value_states = torch.cat([past_value, value_states], dim=2) if use_cache: past_key_value_out = (key_states, value_states) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / ( self.head_dim**0.5 ) if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax( attn_weights, dim=-1, dtype=torch.float32 ).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, self.hidden_size) attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value_out class Eagle3DecoderLayer(nn.Module): """ Eagle-3 decoder layer that processes concatenated embeddings and hidden states. Accepts 2x hidden_size input from concatenated embeddings and fused hidden states. Uses Eagle3Attention for the self-attention computation. """ def __init__( self, config: PretrainedConfig, layer_idx: int, norm_before_residual: bool = False, ): super().__init__() self.hidden_size = config.hidden_size self.norm_before_residual = norm_before_residual self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hidden_norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = LlamaRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.self_attn = Eagle3Attention(config, layer_idx) self.mlp = LlamaMLP(config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, # noqa: ARG002 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, # noqa: ARG002 ) -> tuple: """ Process concatenated embeddings and hidden states through modified decoder layer. :param hidden_states: Input tensor of shape [batch, seq_len, 2*hidden_size] :return: Tuple of layer outputs """ embeds = hidden_states[:, :, : self.hidden_size] hidden = hidden_states[:, :, self.hidden_size : 2 * self.hidden_size] if self.norm_before_residual: hidden = self.hidden_norm(hidden) residual = hidden else: residual = hidden hidden = self.hidden_norm(hidden) embeds = self.input_layernorm(embeds) attn_input = torch.cat([embeds, hidden], dim=-1) attn_output, attn_weights, past_key_value_out = self.self_attn( hidden_states=attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, position_embeddings=position_embeddings, ) hidden_states = residual + attn_output residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (attn_weights,) # type: ignore[assignment] if use_cache: outputs += (past_key_value_out,) # type: ignore[assignment] return outputs @SpeculatorModel.register("eagle3") class Eagle3Speculator(SpeculatorModel): """ EAGLE-3 speculator with vocabulary mapping and multi-layer fusion. EAGLE-3 processes concatenated hidden states from multiple verifier layers through a fusion layer, then combines with embeddings for a custom decoder layer that accepts 2x hidden_size input. """ config_class: ClassVar[type[Eagle3SpeculatorConfig]] = Eagle3SpeculatorConfig # type: ignore[misc] _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [ # type: ignore[misc] "verifier*", ] _keys_to_ignore_on_save: ClassVar[list[str]] = [] # type: ignore[misc,assignment] def __init__( self, config: Eagle3SpeculatorConfig, verifier: Optional[Union[str, os.PathLike, PreTrainedModel]] = None, verifier_attachment_mode: Optional[ Literal["detached", "full", "train_only"] ] = None, ): """ Initialize Eagle3 speculator. :param config: Eagle3SpeculatorConfig instance :param verifier: Optional verifier model :param verifier_attachment_mode: How to attach the verifier """ if not isinstance(config, Eagle3SpeculatorConfig): raise ValueError( f"config must be Eagle3SpeculatorConfig, got {type(config)}" ) self.config: Eagle3SpeculatorConfig = config self.hidden_size = config.transformer_layer_config.hidden_size self.draft_vocab_size = config.draft_vocab_size self.target_vocab_size = config.target_vocab_size # Use target_hidden_size if specified, otherwise use draft model's hidden_size self.target_hidden_size = ( config.target_hidden_size if config.target_hidden_size is not None else self.hidden_size ) super().__init__( config=config, verifier=verifier, verifier_attachment_mode=verifier_attachment_mode, ) self.embed_tokens = nn.Embedding( self.target_vocab_size, self.hidden_size, padding_idx=config.transformer_layer_config.pad_token_id if hasattr(config.transformer_layer_config, "pad_token_id") else None, ) self.fc = nn.Linear( 3 * self.target_hidden_size, # Use target model's hidden size self.hidden_size, bias=False, ) self.layers = nn.ModuleList( [ Eagle3DecoderLayer( config.transformer_layer_config, layer_idx=0, norm_before_residual=config.norm_before_residual, ) ] ) self.norm = LlamaRMSNorm( self.hidden_size, eps=config.transformer_layer_config.rms_norm_eps, ) self.lm_head = nn.Linear( self.hidden_size, self.draft_vocab_size, bias=False, ) self.register_buffer( "d2t", torch.zeros(self.draft_vocab_size, dtype=torch.long), ) self.register_buffer( "t2d", torch.zeros(self.target_vocab_size, dtype=torch.bool), ) # Type hints for buffers self.d2t: torch.Tensor self.t2d: torch.Tensor self.post_init() def forward( self, input_ids: torch.LongTensor, hidden_states: torch.FloatTensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, # noqa: ARG002 return_dict: Optional[bool] = None, ) -> Union[torch.FloatTensor, CausalLMOutputWithPast]: """ Forward pass for EAGLE-3 speculation. :param input_ids: Input token IDs from draft vocabulary :param hidden_states: Concatenated hidden states from 3 verifier layers [B, L, 3*target_H] where target_H is the target model's hidden size :param attention_mask: Optional attention mask :param position_ids: Optional position IDs :param past_key_values: Optional cached key-values :param use_cache: Whether to cache key-values :param output_attentions: Return attention weights :param output_hidden_states: Return hidden states :param return_dict: Return dict output :return: Model outputs with draft vocabulary logits """ return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) inputs_embeds = self.embed_tokens(input_ids) fused_hidden = self.fc(hidden_states) layer_input = torch.cat([inputs_embeds, fused_hidden], dim=-1) batch_size, seq_length = layer_input.shape[:2] if attention_mask is not None and attention_mask.dim() == 2: # noqa: PLR2004 past_key_values_length = ( past_key_values[0][0].shape[2] if past_key_values else 0 ) attention_mask = _prepare_4d_causal_attention_mask( attention_mask, (batch_size, seq_length), hidden_states, past_key_values_length, ) if position_ids is None: device = hidden_states.device position_ids = ( torch.arange( # type: ignore[assignment] seq_length, dtype=torch.long, device=device ) .unsqueeze(0) .expand(batch_size, -1) ) layer_outputs = self.layers[0]( layer_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values[0] if past_key_values else None, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = layer_outputs[0] hidden_states = self.norm(hidden_states) logits = self.compute_logits(hidden_states, map_to_target_vocab=True) if not return_dict: return logits return CausalLMOutputWithPast( logits=logits, past_key_values=[layer_outputs[1]] if use_cache else None, # type: ignore[arg-type] hidden_states=None, attentions=None, ) def compute_logits( self, hidden_states: torch.FloatTensor, map_to_target_vocab: bool = True, ) -> torch.FloatTensor: """ Compute logits with optional vocabulary mapping. :param hidden_states: Hidden states from the model :param map_to_target_vocab: Whether to map draft logits to target vocabulary :return: Logits tensor """ logits = self.lm_head(hidden_states) if not map_to_target_vocab: return logits batch_size, seq_length, _ = logits.shape draft_indices = torch.arange(self.draft_vocab_size, device=logits.device) target_indices = draft_indices + self.d2t mapped_logits = logits.new_full( (batch_size, seq_length, self.target_vocab_size), float("-inf") ) mapped_logits[:, :, target_indices] = logits return mapped_logits def map_draft_to_target_tokens( self, draft_tokens: torch.LongTensor ) -> torch.LongTensor: """ Map draft token IDs to target token IDs. :param draft_tokens: Draft vocabulary token IDs :return: Target vocabulary token IDs """ return draft_tokens + self.d2t[draft_tokens] # type: ignore[return-value] def check_target_token_availability( self, target_tokens: torch.LongTensor ) -> torch.BoolTensor: """ Check if target tokens have draft equivalents. :param target_tokens: Target vocabulary token IDs :return: Boolean mask indicating availability in draft vocabulary """ return self.t2d[target_tokens] # type: ignore[return-value]