From 2751a194bec65bd675ea2cb9003730339b69e721 Mon Sep 17 00:00:00 2001
From: Jan Wassenberg <janwas@google.com>
Date: Thu, 5 Feb 2026 05:59:14 -0800
Subject: [PATCH] Fix paligemma: must subtract image tokens from prompt length

PiperOrigin-RevId: 865905454
---
 gemma/run.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gemma/run.cc b/gemma/run.cc
index 6c6f4d0e..8d005304 100644
--- a/gemma/run.cc
+++ b/gemma/run.cc
@@ -200,7 +200,7 @@ void ReplGemma(const GemmaArgs& args, const Gemma& gemma, KVCache& kv_cache,
                                config.wrapping, abs_pos, prompt_string,
                                image_tokens.Rows());
       runtime_config.image_tokens = &image_tokens;
-      prompt_size = prompt.size();
+      prompt_size = prompt.size() - image_tokens.Rows();
       if (config.wrapping == PromptWrapping::PALIGEMMA) {
         // The end of the prefix for prefix-LM style attention in Paligemma.
         // See Figure 2 of https://arxiv.org/abs/2407.07726.