> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.speculative.eagle.draft_gemma

EAGLE-3 draft model for Gemma4 targets (`Gemma4ForConditionalGeneration`).

Gemma4 is a multimodal decoder whose *text* backbone differs from a Llama-style
dense LLM in a handful of ways: a scaled word embedding (`* sqrt(hidden)`),
zero-centred `(1 + w)` RMSNorm, per-head Q/K norm, GeGLU (`gelu_pytorch_tanh`)
MLPs, alternating sliding-window / full attention with two separate RoPE
schedules, and (on the E2B/E4B checkpoints) Gemma-3n-style per-layer inputs and
AltUp. Almost none of that reaches the EAGLE-3 *draft*: the draft is a single
from-scratch decoder layer that consumes only the post-block auxiliary hidden
states emitted by the frozen target (via `register_forward_hook`) and
re-projects its own Q/K/V, so it never sees the target's sliding mask, per-layer
inputs, AltUp, or experts -- structurally it is the same Llama-style dense draft
used for every other registry entry. The draft's own RMSNorms are trained from
scratch, so the target's zero-centred norm form does not need to be reproduced,
and the constant embedding scale is normalised away by the draft's
`input_layernorm` before the fused `[embed, hidden]` attention input.

Three config quirks *do* have to be reconciled before the shared Llama draft can
build, and that is all this module does (see :func:`_normalize_gemma4_draft_config`):

1. **Activation key.** Gemma text configs name the MLP activation
   `hidden_activation` (`gelu_pytorch_tanh`); the shared MLP reads
   `config.hidden_act`. The GeGLU structure is identical to the draft's SwiGLU
   wiring once the activation is `gelu_pytorch_tanh`, so we copy the value
   across.
2. **RoPE.** Gemma4 stores a nested, per-attention-type `rope_parameters`
   (`full_attention`: `rope_theta=1e6`, `rope_type="proportional"`,
   `partial_rotary_factor=0.25`; `sliding_attention`: `rope_theta=1e4`,
   `rope_type="default"`) that the shared `LlamaRotaryEmbedding` cannot read
   \-- its `_get_rope_config` expects a flat `rope_theta` and would silently
   fall back to base `10000`. The single draft layer runs *full* causal
   attention over the whole sequence, so it mirrors the target's **global**
   (full-attention) schedule. We flatten that schedule to a standard
   full-rotary Llama RoPE (`rope_theta` = the global theta, `head_dim` from
   the config, no partial rotary). Training and inference are then
   self-consistent: the saved checkpoint keeps the canonical
   `architectures: ["LlamaEagle3DraftModel"]` string and a flat `rope_theta`,
   so SGLang / vLLM reproduce the exact same rotary with their existing EAGLE-3
   Llama head -- no Gemma-specific inference support is required. (Reproducing
   the target's `proportional` + partial-rotary global schedule instead would
   be more faithful to the target's long-context frequencies but no inference
   engine can currently serve it, so it is intentionally not done here.)
3. **MoE FFN width.** On a MoE Gemma4 (`enable_moe_block`) the config
   `intermediate_size` is the *per-expert* FFN width (e.g. 2112 on 26B-A4B,
   *below* `hidden_size` 2816). The dense draft would otherwise build a
   contracting MLP and be starved of capacity, so the draft MLP is sized to the
   target's *active* FFN width (`top_k_experts * intermediate_size`). Dense
   targets are untouched.

Everything else (GQA, the EAGLE-3 TTT cache attention, the `fc` projection,
the draft `lm_head` and vocab mapping) is inherited unchanged from
:class:`LlamaEagle3DraftModel`, and the on-disk state-dict layout is identical.

## Module Contents

### Classes

| Name                                                                                                        | Description                             |
| ----------------------------------------------------------------------------------------------------------- | --------------------------------------- |
| [`Gemma4Eagle3DraftModel`](#nemo_automodel-components-speculative-eagle-draft_gemma-Gemma4Eagle3DraftModel) | EAGLE-3 draft model for Gemma4 targets. |

### Functions

| Name                                                                                                                        | Description                                                                |
| --------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_extract_global_rope_theta`](#nemo_automodel-components-speculative-eagle-draft_gemma-_extract_global_rope_theta)         | Return the full-attention (global) RoPE base for the Gemma4 text config.   |
| [`_normalize_gemma4_draft_config`](#nemo_automodel-components-speculative-eagle-draft_gemma-_normalize_gemma4_draft_config) | Reconcile Gemma4 text-config quirks in place so the Llama draft can build. |

### Data

[`_DEFAULT_GLOBAL_ROPE_THETA`](#nemo_automodel-components-speculative-eagle-draft_gemma-_DEFAULT_GLOBAL_ROPE_THETA)

[`logger`](#nemo_automodel-components-speculative-eagle-draft_gemma-logger)

### API

```python
class nemo_automodel.components.speculative.eagle.draft_gemma.Gemma4Eagle3DraftModel(
    config: transformers.PretrainedConfig
)
```

**Bases:** [LlamaEagle3DraftModel](/nemo-automodel/nemo_automodel/components/speculative/eagle/draft_llama#nemo_automodel-components-speculative-eagle-draft_llama-LlamaEagle3DraftModel)

EAGLE-3 draft model for Gemma4 targets.

Identical to :class:`LlamaEagle3DraftModel` except that the Gemma4 text
config is normalized (activation key + flattened global RoPE) before the
shared draft is constructed. See the module docstring for the rationale;
the on-disk checkpoint is byte-for-byte a standard `LlamaEagle3DraftModel`
export so SGLang / vLLM load it with their existing EAGLE-3 Llama head.

```python
nemo_automodel.components.speculative.eagle.draft_gemma._extract_global_rope_theta(
    config: transformers.PretrainedConfig
) -> float
```

Return the full-attention (global) RoPE base for the Gemma4 text config.

Gemma4 nests RoPE parameters per attention type under
`config.rope_parameters` -- `{"full_attention": {"rope_theta": ...},
"sliding_attention": {"rope_theta": ...}}`. The draft runs full causal
attention, so it uses the `full_attention` base. Falls back to a flat
`rope_theta` attribute and then to :data:`_DEFAULT_GLOBAL_ROPE_THETA`.

```python
nemo_automodel.components.speculative.eagle.draft_gemma._normalize_gemma4_draft_config(
    config: transformers.PretrainedConfig
) -> None
```

Reconcile Gemma4 text-config quirks in place so the Llama draft can build.

Sets `hidden_act` from Gemma's `hidden_activation` and flattens the
nested `rope_parameters` to a standard full-rotary Llama RoPE keyed on the
global (full-attention) `rope_theta` (see the module docstring). Mutates
`config` so the draft's serialized `config.json` carries the flattened,
inference-reproducible RoPE.

```python
nemo_automodel.components.speculative.eagle.draft_gemma._DEFAULT_GLOBAL_ROPE_THETA = 1000000.0
```

```python
nemo_automodel.components.speculative.eagle.draft_gemma.logger = logging.getLogger(__name__)
```