GPU model prefill sequence length minimum is 32

talumbau · copybara-github · commit f87e4fa77208 · 2025-01-30T08:14:24.000-08:00
PiperOrigin-RevId: 721395570
diff --git a/ai_edge_torch/generative/examples/experimental/gemma/convert_gemma2_gpu_to_tflite.py b/ai_edge_torch/generative/examples/experimental/gemma/convert_gemma2_gpu_to_tflite.py
@@ -43,7 +43,7 @@
 )
 _PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
     'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
+    (32, 64, 128, 256, 512, 1024),
     'List of the maximum sizes of prefill input tensors.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(

Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@`
`43`	`43`	`)`
`44`	`44`	`_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(`
`45`	`45`	`'prefill_seq_lens',`
`46`		`- (8, 64, 128, 256, 512, 1024),`
	`46`	`+ (32, 64, 128, 256, 512, 1024),`
`47`	`47`	`'List of the maximum sizes of prefill input tensors.',`
`48`	`48`	`)`
`49`	`49`	`_KV_CACHE_MAX_LEN = flags.DEFINE_integer(`