replace with custom sdpa before export

export_et.py

@@ -24,6 +24,7 @@
 #    XnnpackDynamicallyQuantizedPartitioner,
 # )
 from executorch_portable_utils import export_to_edge
+from export_et_util import replace_attention_with_sdpa_attention
 
 from quantize import get_precision
 from torch._export import capture_pre_autograd_graph
@@ -106,6 +107,9 @@ def export_model(model, device, output_path, args=None) -> str:  # noqa: C901
     else:
         raise ValueError(f"Unsupported dtype for ET export: {target_precision}")
 
+
+
+    replace_attention_with_sdpa_attention(export_model)
     with torch.nn.attention.sdpa_kernel(
         [torch.nn.attention.SDPBackend.MATH]
     ), torch.no_grad():

export_et_util.py

@@ -1,7 +1,8 @@
 from executorch.examples.models.llama2.custom_ops import sdpa_with_kv_cache
 from build.model import Attention
+from torch import nn
 
-class AttentionWithSDPA(nn.Module):
+class SDPAAttention(nn.Module):
     def __init__(self, attention: Attention):
         super().__init__()
 
@@ -51,3 +52,11 @@ def forward(
         )
         output = output.view(bsz, seqlen, self.dim)
         return self.wo(output)
+
+
+def replace_attention_with_sdpa_attention(module: nn.Module):
+    for name, child in module.named_children():
+        if isinstance(child, Attention):
+            setattr(module, name, SDPAAttention(child))
+        else:
+            replace_attention_with_sdpa_attention(child)