From d4a51dea10e629672b4a0435126c2ab36f866d62 Mon Sep 17 00:00:00 2001
From: Lunwen He <lwhecser@gmail.com>
Date: Tue, 29 Oct 2024 15:26:48 -0700
Subject: [PATCH] buckify source_transformation

This PR buckies `source_transformation` and move test files into `test`.

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

[ghstack-poisoned]
---
 .../llama/source_transformation/TARGETS       | 33 +++++++++++++++++++
 .../llama/source_transformation/test/TARGETS  | 19 +++++++++++
 .../{ => test}/test_quantized_kv_cache.py     |  0
 .../test_sdpa_with_quantized_kv_cache.py      |  0
 4 files changed, 52 insertions(+)
 create mode 100644 examples/models/llama/source_transformation/TARGETS
 create mode 100644 examples/models/llama/source_transformation/test/TARGETS
 rename examples/models/llama/source_transformation/{ => test}/test_quantized_kv_cache.py (100%)
 rename examples/models/llama/source_transformation/{ => test}/test_sdpa_with_quantized_kv_cache.py (100%)

diff --git a/examples/models/llama/source_transformation/TARGETS b/examples/models/llama/source_transformation/TARGETS
new file mode 100644
index 00000000000..8ef85b40732
--- /dev/null
+++ b/examples/models/llama/source_transformation/TARGETS
@@ -0,0 +1,33 @@
+load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+
+oncall("executorch")
+
+python_library(
+    name = "lib",
+    srcs = [
+        "__init__.py",
+    ],
+    deps = [
+        ":source_transformation",
+    ],
+)
+
+python_library(
+    name = "source_transformation",
+    srcs = [
+        "apply_spin_quant_r1_r2.py",
+        "lora.py",
+        "pre_quantization.py",
+        "prune_vocab.py",
+        "quantized_kv_cache.py",
+        "quantize.py",
+        "rms_norm.py",
+        "rope.py",
+        "sdpa.py",
+        "spin_quant.py",
+        "vulkan_rope.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
diff --git a/examples/models/llama/source_transformation/test/TARGETS b/examples/models/llama/source_transformation/test/TARGETS
new file mode 100644
index 00000000000..1d743c7c217
--- /dev/null
+++ b/examples/models/llama/source_transformation/test/TARGETS
@@ -0,0 +1,19 @@
+load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
+
+oncall("executorch")
+
+python_unittest(
+    name = "source_transformation",
+    srcs = [
+        "test_quantized_kv_cache.py",
+        "test_sdpa_with_quantized_kv_cache.py",
+    ],
+    preload_deps = [
+        "//executorch/extension/llm/custom_ops:custom_ops_aot_lib",
+    ],
+    deps = [
+        "fbsource//third-party/pypi/pytest:pytest",
+        "//executorch/examples/models/llama/source_transformation:lib",
+        "//executorch/examples/models/llama:llama_transformer",
+    ],
+)
diff --git a/examples/models/llama/source_transformation/test_quantized_kv_cache.py b/examples/models/llama/source_transformation/test/test_quantized_kv_cache.py
similarity index 100%
rename from examples/models/llama/source_transformation/test_quantized_kv_cache.py
rename to examples/models/llama/source_transformation/test/test_quantized_kv_cache.py
diff --git a/examples/models/llama/source_transformation/test_sdpa_with_quantized_kv_cache.py b/examples/models/llama/source_transformation/test/test_sdpa_with_quantized_kv_cache.py
similarity index 100%
rename from examples/models/llama/source_transformation/test_sdpa_with_quantized_kv_cache.py
rename to examples/models/llama/source_transformation/test/test_sdpa_with_quantized_kv_cache.py