From 3a63a12624b8b029629e7c53c779d8861b0d5971 Mon Sep 17 00:00:00 2001
From: Charles Zhao <charleszhao@google.com>
Date: Mon, 3 Nov 2025 03:17:22 -0800
Subject: [PATCH] Allow prefill only run by allowing max_prompt_size == seq_len

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

diff --git a/gemma/gemma.cc b/gemma/gemma.cc
index 7991c35..ae0c6c1 100644
--- a/gemma/gemma.cc
+++ b/gemma/gemma.cc
@@ -511,7 +511,7 @@ static void GenerateT(const ModelConfig& config,
     // We use a single divisor, so all sequence lengths must be the same.
     HWY_ASSERT(qbatch.KV(qi).SeqLen() == seq_len);
   }
-  if (max_prompt_size >= seq_len) {
+  if (max_prompt_size > seq_len) {
     HWY_ABORT("max_prompt_size = %zu, increase --seq_len to at least that.",
               max_prompt_size);
   }