Optional Mask input for LM examples

PiperOrigin-RevId: 713729372

Author: talumbau

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -129,6 +129,7 @@ def forward(
       tokens: torch.Tensor,
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
@@ -175,7 +176,15 @@ def _forward_with_embeds(
       input_embeds = input_embeds * self.config.embedding_scale
     x = input_embeds
     updated_kv_entries = []
+    mask_input = mask is not None
     for i, block in enumerate(self.transformer_blocks):
+      mask = (
+          mask
+          if mask_input
+          else self.get_attention_mask(
+              block.config.attn_config.attn_type, input_pos
+          )
+      )
       kv_entry = kv_cache.caches[i] if kv_cache else None
       x, kv_entry = block(x, rope, mask[i], input_pos, kv_entry)
       if kv_entry:

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -54,6 +54,7 @@ def forward(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
+      mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
       called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
@@ -73,8 +74,9 @@ def forward(
     # The first part of input_embeds are image embeddings. Diagonal causal mask
     # doesn't work here.
     embeds_len = input_embeds.shape[1]
-    mask = torch.zeros(embeds_len, self.config.kv_cache_max)
-    mask[:, embeds_len:] = float("-inf")
+    if mask is None:
+      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask[:, embeds_len:] = float("-inf")
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -57,6 +57,7 @@ def forward(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
+      mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
       called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
@@ -73,17 +74,21 @@ def forward(
         repo_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
     )
 
-    if called_by_generate:
-      # PaliGemma2 generate() use a diagonal causal mask even with image embeds.
-      mask = [self.get_attention_mask(
-          self.config.block_config(i).attn_config.attn_type, input_pos
-      ) for i in range(self.config.num_layers)]
-    else:
-      # By default, don't mask image embeds with a diagonal causal mask.
-      embeds_len = input_embeds.shape[1]
-      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
-      mask[:, embeds_len:] = float("-inf")
-      mask = [mask] * self.config.num_layers
+    if mask is None:
+      if called_by_generate:
+        # PaliGemma2 generate() use a diagonal causal mask even with image embeds.
+        mask = [
+            self.get_attention_mask(
+                self.config.block_config(i).attn_config.attn_type, input_pos
+            )
+            for i in range(self.config.num_layers)
+        ]
+      else:
+        # By default, don't mask image embeds with a diagonal causal mask.
+        embeds_len = input_embeds.shape[1]
+        mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+        mask[:, embeds_len:] = float("-inf")
+        mask = [mask] * self.config.num_layers
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -70,6 +70,7 @@ def forward(
       tokens: torch.Tensor,
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
       export_config: Optional[model_builder.ExportConfig] = None,
       called_by_generate: bool = True,
@@ -79,6 +80,7 @@ def forward(
           tokens=tokens,
           input_pos=input_pos,
           kv_cache=kv_cache,
+          mask=mask,
           input_embeds=None,
           export_config=export_config,
           called_by_generate=called_by_generate,
@@ -111,6 +113,7 @@ def forward(
         tokens=None,
         input_pos=input_pos,
         kv_cache=kv_cache,
+        mask=mask,
         input_embeds=input_embeds,
         export_config=export_config,
         called_by_generate=called_by_generate,

ai_edge_torch/generative/examples/test_models/toy_model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 # A toy example which has a single-layer transformer block.
-from typing import Tuple
+from typing import Optional, Tuple
 
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers.attention import TransformerBlock
@@ -52,14 +52,20 @@ def __init__(self, config: cfg.ModelConfig) -> None:
     self.config = config
 
   @torch.inference_mode
-  def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+  def forward(
+      self,
+      idx: torch.Tensor,
+      input_pos: torch.Tensor,
+      mask: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
     x = self.tok_embedding(idx)
     cos, sin = self.rope_cache
 
     cos = cos.index_select(0, input_pos)
     sin = sin.index_select(0, input_pos)
-    mask = self.mask_cache.index_select(2, input_pos)
-    mask = mask[:, :, :, : self.config.max_seq_len]
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.max_seq_len]
 
     x = self.transformer_block(x, (cos, sin), mask, input_pos)
     x = self.final_norm(x)
@@ -98,7 +104,12 @@ def __init__(self, config: cfg.ModelConfig) -> None:
     self.config = config
 
   @torch.inference_mode
-  def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+  def forward(
+      self,
+      idx: torch.Tensor,
+      input_pos: torch.Tensor,
+      mask: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
     x = self.tok_embedding(idx)
     cos, sin = self.rope_cache
 

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -63,14 +63,16 @@ def forward(
       tokens: torch.Tensor,
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
       export_config: Optional[ExportConfig] = None,
   ) -> Tuple[torch.Tensor, kv_utils.KVCache]:
     x = self.tok_embedding(tokens)
     cos, sin = self.rope_cache
     cos = cos.index_select(0, input_pos)
     sin = sin.index_select(0, input_pos)
-    mask = self.mask_cache.index_select(2, input_pos)
-    mask = mask[:, :, :, : self.config.max_seq_len]
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.max_seq_len]
 
     updated_kv_entries = []
     for i, block in enumerate(self.transformer_blocks):

ai_edge_torch/generative/utilities/model_builder.py

@@ -99,6 +99,7 @@ def forward(
       tokens: torch.Tensor,
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
       lora: Optional[lora_utils.LoRA] = None,
       export_config: Optional[ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
@@ -122,8 +123,9 @@ def forward(
         # input_pos=input_pos, n_elem=n_elem, base=attn_config.rotary_base
     )
 
-    mask = self.mask_cache.index_select(2, input_pos)
-    mask = mask[:, :, :, : self.config.kv_cache_max]
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.kv_cache_max]
 
     return self.forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, lora, export_config