ai-edge-torch-nightly 0.6.0.dev20250601__py3-none-any.whl → 0.6.0.dev20250603__py3-none-any.whl

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -21,7 +21,6 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 from ai_edge_torch.generative.utilities import loader
 
-
 flags = converter.define_conversion_flags('llama')
 
 _MODEL_SIZE = flags.DEFINE_enum(
@@ -44,13 +43,14 @@ def main(_):
       custom_loader=loader.maybe_get_custom_loader(
           checkpoint_path, flags.FLAGS.custom_checkpoint_loader
       ),
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+      mask_cache_size=converter.get_mask_cache_size_from_flags(),
   )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
       output_name_prefix=flags.FLAGS.output_name_prefix,
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
       export_config=export_config.get_from_flags(),

ai_edge_torch/generative/examples/llama/llama.py

@@ -93,22 +93,12 @@ class Llama(model_builder.DecoderOnlyModel):
 
   Llama 3.2 shares the same architecture as TinyLlama except ROPE calculation.
   """
+  pass
 
-  def __init__(self, config: cfg.ModelConfig):
-    super().__init__(config)
-    attn_config = self.config.block_config(0).attn_config
 
+def get_1b_model_config() -> cfg.ModelConfig:
+  """Returns the model config for a Llama 3.2-1B model."""
 
-def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for a Llama 3.2-1B model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for a SmolLM model.
-  """
   attn_config = cfg.AttentionConfig(
       num_heads=32,
       head_dim=64,
@@ -147,7 +137,6 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       num_layers=16,
       max_seq_len=max_seq_len,
       embedding_dim=2048,
-      kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
       build_rope=build_rope,
@@ -155,9 +144,9 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+def get_3b_model_config() -> cfg.ModelConfig:
   """Returns the model config for a Llama 3.2-3B model."""
-  config = get_1b_model_config(kv_cache_max_len)
+  config = get_1b_model_config()
   # Llama 3.2 has only one block config.
   attn_config = config.block_config(0).attn_config
   attn_config.num_heads = 24
@@ -167,8 +156,8 @@ def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
-  config = get_1b_model_config(**kwargs)
+def get_fake_model_config() -> cfg.ModelConfig:
+  config = get_1b_model_config()
   config.vocab_size = 128
   config.num_layers = 2
   # SmolLM has only one block config.
@@ -180,6 +169,7 @@ def _build_model(
     checkpoint_path: str,
     config: cfg.ModelConfig,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    mask_cache_size: int = 0,
 ) -> torch.nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
@@ -187,28 +177,25 @@ def _build_model(
       tensor_names=TENSOR_NAMES,
       model_class=Llama,
       custom_loader=custom_loader,
+      mask_cache_size=mask_cache_size,
   )
 
 
 def build_1b_model(
     checkpoint_path: str,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs
+    mask_cache_size: int = 0,
 ) -> torch.nn.Module:
   return _build_model(
-      checkpoint_path,
-      get_1b_model_config(**kwargs),
-      custom_loader=custom_loader,
+      checkpoint_path, get_1b_model_config(), custom_loader, mask_cache_size
   )
 
 
 def build_3b_model(
     checkpoint_path: str,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs
+    mask_cache_size: int = 0,
 ) -> torch.nn.Module:
   return _build_model(
-      checkpoint_path,
-      get_3b_model_config(**kwargs),
-      custom_loader=custom_loader,
+      checkpoint_path, get_3b_model_config(), custom_loader, mask_cache_size
   )

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -31,13 +31,14 @@ def main(_):
       custom_loader=loader.maybe_get_custom_loader(
           checkpoint_path, flags.FLAGS.custom_checkpoint_loader
       ),
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+      mask_cache_size=converter.get_mask_cache_size_from_flags(),
   )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
       output_name_prefix=flags.FLAGS.output_name_prefix,
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
       export_config=export_config.get_from_flags(),

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -42,16 +42,8 @@ class OpenELM(model_builder.DecoderOnlyModel):
   pass
 
 
-def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for an OpenELM model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for an OpenELM model.
-  """
+def get_model_config() -> cfg.ModelConfig:
+  """Returns the model config for an OpenELM model."""
   norm_config = cfg.NormalizationConfig(
       type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6
   )
@@ -98,18 +90,17 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       num_layers=num_layers,
       max_seq_len=2048,
       embedding_dim=3072,
-      kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
   )
   return config
 
 
-def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
-  config = get_model_config(kv_cache_max_len)
+def get_fake_model_config() -> cfg.ModelConfig:
+  config = get_model_config()
   config.vocab_size = 128
   config.num_layers = 2
-  config.max_seq_len = 2 * kv_cache_max_len
+  config.max_seq_len = 256
   config.embedding_dim = 128
   config.block_configs = config.block_configs[: config.num_layers]
   for block_config in config.block_configs:
@@ -122,12 +113,13 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
 def build_model(
     checkpoint_path: str,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs
+    mask_cache_size: int = 0,
 ) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
-      config=get_model_config(**kwargs),
+      config=get_model_config(),
       tensor_names=TENSOR_NAMES,
       model_class=OpenELM,
       custom_loader=custom_loader,
+      mask_cache_size=mask_cache_size,
   )

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -40,7 +40,7 @@ def main(_):
       custom_loader=loader.maybe_get_custom_loader(
           checkpoint_path, flags.FLAGS.custom_checkpoint_loader
       ),
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+      mask_cache_size=converter.get_mask_cache_size_from_flags(),
   )
 
   config = pytorch_model.image_encoder.config.image_embedding
@@ -49,6 +49,7 @@ def main(_):
       output_path=flags.FLAGS.output_path,
       output_name_prefix=f'{flags.FLAGS.output_name_prefix}_{_VERSION.value}',
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       pixel_values_size=torch.Size(
           [1, config.channels, config.image_size, config.image_size]
       ),

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -73,8 +73,9 @@ class Decoder(model_builder.DecoderOnlyModel):
     # The first part of input_embeds are image embeddings. Diagonal causal mask
     # doesn't work here.
     if mask is None:
+      assert kv_cache is not None, "KV cache must be provided."
       embeds_len = input_embeds.shape[1]
-      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask = torch.zeros(embeds_len, kv_cache.get_max_seq_len())
       mask[:, embeds_len:] = attn_config.causal_mask_value
 
     return self._forward_with_embeds(
@@ -87,16 +88,8 @@ class Decoder(model_builder.DecoderOnlyModel):
     )
 
 
-def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for the decoder of a PaliGemma 3B model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for the decoder of a PaliGemma 3B model.
-  """
+def get_decoder_config() -> cfg.ModelConfig:
+  """Returns the model config for the decoder of a PaliGemma 3B model."""
   attn_config = cfg.AttentionConfig(
       num_heads=8,
       head_dim=256,
@@ -125,7 +118,6 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       max_seq_len=8192,
       embedding_dim=embedding_dim,
       embedding_scale=embedding_dim**0.5,
-      kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_use_bias=False,
@@ -133,22 +125,25 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_decoder_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
-  config = get_decoder_config(kv_cache_max_len)
+def get_fake_decoder_config() -> cfg.ModelConfig:
+  config = get_decoder_config()
   # PaliGemma decoder has only one block config.
   config.block_config(0).ff_config.intermediate_size = 128
   config.vocab_size = 128
   config.num_layers = 2
-  config.max_seq_len = 2 * kv_cache_max_len
+  config.max_seq_len = 256
   config.embedding_dim = 128
   config.embedding_scale = 128**0.5
   return config
 
 
-def build_decoder(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+def build_decoder(
+    checkpoint_path: str, mask_cache_size: int = 0
+) -> torch.nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
-      config=get_decoder_config(**kwargs),
+      config=get_decoder_config(),
       tensor_names=TENSOR_NAMES,
       model_class=Decoder,
+      mask_cache_size=mask_cache_size,
   )

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -73,8 +73,9 @@ class Decoder2(gemma2.Gemma2):
 
     if mask is None:
       # By default, don't mask image embeds with a diagonal causal mask.
+      assert kv_cache is not None, "KV cache must be provided."
       embeds_len = input_embeds.shape[1]
-      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask = torch.zeros(embeds_len, kv_cache.get_max_seq_len())
       mask[:, embeds_len:] = attn_config.causal_mask_value
 
     return self._forward_with_embeds(
@@ -82,16 +83,8 @@ class Decoder2(gemma2.Gemma2):
     )
 
 
-def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for the decoder of a PaliGemma 3B model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for the decoder of a PaliGemma 3B model.
-  """
+def get_decoder2_config() -> cfg.ModelConfig:
+  """Returns the model config for the decoder of a PaliGemma 3B model."""
   norm_config = cfg.NormalizationConfig(
       type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
@@ -133,7 +126,6 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       max_seq_len=8192,
       embedding_dim=embedding_dim,
       embedding_scale=embedding_dim**0.5,
-      kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
       lm_head_use_bias=False,
@@ -142,22 +134,25 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_decoder2_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
-  config = get_decoder2_config(kv_cache_max_len)
+def get_fake_decoder2_config() -> cfg.ModelConfig:
+  config = get_decoder2_config()
   # PaliGemma2 decoder has only one block config.
   config.block_config(0).ff_config.intermediate_size = 128
   config.vocab_size = 128
   config.num_layers = 2
-  config.max_seq_len = 2 * kv_cache_max_len
+  config.max_seq_len = 256
   config.embedding_dim = 128
   config.embedding_scale = 128**0.5
   return config
 
 
-def build_decoder2(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+def build_decoder2(
+    checkpoint_path: str, mask_cache_size: int = 0
+) -> torch.nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
-      config=get_decoder2_config(**kwargs),
+      config=get_decoder2_config(),
       tensor_names=TENSOR_NAMES,
       model_class=Decoder2,
+      mask_cache_size=mask_cache_size,
   )

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -45,7 +45,12 @@ class PaliGemmaConfig:
 class PaliGemma(nn.Module):
   """PaliGemma model from the Edge Generative API."""
 
-  def __init__(self, config: PaliGemmaConfig, decoder_class: nn.Module):
+  def __init__(
+      self,
+      config: PaliGemmaConfig,
+      decoder_class: nn.Module,
+      mask_cache_size: int = 0,
+  ):
     super().__init__()
 
     self.image_encoder = image_encoder.SiglipVisionEncoder(
@@ -56,7 +61,7 @@ class PaliGemma(nn.Module):
         config.decoder_config.embedding_dim,
         bias=config.image_projection_use_bias,
     )
-    self.decoder = decoder_class(config.decoder_config)
+    self.decoder = decoder_class(config.decoder_config, mask_cache_size)
     image_embedding_config = config.image_encoder_config.image_embedding
     self.num_patches = (
         image_embedding_config.image_size // image_embedding_config.patch_size
@@ -116,7 +121,7 @@ class PaliGemma(nn.Module):
     )
 
 
-def get_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
+def get_model_config(get_decoder_config) -> PaliGemmaConfig:
   """Returns the model config for a PaliGemma 3B-224 model.
 
   Returns:
@@ -124,16 +129,16 @@ def get_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
   """
   return PaliGemmaConfig(
       image_encoder_config=image_encoder.get_image_encoder_config(),
-      decoder_config=get_decoder_config(**kwargs),
+      decoder_config=get_decoder_config(),
       image_token_id=257152,
       image_projection_use_bias=True,
   )
 
 
-def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
+def get_fake_model_config(get_decoder_config) -> PaliGemmaConfig:
   return PaliGemmaConfig(
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),
-      decoder_config=get_decoder_config(**kwargs),
+      decoder_config=get_decoder_config(),
       image_token_id=127,
       image_projection_use_bias=True,
   )
@@ -143,7 +148,7 @@ def build_model(
     checkpoint_path: str,
     version: int = 2,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs,
+    mask_cache_size: int = 0,
 ) -> PaliGemma:
   if version == 1:
     decoder_class = decoder.Decoder
@@ -154,8 +159,8 @@ def build_model(
     decoder_tensor_names = decoder2.TENSOR_NAMES
     get_decoder_config = decoder2.get_decoder2_config
 
-  config = get_model_config(get_decoder_config, **kwargs)
-  model = PaliGemma(config, decoder_class)
+  config = get_model_config(get_decoder_config)
+  model = PaliGemma(config, decoder_class, mask_cache_size)
   # Load the parameters of image encoder.
   loader = loading_utils.ModelLoader(
       checkpoint_path, image_encoder.TENSOR_NAMES, custom_loader

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -31,13 +31,14 @@ def main(_):
       custom_loader=loader.maybe_get_custom_loader(
           checkpoint_path, flags.FLAGS.custom_checkpoint_loader
       ),
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+      mask_cache_size=converter.get_mask_cache_size_from_flags(),
   )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
       output_name_prefix=flags.FLAGS.output_name_prefix,
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
       export_config=export_config.get_from_flags(),

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -31,13 +31,14 @@ def main(_):
       custom_loader=loader.maybe_get_custom_loader(
           checkpoint_path, flags.FLAGS.custom_checkpoint_loader
       ),
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+      mask_cache_size=converter.get_mask_cache_size_from_flags(),
   )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
       output_name_prefix=flags.FLAGS.output_name_prefix,
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
       export_config=export_config.get_from_flags(),

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -32,13 +32,14 @@ def main(_):
       custom_loader=loader.maybe_get_custom_loader(
           checkpoint_path, flags.FLAGS.custom_checkpoint_loader
       ),
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
+      mask_cache_size=converter.get_mask_cache_size_from_flags(),
   )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
       output_name_prefix=flags.FLAGS.output_name_prefix,
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
       export_config=export_config.get_from_flags(),

ai_edge_torch/generative/examples/phi/phi2.py

@@ -41,16 +41,8 @@ class Phi2(model_builder.DecoderOnlyModel):
   pass
 
 
-def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for a Phi-2 model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for a Phi-2 model.
-  """
+def get_model_config() -> cfg.ModelConfig:
+  """Returns the model config for a Phi-2 model."""
   attn_config = cfg.AttentionConfig(
       num_heads=32,
       head_dim=80,
@@ -77,7 +69,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       vocab_size=51200,
       num_layers=32,
       max_seq_len=2048,
-      kv_cache_max_len=kv_cache_max_len,
       embedding_dim=2560,
       block_configs=block_config,
       final_norm_config=norm_config,
@@ -87,11 +78,11 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
-  config = get_model_config(kv_cache_max_len)
+def get_fake_model_config() -> cfg.ModelConfig:
+  config = get_model_config()
   config.vocab_size = 128
   config.num_layers = 2
-  config.max_seq_len = 2 * kv_cache_max_len
+  config.max_seq_len = 256
   # Phi-2 has only one block config.
   config.block_config(0).ff_config.intermediate_size = 128
   return config
@@ -100,12 +91,13 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
 def build_model(
     checkpoint_path: str,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs
+    mask_cache_size: int = 0,
 ) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
-      config=get_model_config(**kwargs),
+      config=get_model_config(),
       tensor_names=TENSOR_NAMES,
       model_class=Phi2,
       custom_loader=custom_loader,
+      mask_cache_size=mask_cache_size,
   )

ai_edge_torch/generative/examples/phi/phi3.py

@@ -139,16 +139,8 @@ class Phi3_5Mini(model_builder.DecoderOnlyModel):
   pass
 
 
-def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for a Phi-3.5 model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for a Phi-3.5 model.
-  """
+def get_model_config() -> cfg.ModelConfig:
+  """Returns the model config for a Phi-3.5 model."""
   attn_config = cfg.AttentionConfig(
       num_heads=32,
       head_dim=96,
@@ -185,7 +177,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       vocab_size=32064,
       num_layers=32,
       max_seq_len=max_seq_len,
-      kv_cache_max_len=kv_cache_max_len,
       embedding_dim=3072,
       block_configs=block_config,
       final_norm_config=norm_config,
@@ -195,11 +186,11 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
-  config = get_model_config(kv_cache_max_len)
+def get_fake_model_config() -> cfg.ModelConfig:
+  config = get_model_config()
   config.vocab_size = 128
   config.num_layers = 2
-  config.max_seq_len = 2 * kv_cache_max_len
+  config.max_seq_len = 256
   # Phi-3.5 has only one block config.
   config.block_config(0).ff_config.intermediate_size = 128
   return config
@@ -208,13 +199,14 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
 def build_model(
     checkpoint_path: str,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs
+    mask_cache_size: int = 0,
 ) -> torch.nn.Module:
   """Instantiates the model instance and load checkpoint if provided."""
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
-      config=get_model_config(**kwargs),
+      config=get_model_config(),
       tensor_names=TENSOR_NAMES,
       model_class=Phi3_5Mini,
       custom_loader=custom_loader,
+      mask_cache_size=mask_cache_size,
   )

ai_edge_torch/generative/examples/phi/phi4.py

@@ -89,16 +89,8 @@ class Phi4Mini(model_builder.DecoderOnlyModel):
   pass
 
 
-def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
-  """Returns the model config for a Phi-4 model.
-
-  Args:
-    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
-      is 1024.
-
-  Returns:
-    The model config for a Phi-4 model.
-  """
+def get_model_config() -> cfg.ModelConfig:
+  """Returns the model config for a Phi-4 model."""
   attn_config = cfg.AttentionConfig(
       num_heads=24,
       head_dim=128,
@@ -135,7 +127,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       vocab_size=200064,
       num_layers=32,
       max_seq_len=max_seq_len,
-      kv_cache_max_len=kv_cache_max_len,
       embedding_dim=3072,
       block_configs=block_config,
       final_norm_config=norm_config,
@@ -144,11 +135,11 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   return config
 
 
-def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
-  config = get_model_config(kv_cache_max_len)
+def get_fake_model_config() -> cfg.ModelConfig:
+  config = get_model_config()
   config.vocab_size = 128
   config.num_layers = 2
-  config.max_seq_len = 2 * kv_cache_max_len
+  config.max_seq_len = 256
   # Phi-4 has only one block config.
   config.block_config(0).ff_config.intermediate_size = 128
   return config
@@ -157,13 +148,14 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
 def build_model(
     checkpoint_path: str,
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
-    **kwargs
+    mask_cache_size: int = 0,
 ) -> torch.nn.Module:
   """Instantiates the model instance and load checkpoint if provided."""
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
-      config=get_model_config(**kwargs),
+      config=get_model_config(),
       tensor_names=TENSOR_NAMES,
       model_class=Phi4Mini,
       custom_loader=custom_loader,
+      mask_cache_size=mask_cache_size,
   )

ai_edge_torch/generative/examples/phi/verify_util.py

@@ -15,7 +15,6 @@
 """Utils for verifying the Phi model."""
 
 import logging
-import os
 import pathlib
 from typing import Callable, Dict
 
@@ -39,7 +38,6 @@ _BUILDER = {
 def verify_phi(
     version: str,
     checkpoint_dir: str,
-    weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
     prompts: list[str] | None = None,
     atol: float = 1e-04,
@@ -63,7 +61,7 @@ def verify_phi(
     )
     reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   else:
-    reauthored_checkpoint = os.path.join(checkpoint_dir, weight_filename)
+    reauthored_checkpoint = checkpoint_dir
 
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = _BUILDER[version](