Add Gemma 4 attention features: QK/value norms, shared KV, partial RoPE, embedding scale, logit softcap

[jlamypoirier]

Summary

Language model

• embedding_scale (LanguageModelEmbeddingsConfig): multiplicative scale applied to word embeddings after lookup. Gemma 4 uses sqrt(hidden_size). Zero overhead for the default value of 1.0 via a runtime branch. Necessary as a runtime op (not a weight init change) because tied embeddings share the weight with the LM head — baking the scale into weights would also scale logits.
• final_logit_softcap (LanguageModelHeadConfig): applies tanh(logits / cap) * cap before the loss. Gemma 4 uses cap=30. Forward and backward are each @torch.compile-decorated for op fusion. Gradient propagates through the Jacobian (1 - (softcapped / cap)²) before the output-linear backward.

Normalization

• FixedRMSNormConfig / FixedRMSNormalization: no-weight RMS normalization. Triton kernel extended with has_weight: tl_constexpr to skip weight load/multiply when weight=None; torch fallback uses torch.rms_norm(..., weight=None).
• post_mixer_normalization / post_mlp_normalization (DecoderBlockConfig): optional normalization applied to mixer/MLP outputs before the residual add. Gemma 4 applies RMSNorm at both positions.
• pre_mixer_normalization / pre_mlp_normalization (DecoderBlockConfig): independent overrides for the pre-norm at each sub-layer. Both default to normalization when unset, enabling different norm types per sub-layer (e.g. none for pre-MLP in certain Gemma 4 block variants).

Attention

• query_norm / key_norm (AttentionConfig): optional per-head RMSNorm applied to query and key vectors before RoPE. Gradient handled via a local autograd subgraph inside wrap_forward_backward.
• value_norm (AttentionConfig): optional per-head normalization applied to value projections before attention. Gemma 4 uses fixed_rms_norm (no learnable weight).
• shared_key_value (AttentionConfig): single key projection reused as value. Gradients from both key and value paths are summed back to the projection in the backward pass.
• In-place rotary fix: triton_rotary_ wrote results in-place, silently corrupting the saved norm output when query_norm was active (both tensors shared storage via .detach()). Added output_ptr to the Triton kernel and inplace_query flag through the rotary layer so the query gets a fresh allocation when a query_norm context is live.
• ProportionalRotaryConfig / ProportionalRotary: partial RoPE where only the first partial_rotary_factor fraction of head dimensions receive positional encoding (NoPE for the rest, via zero angle scales). Gemma 4 global-attention layers use partial_rotary_factor=0.5.

MoE

• HybridMoEMLPConfig / HybridMoEMLP: new MLP variant combining an always-active dense MLP with top-K routed experts. Each branch has optional pre/post norms (dense_pre_norm, dense_post_norm, moe_pre_norm, moe_post_norm). Gemma 4 uses this layout with separate intermediate sizes for the dense and expert paths.

Checkpoint converter

• Gemma 4 HuggingFace checkpoint converter (fast_llm/models/gpt/conversion/gemma4.py): full import/export support for the Gemma 4 text model family including sliding-window and full-attention pattern blocks, per-head norms, partial RoPE, hybrid MoE blocks, and tied embeddings.
• attention_k_eq_v → shared_key_value: Gemma 4 26B-A4B sets attention_k_eq_v=True for full-attention layers; the converter maps this to AttentionConfig.shared_key_value=True and routes through a single k_proj weight (no v_proj).
• MoE weight layout: HF stores expert weights as [num_experts, out, in] batched tensors; the converter reshapes to Fast-LLM's flat [num_experts * out, in] layout for gate_up_proj and handles the additional transpose for down_proj.
• use_bidirectional_attention: exported as None (Fast-LLM is text-only; bidirectional attention for vision tokens is not implemented).

Not yet implemented

• Per-Layer Embeddings (PLE): Gemma 4 feeds an auxiliary per-layer embedding signal (from a separate 262k-entry table) into each decoder block. Exported as hidden_size_per_layer_input: 0 to disable the feature until it is implemented in Fast-LLM. Round-tripping a real Gemma 4 checkpoint will lose PLE weights. Follow-up work needed.

Tests

• tests/layers/test_embedding.py: 3 base cases × 4 variants = 12 cases.
• tests/layers/test_lm_head.py: adds final_logit_softcap=2.0 case.
• tests/layers/test_attention.py: rewritten as a single parametrized suite — 5 base cases × 6 norm variants + rotary cases + shared_key_value cases = 160 cases covering forward reference, packing equivalence, and flash equivalence.
• tests/layers/test_decoder_block.py: 4 cases (no post-norms, post-mixer only, post-MLP only, both).
• tests/layers/test_mlp.py: hybrid MoE cases added. Test dimensions bumped to 128 to satisfy Triton block-size compile-time assertions under FAST_LLM_SKIP_TRITON_AUTOTUNE.
• tests/layers/test_rotary.py: rewritten as a single parametrized suite — default, big-theta, llama3, yarn, 2d, and proportional variants across head sizes and sequence lengths = 54 cases. Fixed a bug where triton_rotary_ (in-place) corrupted the reference input; now clones query before forward.
• tests/models/: 17 model tests pass for gemma4 (simple, bf16, fp16, checkpoint, resume, conversion, round-trip, load-pretrained, huggingface, frozen-weights, dtype variants).
• tests/models/test_hf_roundtrip.py: new test_hf_roundtrip[gemma4] using the real google/gemma-4-26B-A4B config with scaled-down dims to exercise the full import/export cycle.

Test plan

• pytest -v -n 8 tests/layers/test_rotary.py — 54 passed
• pytest -v -n 8 tests/layers/test_attention.py — 160 passed
• pytest -v tests/layers/test_embedding.py — 12 passed
• pytest -v tests/layers/test_lm_head.py — passed
• pytest -v tests/layers/test_decoder_block.py — 4 passed
• pytest -v -n 8 tests/layers/test_mlp.py — 8 passed
• pytest -v -n 8 --models gemma4 tests/models/ — 18 passed (17 + roundtrip)
• pytest -v -n 8 tests/ — full suite (running)

🤖 Generated with Claude Code