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 Dnn.Quantization.KvCache.Policy;
import Serialization.Tensor;
import Dnn.CompositeComponent;
import Compute.GqaState;
import Dnn.Component;
import Dnn.Components.Residual;
import Dnn.ComponentType;
import Dnn.Components.Gqa;
import Compute.Device;
import Dnn.Components.Linear;
import Compute.ExecutionContext;
import Dnn.TensorTypes;
import Compute.DeviceId;
import Dnn.TensorOps;
import Dnn.ITensor;
import Serialization.SafeTensors;
import Compute.ExecutionContextFactory;
import Dnn.Components.Rope;
import Dnn.TensorDataType;
import Dnn.TensorDataTypeTraits;
import Dnn.Tensor;
import Dnn.Components.Swiglu;
import Dnn.Quantization.Weight.Policies;
import Dnn.ActivationType;
import Serialization.ModelArchive;
import Dnn.Components.IDecoderLayer;
import Compute.DeviceType;
import Compute.DeviceTypeTraits;
import Dnn.Components.GemmaConfig;
import Compute.IExecutionContext;
import Compute.CpuMemoryResource;
import Dnn.Components.RmsNorm;
import Serialization.Mode;

Classes
struct	Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::BlockBuildContexts
	The per-child build contexts and split-scratch geometry this block implies. More...
class	Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >
	One Gemma 4 decoder block; kGlobal selects the global (full-attention) geometry. More...
struct	Mila::Dnn::GemmaBlockWorkspace< TDeviceType, TPrecision >
	Transformer-owned shared activation workspace for GemmaBlock (pooling). 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 / a single shared KV head / 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).

Inference-only (Gemma is an inference target): implements IDecoderLayer's prefill/decode; no training forward/backward. Gemma RMSNorm is x_norm * (1 + weight) (HF Gemma3RMSNorm): the converter writes the weights RAW (zero-centered) and the +1 is applied at the kernel via RmsNormConfig::withUnitOffset(1.0) on every norm – so the stored weights stay identical to the source checkpoint and the shared RmsNorm kernel stays Llama-safe (offset 0 = raw).

HF reference forward order (Gemma4TextDecoderLayer): 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)