Fix verifier or verify_util of each model

mask_cache_size and kv_cache_max_len should have been passed.

PiperOrigin-RevId: 769912691

Author: protobird-git

ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py

@@ -60,6 +60,7 @@ def verify_amd_llama_135m(
   reauthored_model = amd_llama_135m.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/deepseek/verify_util.py

@@ -60,6 +60,7 @@ def verify_deepseek_r1_distill_1_5b(
   reauthored_model = deepseek.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -143,9 +143,7 @@ def verify_reauthored_gemma_model(
   return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
       reauthored_model=verifier.ReauthoredModelWrapper(
-          reauthored_model,
-          mask_as_input=mask_as_input,
-          kv_layout=kv_layout,
+          reauthored_model, mask_as_input, kv_layout
       ),
       tokenizer=GemmaTokenizerWrapper(original_model.tokenizer),
       generate_prompts=generate_prompts,
@@ -171,7 +169,11 @@ def verify_gemma2(
   """
   checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
   logging.info("Building the reauthored model from: %s", checkpoint_path)
-  reauthored_model = gemma2.build_2b_model(checkpoint_path, custom_loader)
+  reauthored_model = gemma2.build_2b_model(
+      checkpoint_path,
+      custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  )
 
   return verify_reauthored_gemma_model(
       checkpoint=checkpoint_dir,
@@ -193,7 +195,11 @@ def verify_gemma1_with_custom_loader(checkpoint_dir: str) -> bool:
   weight_filename = "gemma-2b-it.ckpt"
   checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
   custom_loader = loader.get_custom_loader(checkpoint_path)
-  reauthored_model = gemma1.build_2b_model(checkpoint_path, custom_loader)
+  reauthored_model = gemma1.build_2b_model(
+      checkpoint_path,
+      custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  )
   return verify_reauthored_gemma_model(
       checkpoint=checkpoint_dir,
       variant="2b",

ai_edge_torch/generative/examples/gemma3/verify_util.py

@@ -93,21 +93,27 @@ def generate(
 class UnifiedGemma3Wrapper(verifier.ReauthoredModelWrapper):
   """Unified Gemma3 model wrapper for verification."""
 
-  def __init__(self, model: torch.nn.Module):
-    super().__init__(model, kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED)
+  def __init__(
+      self,
+      model: torch.nn.Module,
+      kv_cache_max_len: int = verifier.DEFAULT_KV_CACHE_MAX_LEN,
+  ):
+    super().__init__(
+        model,
+        kv_layout=kv_utils.KV_LAYOUT_TRANSPOSED,
+        kv_cache_max_len=kv_cache_max_len,
+    )
 
   def _init_kv_cache(self):
     return kv_utils.KVCache.from_model_config(
-        self.model.model.config, kv_layout=self.kv_layout
+        self.kv_cache_max_len, self.model.model.config, kv_layout=self.kv_layout
     )
 
   def forward(
       self, tokens: torch.Tensor, pixel_values: torch.Tensor = None
   ) -> torch.Tensor:
     """Forwards the model."""
-    mask = attn_utils.build_causal_mask_cache(
-        self.model.model.config.kv_cache_max_len
-    )
+    mask = attn_utils.build_causal_mask_cache(self.kv_cache_max_len)
     input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
     mask = mask.index_select(2, input_pos)
     output = self.model.model.forward(
@@ -127,9 +133,7 @@ def generate(
     tokens = torch.tensor([input_ids])
     input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
     kv_cache = self._init_kv_cache()
-    mask_cache = attn_utils.build_causal_mask_cache(
-        self.model.model.config.kv_cache_max_len
-    )
+    mask_cache = attn_utils.build_causal_mask_cache(self.kv_cache_max_len)
     for _ in range(max_new_tokens):
       mask = mask_cache.index_select(2, input_pos)
       output = self.model.model.forward(
@@ -245,7 +249,11 @@ def verify_gemma3(
 
   if variant == "1b":
     reauthored_model = UnifiedGemma3Wrapper(
-        gemma3.build_model_1b(gemma3_model_path, custom_loader)
+        gemma3.build_model_1b(
+            gemma3_model_path,
+            custom_loader,
+            mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
+        )
     )
   else:
     raise ValueError(f"Unsupported Gemma3 variant: {variant}")

ai_edge_torch/generative/examples/hammer/verify_util.py

@@ -66,6 +66,7 @@ def verify_hammer(
   reauthored_model = _BUILDER[model_size](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/llama/verify_util.py

@@ -65,6 +65,7 @@ def verify_llama_3_2(
   reauthored_model = _BUILDER[model_size](
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/openelm/verify_util.py

@@ -60,6 +60,7 @@ def verify_openelm(
   reauthored_model = openelm.build_model(
       checkpoint_path=reauthored_checkpoint,
       custom_loader=custom_loader,
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
 
   logging.info("Loading the tokenizer from: %s", checkpoint_dir)

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -66,7 +66,9 @@ class ReauthoredPaliGemmaWrapper(verifier.ReauthoredModelWrapper):
   """Reauthored PaliGemma model wrapper."""
 
   def _init_kv_cache(self):
-    return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
+    return kv_cache.KVCache.from_model_config(
+        self.kv_cache_max_len, self.model.config.decoder_config
+    )
 
 
 def main(_):
@@ -88,7 +90,9 @@ def main(_):
 
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = paligemma.build_model(
-      reauthored_checkpoint, version=int(_VERSION.value)
+      reauthored_checkpoint,
+      version=int(_VERSION.value),
+      mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN,
   )
   wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
 

ai_edge_torch/generative/examples/paligemma/verify_decoder.py

@@ -51,7 +51,9 @@ def main(_):
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
-  reauthored_model = decoder.build_decoder(reauthored_checkpoint)
+  reauthored_model = decoder.build_decoder(
+      reauthored_checkpoint, mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN
+  )
 
   logging.info("Loading the tokenizer from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,

ai_edge_torch/generative/examples/paligemma/verify_decoder2.py

@@ -48,7 +48,9 @@ def main(_):
   original_language_model = original_full_model.eval().language_model
 
   logging.info("Building the reauthored model from: %s", checkpoint)
-  reauthored_model = decoder2.build_decoder2(checkpoint)
+  reauthored_model = decoder2.build_decoder2(
+      checkpoint, mask_cache_size=verifier.DEFAULT_KV_CACHE_MAX_LEN
+  )
 
   logging.info("Loading the tokenizer from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,