Gemma 4 E2B text model

gemma/activations.h

@@ -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];
@@ -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.
@@ -360,12 +366,15 @@ 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)),
 
         max_workers(ctx.pools.MaxWorkers()),
         s_ffw_in(config.num_layers, max_workers),
@@ -401,6 +410,9 @@ 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);
+    }
     C1.AllocateAndAttachRowPtrs(row_ptrs);
     C2.AllocateAndAttachRowPtrs(row_ptrs);
     ffw_out.AllocateAndAttachRowPtrs(row_ptrs);
@@ -461,6 +473,9 @@ 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);
+    }
 
     attention_storage.SetBatchSize(batch_size);
     // `AttentionActivationsPtrs` holds `MatPtrT` which also require updating;
@@ -480,6 +495,7 @@ struct Activations {
   MatStorageT<BF16> C1;
   MatStorageT<BF16> C2;
   MatStorageT<float> ffw_out;
+  MatStorageT<float> ple_embeds;
 
   const size_t max_workers;
   TensorStats s_ffw_in;

gemma/attention.cc

@@ -198,6 +198,7 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
   CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w1,
              /*add=*/nullptr, env, activations.q);
 
+  if (flags & kSkipKV) 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.
@@ -294,7 +295,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);
         }
 

gemma/configs.cc

@@ -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;
@@ -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:
@@ -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));
   }
@@ -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));
   }
@@ -784,7 +844,8 @@ 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;
-
+    case 35:
+      return Model::GEMMA4_2B;
     // TODO: detect these.
     /*
     return Model::GEMMA2_772M;

gemma/configs.h

@@ -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
@@ -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) {
@@ -221,6 +223,7 @@ enum class Model {
   GEMMA3_12B_LM,
   GEMMA3_27B_LM,
   GEMMA4_26B_MOE,
+  GEMMA4_2B,
   kSentinel,
 };
 
@@ -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`.
   }
 
@@ -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;
 };
 
@@ -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`.
   }
@@ -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.

gemma/flash_attention.cc

@@ -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,

gemma/gemma-inl.h

@@ -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.