Gemma.Block.ixx File Reference

Gemma 4 decoder block (inference: prefill + decode), one type per attention kind. More...

#include <memory>
#include <vector>
#include <string>
#include <format>
#include <sstream>
#include <stdexcept>
#include <optional>
#include <type_traits>
import Serialization.Tensor;
import Serialization.Mode;
import Dnn.Components.MixtureOfExperts;
import Dnn.Components.Router;
import Dnn.Components.GatedMLP;
import Dnn.Components.Swiglu;
import Dnn.ComponentType;
import Compute.DeviceId;
import Dnn.Quantization.KvCache.Policy;
import Dnn.Component;
import Dnn.Components.GemmaConfig;
import Dnn.Components.Residual;
import Compute.IExecutionContext;
import Serialization.SafeTensors;
import Dnn.ActivationType;
import Dnn.TensorOps;
import Dnn.Components.ITransformerBlock;
import Dnn.TensorDataTypeTraits;
import Dnn.ITensor;
import Dnn.Quantization.Weight.Policies;
import Compute.GqaState;
import Dnn.Components.Rope;
import Dnn.TensorTypes;
import Compute.Device;
import Dnn.Tensor;
import Compute.ExecutionContextFactory;
import Dnn.TensorDataType;
import Compute.CpuMemoryResource;
import Compute.DeviceType;
import Dnn.Components.GemmaBlock:Workspace;
import Compute.DeviceTypeTraits;
import Compute.DeviceAllocation;
import Compute.ExecutionContext;
import Serialization.ModelArchive;
import Compute.Observation;
import Dnn.CompositeComponent;
import Dnn.Components.RmsNorm;
import Dnn.Components.Gqa;
import Dnn.Components.Linear;

Classes
struct	Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >::BlockBuildContexts
	The per-child build contexts and split-scratch geometry this block implies. More...
class	Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >
	One Gemma 4 decoder block; kGlobal selects the global (full-attention) geometry. More...

Namespaces
namespace	Mila
	Mila main API namespace.

Detailed Description

Gemma 4 decoder block (inference: prefill + decode), one type per attention kind.

Modeled on LlamaBlock, with the Gemma 4 deltas:

• Sandwich norm (4 RMSNorms): input -> attn -> post_attn -> +res -> pre_ffn -> ffn -> post_ffn -> +res.
• QK-norm: per-head RMSNorm over head_dim on Q and K, applied BEFORE RoPE.
• GeGLU FFN (Swiglu<..., Gelu>), decoupled head_dim, non-square o_proj.
• Per-layer geometry via the compile-time kGlobal flag: global layers use global_head_dim / num_global_kv_heads / K=V (no separate v_proj) / full attention / proportional partial-rotary; sliding layers use head_dim / num_kv_heads / window / full rotation.
• V is per-head normalized (v_norm, no learnable scale, no RoPE) on every layer: sliding layers normalize the separate V projection; global K=V layers derive V from the RAW key projection – V = v_norm(k_proj), distinct from K = RoPE(k_norm(k_proj)).
• Attention scale 1.0 (QK-norm controls magnitude; GqaConfig::withAttentionScale).
• kDelegatedFeedForward selects the FFN wiring: false (the default, and what every published model loads) keeps the inline fc_gate_up -> geglu -> fc_down children; true delegates to a GatedMLP child named mlp, which renames the FFN tensors. See Gemma4MoE.md Phase 2.
• kMixtureOfExperts (requires kDelegatedFeedForward) adds the routed branch beside mlp: Router and MixtureOfExperts read the post-attention residual, and the two branches' post-norms sum ahead of post_ffn_norm. See Gemma4MoE.md Phase 1.

Inference-only (Gemma is an inference target): implements ITransformerBlock's prefill/decode; no training forward/backward. Gemma 4 RMSNorm is x_norm * weight (HF Gemma4RMSNorm), not Gemma 3's x_norm * (1 + weight): the converter writes the weights raw and every norm runs at RmsNormConfig unit offset 0.

HF reference forward order (Gemma4TextDecoderLayer, dense): res0 = x a = self_attn( input_layernorm(x) ) [qkv_proj, q_norm/k_norm, RoPE, GQA, o_proj] x = res0 + post_attention_layernorm(a) res1 = x f = mlp( pre_feedforward_layernorm(x) ) [GeGLU] x = ( res1 + post_feedforward_layernorm(f) ) * layer_scalar