Gemma.ixx File Reference

Gemma 4 decoder-only transformer network (inference: prefill + decode). More...

#include <string>
#include <vector>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <cstdint>
#include <format>
#include <algorithm>
#include <type_traits>
#include <cmath>
import Serialization.PretrainedReader;
import Dnn.ModelType;
import Dnn.LanguageNetwork;
import Compute.DeviceId;
import Dnn.TensorDataTypeTraits;
import Dnn.ComponentType;
import Dnn.Tensor;
import Compute.DeviceType;
import Serialization.Tensor;
import Dnn.TensorTypes;
import Compute.ExecutionContextFactory;
import Dnn.Components.IDecoderLayer;
import Dnn.Quantization.KvCache.Policy;
import Dnn.Component;
import Dnn.Components.GemmaConfig;
import Dnn.Components.TokenEmbedding;
import Serialization.ModelArchive;
import Dnn.Components.GemmaBlock;
import Dnn.Quantization.Weight.Policies;
import Logging.Logger;
import Compute.DeviceTypeTraits;
import Compute.GqaState;
import Dnn.ITensor;
import Dnn.Components.Linear;
import Serialization.Mode;
import Dnn.Components.RmsNorm;
import Compute.Device;
import Compute.CpuMemoryResource;
import Dnn.TensorDataType;
import Compute.ExecutionContext;
import Serialization.Metadata;

Classes
class	Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >
	Gemma 4 transformer (decoder-only) for autoregressive inference. More...

Namespaces
namespace	Mila
	Mila main API namespace.

Variables
constexpr int64_t	Mila::Dnn::kGemmaFlashPrefillMinContext = 16384
constexpr int64_t	Mila::Dnn::kGemmaPrefillActivationBudgetBytes = int64_t{ 1536 } * 1024 * 1024
constexpr int64_t	Mila::Dnn::kGemmaPrefillChunkOverride = 0

Detailed Description

Gemma 4 decoder-only transformer network (inference: prefill + decode).

Device-templated Gemma 4 autoregressive decoder. Modeled on LlamaTransformer, with the two structural deltas that make Gemma heterogeneous:

• The layer list is NOT homogeneous. Gemma interleaves sliding (local) and full-attention (global) blocks 5:1 over 48 layers (final layer global), and the two are distinct GemmaBlock instantiations (kGlobal false/true) that differ in head_dim / KV-head count / K=V / window / RoPE. The transformer drives them polymorphically through IDecoderLayer (one virtual call per layer per token step, negligible against the per-layer GEMMs). See Gemma.md section 8.
• One shared GQA transient workspace serves both geometries. CudaGqaOp::setState takes only the raw scratch pointer and indexes it with its own HS_, so sizing q_permute / v_out at the MAX head_dim (global 512) lets the local layers (head_dim 256) use a prefix of the same buffer; preatt / att are head_dim- independent ([B, NH, chunk, T], NH shared at 16).

Inference-only (Gemma is an inference target): forward()/backward() are not implemented; the generation loop drives prefill()/decode().

Two Gemma deltas are handled by deliberate design decision:

• Embedding scale (x sqrt(hidden_size)) is applied at runtime in TokenEmbedding::forward via TokenEmbeddingConfig::embedding_scale (set in createGraph). The table is stored raw so it can be shared with the tied lm_head; see WeightTying.md D5. (Superseded the earlier converter-fold decision, BACKLOG Step 5d 2026-06-20, when weight tying landed.)
• Final logit softcap (30 * tanh(logits / 30)) is applied host-side at the sampler: it is strictly monotonic, so it does not change greedy argmax, and GemmaConfig::getFinalLogitSoftcapping() carries the scalar for samplers that need it.