Update on "[Executorch][llm] Enable leveraging ring kv cache via module swap"

This allows us to make some of the attention modules to use sliding window kv cache. Will help enable models like gemma3.

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

[ghstack-poisoned]

Author: kimishpatel

.ci/scripts/test_model.sh

@@ -87,10 +87,6 @@ test_model() {
     bash examples/models/llava/install_requirements.sh
     STRICT="--no-strict"
   fi
-  if [[ "$MODEL_NAME" == "llama3_2_vision_encoder" || "$MODEL_NAME" == "llama3_2_text_decoder" ]]; then
-    # Install requirements for llama vision.
-    bash examples/models/llama3_2_vision/install_requirements.sh
-  fi
   if [[ "${MODEL_NAME}" == "qwen2_5" ]]; then
       # Install requirements for export_llama
       bash examples/models/llama/install_requirements.sh

.ci/scripts/unittest-linux.sh

@@ -24,9 +24,6 @@ if [[ "$BUILD_TOOL" == "cmake" ]]; then
     CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
     .ci/scripts/setup-linux.sh "$@"
 
-    # Install llama3_2_vision dependencies.
-    PYTHON_EXECUTABLE=python ./examples/models/llama3_2_vision/install_requirements.sh
-
     .ci/scripts/unittest-linux-cmake.sh
 elif [[ "$BUILD_TOOL" == "buck2" ]]; then
     # Removing this breaks sccache in the Buck build, apparently

.ci/scripts/unittest-macos.sh

@@ -29,7 +29,6 @@ if [[ "$BUILD_TOOL" == "cmake" ]]; then
     # Install llama3_2_vision dependencies.
     PYTHON_EXECUTABLE=python \
     ${CONDA_RUN} --no-capture-output \
-    ./examples/models/llama3_2_vision/install_requirements.sh
 
     .ci/scripts/unittest-macos-cmake.sh
 elif [[ "$BUILD_TOOL" == "buck2" ]]; then

.lintrunner.toml

@@ -220,6 +220,13 @@ exclude_patterns = [
     'extension/**',
     'kernels/optimized/**',
     # Justified <functional> include.
+    'kernels/portable/cpu/op_bitwise*.cpp',
+    'kernels/portable/cpu/op_eq.cpp',
+    'kernels/portable/cpu/op_ge.cpp',
+    'kernels/portable/cpu/op_gt.cpp',
+    'kernels/portable/cpu/op_le.cpp',
+    'kernels/portable/cpu/op_lt.cpp',
+    'kernels/portable/cpu/op_ne.cpp',
     'runtime/kernel/thread_parallel_interface.h',
     'scripts/**',
     'third-party/**',

CMakeLists.txt

@@ -242,6 +242,8 @@ option(EXECUTORCH_USE_DL "Use libdl library" ON)
 
 option(EXECUTORCH_BUILD_CADENCE "Build the Cadence DSP backend" OFF)
 
+option(EXECUTORCH_BUILD_CORTEX_M "Build the Cortex-M backend" OFF)
+
 #
 # pthreadpool: build pthreadpool library. Disable on unsupported platforms
 #
@@ -715,6 +717,10 @@ if(EXECUTORCH_BUILD_XNNPACK)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/xnnpack)
 endif()
 
+if(EXECUTORCH_BUILD_CORTEX_M)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cortex_m)
+endif()
+
 if(EXECUTORCH_BUILD_DEVTOOLS)
   if(NOT EXECUTORCH_BUILD_ARM_BAREMETAL)
     set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER

backends/arm/scripts/build_executorch.sh

@@ -129,6 +129,7 @@ cmake                                                 \
     -DEXECUTORCH_BUILD_ARM_BAREMETAL=ON               \
     -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON           \
     -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON       \
+    -DEXECUTORCH_BUILD_CORTEX_M=ON                    \
     -DEXECUTORCH_ENABLE_LOGGING=ON                    \
     ${build_devtools_flags}                           \
     ${build_with_etdump_flags}                        \

backends/arm/scripts/pre-push

@@ -8,11 +8,13 @@
 # non-interactive mode. "$#" gives the number of positional arguments.
 [ "$#" -eq 0 ] && is_script_interactive=1 || is_script_interactive=0
 
-RESET='\e[0m'
-RED='\e[31m'
-GREEN='\e[32m'
-YELLOW='\e[33m'
-BLUE='\e[34m'
+if [ $is_script_interactive -eq 1 ]; then
+    RESET='\e[0m'
+    RED='\e[31m'
+    GREEN='\e[32m'
+    YELLOW='\e[33m'
+    BLUE='\e[34m'
+fi
 
 INFO="${BLUE}[INFO]${RESET}"
 WARNING="${YELLOW}[WARNING]${RESET}"

backends/arm/test/test_arm_baremetal.sh

@@ -83,8 +83,6 @@ test_pytest_ops() { # Test ops and other things
 test_pytest_models() { # Test ops and other things
     echo "${TEST_SUITE_NAME}: Run pytest"
 
-    examples/models/llama3_2_vision/install_requirements.sh
-
     # Prepare for pytest
     backends/arm/scripts/build_executorch.sh
 
@@ -117,8 +115,6 @@ test_pytest_ops_ethosu_fvp() { # Same as test_pytest but also sometime verify us
 test_pytest_models_ethosu_fvp() { # Same as test_pytest but also sometime verify using Corstone FVP
     echo "${TEST_SUITE_NAME}: Run pytest with fvp"
 
-    examples/models/llama3_2_vision/install_requirements.sh
-
     # Prepare Corstone-3x0 FVP for pytest
     backends/arm/scripts/build_executorch.sh
     backends/arm/scripts/build_portable_kernels.sh
@@ -154,6 +150,13 @@ test_run_ethosu_fvp() { # End to End model tests using run.sh
     echo "${TEST_SUITE_NAME}: Test ethos-u target Ethos-U85"
     examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=add
     examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul
+
+    # Cortex-M op tests
+    examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qadd --bundleio
+    examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio
+    examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio --no_delegate --portable_kernels="aten::sub.out,aten::add.out,aten::mul.out"
+    examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=qops --bundleio
+
     echo "${TEST_SUITE_NAME}: PASS"
     }
 

backends/cortex_m/CMakeLists.txt

@@ -0,0 +1,61 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Kernel library for Cortex-M operators. Please keep this file formatted by running:
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+cmake_minimum_required(VERSION 3.19)
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+if(NOT CMAKE_CXX_STANDARD)
+  set(CMAKE_CXX_STANDARD 17)
+endif()
+
+# Source root directory for executorch.
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
+endif()
+
+include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
+include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake)
+
+if(NOT PYTHON_EXECUTABLE)
+  resolve_python_executable()
+endif()
+
+# Cortex-M ops kernel sources
+set(_cortex_m_kernels__srcs
+    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantize_per_tensor.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_dequantize_per_tensor.cpp
+)
+
+# Generate C++ bindings to register kernels into Executorch (for runtime).
+# Here select all ops in operators.yaml
+set(_yaml_file ${CMAKE_CURRENT_LIST_DIR}/ops/operators.yaml)
+gen_selected_ops(LIB_NAME "cortex_m_ops_lib" OPS_SCHEMA_YAML "${_yaml_file}")
+
+# Generate bindings for the kernels
+generate_bindings_for_kernels(
+  LIB_NAME "cortex_m_ops_lib" CUSTOM_OPS_YAML "${_yaml_file}"
+)
+message("Generated files ${gen_command_sources}")
+
+# Build a library for _cortex_m_kernels_srcs
+add_library(cortex_m_kernels ${_cortex_m_kernels__srcs})
+target_link_libraries(cortex_m_kernels PRIVATE executorch)
+target_compile_options(cortex_m_kernels PUBLIC ${_common_compile_options})
+
+# cortex_m_ops_lib: Register Cortex-M ops kernels into Executorch runtime
+gen_operators_lib(
+  LIB_NAME "cortex_m_ops_lib" KERNEL_LIBS cortex_m_kernels DEPS executorch
+)
+
+install(
+  TARGETS cortex_m_kernels cortex_m_ops_lib
+  DESTINATION lib
+  PUBLIC_HEADER DESTINATION include/executorch/backends/cortex_m/ops/
+)

backends/cortex_m/ops/op_dequantize_per_tensor.cpp

@@ -29,6 +29,7 @@ namespace {
  */
 void check_dequantize_args(
     const Tensor& input,
+    int64_t zero_point,
     int64_t quant_min,
     int64_t quant_max,
     ScalarType dtype,
@@ -39,6 +40,18 @@ void check_dequantize_args(
       "input.scalar_type() %" PRId8 " is not char type",
       static_cast<int8_t>(input.scalar_type()));
 
+  // Check zp range
+  ET_CHECK_MSG(
+      zero_point >= quant_min,
+      "zero_point must be %" PRId64 " <= quant_min %" PRId64,
+      zero_point,
+      quant_min);
+  ET_CHECK_MSG(
+      zero_point <= quant_max,
+      "zero_point must be %" PRId64 " >= quant_max %" PRId64,
+      zero_point,
+      quant_max);
+
   // Check output dtype is float
   ET_CHECK_MSG(
       out.scalar_type() == ScalarType::Float,
@@ -73,18 +86,10 @@ void check_dequantize_args(
 /**
  * Scalar implementation of quantization for a single value.
  */
-template <typename K, typename T>
-T dequantize_val(
-    float scale,
-    int32_t zero_point,
-    K value,
-    int64_t quant_min,
-    int64_t quant_max) {
-  (void)quant_min;
-  (void)quant_max;
-  return static_cast<T>((static_cast<int32_t>(value) - zero_point) * scale);
+template <typename Q, typename F>
+F dequantize_val(float scale, int32_t zero_point, Q qvalue) {
+  return static_cast<F>((static_cast<int32_t>(qvalue) - zero_point) * scale);
 }
-
 } // namespace
 
 Tensor& dequantize_per_tensor_out(
@@ -106,29 +111,71 @@ Tensor& dequantize_per_tensor_out(
       "Failed to resize out Tensor in dequantize_per_tensor_out");
 
   // Validate input parameters
-  check_dequantize_args(input, quant_min, quant_max, dtype, out);
+  check_dequantize_args(input, zero_point, quant_min, quant_max, dtype, out);
 
-  // Pre-compute inverse scale for better performance
   int32_t zp = static_cast<int32_t>(zero_point);
-  int32_t qmin = static_cast<int32_t>(quant_min);
-  int32_t qmax = static_cast<int32_t>(quant_max);
 
   // Get pointers to input and output data
   const int8_t* input_data = input.const_data_ptr<int8_t>();
   float* out_data = out.mutable_data_ptr<float>();
   const size_t numel = input.numel();
 
+  size_t i = 0;
 #if defined(HAS_HELIUM_SIMD)
-// Helium MVE implementation for float32 to int8 quantization
-#Error "Implement MVE version!"
-#else
-  // Scalar implementation for float32 to int8 quantization
-  for (size_t i = 0; i < numel; i++) {
-    out_data[i] =
-        dequantize_val<int8_t, float>(scale, zp, input_data[i], qmin, qmax);
+  // Helium MVE implementation for int8 to float quantization
+  static uint8x16_t voffset{
+      0x0,
+      0x8,
+      0x4,
+      0xC,
+      0x1,
+      0x9,
+      0x5,
+      0xD,
+      0x2,
+      0xA,
+      0x6,
+      0xE,
+      0x3,
+      0xB,
+      0x7,
+      0xF};
+
+  int16x8_t vzp = vdupq_n_s16(static_cast<int16_t>(zp));
+  float32x4_t vscale = vdupq_n_f32(static_cast<float>(scale));
+
+  for (; i + 15 < numel; i += 16) {
+    int8x16_t in_084C195D2A6E3B7F =
+        vldrbq_gather_offset_s8(input_data, voffset);
+
+    int16x8_t in_04152637 = vsubq_s16(vmovlbq_s8(in_084C195D2A6E3B7F), vzp);
+    int16x8_t in_8C9DAEBF = vsubq_s16(vmovltq_s8(in_084C195D2A6E3B7F), vzp);
+
+    float32x4_t inf_0123 = vcvtq_f32_s32(vmovlbq_s16(in_04152637));
+    float32x4_t inf_4567 = vcvtq_f32_s32(vmovltq_s16(in_04152637));
+    float32x4_t inf_89AB = vcvtq_f32_s32(vmovlbq_s16(in_8C9DAEBF));
+    float32x4_t inf_CDEF = vcvtq_f32_s32(vmovltq_s16(in_8C9DAEBF));
+
+    float32x4_t out_0123 = vmulq_f32(inf_0123, vscale);
+    float32x4_t out_4567 = vmulq_f32(inf_4567, vscale);
+    float32x4_t out_89AB = vmulq_f32(inf_89AB, vscale);
+    float32x4_t out_CDEF = vmulq_f32(inf_CDEF, vscale);
+
+    vstrwq_f32(out_data + 0, out_0123);
+    vstrwq_f32(out_data + 4, out_4567);
+    vstrwq_f32(out_data + 8, out_89AB);
+    vstrwq_f32(out_data + 12, out_CDEF);
+
+    input_data += 16;
+    out_data += 16;
   }
-#endif
+#endif // defined(HAS_HELIUM_SIMD)
 
+  for (; i < numel; i++) {
+    *out_data = dequantize_val<int8_t, float>(scale, zp, *input_data);
+    input_data++;
+    out_data++;
+  }
   return out;
 }