nemo_automodel.components.speculative.eagle.draft_llama

Bases: _PeagleAttentionMixin, Module

EAGLE-3 draft attention over [input_emb, hidden] 2H features.

Driven through a shared cache_hidden = [K_list, V_list] pair. At step k (0-indexed), with K_list and V_list already holding entries from steps 0..k-1:

1. step_idx = len(K_list) (equal to k) gives the rotary phase shift, so the draft’s K_k encodes “this is k tokens into the future”. The shifted cos / sin are computed from position_ids + step_idx.

2. The freshly projected K, V (after GQA expansion) are appended to the cache lists in place.

3. The attention output is the EAGLE-3 mixed pattern:

   attn_weights = [ Q @ K_0^T / sqrt(d) + mask ] || diag_1 || ... || diag_k

   where diag_i[t] = (Q_t * K_i_t).sum(-1) / sqrt(d). The softmax is taken over the full extended column axis of length T + k. Output is

   out = attn_probs[..., :T] @ V_0 + sum_{i=1..k} attn_probs[..., T+i-1, None] * V_i.

   In English: Q at position t attends to all K_0 positions (the regular T x T causal block), and additionally to the same position t in each previous draft step i >= 1. Implementation-wise we replace SpecForge llama3_eagle.py’s two O(k^2) cat / add Python loops with single vectorized einsum calls.

cache_hidden is mutated in place; callers are responsible for re-initializing it to [[], []] at the start of each training batch.

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaAttention._eager_attention_forward(
    q: torch.Tensor,
    cache_k: list[torch.Tensor],
    cache_v: list[torch.Tensor],
    attention_mask: torch.Tensor,
    step_idx: int,
    batch_size: int,
    seq_len: int
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaAttention._flash_attention_forward(
    q: torch.Tensor,
    cache_k: list[torch.Tensor],
    cache_v: list[torch.Tensor],
    step_idx: int,
    batch_size: int,
    seq_len: int,
    cu_seqlens: torch.Tensor | None = None,
    max_seqlen: int | None = None
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaAttention._flash_block1_varlen(
    q_fa: torch.Tensor,
    k0_fa: torch.Tensor,
    v0_fa: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: int,
    batch_size: int,
    seq_len: int
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaAttention._project_qkv(
    combined_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaAttention._repeat_kv(
    k: torch.Tensor,
    v: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaAttention.forward(
    combined_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    cache_hidden: list[list[torch.Tensor]],
    cu_seqlens: torch.Tensor | None = None,
    max_seqlen: int | None = None
) -> torch.Tensor

Bases: _PeagleDecoderLayerMixin, Module

Single decoder layer used by the minimal EAGLE-3 draft model.

Attribute names mirror SGLang’s LlamaDecoderLayer in sglang/srt/models/llama_eagle3.py: input_layernorm is applied to the per-step token embeddings (embeds in SGLang), hidden_norm is applied to the carried hidden state. is_input_layer is the layer-0 flag that gates the [embeds, hidden] concatenation (always true for our single-layer draft).

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaDecoderLayer.forward(
    input_embeds: torch.Tensor,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: torch.Tensor,
    cache_hidden: list[list[torch.Tensor]],
    cu_seqlens: torch.Tensor | None = None,
    max_seqlen: int | None = None
) -> torch.Tensor

Bases: Module

Standard Llama-style SwiGLU MLP on hidden-size activations.

nemo_automodel.components.speculative.eagle.draft_llama.Eagle3LlamaMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Bases: Module

Inner backbone matching SGLang’s LlamaModel in llama_eagle3.py.

Owns embed_tokens, the fc projection from concatenated target aux hidden states to draft hidden size, the (single-element) draft layers ModuleList, and the final norm. The LlamaEagle3DraftModel wrapper around this module adds the top-level lm_head and the training-facing public API.

Bases: _PeagleVanillaLayerMixin, Module

Vanilla Llama decoder layer for P-EAGLE depths >= 1.

The EAGLE-3 first layer (:class:Eagle3LlamaDecoderLayer) fuses the token embedding and the projected target hidden state (2H attention input). P-EAGLE stacks num_hidden_layers layers; every layer after the first is a standard Llama block operating on plain hidden states (H), matching speculators’ decoder_layer_class (a vanilla LlamaDecoderLayer). Only the P-EAGLE flex-attention path is implemented (these deeper layers do not participate in the EAGLE-3 cache_hidden TTT recurrence).

Bases: _PeagleDraftMixin, PreTrainedModel

Llama-style dense EAGLE-3 draft model (Llama, Phi-3, Qwen3).

State dict keys match SGLang’s LlamaForCausalLMEagle3 so the saved checkpoint can be loaded by SGLang’s inference engine without any remapping (SGLang’s load_weights fuses q/k/v_proj into qkv_proj and gate/up_proj into gate_up_proj via its standard stacked_params_mapping).

The class name is retained for checkpoint-architectures compatibility; the implementation is config-driven and works for any HF dense decoder-only config that exposes hidden_size, num_attention_heads, num_key_value_heads, attention_bias, mlp_bias, rope_theta, and rms_norm_eps. A decoupled head_dim is read via getattr(config, "head_dim", ...) in the attention layer.

Scope:

• single draft decoder layer
• no KV-cache optimization
• no speculative runtime integration

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.compute_logits(
    hidden_states: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.copy_embeddings_from_target(
    target_embedding: torch.nn.Embedding
) -> None

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.embed_input_ids(
    input_ids: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.forward(
    input_ids: torch.Tensor,
    projected_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    position_ids: typing.Optional[torch.Tensor] = None,
    cache_hidden: typing.Optional[list[list[torch.Tensor]]] = None,
    seq_lens: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.freeze_embeddings() -> None

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.gradient_checkpointing_disable() -> None

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs = None
) -> None

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.project_hidden_states(
    aux_hidden_states: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.eagle.draft_llama.LlamaEagle3DraftModel.set_vocab_mapping(
    selected_token_ids: torch.Tensor
) -> None

Build a standard causal + padding mask for SDPA/eager attention.

Return True when each row is a contiguous valid-prefix followed by padding.

Best-effort load of flash-attn without breaking eager-only users.

safe_import_from already handles missing modules and missing symbols, but some broken flash-attn installs fail with lower-level loader errors (e.g. ABI / shared-library issues) that should not prevent importing this module for the eager path. Returns the dense flash_attn_func and the flash_attn_varlen_func (used by the packed block-causal path).

Build FlashAttention varlen cu_seqlens (int32) from packed seq_lens.

Documents are flattened row-major to match the varlen attention’s reshape(B*T, ...) token order. Returns (cu_seqlens, max_seqlen).