Skip to content

Gemma 4 E2B text model #951

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jul 16, 2026
Merged

Gemma 4 E2B text model #951

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 28 additions & 5 deletions gemma/activations.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,14 @@ static inline size_t MaxQkvDim(const ModelConfig& config) {
}
return max_dim;
}
static inline size_t MaxFFHiddenDim(const ModelConfig& config) {
size_t max_dim = 0;
for (const auto& lc : config.layer_configs) {
max_dim = HWY_MAX(max_dim, static_cast<size_t>(lc.ff_hidden_dim));
}
return max_dim;
}


static inline float ChooseQueryScale(const ModelConfig& config) {
const LayerConfig& layer_config = config.layer_configs[0];
Expand Down Expand Up @@ -115,9 +123,7 @@ struct AttentionActivations {
layer_config.post_qk == PostQKType::HalfRope)),
inv_timescale_global(CreateInvTimescale(
allocator,
config.partial_rotary_factor < 1.0f
? max_qkv_dim
: max_qkv_dim / 4,
max_qkv_dim,
layer_config.post_qk == PostQKType::HalfRope, 1000000.0,
config.partial_rotary_factor)) {
// Batch size can be 0 in experimental code so do not assert.
Expand Down Expand Up @@ -360,12 +366,18 @@ struct Activations {

pre_ffw_rms_out(MatFactory("pre_ffw_rms_out", batch_size,
config.model_dim, ctx.allocator)),
C1(MatFactory("C1", batch_size, layer_config.ff_hidden_dim,
C1(MatFactory("C1", batch_size, MaxFFHiddenDim(config),
ctx.allocator)),
C2(MatFactory("C2", batch_size, layer_config.ff_hidden_dim,
C2(MatFactory("C2", batch_size, MaxFFHiddenDim(config),
ctx.allocator)),
ffw_out(
MatFactory("ffw_out", batch_size, config.model_dim, ctx.allocator)),
ple_embeds(
MatFactory("ple_embeds", batch_size,
config.num_layers * config.ple_dim, ctx.allocator)),
gate_out(
MatFactory("gate_out", batch_size, config.ple_dim, ctx.allocator)),
ple_token_emb(config.num_layers * config.ple_dim),

max_workers(ctx.pools.MaxWorkers()),
s_ffw_in(config.num_layers, max_workers),
Expand Down Expand Up @@ -401,6 +413,10 @@ struct Activations {
x.AllocateAndAttachRowPtrs(row_ptrs);
x_bf.AllocateAndAttachRowPtrs(row_ptrs);
logits.AllocateAndAttachRowPtrs(row_ptrs);
if (config.ple_dim > 0) {
ple_embeds.AllocateAndAttachRowPtrs(row_ptrs);
gate_out.AllocateAndAttachRowPtrs(row_ptrs);
}
C1.AllocateAndAttachRowPtrs(row_ptrs);
C2.AllocateAndAttachRowPtrs(row_ptrs);
ffw_out.AllocateAndAttachRowPtrs(row_ptrs);
Expand Down Expand Up @@ -461,6 +477,10 @@ struct Activations {
C1.OverrideRows(batch_size);
C2.OverrideRows(batch_size);
ffw_out.OverrideRows(batch_size);
if (layer_config.ple_dim > 0) {
ple_embeds.OverrideRows(batch_size);
gate_out.OverrideRows(batch_size);
}

attention_storage.SetBatchSize(batch_size);
// `AttentionActivationsPtrs` holds `MatPtrT` which also require updating;
Expand All @@ -480,6 +500,9 @@ struct Activations {
MatStorageT<BF16> C1;
MatStorageT<BF16> C2;
MatStorageT<float> ffw_out;
MatStorageT<float> ple_embeds;
MatStorageT<BF16> gate_out;
std::vector<float> ple_token_emb;

const size_t max_workers;
TensorStats s_ffw_in;
Expand Down
25 changes: 16 additions & 9 deletions gemma/attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,13 +191,20 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,

const size_t qkv_dim = layer_config.qkv_dim;
const size_t kv_heads = layer_config.kv_heads;
const size_t cache_layer_size = layer_config.CacheLayerSize();

// The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim,
// Resolve KV cache layer index and skip flag
const size_t kv_layer_idx = (layer_config.kv_share_layer_idx >= 0)
? static_cast<size_t>(layer_config.kv_share_layer_idx)
: layer_idx;
const bool skip_kv = (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV);
const size_t cache_layer_size = activations.config.layer_configs[kv_layer_idx].CacheLayerSize();

// The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim,
// model_dim], which we reshaped to (heads + kv_heads * 2) * qkv_dim rows.
CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w1,
/*add=*/nullptr, env, activations.q);

if (skip_kv) return;
// Set up MatMul row pointers for writing to KV, which consists of
// `kv_heads` pairs of (k, v) vectors. This safely handles wraparound
// because rows are computed modulo seq_len.
Expand All @@ -213,8 +220,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
HWY_DASSERT(cache_pos < activations.SeqLen());

const size_t layer_offset = qbatch.KV(qi).cache->layer_flat_offsets.empty()
? layer_idx * cache_layer_size
: qbatch.KV(qi).cache->layer_flat_offsets[layer_idx];
? kv_layer_idx * cache_layer_size
: qbatch.KV(qi).cache->layer_flat_offsets[kv_layer_idx];

env.row_ptrs[0][interleaved_idx] = reinterpret_cast<uint8_t*>(
qbatch.KV(qi).kv_cache.Row(cache_pos) + layer_offset);
Expand Down Expand Up @@ -256,12 +263,12 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
auto& k_cache = qbatch.KV(qi).k_cache;
KV_t* HWY_RESTRICT k =
k_cache.Row(cache_pos / (2 * kFloatsPerVector)) +
qbatch.KV(qi).cache->KOffset(layer_idx, head, kFloatsPerVector,
qbatch.KV(qi).cache->KOffset(kv_layer_idx, head, kFloatsPerVector,
cache_pos);
auto& v_cache = qbatch.KV(qi).v_cache;
KV_t* HWY_RESTRICT v =
v_cache.Row(cache_pos / (2 * kFloatsPerVector)) +
qbatch.KV(qi).cache->VOffset(layer_idx, head, kFloatsPerVector,
qbatch.KV(qi).cache->VOffset(kv_layer_idx, head, kFloatsPerVector,
cache_pos);
if (token_idx >= num_tokens) {
// Create a zero-filled K/V pair for padding for out-of-sequence
Expand All @@ -273,8 +280,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
HWY_DASSERT(cache_pos < activations.SeqLen());
auto& kv_cache = qbatch.KV(qi).kv_cache;
const size_t layer_offset = qbatch.KV(qi).cache->layer_flat_offsets.empty()
? layer_idx * cache_layer_size
: qbatch.KV(qi).cache->layer_flat_offsets[layer_idx];
? kv_layer_idx * cache_layer_size
: qbatch.KV(qi).cache->layer_flat_offsets[kv_layer_idx];
KV_t* HWY_RESTRICT kv = kv_cache.Row(cache_pos) +
layer_offset +
head * qkv_dim * 2;
Expand All @@ -294,7 +301,7 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0, kv_f32,
qkv_dim, env.ctx, worker);
});
} else if (layer_config.post_qk == PostQKType::NormLocalRope) {
} else if (layer_config.post_qk == PostQKType::NormLocalRope || layer_config.use_qk_norm) {
RMSNormNoScaleInplace(kv_f32, qkv_dim, env.ctx, worker);
}

Expand Down
65 changes: 63 additions & 2 deletions gemma/configs.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -470,7 +470,7 @@ static LayerConfig LayerConfigGemma4_26B_MoE_LM(size_t model_dim) {
static ModelConfig ConfigGemma4_26B_MoE() {
ModelConfig config = ConfigBaseGemmaV4();
config.display_name = "Gemma4_26B_MoE";
config.final_cap = 30.0f;
config.final_cap = 0.0f;
config.att_cap = 0.0f;
config.model = Model::GEMMA4_26B_MOE;
config.wrapping = PromptWrapping::GEMMA_IT;
Expand All @@ -495,6 +495,62 @@ static ModelConfig ConfigGemma4_26B_MoE() {
return config;
}

static LayerConfig LayerConfigGemma4_2B_Local(size_t model_dim) {
LayerConfig config;
config.model_dim = model_dim;
config.ff_hidden_dim = 6144;
config.heads = 8;
config.kv_heads = 1;
config.qkv_dim = 256;
config.optimized_gating = true;
config.post_norm = PostNormType::Scale;
config.activation = ActivationType::Gelu;
config.post_qk = PostQKType::NormLocalRope;
config.use_qk_norm = true;
config.norm_v = true;
config.ple_dim = 256;
return config;
}

static LayerConfig LayerConfigGemma4_2B_Global(size_t model_dim) {
LayerConfig config = LayerConfigGemma4_2B_Local(model_dim);
config.qkv_dim = 512;
return config;
}

// Until we have the audio checkpoints included, we use the LM config directly.
static ModelConfig ConfigGemma4_2B() {
ModelConfig config = ConfigBaseGemmaV4();
config.display_name = "Gemma4_2B";
config.model = Model::GEMMA4_2B;
config.wrapping = PromptWrapping::GEMMA_IT;
config.model_dim = 1536;
config.vocab_size = kGemmaV3VocabSize; // 262144
config.max_seq_len = 128 * 1024;
config.final_cap = 0.0f;
config.ple_dim = 256;
config.num_layers = 35;
config.use_global_timescale = true;
config.partial_rotary_factor = 0.25f;
config.query_scale = QueryScaleType::One;
LayerConfig local_config = LayerConfigGemma4_2B_Local(config.model_dim);
config.layer_configs = {config.num_layers, local_config};
// Global attention layers: [4, 9, 14, 19, 24, 29, 34] (stride 5)
for (size_t i = 0; i < config.num_layers; ++i) {
if (i % 5 == 4) {
config.layer_configs[i] = LayerConfigGemma4_2B_Global(config.model_dim);
}
}
// Double-wide MLP for last 20 layers (KV-shared layers 15-34)
for (size_t i = 15; i < config.num_layers; ++i) {
config.layer_configs[i].ff_hidden_dim = 12288;
config.layer_configs[i].kv_share_layer_idx = (i % 5 == 4) ? 14 : 13;
}
config.attention_window_sizes = RepeatedAttentionWindowSizes<35, 5>(
{512, 512, 512, 512, config.max_seq_len});
return config;
}

static ModelConfig ConfigFromModel(Model model) {
switch (model) {
case Model::GEMMA2_2B:
Expand Down Expand Up @@ -529,6 +585,8 @@ static ModelConfig ConfigFromModel(Model model) {
return ConfigGemma3_27B_LM();
case Model::GEMMA4_26B_MOE:
return ConfigGemma4_26B_MoE();
case Model::GEMMA4_2B:
return ConfigGemma4_2B();
default:
HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
}
Expand Down Expand Up @@ -570,6 +628,8 @@ const char* ModelPrefix(Model model) {
return "gemma3-27b-lm";
case Model::GEMMA4_26B_MOE:
return "gemma4-26b-moe";
case Model::GEMMA4_2B:
return "gemma4-2b";
default:
HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
}
Expand Down Expand Up @@ -770,6 +830,8 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
case 34:
return (layer_types & kDeducedViT) ? Model::GEMMA3_4B
: Model::GEMMA3_4B_LM;
case 35:
return Model::GEMMA4_2B;
case 42:
if (layer_types & kDeducedViT) {
return (layer_types & kDeduced448) ? Model::PALIGEMMA2_10B_448
Expand All @@ -784,7 +846,6 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
case 62:
return (layer_types & kDeducedViT) ? Model::GEMMA3_27B
: Model::GEMMA3_27B_LM;

// TODO: detect these.
/*
return Model::GEMMA2_772M;
Expand Down
14 changes: 12 additions & 2 deletions gemma/configs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ namespace gcpp {

constexpr size_t kMaxBF16PerVector = HWY_ARCH_MAX_BYTES / sizeof(BF16);

HWY_INLINE_VAR constexpr int kSkipKV = 1;

HWY_INLINE_VAR constexpr size_t kMaxQKVDim = 1024;

#ifndef GEMMA_FUSED_FFN
Expand Down Expand Up @@ -131,7 +133,7 @@ enum class PostQKType {
Rope,
HalfRope,
NormLocalRope = 8, // Norm without scale, and rope for local attention layers
kSentinel // must be last
kSentinel // must be last
};

static inline bool EnumValid(PostQKType type) {
Expand Down Expand Up @@ -221,6 +223,7 @@ enum class Model {
GEMMA3_12B_LM,
GEMMA3_27B_LM,
GEMMA4_26B_MOE,
GEMMA4_2B,
kSentinel,
};

Expand Down Expand Up @@ -303,6 +306,8 @@ struct LayerConfig : public IFields {
visitor(norm_v);
visitor(num_experts);
visitor(num_experts_per_datapoint);
visitor(ple_dim);
visitor(kv_share_layer_idx);
// Append new fields here, then update `python/configs.cc`.
}

Expand Down Expand Up @@ -338,6 +343,8 @@ struct LayerConfig : public IFields {
bool norm_v = false; // Normalize V projections before caching.
uint32_t num_experts = 0;
uint32_t num_experts_per_datapoint = 0;
uint32_t ple_dim = 0; // Per-Layer Embedding dimension (0 = disabled).
int kv_share_layer_idx = -1;
InternalLayerConfig internal;
};

Expand Down Expand Up @@ -437,6 +444,7 @@ struct ModelConfig : public IFields {

visitor(use_global_timescale);
visitor(partial_rotary_factor);
visitor(ple_dim);

// Append new fields here, then update `python/configs.cc`.
}
Expand Down Expand Up @@ -552,7 +560,9 @@ struct ModelConfig : public IFields {

InternalModelConfig internal;
bool use_global_timescale = false; // for Gemma 3
float partial_rotary_factor = 1.0f; // Fraction of dims with RoPE (0.25 for Gemma4 MoE).
float partial_rotary_factor =
1.0f; // Fraction of dims with RoPE (0.25 for Gemma4 MoE).
uint32_t ple_dim = 0; // Per-Layer Embedding dimension (0 = disabled).
};

// Returns the sub-config for the ViT model of the PaliGemma model.
Expand Down
13 changes: 9 additions & 4 deletions gemma/flash_attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ void RMSNormAndPositionalEncoding(const size_t num_tokens, const QBatch& qbatch,
RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0, q_row,
layer_config.qkv_dim, ctx, worker);
});
} else if (layer_config.post_qk == PostQKType::NormLocalRope) {
} else if (layer_config.post_qk == PostQKType::NormLocalRope || layer_config.use_qk_norm) {
RMSNormNoScaleInplace(q_row, layer_config.qkv_dim, ctx, worker);
}
PositionalEncodingQK(q_row, layer_idx, activations, ctx, worker, pos,
Expand Down Expand Up @@ -2610,6 +2610,11 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism,
const size_t seq_len =
static_cast<size_t>(activations.div_seq_len.GetDivisor());

// Resolve KV cache layer index
const size_t kv_layer_idx = (layer_config.kv_share_layer_idx >= 0)
? static_cast<size_t>(layer_config.kv_share_layer_idx)
: layer_idx;

using DF = hn::ScalableTag<float>;
const DF df;
const size_t kNF = hn::Lanes(df);
Expand All @@ -2618,7 +2623,7 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism,
ParallelFor(
Parallelism::kWithinCluster, activations.q.Rows(), ctx,
/*cluster_idx=*/0, Callers::kFlashAttention,
[&](size_t row, size_t worker) {
[&](size_t row, size_t worker) HWY_ATTR {
CompressPerThread tls;
const hn::ScalableTag<float> df;
CompressTraits<BF16>::Compress(
Expand All @@ -2642,14 +2647,14 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism,
kRoundedQkvDim * 2 * kNF));
kT.SetPtr(
kT_cache.Row(0) + qbatch.KV(param.qi_index)
.cache->KOrVOffset(layer_idx, param.kv_head, kNF),
.cache->KOrVOffset(kv_layer_idx, param.kv_head, kNF),
kT_cache.Stride());
auto& vT_cache = qbatch.KV(param.qi_index).v_cache;
MatPtrT<KV_t> vT("v_T_view", Extents2D(hwy::DivCeil(seq_len, 2 * kNF),
kRoundedQkvDim * 2 * kNF));
vT.SetPtr(
vT_cache.Row(0) + qbatch.KV(param.qi_index)
.cache->KOrVOffset(layer_idx, param.kv_head, kNF),
.cache->KOrVOffset(kv_layer_idx, param.kv_head, kNF),
vT_cache.Stride());
MatPtrT<float>& att_out =
param.i_of_n == 0 ? activations.att_out : activations.att_out_reps;
Expand Down
2 changes: 2 additions & 0 deletions gemma/gemma-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,8 @@ static inline void FFWNoVit(const LayerWeightsPtrs& layer,
Activations& activations, MatMulEnv& env) {
GCPP_ZONE(env.ctx, hwy::Profiler::GlobalIdx(), Zones::kGenFFW);
const LayerConfig& layer_config = layer.layer_config;
activations.C1.OverrideCols(layer_config.ff_hidden_dim);
activations.C2.OverrideCols(layer_config.ff_hidden_dim);

HWY_DASSERT(!layer_config.ff_biases); // Only used in Vit.

Expand Down
Loading
Loading