GEMMA 4 ARCHITECTURE

self.scaling = 1.0

# v_proj is None when attention_k_eq_v=True
value_states = self.v_proj(hidden_states).view(hidden_shape) \
    if self.v_proj is not None else key_states

# K path: scaled norm + RoPE
key_states = self.k_norm(key_states)          # with_scale=True
key_states = apply_rotary_pos_emb(key_states, cos, sin)

# V path: unscaled norm, NO RoPE
value_states = self.v_norm(value_states)      # with_scale=False

class Gemma4RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6, with_scale=True):
        super().__init__()
        self.eps = eps
        self.with_scale = with_scale
        if self.with_scale:
            self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, hidden_states):
        normed = self._norm(hidden_states.float())
        if self.with_scale:
            normed = normed * self.weight.float()
        return normed.type_as(hidden_states)

class Gemma4TextScaledWordEmbedding(nn.Embedding):
    def __init__(self, num_embeddings, embedding_dim,
                 padding_idx, embed_scale=1.0):
        super().__init__(num_embeddings, embedding_dim, padding_idx)
        self.register_buffer(
            "embed_scale", torch.tensor(embed_scale))

    def forward(self, input_ids):
        return super().forward(input_ids) * self.embed_scale

# Main embedding: scale = sqrt(5376) ≈ 73.3
# Per-layer embedding: scale = sqrt(256) = 16.0

# After attention + FFN residuals:
if self.hidden_size_per_layer_input:
    residual = hidden_states
    hidden_states = self.per_layer_input_gate(hidden_states)  # d → 256
    hidden_states = self.act_fn(hidden_states)
    hidden_states = hidden_states * per_layer_input  # gated element-wise
    hidden_states = self.per_layer_projection(hidden_states)  # 256 → d
    hidden_states = self.post_per_layer_input_norm(hidden_states)
    hidden_states = residual + hidden_states

# Dense path (standard MLP)
hidden_states = self.pre_feedforward_layernorm(residual)
hidden_states = self.mlp(hidden_states)

if self.enable_moe_block:
    hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)

    # MoE path (parallel, from pre-MLP residual)
    hidden_states_flat = residual.reshape(-1, residual.shape[-1])
    _, top_k_weights, top_k_index = self.router(hidden_states_flat)
    hidden_states_2 = self.pre_feedforward_layernorm_2(hidden_states_flat)
    hidden_states_2 = self.experts(hidden_states_2, top_k_index, top_k_weights)
    hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2)

    # Sum both paths, then final post-norm
    hidden_states = hidden_states_1 + hidden_states_2
    hidden_states = self.post_feedforward_layernorm(hidden_states)

hidden_states = residual + hidden_states  # residual add

def forward(self, hidden_states):
    hidden_states = self.norm(hidden_states)        # RMSNorm, no scale
    hidden_states = hidden_states * self.scale * self.scalar_root_size

    expert_scores = self.proj(hidden_states)        # Linear → num_experts
    router_probs = softmax(expert_scores, dim=-1)

    top_k_weights, top_k_index = torch.topk(
        router_probs, k=self.config.top_k_experts)

    top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True)
    top_k_weights = top_k_weights * self.per_expert_scale[top_k_index]

    return router_probs, top_k_weights, top_k_index

hidden_states *= self.layer_scalar  # buffer, not nn.Parameter

if self.config.final_logit_softcapping is not None:
    logits = logits / self.config.final_logit_softcapping
    logits = torch.tanh(logits)
    logits = logits * self.config.final_logit_softcapping

# From Gemma4TextConfig:
rope_parameters = {
    "sliding_attention": {
        "rope_type": "default",
        "rope_theta": 10_000.0
    },
    "full_attention": {
        "rope_type": "proportional",
        "partial_rotary_factor": 0.25,
        "rope_theta": 1_000_000.0
    }
}

# Last layer forced to full_attention regardless of pattern:
if self.layer_types[-1] != "full_attention":
    self.layer_types[-1] = "full_attention"

# In shared layer forward:
if self.is_kv_shared_layer and past_key_values is not None:
    key_states, value_states = \
        past_key_values.shared_layers[self.kv_shared_layer_index]

# In non-shared layer: store KV for later reuse
if self.store_full_length_kv:
    past_key_values.shared_layers[self.layer_idx] = \
        key_states, value_states