Gemma 4 E2B Vision

BUILD.bazel

@@ -801,6 +801,7 @@ cc_library(
         "gemma/gemma4_moe.cc",
         "gemma/tiled_attention.cc",
         "gemma/vit.cc",
+        "gemma/vit_gemma4.cc",
     ],
     hdrs = [
         "deepseek/deepseek.h",

CMakeLists.txt

@@ -162,6 +162,7 @@ set(SOURCES
   gemma/tokenizer.cc
   gemma/tokenizer.h
   gemma/vit.cc
+  gemma/vit_gemma4.cc
   gemma/vit.h
   gemma/weights.cc
   gemma/weights.h

gemma/configs.cc

@@ -519,10 +519,10 @@ static LayerConfig LayerConfigGemma4_2B_Global(size_t model_dim) {
 }
 
 // Until we have the audio checkpoints included, we use the LM config directly.
-static ModelConfig ConfigGemma4_2B() {
+static ModelConfig ConfigGemma4_2B_LM() {
   ModelConfig config = ConfigBaseGemmaV4();
-  config.display_name = "Gemma4_2B";
-  config.model = Model::GEMMA4_2B;
+  config.display_name = "Gemma4_2B_LM";
+  config.model = Model::GEMMA4_2B_LM;
   config.wrapping = PromptWrapping::GEMMA_IT;
   config.model_dim = 1536;
   config.vocab_size = kGemmaV3VocabSize;  // 262144
@@ -551,6 +551,35 @@ static ModelConfig ConfigGemma4_2B() {
   return config;
 }
 
+static ModelConfig ConfigGemma4_2B() {
+  ModelConfig config = ConfigGemma4_2B_LM();
+  config.display_name = "Gemma4_2B";
+  config.model = Model::GEMMA4_2B;
+  config.wrapping = PromptWrapping::GEMMA_IT;
+  config.use_global_timescale = true;
+
+  config.vit_config.model_dim = 768;
+  config.vit_config.patch_width = 16;
+  config.vit_config.seq_len = 2520;
+  config.vit_config.pool_dim = 3;
+  config.vit_config.image_size = 896;
+
+  LayerConfig vit_layer;
+  vit_layer.model_dim = 768;
+  vit_layer.ff_hidden_dim = 3072;
+  vit_layer.heads = 12;
+  vit_layer.kv_heads = 12;
+  vit_layer.qkv_dim = 64;
+  vit_layer.type = LayerAttentionType::kVitGemma4;
+  vit_layer.use_qk_norm = true;
+  vit_layer.post_norm = PostNormType::Scale;
+
+  config.vit_config.layer_configs = {16, vit_layer};
+  return config;
+}
+
+
+
 static LayerConfig LayerConfigDeepSeek4_Flash(size_t model_dim) {
   LayerConfig config;
   config.model_dim = model_dim;
@@ -704,6 +733,8 @@ static ModelConfig ConfigFromModel(Model model) {
       return ConfigGemma4_2B();
     case Model::DEEPSEEK4_FLASH:
       return ConfigDeepSeek4_Flash();
+    case Model::GEMMA4_2B_LM:
+      return ConfigGemma4_2B_LM();
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -749,6 +780,8 @@ const char* ModelPrefix(Model model) {
       return "gemma4-2b";
     case Model::DEEPSEEK4_FLASH:
       return "deepseek4-flash";
+    case Model::GEMMA4_2B_LM:
+      return "gemma4-2b-lm";
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -780,7 +813,10 @@ ModelConfig::ModelConfig(const Model model, Type weight,
   if (model != Model::UNKNOWN) *this = ConfigFromModel(model);
   HWY_ASSERT(this->model == model);
   this->weight = weight;
-  this->wrapping = wrapping;
+  if (this->wrapping != PromptWrapping::PALIGEMMA &&
+      this->wrapping != PromptWrapping::GEMMA_VLM) {
+    this->wrapping = wrapping;
+  }
 }
 
 static Model FindModel(const std::string& specifier) {
@@ -950,7 +986,8 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
       return (layer_types & kDeducedViT) ? Model::GEMMA3_4B
                                          : Model::GEMMA3_4B_LM;
     case 35:
-      return Model::GEMMA4_2B;
+      return (layer_types & kDeducedViT) ? Model::GEMMA4_2B
+                                         : Model::GEMMA4_2B_LM;
     case 42:
       if (layer_types & kDeducedViT) {
         return (layer_types & kDeduced448) ? Model::PALIGEMMA2_10B_448

gemma/configs.h

@@ -93,11 +93,13 @@ enum class LayerAttentionType {
   // compressed KV latent plus a decoupled RoPE key, optionally with
   // sequence-axis compression (CSA/HCA) selected via `AttentionVariant`.
   kDeepSeekMLA,
+  kVitGemma4,
 };
 
 static inline bool EnumValid(LayerAttentionType type) {
   return type == LayerAttentionType::kGemma ||
          type == LayerAttentionType::kVit ||
+         type == LayerAttentionType::kVitGemma4 ||
          type == LayerAttentionType::kDeepSeekMLA;
 }
 
@@ -273,6 +275,7 @@ enum class Model {
   GEMMA4_26B_MOE,
   GEMMA4_2B,
   DEEPSEEK4_FLASH,
+  GEMMA4_2B_LM,
   kSentinel,
 };
 

gemma/gemma.cc

@@ -22,6 +22,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <optional>
+#include <vector>
 
 #include "compression/types.h"  // GEMMA_DISABLED_TARGETS
 #ifndef HWY_DISABLED_TARGETS
@@ -81,6 +82,9 @@ HWY_BEFORE_NAMESPACE();
 namespace gcpp {
 namespace HWY_NAMESPACE {
 
+using gcpp::BF16;
+using gcpp::MatPtr;
+
 void Attention(LayerAttentionType type, const size_t num_tokens,
                const size_t layer_idx, const LayerWeightsPtrs& layer,
                Activations& activations, QBatch& qbatch, MatMulEnv& env) {
@@ -128,7 +132,6 @@ HWY_NOINLINE void TransformerLayer(const size_t num_tokens,
 
   PostNorm(layer_config.post_norm, layer.post_attention_norm_scale,
            activations.attention.att_sums, env.ctx);
-
   ResidualConnection(activations.attention.att_sums, activations.x, layer,
                      /*is_attention=*/true, env.ctx);
 
@@ -196,7 +199,8 @@ HWY_NOINLINE void TransformerLayer(const size_t num_tokens,
                 });
   }
 
-  if (layer.skip_scale.HasPtr()) {
+  const Model m = activations.attention.config.model;
+  if (layer.skip_scale.HasPtr() && m != Model::GEMMA4_26B_MOE) {
     float skip_scale_val = 1.0f;
     if (layer.skip_scale.GetType() == Type::kF32) {
       skip_scale_val = static_cast<const float*>(layer.skip_scale.Packed())[0];
@@ -240,12 +244,18 @@ HWY_NOINLINE size_t EmbedMMToken(int token, size_t x_row, size_t pos,
   GCPP_ZONE(ctx, hwy::Profiler::GlobalIdx(), Zones::kGenEmbed);
 
   // Image tokens just need to be copied.
-  if (model_config.wrapping == PromptWrapping::GEMMA_VLM &&
-      image_tokens != nullptr && token == -2 &&
-      image_token_position < image_tokens->Rows()) {
-    hwy::CopyBytes(image_tokens->Row(image_token_position), x.Row(x_row),
-                   x.Cols() * x.ElementBytes());
-    return image_token_position + 1;
+  if ((model_config.wrapping == PromptWrapping::GEMMA_VLM ||
+       !model_config.vit_config.layer_configs.empty()) &&
+      image_tokens != nullptr && token == -2) {
+    if (image_token_position < image_tokens->Rows()) {
+      hwy::CopyBytes(image_tokens->Row(image_token_position), x.Row(x_row),
+                     x.Cols() * x.ElementBytes());
+      return image_token_position + 1;
+    } else {
+      HWY_WARN(
+          "EmbedMMToken: token == -2 but image_token_position %zu >= Rows %zu",
+          image_token_position, image_tokens->Rows());
+    }
   }
 
   if (model_config.wrapping == PromptWrapping::PALIGEMMA &&
@@ -294,7 +304,7 @@ static HWY_NOINLINE void ComputePLEEmbeddings(size_t tbatch_size,
   // 1. Convert activations.x (float) to activations.x_bf (BF16)
   for (size_t r = 0; r < tbatch_size; ++r) {
     for (size_t c = 0; c < config.model_dim; ++c) {
-      activations.x_bf.Row(r)[c] = BF16(activations.x.Row(r)[c]);
+      activations.x_bf.Row(r)[c] = gcpp::BF16(activations.x.Row(r)[c]);
     }
   }
 
@@ -329,8 +339,12 @@ static HWY_NOINLINE void ComputePLEEmbeddings(size_t tbatch_size,
   float* token_emb = activations.ple_token_emb.data();
   for (size_t r = 0; r < tbatch_size; ++r) {
     int token = tokens[r];
+    float* out_row = activations.ple_embeds.Row(r);
+    // Use pad_token_id (0) for multimodal placeholder tokens (which are
+    // negative) as done in JAX/PyTorch implementation.
+    const int ple_token = token < 0 ? 0 : token;
     CallUpcasted(&weights.ple_embeddings, [&](const auto* weights_t) HWY_ATTR {
-      const size_t embedding_ofs = token * weights_t->Stride();
+      const size_t embedding_ofs = ple_token * weights_t->Stride();
       const auto embedding_span =
           MakeSpan(weights_t->Row(0), embedding_ofs + ple_total_dim);
       const hn::ScalableTag<float> df;
@@ -340,7 +354,6 @@ static HWY_NOINLINE void ComputePLEEmbeddings(size_t tbatch_size,
       const float token_scale =
           sqrtf(static_cast<float>(config.ple_dim)) * weights_t->Scale();
       const float scaled_token_scale = token_scale * scale_input;
-      float* out_row = activations.ple_embeds.Row(r);
 
       // Vectorized embedding loop (aligned with precomputed scale intent)
       using DF = hn::ScalableTag<float>;
@@ -714,9 +727,8 @@ ChooseSampleFunc(const RuntimeConfig& runtime_config,
   // If user provided a sample_func, use it.
   if (runtime_config.sample_func) return runtime_config.sample_func;
 
-  // Fast path for top-1 with no accept_token.
   if (runtime_config.top_k == 1 && !runtime_config.accept_token) {
-    return [&](size_t /*qi*/, size_t /*pos*/, Logits logits, size_t worker)
+    return [&](size_t qi, size_t pos, Logits logits, size_t worker)
                HWY_ATTR -> TokenAndProb {
                  GCPP_ZONE(ctx, worker, Zones::kGenSampleTop1);
                  return Top1OfSoftmax(logits);
@@ -1002,8 +1014,14 @@ void GenerateImageTokensT(const ModelConfig& config,
     Activations prefill_activations(runtime_config, vit_config, num_tokens,
                                     num_tokens, env.ctx, env.row_ptrs);
     // Weights are for the full PaliGemma model, not just the ViT part.
-    PrefillVit(config, weights, prefill_runtime_config, image, image_tokens,
-               prefill_activations, env);
+    if (config.model == Model::GEMMA4_2B) {
+      PrefillVitGemma4(config, weights, prefill_runtime_config, image,
+                       image_tokens, prefill_activations, env);
+    } else {
+      PrefillVit(config, weights, prefill_runtime_config, image, image_tokens,
+                 prefill_activations, env);
+    }
+
   }  // end GCPP_ZONE before we print results.
 
   // No-op if the profiler is disabled. Printing now ensures that the
@@ -1037,9 +1055,11 @@ Gemma::Gemma(const GemmaArgs& args, ThreadingContext& ctx)
         model_.Config().max_seq_len, args.inference.seq_len);
     model_.MutableConfig().SetMaxSeqLen(args.inference.seq_len);
   }
+
   // Negligible CPU time in the ctor body (except ReadFromBlobs).
   weight_read_mode_ = weights_.ReadFromBlobs(model_, reader_, args.loader,
                                              args.inference, mat_owners_, ctx);
+
   // Read everything into memory, or `weights_.mapped_` keeps the mapping alive.
   reader_.CloseFile();
 }

gemma/model_store.cc

@@ -247,7 +247,8 @@ static int DeduceLayerTypes(const BlobReader& reader) {
   bool has_query_norm = false;
   for (size_t key_idx = 0; key_idx < reader.Keys().size(); ++key_idx) {
     const std::string& key = reader.Keys()[key_idx];
-    if (key.find("qkv_ein_w") != std::string::npos) {  // NOLINT
+    if (key.find("qkv_ein_w") != std::string::npos ||
+        key.find("vit_qkv") != std::string::npos) {  // NOLINT
       layer_types |= kDeducedViT;
     }
     if (key.find("img_pos_emb") != std::string::npos) {  // NOLINT
@@ -290,6 +291,18 @@ static ModelConfig ReadOrDeduceConfig(BlobReader& reader,
           WarnIfExtra(result, kConfigName);
           HWY_ASSERT_M(result.pos != 0, "Error deserializing config");
         }));
+    const PromptWrapping expected_wrapping = ChooseWrapping(config.model);
+    if (expected_wrapping == PromptWrapping::PALIGEMMA ||
+        expected_wrapping == PromptWrapping::GEMMA_VLM) {
+      if (config.wrapping != expected_wrapping) {
+        HWY_WARN("Overriding config.wrapping=%d with expected %d for model %s",
+                 static_cast<int>(config.wrapping),
+                 static_cast<int>(expected_wrapping),
+                 ModelPrefix(config.model));
+        config.wrapping = expected_wrapping;
+      }
+    }
+    // KV sharing is supported, do not force -1.
   }
   // Optionally deduce so we can verify it against the config we read.
   std::optional<Model> deduced_model;

gemma/run.cc

@@ -113,10 +113,15 @@ void ReplGemma(const GemmaArgs& args, const Gemma& gemma, KVCache& kv_cache,
   image_tokens.AllocateAndAttachRowPtrs(env.row_ptrs);
   if (have_image) {
     HWY_ASSERT(config.wrapping == PromptWrapping::PALIGEMMA ||
-               config.wrapping == PromptWrapping::GEMMA_VLM);
+               config.wrapping == PromptWrapping::GEMMA_VLM ||
+               !config.vit_config.layer_configs.empty());
     HWY_ASSERT(image.ReadPPM(inference.image_file.path));
-    const size_t image_size = config.vit_config.image_size;
-    image.Resize(image_size, image_size);
+    // Gemma 4 ViT does its own aspect-ratio-preserving resize, so skip
+    // the forced square resize for Gemma 4 VLM.
+    if (config.model != Model::GEMMA4_2B) {
+      const size_t image_size = config.vit_config.image_size;
+      image.Resize(image_size, image_size);
+    }
     RuntimeConfig runtime_config = {.verbosity = verbosity,
                                     .use_spinning = args.threading.spin};
     gemma.GenerateImageTokens(runtime_config, kv_cache.SeqLen(), image,