Refactor LoRAFeedForward.forward to use shared lora_call (#20831)

Differential Revision: D110890567

Pull Request resolved: #20831

Author: derekdixu

examples/models/llama/feed_forward.py

@@ -1,6 +1,6 @@
 import torch.nn.functional as F
 
-from executorch.examples.models.llama.lora import LoRALinear
+from executorch.examples.models.llama.lora import lora_call, LoRALinear
 from executorch.examples.models.llama.model_args import ModelArgs
 from torch import nn
 
@@ -65,15 +65,8 @@ def __init__(self, dim: int, hidden_dim: int, args: ModelArgs):
         )
 
     def forward(self, x, lora_blob=None):
-        # CoreML LoRA-as-IO Path-2: when `lora_blob` is provided, route per-
-        # projection slices to LoRALinear instances tagged with `_lora_key`.
-        # Default behavior (lora_blob=None) is unchanged.
-        def _call(linear, x_in):
-            if lora_blob is not None:
-                key = getattr(linear, "_lora_key", None)
-                if key is not None and key in lora_blob:
-                    a, b = lora_blob[key]
-                    return linear(x_in, a, b)
-            return linear(x_in)
-
-        return _call(self.w2, F.silu(_call(self.w1, x)) * _call(self.w3, x))
+        return lora_call(
+            self.w2,
+            F.silu(lora_call(self.w1, x, lora_blob)) * lora_call(self.w3, x, lora_blob),
+            lora_blob,
+        )