Update on "[Executorch] Add quantized kv cache to oss ci"

Fixes to make sure quantized kv cache works in oss

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

[ghstack-poisoned]

Author: kimishpatel

examples/models/llama/source_transformation/quantized_kv_cache.py

@@ -4,27 +4,34 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from pathlib import Path
 import logging
 from enum import Enum
 
-from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
-
 import torch
 import torch.nn as nn
 from executorch.examples.models.llama.llama_transformer import KVCache
+
+from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
 from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib  # noqa: F401
 
 
 try:
-    op = torch.ops.quantized_decomposed.quantize_per_token
+    op = torch.ops.quantized_decomposed.quantize_per_token.out
     assert op is not None
 except:
-    libs = list(Path(__file__).parent.resolve().glob("libquantized_ops_aot_lib.*"))
+    import executorch
+    import glob
+
+    executorch_package_path = executorch.__path__[0]
+    libs = list(
+        glob.glob(
+            f"{executorch_package_path}/**/libquantized_ops_aot_lib.*", recursive=True
+        )
+    )
     assert len(libs) == 1, f"Expected 1 library but got {len(libs)}"
     logging.info(f"Loading custom ops library: {libs[0]}")
     torch.ops.load_library(libs[0])
-    op = torch.ops.quantized_decomposed.quantize_per_token
+    op = torch.ops.quantized_decomposed.quantize_per_token.out
     assert op is not None
 
 """
@@ -204,7 +211,6 @@ def update(self, input_pos, k_val, v_val):
                 seq_length = k_val.size(dim_to_slice)
                 narrowed_k = k_out.narrow(dim_to_slice, start_pos, seq_length)
                 narrowed_k.copy_(k_val)
-                # pyre-ignore: Incompatible parameter type [6]
                 narrowed_v = v_out.narrow(dim_to_slice, start_pos, seq_length)
                 narrowed_v.copy_(v_val)
             else:

kernels/quantized/CMakeLists.txt

@@ -60,14 +60,17 @@ if(NOT CMAKE_GENERATOR STREQUAL "Xcode"
     set(_quantized_aot_ops
         "quantized_decomposed::add.out"
         "quantized_decomposed::choose_qparams.Tensor_out"
+        "quantized_decomposed::choose_qparams_per_token_asymmetric.out"
         "quantized_decomposed::dequantize_per_channel.out"
         "quantized_decomposed::dequantize_per_tensor.out"
         "quantized_decomposed::dequantize_per_tensor.Tensor_out"
+        "quantized_decomposed::dequantize_per_token.out"
         "quantized_decomposed::mixed_linear.out"
         "quantized_decomposed::mixed_mm.out"
         "quantized_decomposed::quantize_per_channel.out"
         "quantized_decomposed::quantize_per_tensor.out"
         "quantized_decomposed::quantize_per_tensor.Tensor_out"
+        "quantized_decomposed::quantize_per_token.out"
     )
     gen_selected_ops(
       LIB_NAME "quantized_ops_aot_lib" ROOT_OPS ${_quantized_aot_ops}