Add mask as optional input for converter

PiperOrigin-RevId: 718980201

Author: talumbau

ai_edge_torch/generative/utilities/converter.py

@@ -174,6 +174,9 @@ def _export_helper(
           'input_pos': prefill_input_pos,
           'kv_cache': kv,
       }
+      if export_config.prefill_mask is not None:
+        sample_kwargs['mask'] = export_config.prefill_mask
+
       if lora is not None:
         prefill_signature_name += f'_lora_r{lora.get_rank()}'
         sample_kwargs['lora'] = lora
@@ -199,6 +202,8 @@ def _export_helper(
         'input_pos': decode_input_pos,
         'kv_cache': kv,
     }
+    if export_config.decode_mask is not None:
+      sample_kwargs['mask'] = export_config.decode_mask
     if lora is not None:
       sample_kwargs['lora'] = lora
 

ai_edge_torch/generative/utilities/model_builder.py

@@ -55,6 +55,9 @@ class ExportConfig:
   # On prefill signatures, should the model produce logit output?
   # When False, only decode signatures will produce output.
   output_logits_on_prefill: bool = False
+  # Attention masks given as inputs to the model.
+  prefill_mask: Optional[torch.Tensor] = None
+  decode_mask: Optional[torch.Tensor] = None
 
 
 class DecoderOnlyModel(nn.Module):