[Executorch][BE] Rename sdpa_with_kv_cache.py to custom_ops.py

Pull Request resolved: #6996

Because now we have more than sdpa_with_kv_cache in it
ghstack-source-id: 254687229
@exported-using-ghexport

Differential Revision: [D66269486](https://our.internmc.facebook.com/intern/diff/D66269486/)

Author: kimishpatel

examples/models/llama/runner/native.py

@@ -23,7 +23,7 @@
 from executorch.examples.models.llama.runner.generation import LlamaRunner
 
 # Note: import this after portable_lib
-from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
+import executorch.extension.llm.custom_ops  # noqa # usort: skip
 from executorch.kernels import quantized  # noqa
 
 

examples/models/llama/source_transformation/sdpa.py

@@ -99,7 +99,7 @@ def _replace_sdpa_with_custom_op(module: torch.nn.Module):
 
 
 def replace_sdpa_with_custom_op(module: torch.nn.Module) -> torch.nn.Module:
-    from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa
+    import executorch.extension.llm.custom_ops  # noqa
 
     _replace_sdpa_with_custom_op(module)
     return module

examples/models/llava/test/test_llava.py

@@ -18,7 +18,7 @@
 from executorch.extension.pybindings.portable_lib import (
     _load_for_executorch_from_buffer,
 )
-from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
+import executorch.extension.llm.custom_ops  # noqa # usort: skip
 from executorch.kernels import quantized  # noqa # usort: skip
 
 logging.basicConfig(level=logging.INFO)

examples/models/llava/test/test_pte.py

@@ -14,7 +14,7 @@
 from PIL import Image
 
 # Custom ops has to be loaded after portable_lib.
-from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
+import executorch.extension.llm.custom_ops  # noqa # usort: skip
 from executorch.kernels import quantized  # noqa # usort: skip
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"

extension/llm/README.md

@@ -38,7 +38,7 @@ A sampler class in C++ to sample the logistics given some hyperparameters.
 ## custom_ops
 Contains custom op, such as:
 - custom sdpa: implements CPU flash attention and avoids copies by taking the kv cache as one of its arguments.
-  - _sdpa_with_kv_cache.py_, _op_sdpa_aot.cpp_: custom op definition in PyTorch with C++ registration.
+  - _custom_ops.py_, _op_sdpa_aot.cpp_: custom op definition in PyTorch with C++ registration.
   - _op_sdpa.cpp_: the optimized operator implementation and registration of _sdpa_with_kv_cache.out_.
 
 ## runner

extension/llm/custom_ops/__init__.py

@@ -0,0 +1 @@
+from .custom_ops import *

extension/llm/custom_ops/sdpa_with_kv_cache.py renamed to extension/llm/custom_ops/custom_ops.py

@@ -81,7 +81,7 @@ def define_common_targets():
     runtime.python_library(
         name = "custom_ops_aot_py",
         srcs = [
-            "sdpa_with_kv_cache.py",
+            "custom_ops.py",
         ],
         visibility = [
             "//executorch/...",

extension/llm/custom_ops/targets.bzl

@@ -11,7 +11,7 @@
 import torch
 import torch.nn.functional as F
 
-from .sdpa_with_kv_cache import custom_ops_lib  # noqa
+from .custom_ops import custom_ops_lib  # noqa
 
 
 def _sdpa_with_kv_cache_ref(q, k, v, k_cache, v_cache, attn_mask, start_pos, seq_len):