Fix kv cache pyre and build

Differential Revision: D68132579

Pull Request resolved: #7638

Author: jackzhxng

examples/models/llama/TARGETS

@@ -109,6 +109,7 @@ runtime.python_library(
         "//executorch/backends/vulkan/_passes:vulkan_passes",
         "//executorch/examples/models:model_base",
         "//executorch/examples/models:models",
+        "//executorch/exir/passes:init_mutable_pass",
         "//executorch/extension/llm/custom_ops:custom_ops_aot_py",
         "//executorch/extension/llm/export:export_lib",
         # one definition has to be included in the user of the libarary

exir/emit/test/TARGETS

@@ -20,8 +20,10 @@ python_unittest(
         "//executorch/exir/emit:lib",
         "//executorch/exir/passes:const_prop_pass",
         "//executorch/exir/passes:constant_prop_pass",
+        "//executorch/exir/passes:init_mutable_pass",
         "//executorch/exir/tests:lib",
         "//executorch/exir/tests:models",
         "//executorch/extension/pybindings:portable_lib",
+        "//executorch/runtime:runtime",
     ],
 )

exir/passes/TARGETS

@@ -220,6 +220,16 @@ python_library(
     ],
 )
 
+python_library(
+    name = "init_mutable_pass",
+    srcs = [
+        "init_mutable_pass.py",
+    ],
+    deps = [
+        "//executorch/exir:delegate",
+    ],
+)
+
 python_library(
     name = "sym_shape_eval_pass",
     srcs = [

extension/llm/export/TARGETS

@@ -33,6 +33,7 @@ runtime.python_library(
         "//executorch/backends/vulkan/partitioner:vulkan_partitioner",
         "//executorch/backends/vulkan/quantizer:vulkan_quantizer",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
+        "//executorch/exir:delegate",
         "//executorch/exir:lib",
         "//executorch/exir/backend:backend_details",
         "//executorch/extension/export_util:export_util",

extension/llm/export/builder.py

@@ -27,7 +27,7 @@
 from executorch.exir.backend.utils import format_delegated_graph
 from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig
 
-from executorch.exir.pass_manager import PassType
+from executorch.exir.pass_base import ExportPass
 from executorch.exir.passes import MemoryPlanningPass
 from executorch.exir.passes.quant_fusion_pass import QuantFusionPass
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
@@ -417,12 +417,11 @@ def to_backend(self, partitioners: Optional[List[Partitioner]]) -> "LLMEdgeManag
         return self
 
     def to_executorch(
-        self, passes: Optional[List[PassType]] = None
+        self, passes: Optional[List[ExportPass]] = None
     ) -> "LLMEdgeManager":
         """
         Lower the model to executorch and get an ExecutorchProgram.
         """
-        assert self.edge_manager, "Need to run export_to_edge() first"
         to_executorch_passes = [
             # If there are Linear operations left in the graph, let's execute
             # them with the optimized op_linear rather than materializing a
@@ -431,11 +430,19 @@ def to_executorch(
             QuantFusionPass(),
         ]
         if passes:
+            # pyre-fixme[6]: In call `list.extend`, for 1st positional argument,
+            # expected `Iterable[Union[ConvertToLinearPass, QuantFusionPass]]` but
+            # got `List[ExportPass]
             to_executorch_passes.extend(passes)
 
+        assert self.edge_manager, "Need to run export_to_edge() first"
         self.export_program = self.edge_manager.to_executorch(
             ExecutorchBackendConfig(
                 extract_delegate_segments=True,
+                # pyre-fixme[6]: In call `ExecutorchBackendConfig.__init__`, for
+                # argument `passes`, expected `List[typing.Callable[[GraphModule],
+                # Optional[PassResult]]]` but got `List[Union[ConvertToLinearPass,
+                # QuantFusionPass]]`.
                 passes=to_executorch_passes,
                 memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
                 sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(),

runtime/TARGETS

@@ -10,5 +10,6 @@ runtime.python_library(
     ],
     visibility = [
         "//executorch/runtime/...",
+        "//executorch/exir/emit/test/...",
     ],
 )