Allow prefill only run by allowing max_prompt_size == seq_len

Charles Zhao · copybara-github · commit 3a63a12624b8 · 2025-11-03T03:17:54.000-08:00
PiperOrigin-RevId: 827415258
diff --git 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);
   }

Original file line number	Diff line number	Diff line change
`@@ -511,7 +511,7 @@ static void GenerateT(const ModelConfig& config,`
`511`	`511`	`// We use a single divisor, so all sequence lengths must be the same.`
`512`	`512`	`HWY_ASSERT(qbatch.KV(qi).SeqLen() == seq_len);`
`513`	`513`	`}`
`514`		`- if (max_prompt_size >= seq_len) {`
	`514`	`+ if (max_prompt_size > seq_len) {`
`515`	`515`	`HWY_ABORT("max_prompt_size = %zu, increase --seq_len to at least that.",`
`516`	`516`	`max_prompt_size);`
`517`	`517`	`}`