diff --git a/BUILD.bazel b/BUILD.bazel
index 982b47cd..3f00d2d9 100644
--- a/BUILD.bazel
+++ b/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",
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 827f0fc4..688c4fb1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -168,6 +168,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
diff --git a/gemma/configs.cc b/gemma/configs.cc
index 997e76f4..6c307d08 100644
--- a/gemma/configs.cc
+++ b/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,33 @@ 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;
@@ -756,6 +783,8 @@ static ModelConfig ConfigFromModel(Model model) {
       return ConfigDeepSeek4_Flash();
     case Model::T5GEMMA_S_S:
       return ConfigT5Gemma_S_S();
+    case Model::GEMMA4_2B_LM:
+      return ConfigGemma4_2B_LM();
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -803,6 +832,8 @@ const char* ModelPrefix(Model model) {
       return "deepseek4-flash";
     case Model::T5GEMMA_S_S:
       return "t5gemma-s-s";
+    case Model::GEMMA4_2B_LM:
+      return "gemma4-2b-lm";
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -834,7 +865,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) {
@@ -1011,7 +1045,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
diff --git a/gemma/configs.h b/gemma/configs.h
index 26096291..df13e23d 100644
--- a/gemma/configs.h
+++ b/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;
 }
 
@@ -274,6 +276,7 @@ enum class Model {
   DEEPSEEK4_FLASH,
   // T5Gemma family - starting with S/S.
   T5GEMMA_S_S,
+  GEMMA4_2B_LM,
   kSentinel,
 };
 
diff --git a/gemma/gemma.cc b/gemma/gemma.cc
index fb736789..5b9823cf 100644
--- a/gemma/gemma.cc
+++ b/gemma/gemma.cc
@@ -23,6 +23,7 @@
 #include <cstdint>
 #include <limits>
 #include <optional>
+#include <vector>
 
 #include "compression/types.h"  // GEMMA_DISABLED_TARGETS
 #ifndef HWY_DISABLED_TARGETS
@@ -84,6 +85,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) {
@@ -131,7 +135,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);
 
@@ -199,7 +202,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];
@@ -1019,12 +1023,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 &&
@@ -1054,7 +1064,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]);
     }
   }
 
@@ -1089,8 +1099,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;
@@ -1100,7 +1114,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>;
@@ -1474,9 +1487,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);
@@ -1796,8 +1808,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
@@ -1831,9 +1849,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();
 }
diff --git a/gemma/model_store.cc b/gemma/model_store.cc
index b03f413b..40ff101c 100644
--- a/gemma/model_store.cc
+++ b/gemma/model_store.cc
@@ -249,7 +249,8 @@ static int DeduceLayerTypes(const BlobReader& reader) {
   bool has_t5gemma_decoder = 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
@@ -306,6 +307,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;
diff --git a/gemma/run.cc b/gemma/run.cc
index 5a3ecd8d..63eeafa5 100644
--- a/gemma/run.cc
+++ b/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,
diff --git a/gemma/tensor_info.cc b/gemma/tensor_info.cc
index f444e418..62e3fae5 100644
--- a/gemma/tensor_info.cc
+++ b/gemma/tensor_info.cc
@@ -44,6 +44,21 @@ void TensorInfoRegistry::Add(const std::string& suffix,
 // Non-layer tensors.
 void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) {
   const std::string no_suffix;
+  size_t pos_emb_rows = config.vit_config.seq_len;
+  std::vector<size_t> img_emb_axes = {3, 0, 1, 2};
+  std::vector<size_t> img_emb_shape = {config.vit_config.model_dim,
+                                       config.vit_config.patch_width,
+                                       config.vit_config.patch_width, 3};
+  bool cols_take_extra_dims = true;
+
+  if (!config.vit_config.layer_configs.empty() &&
+      config.vit_config.layer_configs[0].type == LayerAttentionType::kVitGemma4) {
+    pos_emb_rows = 20480;
+    img_emb_axes = {1, 0};
+    img_emb_shape = {config.vit_config.model_dim,
+                     3 * config.vit_config.patch_width * config.vit_config.patch_width};
+    cols_take_extra_dims = false;
+  }
   Add(no_suffix, {
                      .base_name = "c_embedding",
                      .source_names = {"embedder/input_embedding"},
@@ -83,12 +98,12 @@ void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) {
   Add(no_suffix,
       {
           .base_name = "img_emb_kernel",
-          .source_names = {"img/embedding/kernel"},
-          .axes = {3, 0, 1, 2},
-          .shape = {config.vit_config.model_dim, config.vit_config.patch_width,
-                    config.vit_config.patch_width, 3},
+          .source_names = {"img/embedding/kernel",
+                           "vit/entry/input_projection/w"},
+          .axes = img_emb_axes,
+          .shape = img_emb_shape,
           .min_size = Type::kBF16,
-          .cols_take_extra_dims = true,
+          .cols_take_extra_dims = cols_take_extra_dims,
       });
   Add(no_suffix,
       {
@@ -108,9 +123,10 @@ void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) {
       });
   Add(no_suffix, {
                      .base_name = "img_pos_emb",
-                     .source_names = {"img/pos_embedding"},
+                     .source_names = {"img/pos_embedding",
+                                      "vit/entry/pos_emb"},
                      .axes = {0, 1},
-                     .shape = {/*1,*/ config.vit_config.seq_len,
+                     .shape = {/*1,*/ pos_emb_rows,
                                config.vit_config.model_dim},
                      .min_size = Type::kF32,
                  });
@@ -195,6 +211,103 @@ void TensorInfoRegistry::AddImageLayerTensors(const ModelConfig& config,
                                               const size_t img_layer_idx) {
   const std::string suffix = LayerSuffix(img_layer_idx);
 
+  if (layer_config.type == LayerAttentionType::kVitGemma4) {
+    Add(suffix, {
+        .base_name = "vit_qkv1_w",
+        .source_names = {"vit/transformer/stacked_layers/block/attn/q_einsum/w"},
+        .axes = {0, 2, 1},
+        .shape = {layer_config.heads * layer_config.qkv_dim, config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_qkv2_w",
+        .source_names = {"vit/transformer/stacked_layers/block/attn/kv_einsum/w"},
+        .axes = {1, 0, 3, 2},
+        .shape = {2 * layer_config.kv_heads * layer_config.qkv_dim, config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_gating1_w",
+        .source_names = {"vit/transformer/stacked_layers/block/mlp/gating_einsum1/w"},
+        .axes = {0, 2, 1},
+        .shape = {layer_config.ff_hidden_dim, config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_gating2_w",
+        .source_names = {"vit/transformer/stacked_layers/block/mlp/gating_einsum2/w"},
+        .axes = {0, 2, 1},
+        .shape = {layer_config.ff_hidden_dim, config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_att_ein",
+        .source_names = {"vit/transformer/stacked_layers/block/attn/attn_vec_einsum/w"},
+        .axes = {2, 0, 1},
+        .shape = {config.vit_config.model_dim, layer_config.heads,
+                  layer_config.qkv_dim},
+        .min_size = Type::kBF16,
+        .cols_take_extra_dims = true,
+    });
+    Add(suffix, {
+        .base_name = "vit_pr_att",
+        .source_names = {"vit/transformer/stacked_layers/block/pre_attention_norm/scale"},
+        .axes = {0},
+        .shape = {config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_po_att",
+        .source_names = {"vit/transformer/stacked_layers/block/post_attention_norm/scale"},
+        .axes = {0},
+        .shape = {config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_pr_ff",
+        .source_names = {"vit/transformer/stacked_layers/block/pre_ffw_norm/scale"},
+        .axes = {0},
+        .shape = {config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_po_ff",
+        .source_names = {"vit/transformer/stacked_layers/block/post_ffw_norm/scale"},
+        .axes = {0},
+        .shape = {config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_gate_ein",
+        .source_names = {"vit/transformer/stacked_layers/block/mlp/gating_einsum/w"},
+        .axes = {0, 1, 2},
+        .shape = {2, layer_config.ff_hidden_dim, config.vit_config.model_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_linear_w",
+        .source_names = {"vit/transformer/stacked_layers/block/mlp/linear/w"},
+        .axes = {1, 0},
+        .shape = {config.vit_config.model_dim, layer_config.ff_hidden_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_q_norm",
+        .source_names = {"vit/transformer/stacked_layers/block/attn/query_norm/scale"},
+        .axes = {0},
+        .shape = {layer_config.qkv_dim},
+        .min_size = Type::kBF16,
+    });
+    Add(suffix, {
+        .base_name = "vit_key_norm",
+        .source_names = {"vit/transformer/stacked_layers/block/attn/key_norm/scale"},
+        .axes = {0},
+        .shape = {layer_config.qkv_dim},
+        .min_size = Type::kBF16,
+    });
+    return;
+  }
+
   // Vit layers.
   Add(suffix, {
                   .base_name = "attn_out_w",
@@ -1017,17 +1130,25 @@ TensorInfoRegistry::TensorInfoRegistry(const ModelConfig& config) {
 
 TensorInfo TensorInfoRegistry::TensorInfoFromSourcePath(const std::string& path,
                                                         int layer_idx) const {
+  const TensorInfo* best_tensor = nullptr;
+  size_t longest_match_len = 0;
   for (const TensorInfo& tensor : tensors_) {
     for (const std::string& source_name : tensor.source_names) {
       // path ends with source_name?
       const size_t pos = path.rfind(source_name);
       if (pos != std::string::npos && path.size() == pos + source_name.size()) {
-        std::string name = tensor.base_name;
-        if (layer_idx >= 0) name += LayerSuffix(static_cast<size_t>(layer_idx));
-        return TensorInfoFromName(name);
+        if (source_name.size() > longest_match_len) {
+          best_tensor = &tensor;
+          longest_match_len = source_name.size();
+        }
       }
     }
   }
+  if (best_tensor != nullptr) {
+    std::string name = best_tensor->base_name;
+    if (layer_idx >= 0) name += LayerSuffix(static_cast<size_t>(layer_idx));
+    return TensorInfoFromName(name);
+  }
   return TensorInfo();
 }
 
diff --git a/gemma/tokenizer.cc b/gemma/tokenizer.cc
index a6b7b273..fc27fc86 100644
--- a/gemma/tokenizer.cc
+++ b/gemma/tokenizer.cc
@@ -71,7 +71,7 @@ GemmaChatTemplate::GemmaChatTemplate(const GemmaTokenizer& tokenizer,
     sot_user_ = {105, 2364, 107};
     sot_model_ = {105, 4368, 107, 100, 45518, 107, 101};
     eot_ = {106, 107};
-  } else if (model == Model::GEMMA4_2B) {
+  } else if (model == Model::GEMMA4_2B || model == Model::GEMMA4_2B_LM) {
     sot_user_ = {105, 2364, 107};
     sot_model_ = {105, 4368, 107};
     eot_ = {106, 107};
@@ -85,10 +85,21 @@ GemmaChatTemplate::GemmaChatTemplate(const GemmaTokenizer& tokenizer,
   }
 
   HWY_ASSERT(tokenizer.Encode("\n", &pali_sep_));
-  vlm_soi_.reserve(2);
-  HWY_ASSERT(tokenizer.Encode("\n\n<start_of_image>", &vlm_soi_));
-  vlm_eoi_.reserve(2);
-  HWY_ASSERT(tokenizer.Encode("<end_of_image>\n\n", &vlm_eoi_));
+
+  // Gemma 4 uses different image boundary tokens than Gemma 3.
+  // Gemma 4: <|image> (255999) ... <image|> (258882)
+  // Gemma 3: \n\n<start_of_image> ... <end_of_image>\n\n
+  if (model == Model::GEMMA4_2B) {
+    // HF format: \n\n<|image><|image|>...<image|>\n\n
+    // Token 108 = \n in Gemma 4 tokenizer.
+    vlm_soi_ = {108, 108, 255999};  // \n\n<|image>
+    vlm_eoi_ = {258882, 108, 108};  // <image|>\n\n
+  } else {
+    vlm_soi_.reserve(2);
+    HWY_ASSERT(tokenizer.Encode("\n\n<start_of_image>", &vlm_soi_));
+    vlm_eoi_.reserve(2);
+    HWY_ASSERT(tokenizer.Encode("<end_of_image>\n\n", &vlm_eoi_));
+  }
 }
 
 std::vector<int> GemmaChatTemplate::Apply(size_t pos,
@@ -172,6 +183,7 @@ std::vector<int> WrapAndTokenize(const GemmaTokenizer& tokenizer,
     case PromptWrapping::PALIGEMMA:
       HWY_ASSERT(pos == 0);
       return chat_template.WrapPali(text_part, image_batch_size);
+    case PromptWrapping::GEMMA_IT:  // Gemma 4 VLM uses IT wrapping
     case PromptWrapping::GEMMA_VLM:
       return chat_template.Apply(
           pos, chat_template.WrapVLM(text_part, image_batch_size));
diff --git a/gemma/vit.h b/gemma/vit.h
index d6562f6c..6aa72b9f 100644
--- a/gemma/vit.h
+++ b/gemma/vit.h
@@ -35,6 +35,11 @@ namespace gcpp {
                   const RuntimeConfig& runtime_config, const Image& image,     \
                   ImageTokens& image_tokens, Activations& activations,         \
                   MatMulEnv& env);                                             \
+                                                                               \
+  void PrefillVitGemma4(const ModelConfig& model_config, const WeightsPtrs& weights, \
+                        const RuntimeConfig& runtime_config, const Image& image,     \
+                        ImageTokens& image_tokens, Activations& activations,         \
+                        MatMulEnv& env);                                             \
   /* NOLINTNEXTLINE(google-readability-namespace-comments) */                  \
   }  // namespace NAMESPACE
 
diff --git a/gemma/vit_gemma4.cc b/gemma/vit_gemma4.cc
new file mode 100644
index 00000000..0967e597
--- /dev/null
+++ b/gemma/vit_gemma4.cc
@@ -0,0 +1,654 @@
+// Copyright 2026 Google LLC
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <vector>
+
+#include "compression/types.h"  // GEMMA_DISABLED_TARGETS
+#include "gemma/flash_structs.h"
+#ifndef HWY_DISABLED_TARGETS
+#define HWY_DISABLED_TARGETS GEMMA_DISABLED_TARGETS
+#endif  // HWY_DISABLED_TARGETS
+
+#include "gemma/activations.h"
+#include "gemma/gemma.h"
+#include "gemma/gemma_args.h"
+#include "gemma/weights.h"
+#include "paligemma/image.h"
+#include "hwy/contrib/thread_pool/thread_pool.h"
+#include "hwy/profiler.h"
+
+// Compiles this file for multiple architectures via "foreach_target.h", to
+// which we pass the filename via macro 'argument'.
+// clang-format off
+#undef HWY_TARGET_INCLUDE
+#define HWY_TARGET_INCLUDE "gemma/vit_gemma4.cc" // NOLINT
+// clang-format on
+#include "hwy/foreach_target.h"  // IWYU pragma: keep
+#include "hwy/highway.h"
+// After highway.h
+#include "gemma/attention.h"
+#include "gemma/flash_attention.h"
+#include "gemma/gemma-inl.h"
+#include "ops/ops-inl.h"
+
+HWY_BEFORE_NAMESPACE();
+namespace gcpp {
+namespace HWY_NAMESPACE {
+
+void GetAspectRatioPreservingSize(int height, int width, int patch_size,
+                                  int max_patches, int pooling_kernel_size,
+                                  int& target_height, int& target_width) {
+  float total_px = height * width;
+  float target_px = max_patches * (patch_size * patch_size);
+  float factor = std::sqrt(target_px / total_px);
+  float ideal_height = factor * height;
+  float ideal_width = factor * width;
+  int side_mult = pooling_kernel_size * patch_size;  // 3 * 16 = 48
+
+  target_height =
+      static_cast<int>(std::floor(ideal_height / side_mult)) * side_mult;
+  target_width =
+      static_cast<int>(std::floor(ideal_width / side_mult)) * side_mult;
+
+  if (target_height == 0 && target_width == 0) {
+    HWY_ABORT("Attempting to resize to a 0 x 0 image.");
+  }
+
+  int max_side_length =
+      (max_patches / (pooling_kernel_size * pooling_kernel_size)) * side_mult;
+  if (target_height == 0) {
+    target_height = side_mult;
+    target_width = std::min(
+        static_cast<int>(std::floor(static_cast<float>(width) / height)) *
+            side_mult,
+        max_side_length);
+  } else if (target_width == 0) {
+    target_width = side_mult;
+    target_height = std::min(
+        static_cast<int>(std::floor(static_cast<float>(height) / width)) *
+            side_mult,
+        max_side_length);
+  }
+
+  if (target_height * target_width > target_px) {
+    HWY_ABORT("Resized image exceeds max patches.");
+  }
+}
+
+static HWY_NOINLINE void EmbedImagePatchesGemma4(
+    const Image& image, const ModelConfig& model_config,
+    const WeightsPtrs& weights, Activations& activations, MatMulEnv& env) {
+  const size_t model_dim = model_config.vit_config.model_dim;      // 768
+  const size_t patch_width = model_config.vit_config.patch_width;  // 16
+  const size_t max_patches = model_config.vit_config.seq_len;      // 2520
+  const size_t patch_area = patch_width * patch_width * 3;         // 768
+  const hwy::Divisor div_patch_dim(patch_width);
+
+  const size_t num_real_patches =
+      (image.height() / patch_width) * (image.width() / patch_width);
+  HWY_ASSERT(num_real_patches <= max_patches);
+
+  MatStorageT<float> image_patches("patches",
+                                   Extents2D(max_patches, patch_area),
+                                   env.ctx.allocator, MatPadding::kOdd);
+  gcpp::ZeroInit(image_patches);
+
+  for (size_t i = 0; i < num_real_patches; ++i) {
+    float* row = image_patches.Row(i);
+    image.GetPatch(i, div_patch_dim, row);
+  }
+
+  CallMatMul(image_patches, weights.vit_img_embedding_kernel,
+             weights.vit_img_embedding_bias.PackedScale1(), env, activations.x);
+
+  const size_t real_patch_width = image.width() / patch_width;
+
+  CallUpcastedActivation(
+      &weights.vit_img_pos_embedding, [&](const auto* pos_emb_t) {
+        for (size_t i = 0; i < num_real_patches; ++i) {
+          size_t y = i / real_patch_width;
+          size_t x = i % real_patch_width;
+
+          HWY_DASSERT(x < 10240);
+          HWY_DASSERT(y < 10240);
+
+          const auto* x_emb = pos_emb_t->Row(2 * x);
+          const auto* y_emb = pos_emb_t->Row(2 * y + 1);
+          float* act_row = activations.x.Row(i);
+
+          for (size_t d = 0; d < model_dim; ++d) {
+            act_row[d] +=
+                static_cast<float>(x_emb[d]) + static_cast<float>(y_emb[d]);
+          }
+        }
+      });
+}
+
+class VitGemma4Attention {
+ public:
+  VitGemma4Attention(size_t num_tokens, size_t num_real_patches,
+                     size_t layer_idx, Activations& activations,
+                     const LayerWeightsPtrs& layer, MatMulEnv& env,
+                     size_t real_patch_width)
+      : num_tokens_(num_tokens),
+        num_real_patches_(num_real_patches),
+        activations_(activations),
+        layer_(layer),
+        layer_config_(layer.layer_config),
+        env_(env),
+        real_patch_width_(real_patch_width),
+        pool_(env_.ctx.pools.Pool(0)),
+        caller1_(env_.ctx.pool_callers.Get(Callers::kVitDotSoftmax1)),
+        caller2_(env_.ctx.pool_callers.Get(Callers::kVitDotSoftmax2)),
+        caller3_(env_.ctx.pool_callers.Get(Callers::kVitDotSoftmax3)),
+        caller4_(env_.ctx.pool_callers.Get(Callers::kVitDotSoftmax4)) {}
+
+  HWY_INLINE void operator()() {
+    ComputeQKV();
+    ApplyQKVNorm();
+    FlashAttention();
+    SumHeads();
+  }
+
+ private:
+  HWY_NOINLINE void ApplyQKVNorm() {
+    PROFILER_ZONE("Gen.VitGemma4Attention.QKVNorm");
+    const size_t heads = layer_config_.heads;
+    const size_t kv_heads = layer_config_.kv_heads;
+    const size_t qkv_dim = layer_config_.qkv_dim;
+    auto& qkv = activations_.attention.q;
+
+    float inv_timescale[16];
+    for (size_t dim = 0; dim < 16; ++dim) {
+      inv_timescale[dim] =
+          1.0f / std::pow(100.0f, static_cast<float>(2 * dim) / 32.0f);
+    }
+
+    // Q-norm
+    if (layer_.query_norm_scale.HasPtr()) {
+      CallUpcasted(&layer_.query_norm_scale, [&](const auto* weights_t) {
+        ParallelFor(Parallelism::kWithinCluster, num_real_patches_, env_.ctx,
+                    /*cluster_idx=*/0, Callers::kFlashAttention,
+                    [&](size_t token, size_t worker) {
+                      const size_t y = token / real_patch_width_;
+                      const size_t x = token % real_patch_width_;
+                      for (size_t h = 0; h < heads; ++h) {
+                        float* q_ptr = qkv.Row(token) + h * qkv_dim;
+                        RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0,
+                                       q_ptr, qkv_dim, env_.ctx, worker);
+                        Rope(q_ptr, 32, inv_timescale, static_cast<int>(x),
+                             env_.ctx, worker);
+                        Rope(q_ptr + 32, 32, inv_timescale, static_cast<int>(y),
+                             env_.ctx, worker);
+                      }
+                    });
+      });
+    } else if (layer_config_.use_qk_norm) {
+      ParallelFor(Parallelism::kWithinCluster, num_real_patches_, env_.ctx,
+                  /*cluster_idx=*/0, Callers::kFlashAttention,
+                  [&](size_t token, size_t worker) {
+                    const size_t y = token / real_patch_width_;
+                    const size_t x = token % real_patch_width_;
+                    for (size_t h = 0; h < heads; ++h) {
+                      float* q_ptr = qkv.Row(token) + h * qkv_dim;
+                      RMSNormNoScaleInplace(q_ptr, qkv_dim, env_.ctx, worker);
+                      Rope(q_ptr, 32, inv_timescale, static_cast<int>(x),
+                           env_.ctx, worker);
+                      Rope(q_ptr + 32, 32, inv_timescale, static_cast<int>(y),
+                           env_.ctx, worker);
+                    }
+                  });
+    }
+
+    // K-norm
+    if (layer_.key_norm_scale.HasPtr()) {
+      CallUpcasted(&layer_.key_norm_scale, [&](const auto* weights_t) {
+        ParallelFor(Parallelism::kWithinCluster, num_real_patches_, env_.ctx,
+                    /*cluster_idx=*/0, Callers::kFlashAttention,
+                    [&](size_t token, size_t worker) {
+                      const size_t y = token / real_patch_width_;
+                      const size_t x = token % real_patch_width_;
+                      for (size_t h = 0; h < kv_heads; ++h) {
+                        float* k_ptr =
+                            qkv.Row(token) + heads * qkv_dim + 2 * h * qkv_dim;
+                        RMSNormInplace(weights_t->PackedScale1(), /*w_ofs=*/0,
+                                       k_ptr, qkv_dim, env_.ctx, worker);
+                        Rope(k_ptr, 32, inv_timescale, static_cast<int>(x),
+                             env_.ctx, worker);
+                        Rope(k_ptr + 32, 32, inv_timescale, static_cast<int>(y),
+                             env_.ctx, worker);
+                      }
+                    });
+      });
+    } else if (layer_config_.use_qk_norm) {
+      ParallelFor(Parallelism::kWithinCluster, num_real_patches_, env_.ctx,
+                  /*cluster_idx=*/0, Callers::kFlashAttention,
+                  [&](size_t token, size_t worker) {
+                    const size_t y = token / real_patch_width_;
+                    const size_t x = token % real_patch_width_;
+                    for (size_t h = 0; h < kv_heads; ++h) {
+                      float* k_ptr =
+                          qkv.Row(token) + heads * qkv_dim + 2 * h * qkv_dim;
+                      RMSNormNoScaleInplace(k_ptr, qkv_dim, env_.ctx, worker);
+                      Rope(k_ptr, 32, inv_timescale, static_cast<int>(x),
+                           env_.ctx, worker);
+                      Rope(k_ptr + 32, 32, inv_timescale, static_cast<int>(y),
+                           env_.ctx, worker);
+                    }
+                  });
+    }
+
+    ParallelFor(Parallelism::kWithinCluster, num_real_patches_, env_.ctx,
+                /*cluster_idx=*/0, Callers::kFlashAttention,
+                [&](size_t token, size_t worker) {
+                  for (size_t h = 0; h < kv_heads; ++h) {
+                    float* v_ptr = qkv.Row(token) + heads * qkv_dim +
+                                   (2 * h + 1) * qkv_dim;
+                    RMSNormNoScaleInplace(v_ptr, qkv_dim, env_.ctx, worker);
+                  }
+                });
+  }
+
+  HWY_NOINLINE void ComputeQKV() {
+    PROFILER_ZONE("Gen.VitGemma4Attention.QKV");
+    auto& qkv = activations_.attention.q;
+    const size_t heads = layer_config_.heads;
+    const size_t kv_heads = layer_config_.kv_heads;
+    const size_t qkv_dim = layer_config_.qkv_dim;
+
+    HWY_ASSERT(qkv.Rows() == num_tokens_);
+    HWY_ASSERT(qkv.Cols() == (heads + 2 * kv_heads) * qkv_dim);
+
+    // Q part:
+    MatPtrT<float> q_part("q_part", Extents2D(num_tokens_, heads * qkv_dim));
+    q_part.SetPtr(qkv.Row(0), qkv.Stride());
+    CallMatMul(activations_.attention.pre_att_rms_out, layer_.qkv_einsum_w1,
+               /*bias=*/nullptr, env_, q_part);
+
+    // KV part:
+    MatPtrT<float> kv_part("kv_part",
+                           Extents2D(num_tokens_, 2 * kv_heads * qkv_dim));
+    kv_part.SetPtr(qkv.Row(0) + heads * qkv_dim, qkv.Stride());
+    CallMatMul(activations_.attention.pre_att_rms_out, layer_.qkv_einsum_w2,
+               /*bias=*/nullptr, env_, kv_part);
+  }
+
+  // Applies the query scale to the query and converts to QType.
+  template <typename QKVType, typename QType>
+  void ScaleQuery(const MatPtrT<QKVType>& qkv, const size_t num_tokens,
+                  const size_t heads, const size_t qkv_dim,
+                  const float query_scale, MatPtrT<QType>& q_output) {
+    ParallelFor(
+        Parallelism::kWithinCluster, heads, env_.ctx,
+        /*cluster_idx=*/0, Callers::kFlashAttention,
+        [&](size_t head, size_t worker) {
+          size_t src_q_offset = head * qkv_dim;
+          size_t dst_q_offset = head * qkv_dim;
+          for (size_t token = 0; token < num_tokens; ++token) {
+            const float* HWY_RESTRICT src_q = qkv.Row(token) + src_q_offset;
+            QType* HWY_RESTRICT dst_q = q_output.Row(token) + dst_q_offset;
+            if (token < num_real_patches_) {
+              for (size_t i = 0; i < qkv_dim; ++i) {
+                dst_q[i] = hwy::ConvertScalarTo<QType>(
+                    hwy::ConvertScalarTo<float>(src_q[i]) * query_scale);
+              }
+            } else {
+              for (size_t i = 0; i < qkv_dim; ++i) {
+                dst_q[i] = hwy::ConvertScalarTo<QType>(0.0f);
+              }
+            }
+          }
+        });
+  }
+
+  // Transposes K and V and converts to KVType.
+  template <typename QKVType, typename KVType>
+  void TransposeKAndV(const MatPtrT<QKVType>& qkv, const size_t num_tokens,
+                      const size_t heads, const size_t qkv_dim,
+                      MatPtrT<KVType>& k_output, MatPtrT<KVType>& v_output) {
+    using DF = hn::ScalableTag<float>;
+    const DF df;
+    const size_t kNF = hn::Lanes(df);
+    const size_t kNumTokensH = hwy::DivCeil(num_tokens, 2 * kNF);
+    const size_t kRoundedKVDim = hwy::RoundUpTo(qkv_dim, 2 * kNF);
+    ParallelFor(
+        Parallelism::kWithinCluster, heads, env_.ctx,
+        /*cluster_idx=*/0, Callers::kFlashAttention,
+        [&](size_t head, size_t worker) {
+          const size_t k_offset = heads * qkv_dim + 2 * head * qkv_dim;
+          const size_t v_offset = heads * qkv_dim + (2 * head + 1) * qkv_dim;
+
+          const size_t k_or_v_output_offset = head * 2 * kNF * kRoundedKVDim;
+          for (size_t token_h = 0; token_h < kNumTokensH; ++token_h) {
+            KVType* HWY_RESTRICT dst_k = k_output.Row(token_h);
+            KVType* HWY_RESTRICT dst_v = v_output.Row(token_h);
+            size_t dst_k_index = k_or_v_output_offset;
+            for (size_t q = 0; q < qkv_dim; q += 2) {
+              for (size_t token_l = 0; token_l < 2 * kNF;
+                   ++token_l, dst_k_index += 2) {
+                const size_t token_idx = token_h * 2 * kNF + token_l;
+                if (token_idx < num_real_patches_) {
+                  const QKVType* HWY_RESTRICT src_k =
+                      qkv.Row(token_idx) + k_offset;
+                  dst_k[dst_k_index] = hwy::ConvertScalarTo<KVType>(src_k[q]);
+                  dst_k[dst_k_index + 1] =
+                      hwy::ConvertScalarTo<KVType>(src_k[q + 1]);
+                } else {
+                  dst_k[dst_k_index] = hwy::ConvertScalarTo<KVType>(0.0f);
+                  dst_k[dst_k_index + 1] = hwy::ConvertScalarTo<KVType>(0.0f);
+                }
+              }
+            }
+            size_t dst_v_index = k_or_v_output_offset;
+            for (size_t q = 0; q < qkv_dim; q += 2 * kNF) {
+              for (size_t token_l = 0; token_l < 2 * kNF; ++token_l) {
+                const size_t token_idx = token_h * 2 * kNF + token_l;
+                if (token_idx < num_real_patches_) {
+                  const QKVType* HWY_RESTRICT src_v =
+                      qkv.Row(token_idx) + v_offset;
+                  if (q + 2 * kNF <= qkv_dim) {
+                    for (size_t q_l = 0; q_l < 2 * kNF; ++q_l) {
+                      dst_v[dst_v_index++] =
+                          hwy::ConvertScalarTo<KVType>(src_v[q + q_l]);
+                    }
+                  } else {
+                    for (size_t q_l = 0; q_l < qkv_dim - q; ++q_l) {
+                      dst_v[dst_v_index++] =
+                          hwy::ConvertScalarTo<KVType>(src_v[q + q_l]);
+                    }
+                  }
+                } else {
+                  const size_t cnt =
+                      (q + 2 * kNF <= qkv_dim) ? (2 * kNF) : (qkv_dim - q);
+                  for (size_t q_l = 0; q_l < cnt; ++q_l) {
+                    dst_v[dst_v_index++] = hwy::ConvertScalarTo<KVType>(0.0f);
+                  }
+                }
+              }
+            }
+            // Zero out the padding area.
+            // In the loops above, the dst_k loop has written 2kNF x 2
+            // consecutive elements for each q +=2, and the dst_v loop has
+            // written 2kNF x 2kNF consecutive elements for each q += 2 * kNF.
+            // Both of them therefore write 2kNF elements for each increment of
+            // q, so we can combine both into a single loop for the padding.
+            // This could be further simplified by writing a zero vector.
+            for (size_t q = qkv_dim; q < kRoundedKVDim; ++q) {
+              for (size_t token_l = 0; token_l < 2 * kNF; ++token_l) {
+                dst_k[dst_k_index++] = hwy::ConvertScalarTo<KVType>(0.0f);
+                dst_v[dst_v_index++] = hwy::ConvertScalarTo<KVType>(0.0f);
+              }
+            }
+          }
+        });
+  }
+
+  // Computes the flash attention parameters. This is mostly about deciding on
+  // the tile sizes and filling the param structs with the correct offsets.
+  template <typename QType, typename KVType>
+  void ComputeParams(const uint32_t num_tokens, const size_t seq_len,
+                     const size_t heads, const uint32_t qkv_dim,
+                     const MatPtrT<QType>& q, const MatPtrT<KVType>& k,
+                     const MatPtrT<KVType>& v, const MatPtrT<float>& att_out,
+                     std::vector<Tile148Params>& flash_params) {
+    flash_params.clear();
+    for (uint32_t head = 0; head < heads; ++head) {
+      uint32_t token = 0;
+      while (token + k8xNFVTileSize <= num_tokens) {
+        flash_params.push_back(Tile148Params{
+            .v_tile_size = k8xNFVTileSize,
+            .qi_index = token,
+            .kv_head = head,
+        });
+        token += k8xNFVTileSize;
+      }
+      if (token + k4xNFVTileSize <= num_tokens) {
+        flash_params.push_back(Tile148Params{
+            .v_tile_size = k4xNFVTileSize,
+            .qi_index = token,
+            .kv_head = head,
+        });
+        token += k4xNFVTileSize;
+      }
+      while (token < num_tokens) {
+        flash_params.push_back(Tile148Params{
+            .v_tile_size = 1,
+            .qi_index = token,
+            .kv_head = head,
+        });
+        token += 1;
+      }
+    }
+    for (auto& param : flash_params) {
+      param.min_start_pos = 0;
+      param.max_last_pos = num_tokens - 1;
+      for (size_t i = 0; i < param.v_tile_size; ++i) {
+        param.q_offsets[i] =
+            q.Row(param.qi_index + i) + param.kv_head * qkv_dim - q.Row(0);
+        param.out_offsets[i] = att_out.Row(param.qi_index + i) +
+                               param.kv_head * qkv_dim - att_out.Row(0);
+        param.start_pos[i] = 0;
+        param.last_pos[i] = num_tokens - 1;
+      }
+    }
+  }
+
+  // Runs the flash attention algorithm on Q, K, V.
+  HWY_NOINLINE void FlashAttention() {
+    GCPP_ZONE(env_.ctx, 0, Zones::kVitFlashAttentionInclusive);
+    const size_t qkv_dim = layer_config_.qkv_dim;
+    const size_t heads = layer_config_.heads;
+    HWY_ASSERT_M(heads == layer_config_.kv_heads, "Vit expects MHA");
+    const size_t kNF = FloatsPerVector();
+    const size_t kRoundedKVDim = hwy::RoundUpTo(qkv_dim, 2 * kNF);
+    auto& attn = activations_.attention;
+    const size_t seq_len = static_cast<size_t>(attn.div_seq_len.GetDivisor());
+    if (attn.vit_K_T.Rows() >= seq_len) {
+      attn.vit_K_T.ReshapePackedRowsToCols(2 * kNF);
+      attn.vit_V_T.ReshapePackedRowsToCols(2 * kNF);
+    }
+    const float query_scale = 1.0f;
+    ScaleQuery(attn.q, num_tokens_, heads, qkv_dim, query_scale, attn.q_bf);
+    TransposeKAndV(attn.q, num_tokens_, heads, qkv_dim, attn.vit_K_T,
+                   attn.vit_V_T);
+    // Zero-out the remaining query outputs to prevent garbage propagation
+    for (size_t token = num_real_patches_; token < num_tokens_; ++token) {
+      float* out_row = attn.att_out.Row(token);
+      for (size_t d = 0; d < heads * qkv_dim; ++d) {
+        out_row[d] = 0.0f;
+      }
+    }
+    ComputeParams(num_real_patches_, num_real_patches_, heads, qkv_dim,
+                  attn.q_bf, attn.vit_K_T, attn.vit_V_T, attn.att_out,
+                  attn.flash_params);
+    size_t num_tasks = attn.flash_params.size();
+
+    // For each param, compute fused flash Q.K, softmax and weighted V.
+    const auto func = [&, &ctx = env_.ctx](const size_t task,
+                                           size_t worker) HWY_ATTR {
+      GCPP_ZONE(ctx, worker, Zones::kFlashAttentionFlashAttention);
+      auto& param = attn.flash_params[task];
+      MatPtrT<KV_t> kT("k_T_view", Extents2D(hwy::DivCeil(seq_len, 2 * kNF),
+                                             kRoundedKVDim * 2 * kNF));
+      kT.SetPtr(attn.vit_K_T.Row(0) + param.kv_head * kRoundedKVDim * 2 * kNF,
+                attn.vit_K_T.Stride());
+      MatPtrT<KV_t> vT("v_T_view", Extents2D(hwy::DivCeil(seq_len, 2 * kNF),
+                                             kRoundedKVDim * 2 * kNF));
+      vT.SetPtr(attn.vit_V_T.Row(0) + param.kv_head * kRoundedKVDim * 2 * kNF,
+                attn.vit_V_T.Stride());
+      DispatchDispatchTileFlashAttention148(
+          param, attn.q_bf, kT, vT, /*layer_idx=*/0, attn, attn.att_out,
+          qkv_dim, ctx, worker, /*attention_impl=*/AttentionImpl::kFlash);
+    };
+
+    {
+      PROFILER_ZONE("Gen.VitGemma4FlashAttention.ForkJoin");
+      // Full parallelism is helpful, SmallParallelFor is insufficient.
+      HierarchicalParallelFor(num_tasks, env_.ctx, Callers::kFlashAttention,
+                              func);
+    }
+  }
+
+  // Sums encoded (`att_out`) over num_heads (`layer_config_.heads`) and
+  // head_dim (`qkv_dim`) into output (`att_sums`).
+  HWY_NOINLINE void SumHeads() {
+    CallMatMul(activations_.attention.att_out, layer_.attn_vec_einsum_w,
+               /*bias=*/nullptr, env_, activations_.attention.att_sums);
+  }
+
+ private:
+  const size_t num_tokens_;
+  const size_t num_real_patches_;
+  Activations& activations_;
+  const LayerWeightsPtrs& layer_;
+  const LayerConfig& layer_config_;
+  MatMulEnv& env_;
+  const size_t real_patch_width_;
+  hwy::ThreadPool& pool_;
+  hwy::pool::Caller caller1_;
+  hwy::pool::Caller caller2_;
+  hwy::pool::Caller caller3_;
+  hwy::pool::Caller caller4_;
+};
+
+void VitGemma4TransformerLayer(size_t num_tokens, size_t num_real_patches,
+                               const size_t layer_idx,
+                               const LayerWeightsPtrs& layer,
+                               Activations& activations, MatMulEnv& env,
+                               size_t real_patch_width) {
+  RMSNormBatched(activations.x, layer.pre_attention_norm_scale,
+                 activations.attention.pre_att_rms_out, env.ctx);
+
+  VitGemma4Attention(num_tokens, num_real_patches, layer_idx, activations,
+                     layer, env, real_patch_width)();
+
+  PostNorm(layer.layer_config.post_norm, layer.post_attention_norm_scale,
+           activations.attention.att_sums, env.ctx);
+
+  AddFromBatched(activations.attention.att_sums, activations.x, env.ctx);
+
+  RMSNormBatched(activations.x, layer.pre_ffw_norm_scale,
+                 activations.pre_ffw_rms_out, env.ctx);
+
+  FFWNoVit(layer, activations, env);
+
+  PostNorm(layer.layer_config.post_norm, layer.post_ffw_norm_scale,
+           activations.ffw_out, env.ctx);
+
+  AddFromBatched(activations.ffw_out, activations.x, env.ctx);
+}
+
+// Prefills the image tokens with the ViT encoder.
+void PrefillVitGemma4(const ModelConfig& model_config,
+                      const WeightsPtrs& weights,
+                      const RuntimeConfig& runtime_config, const Image& image,
+                      ImageTokens& image_tokens, Activations& activations,
+                      MatMulEnv& env) {
+  PROFILER_ZONE("Gen.PrefillVitGemma4");
+  const size_t max_patches = model_config.vit_config.seq_len;      // 2520
+  const size_t vit_model_dim = model_config.vit_config.model_dim;  // 768
+  HWY_ASSERT(max_patches == activations.x.Rows());
+
+  // 1. Aspect-ratio-preserving resize
+  Image resized_image = image;
+  int target_height, target_width;
+  GetAspectRatioPreservingSize(image.height(), image.width(),
+                               model_config.vit_config.patch_width,  // 16
+                               max_patches,
+                               3,  // pooling_kernel_size
+                               target_height, target_width);
+  resized_image.ResizeBilinear(target_width, target_height);
+
+  // 2. Embed patches & add position embeddings
+  EmbedImagePatchesGemma4(resized_image, model_config, weights, activations,
+                          env);
+
+  const size_t patch_width = model_config.vit_config.patch_width;  // 16
+  const size_t real_patch_width = target_width / patch_width;
+
+  const size_t num_real_patches = (resized_image.height() / patch_width) *
+                                  (resized_image.width() / patch_width);
+
+  // 3. Go through all layers.
+  for (size_t layer_idx = 0;
+       layer_idx < model_config.vit_config.layer_configs.size(); ++layer_idx) {
+    VitGemma4TransformerLayer(max_patches, num_real_patches, layer_idx,
+                              *weights.VitLayer(layer_idx), activations, env,
+                              real_patch_width);
+  }
+
+  // 4. Final Norm (RMSNorm in Gemma 4? Yes, post_attention_norm etc are
+  // RMSNorm). JAX has vit/transformer/norm_out/scale. In C++ weights.h we
+  // registered: vit_encoder_norm_bias(finder_("enc_norm_bias")) (not used if
+  // RMSNorm) vit_encoder_norm_scale(finder_("enc_norm_scale")) We should apply
+  // final norm. Wait, does it use bias? Gemma 4 final norm is RMSNorm (no
+  // bias). So vit_encoder_norm_bias will be empty. We can use PostNorm or
+  // RMSNormInplaceBatched.
+  // RMSNormInplaceBatched(weights.vit_encoder_norm_scale, activations.x,
+  // env.ctx); Actually, vit.cc uses LayerNormBatched: Gemma 4 ViT has no final
+  // encoder norm.
+
+  // 5. Pooling (3x3 average pooling)
+  const size_t H_p = target_height / patch_width;
+  const size_t W_p = target_width / patch_width;
+  const size_t H_g = H_p / 3;
+  const size_t W_g = W_p / 3;
+  const size_t num_pooled_tokens = H_g * W_g;
+  const size_t max_soft_tokens = max_patches / 9;  // 280
+
+  HWY_ASSERT(num_pooled_tokens <= max_soft_tokens);
+
+  MatStorageT<float> pooled("pooled", Extents2D(max_soft_tokens, vit_model_dim),
+                            env.ctx.allocator, MatPadding::kOdd);
+  gcpp::ZeroInit(pooled);
+
+  for (size_t gy = 0; gy < H_g; ++gy) {
+    for (size_t gx = 0; gx < W_g; ++gx) {
+      const size_t g_idx = gy * W_g + gx;
+      float* dst_row = pooled.Row(g_idx);
+      for (size_t dy = 0; dy < 3; ++dy) {
+        for (size_t dx = 0; dx < 3; ++dx) {
+          const size_t py = 3 * gy + dy;
+          const size_t px = 3 * gx + dx;
+          const size_t p_idx = py * W_p + px;
+          const float* src_row = activations.x.Row(p_idx);
+          for (size_t d = 0; d < vit_model_dim; ++d) {
+            dst_row[d] += src_row[d];
+          }
+        }
+      }
+      for (size_t d = 0; d < vit_model_dim; ++d) {
+        dst_row[d] /= 9.0f;
+      }
+    }
+  }
+
+  // 6. Unscaled RMSNorm (embedding_pre_projection_norm in HF
+  // Gemma4MultimodalEmbedder)
+  RMSNormNoScaleInplaceBatched(pooled, env.ctx);
+
+  // 7. Project to LLM space
+  // Apply head embedding into image_tokens of size of the LLM kModelDim.
+  CallMatMul(pooled, weights.vit_img_head_kernel,
+             /*bias=*/nullptr, env, image_tokens);
+}
+
+// NOLINTNEXTLINE(google-readability-namespace-comments)
+}  // namespace HWY_NAMESPACE
+}  // namespace gcpp
+HWY_AFTER_NAMESPACE();
diff --git a/gemma/weights.h b/gemma/weights.h
index 0dddffb2..f770cf54 100644
--- a/gemma/weights.h
+++ b/gemma/weights.h
@@ -105,6 +105,26 @@ struct ClampRange {
 
 // Per-layer weight metadata and pointers. The tensor data is owned by
 // `MatOwner`.
+//
+// Gemma 4 ViT layers use separate tensor names from the Gemma 3 ViT (kVit)
+// layers. VitCond remaps the generic layer tensor names (e.g. "pre_att_ns")
+// to Gemma-4-specific names (e.g. "vit_pr_att") so both can coexist in the
+// same LayerWeightsPtrs struct. For kVit layers the names pass through
+// unchanged.
+static inline std::string VitCond(const LayerConfig& config,
+                                  const std::string& name) {
+  if (config.type == LayerAttentionType::kVitGemma4) {
+    if (name == "pre_att_ns") return "vit_pr_att";
+    if (name == "post_att_ns") return "vit_po_att";
+    if (name == "pre_ff_ns") return "vit_pr_ff";
+    if (name == "post_ff_ns") return "vit_po_ff";
+    if (name == "gating_ein") return "vit_gate_ein";
+    if (name == "query_norm") return "vit_q_norm";
+    return "vit_" + name;
+  }
+  return name;
+}
+
 struct LayerWeightsPtrs {
   // Initializes tensor metadata without allocating.
   // NOTE: do not store layer_idx, TransformerLayer and Attention may use
@@ -113,9 +133,9 @@ struct LayerWeightsPtrs {
                    const TensorInfoRegistry& tensors)
       : layer_idx(layer_idx),
         finder_(LayerSuffix(layer_idx), tensors),
-        qkv_einsum_w(finder_("qkv_ein")),
-        qkv_einsum_w1(finder_("qkv1_w")),
-        qkv_einsum_w2(finder_("qkv2_w")),
+        qkv_einsum_w(finder_(VitCond(config, "qkv_ein"))),
+        qkv_einsum_w1(finder_(VitCond(config, "qkv1_w"))),
+        qkv_einsum_w2(finder_(VitCond(config, "qkv2_w"))),
         attention_output_biases(finder_("attn_ob")),
         // MultiHeadDotProductAttention.
         vit({.attn_out_w = finder_("attn_out_w"),
@@ -130,23 +150,23 @@ struct LayerWeightsPtrs {
              .layer_norm_0_scale = finder_("ln_0_scale"),
              .layer_norm_1_bias = finder_("ln_1_bias"),
              .layer_norm_1_scale = finder_("ln_1_scale")}),
-        gating_einsum_w(finder_("gating_ein")),
-        gating_einsum_w1(finder_("gating1_w")),
-        gating_einsum_w2(finder_("gating2_w")),
-        linear_w(finder_("linear_w")),
-        pre_attention_norm_scale(finder_("pre_att_ns")),
-        pre_ffw_norm_scale(finder_("pre_ff_ns")),
-        post_attention_norm_scale(finder_("post_att_ns")),
-        post_ffw_norm_scale(finder_("post_ff_ns")),
+        gating_einsum_w(finder_(VitCond(config, "gating_ein"))),
+        gating_einsum_w1(finder_(VitCond(config, "gating1_w"))),
+        gating_einsum_w2(finder_(VitCond(config, "gating2_w"))),
+        linear_w(finder_(VitCond(config, "linear_w"))),
+        pre_attention_norm_scale(finder_(VitCond(config, "pre_att_ns"))),
+        pre_ffw_norm_scale(finder_(VitCond(config, "pre_ff_ns"))),
+        post_attention_norm_scale(finder_(VitCond(config, "post_att_ns"))),
+        post_ffw_norm_scale(finder_(VitCond(config, "post_ff_ns"))),
         skip_scale(finder_("skip_scale")),
         ffw_gating_biases(finder_("ffw_gat_b")),
         ffw_output_biases(finder_("ffw_out_b")),
 
-        attn_vec_einsum_w(finder_("att_ein")),
-        att_weights(finder_("att_w")),
+        attn_vec_einsum_w(finder_(VitCond(config, "att_ein"))),
+        att_weights(finder_(VitCond(config, "att_w"))),
 
-        key_norm_scale(finder_("key_norm")),
-        query_norm_scale(finder_("query_norm")),
+        key_norm_scale(finder_(VitCond(config, "key_norm"))),
+        query_norm_scale(finder_(VitCond(config, "query_norm"))),
 
         router_scale(finder_("router_scale")),
         p_expert_sc(finder_("p_expert_sc")),
@@ -336,6 +356,22 @@ struct LayerWeightsPtrs {
       func(TENSOR_ARGS(vit.layer_norm_0_scale, kMustRead));
       func(TENSOR_ARGS(vit.layer_norm_1_bias, kMustRead));
       func(TENSOR_ARGS(vit.layer_norm_1_scale, kMustRead));
+    }
+    if (layer_config.type == LayerAttentionType::kVitGemma4) {
+      func(TENSOR_ARGS(qkv_einsum_w1, kMustRead));
+      func(TENSOR_ARGS(qkv_einsum_w2, kMustRead));
+      func(TENSOR_ARGS(att_weights, kMaybeRead));
+      func(TENSOR_ARGS(attn_vec_einsum_w, kMaybeRead));
+      func(TENSOR_ARGS(gating_einsum_w, kMustRead));
+      func(TENSOR_ARGS(linear_w, kMustRead));
+      func(TENSOR_ARGS(pre_attention_norm_scale, kMustRead));
+      func(TENSOR_ARGS(post_attention_norm_scale, kMustRead));
+      func(TENSOR_ARGS(pre_ffw_norm_scale, kMustRead));
+      func(TENSOR_ARGS(post_ffw_norm_scale, kMustRead));
+      if (layer_config.use_qk_norm) {
+        func(TENSOR_ARGS(key_norm_scale, kMustRead));
+        func(TENSOR_ARGS(query_norm_scale, kMustRead));
+      }
       return;
     }
     if (layer_config.type == LayerAttentionType::kGemma) {
@@ -816,15 +852,34 @@ struct WeightsPtrs {
     }
 
     if (!config_.vit_config.layer_configs.empty()) {  // Vit parts.
-      func(TENSOR_ARGS(vit_encoder_norm_bias, kMustRead));
-      func(TENSOR_ARGS(vit_encoder_norm_scale, kMustRead));
-      func(TENSOR_ARGS(vit_img_embedding_bias, kMustRead));
+      // Gemma 4 ViT does not use encoder norm or embedding/head biases;
+      // mark them kMaybeRead so loading doesn't abort when absent.
+      if (config_.model == Model::GEMMA4_2B) {
+        func(TENSOR_ARGS(vit_encoder_norm_bias, kMaybeRead));
+        func(TENSOR_ARGS(vit_encoder_norm_scale, kMaybeRead));
+      } else {
+        func(TENSOR_ARGS(vit_encoder_norm_bias, kMustRead));
+        func(TENSOR_ARGS(vit_encoder_norm_scale, kMustRead));
+      }
+      if (config_.model == Model::GEMMA4_2B) {
+        func(TENSOR_ARGS(vit_img_embedding_bias, kMaybeRead));
+      } else {
+        func(TENSOR_ARGS(vit_img_embedding_bias, kMustRead));
+      }
       func(TENSOR_ARGS(vit_img_embedding_kernel, kMustRead));
       func(TENSOR_ARGS(vit_img_pos_embedding, kMustRead));
-      func(TENSOR_ARGS(vit_img_head_bias, kMustRead));
+      if (config_.model == Model::GEMMA4_2B) {
+        func(TENSOR_ARGS(vit_img_head_bias, kMaybeRead));
+      } else {
+        func(TENSOR_ARGS(vit_img_head_bias, kMustRead));
+      }
       func(TENSOR_ARGS(vit_img_head_kernel, kMustRead));
 
-      if (config_.wrapping == PromptWrapping::GEMMA_VLM) {
+      // RMS norm applied to soft tokens before projection. Used by all VLM
+      // models; kMaybeRead for Gemma 4 because its export path may omit it.
+      if (config_.model == Model::GEMMA4_2B) {
+        func(TENSOR_ARGS(mm_embed_norm, kMaybeRead));
+      } else {
         func(TENSOR_ARGS(mm_embed_norm, kMustRead));
       }
     }
@@ -844,7 +899,9 @@ struct WeightsPtrs {
     HWY_ASSERT(config_.vit_config.layer_configs.empty() == vit_layers.empty());
     for (size_t layer_idx = 0; layer_idx < vit_layers.size(); ++layer_idx) {
       HWY_ASSERT(vit_layers[layer_idx].layer_config.type ==
-                 LayerAttentionType::kVit);
+                     LayerAttentionType::kVit ||
+                 vit_layers[layer_idx].layer_config.type ==
+                     LayerAttentionType::kVitGemma4);
       other_layer1 = other1 ? other1->VitLayer(layer_idx) : nullptr;
       other_layer2 = other2 ? other2->VitLayer(layer_idx) : nullptr;
       VitLayer(layer_idx)->ForEachTensor(other_layer1, other_layer2, func);
diff --git a/io/fields.cc b/io/fields.cc
index 6b54a465..2773463d 100644
--- a/io/fields.cc
+++ b/io/fields.cc
@@ -208,6 +208,7 @@ class ReadVisitor : public VisitorBase {
     HWY_ASSERT(!end_.empty() && result_.pos <= end_.back());
     // Count extra, which indicates old code and new data.
     result_.extra_u32 += end_.back() - result_.pos;
+    result_.pos = end_.back();
     end_.pop_back();
   }
 
diff --git a/paligemma/image.cc b/paligemma/image.cc
index 5f94a3cb..a34e2850 100644
--- a/paligemma/image.cc
+++ b/paligemma/image.cc
@@ -210,6 +210,51 @@ void Image::Resize(int new_width, int new_height) {
   width_ = new_width;
 }
 
+void Image::ResizeBilinear(int new_width, int new_height) {
+  HWY_ASSERT(new_width > 0 && new_height > 0);
+  HWY_ASSERT(new_width <= INT32_MAX / 3 &&
+             new_width * 3 <= INT32_MAX / new_height);
+  if (width_ == new_width && height_ == new_height) return;
+
+  std::vector<float> new_data(static_cast<size_t>(new_width) * new_height * 3);
+  const float scale_x = static_cast<float>(width_) / new_width;
+  const float scale_y = static_cast<float>(height_) / new_height;
+
+  for (int y = 0; y < new_height; ++y) {
+    const float src_y = (y + 0.5f) * scale_y - 0.5f;
+    const float src_y_clamped = std::clamp(src_y, 0.0f, static_cast<float>(height_ - 1));
+    const int y_l = static_cast<int>(std::floor(src_y_clamped));
+    const int y_h = std::min(height_ - 1, y_l + 1);
+    const float d_y = src_y_clamped - y_l;
+
+    for (int x = 0; x < new_width; ++x) {
+      const float src_x = (x + 0.5f) * scale_x - 0.5f;
+      const float src_x_clamped = std::clamp(src_x, 0.0f, static_cast<float>(width_ - 1));
+      const int x_l = static_cast<int>(std::floor(src_x_clamped));
+      const int x_h = std::min(width_ - 1, x_l + 1);
+      const float d_x = src_x_clamped - x_l;
+
+      const float w00 = (1.0f - d_x) * (1.0f - d_y);
+      const float w10 = d_x * (1.0f - d_y);
+      const float w01 = (1.0f - d_x) * d_y;
+      const float w11 = d_x * d_y;
+
+      for (int k = 0; k < 3; ++k) {
+        const float v00 = data_[(static_cast<size_t>(y_l) * width_ + x_l) * 3 + k];
+        const float v10 = data_[(static_cast<size_t>(y_l) * width_ + x_h) * 3 + k];
+        const float v01 = data_[(static_cast<size_t>(y_h) * width_ + x_l) * 3 + k];
+        const float v11 = data_[(static_cast<size_t>(y_h) * width_ + x_h) * 3 + k];
+
+        new_data[(static_cast<size_t>(y) * new_width + x) * 3 + k] =
+            w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
+      }
+    }
+  }
+  data_ = std::move(new_data);
+  height_ = new_height;
+  width_ = new_width;
+}
+
 bool Image::WriteBinary(const std::string& filename) const {
   // Writes the floating point values as float32 in binary format.
   std::ofstream file(filename);
diff --git a/paligemma/image.h b/paligemma/image.h
index e54bf86c..938acd21 100644
--- a/paligemma/image.h
+++ b/paligemma/image.h
@@ -41,6 +41,8 @@ class Image {
   // Resizes to width x height (nearest-neighbor for now, bilinear or antialias
   // would be better).
   void Resize(int width, int height);
+  // Resizes to width x height using bilinear interpolation.
+  void ResizeBilinear(int width, int height);
   // Writes the file as plain floats in binary. Useful to e.g. load in a colab.
   bool WriteBinary(const std::string& filename) const;
   // Stores the patch for the given patch number in `patch`.
diff --git a/python/configs.cc b/python/configs.cc
index c16924aa..a226ea8f 100644
--- a/python/configs.cc
+++ b/python/configs.cc
@@ -45,6 +45,7 @@ PYBIND11_MODULE(configs, py_module) {
   enum_<gcpp::LayerAttentionType>(py_module, "LayerAttentionType")
       .value("kGemma", gcpp::LayerAttentionType::kGemma)
       .value("kVit", gcpp::LayerAttentionType::kVit)
+      .value("kVitGemma4", gcpp::LayerAttentionType::kVitGemma4)
       .value("kDeepSeekMLA", gcpp::LayerAttentionType::kDeepSeekMLA);
 
   enum_<gcpp::AttentionVariant>(py_module, "AttentionVariant")