CMake support for quantized MoE optimized build

examples/models/llama/README.md

@@ -561,15 +561,14 @@ registered in `executorch.extension.llm.custom_ops.custom_ops`.
 
 The runtime kernel ships in `extension/llm/custom_ops/op_moe.cpp`. It
 always compiles with a portable reference fallback (unpack + dequant +
-`cpublas::gemm`) that works on any platform. `ENABLE_QUANTIZED_MOE_FFN`
-is an **optimization gate**, not a correctness requirement — when
-defined, the kernel uses torchao's fused `linear_operator` (NEON
-i8mm/dotprod on aarch64) instead of the reference path.
+`cpublas::gemm`) that works on any platform. The optimized build option
+uses torchao's fused `linear_operator` (NEON i8mm/dotprod on aarch64)
+instead of the reference path.
 
 In CMake, opt in to the optimized path with:
 
 ```cmake
--DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON
+-DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON
 ```
 
 In Buck, `_get_quantized_moe_deps()` in `targets.bzl` wires:

examples/models/llama/export_llama_lib.py

@@ -591,8 +591,8 @@ def build_args_parser() -> argparse.ArgumentParser:
             "Replace eager MoE feed-forward modules with the "
             "`llama::quantized_moe_ffn` portable-runtime custom op (INT4 "
             "weights, INT8 dyn-quant activations via torchao). On aarch64 "
-            "with ENABLE_QUANTIZED_MOE_FFN the optimized torchao kernel is "
-            "used; otherwise a portable reference fallback runs."
+            "an optimized runtime build uses the torchao kernel; otherwise "
+            "a portable reference fallback runs."
         ),
     )
 

extension/llm/custom_ops/CMakeLists.txt

@@ -20,7 +20,7 @@ set(_common_compile_options
     $<$<CXX_COMPILER_ID:MSVC>:/wd4996>
     $<$<NOT:$<CXX_COMPILER_ID:MSVC>>:-Wno-deprecated-declarations -fPIC>
 )
-if(CMAKE_SYSTEM_PROCESSOR MATCHES "arm64|aarch64")
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64)$")
   list(APPEND _common_compile_options
        "$<$<NOT:$<CXX_COMPILER_ID:MSVC>>:-march=armv8.2-a+dotprod>"
   )
@@ -86,28 +86,72 @@ target_link_libraries(custom_ops PUBLIC ${custom_ops_libs} executorch_core)
 # The MoE kernel always compiles with a reference fallback (unpack + dequant +
 # cpublas::gemm) using the torchao weight_packing headers from third-party/ao
 # (already on the include path).  Pass
-# -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON to additionally link the
-# optimized torchao linear_operator (fused INT8-dyn-act GEMM, aarch64 NEON
+# -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON to additionally link
+# the optimized torchao linear_operator (fused INT8-dyn-act GEMM, aarch64 NEON
 # dotprod).
-option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE
+option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED
        "Link the optimized torchao linear kernel for llama::quantized_moe_ffn"
        OFF
 )
-if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE)
-  if(NOT TARGET torchao_ops_linear_8bit_act_xbit_weight_executorch)
+
+function(_target_enable_quantized_moe_torchao target)
+  target_compile_definitions(
+    ${target}
+    PRIVATE TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_ENABLE_ARM_NEON_DOT=1
+            TORCHAO_PARALLEL_EXECUTORCH=1
+            TORCHAO_SHARED_KERNELS_BUILD_EXECUTORCH=1
+  )
+endfunction()
+
+if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED)
+  if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
     message(
       FATAL_ERROR
-        "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON requires target "
-        "torchao_ops_linear_8bit_act_xbit_weight_executorch, which is not "
-        "defined. Build the torchao ops or set this option OFF."
+        "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON is supported only on "
+        "aarch64 and arm64."
     )
   endif()
-  # Compile definition and link must be gated on the same condition, else the
-  # ENABLE_QUANTIZED_MOE_FFN path compiles without the library that defines it.
-  target_compile_definitions(custom_ops PUBLIC ENABLE_QUANTIZED_MOE_FFN=1)
-  target_link_libraries(
-    custom_ops PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch
+
+  # The selected target is also reused by custom_ops_aot_lib below.
+  if(EXECUTORCH_BUILD_KERNELS_TORCHAO)
+    if(NOT TARGET torchao_ops_executorch)
+      message(FATAL_ERROR "EXECUTORCH_BUILD_KERNELS_TORCHAO=ON but target "
+                          "torchao_ops_executorch is not defined."
+      )
+    endif()
+    set(quantized_moe_torchao_target torchao_ops_executorch)
+  else()
+    set(quantized_moe_torchao_target torchao_moe_linear)
+    add_library(
+      torchao_moe_linear STATIC
+      ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/linear_8bit_act_xbit_weight.cpp
+      ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/quantization/quantize.cpp
+      ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/reduction/compute_sum.cpp
+      ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/reduction/find_min_and_max.cpp
+      ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/valpacking/interleave.cpp
+    )
+    target_include_directories(
+      torchao_moe_linear PRIVATE ${EXECUTORCH_ROOT}/third-party/ao
+    )
+    _target_enable_quantized_moe_torchao(torchao_moe_linear)
+    target_link_libraries(
+      torchao_moe_linear PRIVATE cpuinfo executorch_core extension_threadpool
+    )
+    target_compile_options(
+      torchao_moe_linear PRIVATE ${_common_compile_options}
+    )
+    install(
+      TARGETS torchao_moe_linear
+      EXPORT ExecuTorchTargets
+      DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    )
+  endif()
+
+  target_compile_definitions(
+    custom_ops PRIVATE EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO=1
   )
+  _target_enable_quantized_moe_torchao(custom_ops)
+  target_link_libraries(custom_ops PUBLIC ${quantized_moe_torchao_target})
 endif()
 
 target_compile_options(custom_ops PUBLIC ${_common_compile_options})
@@ -190,21 +234,13 @@ if(EXECUTORCH_BUILD_KERNELS_LLM_AOT)
     custom_ops_aot_lib PUBLIC cpublas torch extension_tensor
                               extension_threadpool
   )
-  if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE)
-    if(NOT TARGET torchao_ops_linear_8bit_act_xbit_weight_executorch)
-      message(
-        FATAL_ERROR
-          "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON requires target "
-          "torchao_ops_linear_8bit_act_xbit_weight_executorch, which is not "
-          "defined. Build the torchao ops or set this option OFF."
-      )
-    endif()
+  if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED)
     target_compile_definitions(
-      custom_ops_aot_lib PUBLIC ENABLE_QUANTIZED_MOE_FFN=1
+      custom_ops_aot_lib PRIVATE EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO=1
     )
+    _target_enable_quantized_moe_torchao(custom_ops_aot_lib)
     target_link_libraries(
-      custom_ops_aot_lib
-      PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch
+      custom_ops_aot_lib PUBLIC ${quantized_moe_torchao_target}
     )
   endif()
   if(WIN32)

extension/llm/custom_ops/op_moe.cpp

@@ -9,19 +9,20 @@
 #include <executorch/extension/llm/custom_ops/op_moe.h>
 
 #include <executorch/extension/kernel_util/make_boxed_from_unboxed_functor.h>
-#include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/kernels/optimized/blas/CPUBlas.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
-#include <executorch/runtime/kernel/thread_parallel_interface.h>
 
 #include <torchao/csrc/cpu/shared_kernels/internal/packed_weights_header.h>
 #include <torchao/csrc/cpu/torch_free_kernels/weight_packing/weight_packing.h>
 
-#ifdef ENABLE_QUANTIZED_MOE_FFN
-#include <torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/kernel_selector.h>
+#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
+#include <cpuinfo.h>
+#include <torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/kernel_config.h>
 #include <torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/linear_8bit_act_xbit_weight.h>
+#include <torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/packed_weights_format.h>
+#include <torchao/csrc/cpu/torch_free_kernels/aarch64/linear/channelwise_8bit_activation_groupwise_lowbit_weight/channelwise_8bit_activation_groupwise_lowbit_weight.h>
 #include <optional> // std::nullopt, used only by the optimized aarch64 path
-#endif // ENABLE_QUANTIZED_MOE_FFN
+#endif // EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
 
 #include <array>
 #include <cmath>
@@ -37,6 +38,49 @@ namespace {
 
 using ::executorch::aten::string_view;
 
+#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
+template <int kWeightNbit>
+const torchao::ops::linear_8bit_act_xbit_weight::UKernelConfig&
+universal_ukernel_config() {
+  using torchao::ops::linear_8bit_act_xbit_weight::UKernelConfig;
+  namespace kernel = torchao::kernels::cpu::aarch64::linear::
+      channelwise_8bit_activation_groupwise_lowbit_weight;
+
+  static const auto config = [] {
+    ET_CHECK_MSG(
+        cpuinfo_initialize() && cpuinfo_has_arm_neon_dot(),
+        "quantized_moe_ffn optimized path requires Arm NEON dot product");
+    auto result = UKernelConfig::make(
+        /*preferred_alignment=*/16,
+        /*n_step=*/8,
+        /*nr=*/8,
+        /*kr=*/16,
+        /*sr=*/2,
+        kWeightNbit,
+        /*has_weight_zeros=*/false,
+        /*has_bias=*/false,
+        &torchao::weight_packing::packed_weights_size,
+        &torchao::weight_packing::packed_weights_offset,
+        &torchao::weight_packing::pack_weights<kWeightNbit, 8, 16, 2>,
+        {});
+    result.linear_configs[0] = UKernelConfig::linear_config_type({
+        /*m_step=*/1,
+        /*mr=*/1,
+        &kernel::packed_activations_size,
+        &kernel::packed_activations_offset,
+        &kernel::pack_activations<1, 16, 2>,
+        &kernel::kernel_1x8x16_f32_neondot<
+            kWeightNbit,
+            /*has_weight_zeros=*/false,
+            /*has_lut=*/false>,
+    });
+    result.validate();
+    return result;
+  }();
+  return config;
+}
+#endif
+
 // Numerically-stable sigmoid. Branching on sign keeps exp()'s argument
 // non-positive on both sides, so it can never overflow.
 inline float stable_sigmoid(float v) {
@@ -189,7 +233,7 @@ inline void reference_linear(
 
 // Dispatch a single per-expert grouped GEMM through torchao's
 // linear_operator (optimized, aarch64) or reference unpack+dequant+gemm.
-#ifdef ENABLE_QUANTIZED_MOE_FFN
+#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
 template <int kWeightNbit>
 inline void torchao_linear(
     const uint8_t* packed_w_blob,
@@ -205,21 +249,22 @@ inline void torchao_linear(
           static_cast<int64_t>(torchao::ops::PackedWeightsHeader::size()),
       "torchao packed blob too small to contain header");
   auto header = torchao::ops::PackedWeightsHeader::read(packed_w_blob);
-  // Select the ukernel from the format declared in the header. This resolves
-  // the universal or kleidi packing automatically; a format whose kernels are
-  // not compiled into this build (e.g. kleidi when TORCHAO_ENABLE_KLEIDI is
-  // unset) throws here instead of being silently mis-read.
-  // TODO: enable KleidiAI here — build this op on xplat arm64 with
-  // -DTORCHAO_ENABLE_KLEIDI (+ -DTORCHAO_ENABLE_ARM_I8MM) and link the kleidi
-  // kernel target so a kleidi header actually resolves to a kleidi ukernel.
-  // Must be coordinated with the AoT packer emitting kleidi headers (see
-  // targets.bzl).
-  auto uk = torchao::ops::linear_8bit_act_xbit_weight::select_ukernel_config<
-      kWeightNbit>(header);
-
-  // Validate the blob against the *selected* format's layout. nr/kr/sr and the
-  // size formula differ between universal and kleidi, so derive them from the
-  // chosen config rather than assuming a fixed layout.
+  ET_CHECK_MSG(
+      header.type ==
+          torchao::ops::PackedWeightsType::
+              linear_8bit_act_xbit_weight_universal,
+      "quantized_moe_ffn requires universal torchao packed weights");
+  const auto format = torchao::ops::linear_8bit_act_xbit_weight::
+      PackedWeightsFormat::from_packed_weights_header(header);
+  ET_CHECK_MSG(
+      format.weight_nbit == kWeightNbit && !format.has_weight_zeros &&
+          !format.has_bias && format.nr == 8 && format.kr == 16 &&
+          format.sr == 2,
+      "quantized_moe_ffn received an unsupported universal weight format");
+  const auto& uk = universal_ukernel_config<kWeightNbit>();
+
+  // Validate the blob against the universal config's layout without
+  // duplicating its packed-weight size formula.
   const int64_t required_bytes =
       static_cast<int64_t>(torchao::ops::PackedWeightsHeader::size()) +
       static_cast<int64_t>(uk.packed_weights_size(
@@ -258,7 +303,7 @@ inline void torchao_linear(
       /*clamp_min=*/0.0f,
       /*clamp_max=*/0.0f);
 }
-#endif // ENABLE_QUANTIZED_MOE_FFN
+#endif // EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
 
 inline void expert_linear_dispatch(
     int64_t weight_nbit,
@@ -270,12 +315,11 @@ inline void expert_linear_dispatch(
     int64_t k,
     int64_t group_size,
     float* out) {
-#ifndef ENABLE_QUANTIZED_MOE_FFN
+#ifndef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
   // Reference path only: it unpacks the universal layout, so validate the blob
   // holds the header plus the universal packed weight-data bytes for the
   // claimed dims before any path dereferences it. The torchao path validates
-  // against its own selected format (universal or kleidi) inside
-  // torchao_linear.
+  // the same required universal format inside torchao_linear.
   constexpr int kNr = 8, kKr = 16, kSr = 2;
   const int64_t required_bytes =
       static_cast<int64_t>(torchao::ops::PackedWeightsHeader::size()) +
@@ -299,10 +343,10 @@ inline void expert_linear_dispatch(
       static_cast<long long>(k),
       static_cast<long long>(group_size),
       static_cast<long long>(weight_nbit));
-#endif // !ENABLE_QUANTIZED_MOE_FFN
+#endif // !EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
   switch (weight_nbit) {
     case 4:
-#ifdef ENABLE_QUANTIZED_MOE_FFN
+#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
       torchao_linear<4>(
           packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out);
 #else
@@ -311,7 +355,7 @@ inline void expert_linear_dispatch(
 #endif
       return;
     case 8:
-#ifdef ENABLE_QUANTIZED_MOE_FFN
+#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO
       torchao_linear<8>(
           packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out);
 #else
@@ -634,26 +678,7 @@ Tensor& quantized_moe_ffn_out(
     }
   };
 
-#ifdef ENABLE_QUANTIZED_MOE_FFN
-  // torchao linear path (perf-sensitive). The kernel threads internally on the
-  // shared pool in the common config, or runs single-threaded when only the
-  // thread-pool-free variant is linked. Distribute experts across the pool
-  // ourselves only when the kernel won't and the pool has more than one thread
-  // -- running both would nest on one pthreadpool and deadlock.
-  const bool parallelize_experts = torchao::ops::linear_8bit_act_xbit_weight::
-                                       linear_operator_num_threads() == 1 &&
-      ::executorch::extension::threadpool::get_threadpool()
-              ->get_thread_count() > 1;
-#else
-  // Portable reference path: prefer simplicity over speed and run experts
-  // serially.
-  const bool parallelize_experts = false;
-#endif
-  if (parallelize_experts) {
-    torch::executor::parallel_for(0, E, /*grain_size=*/1, run_experts);
-  } else {
-    run_experts(0, E);
-  }
+  run_experts(0, E);
 
   // ----- 7. Weighted scatter-add unpermute (cross-expert reduction) -----
   // Each token sums the contributions of its top-k experts; run serially to

extension/llm/custom_ops/targets.bzl

@@ -53,17 +53,13 @@ def _get_quantized_moe_preproc_flags():
     if runtime.is_oss:
         return []
     if is_xplat():
-        # TODO: enable KleidiAI for the runtime here by adding
-        # -DTORCHAO_ENABLE_KLEIDI (+ -DTORCHAO_ENABLE_ARM_I8MM=1) on arm64 and
-        # linking the kleidi kernel target in _get_quantized_moe_deps(). The
-        # runtime (op_moe.cpp) already selects the ukernel from the header, so
-        # kleidi headers resolve automatically once the kernels are compiled.
-        # Must be paired with the AoT packer emitting kleidi headers (see
-        # _get_quantized_moe_aot_packer_deps()).
+        # TODO: enable KleidiAI by adding its runtime config to op_moe.cpp,
+        # compiling and linking its kernels here, and pairing it with an AoT
+        # packer that emits Kleidi headers.
         return select({
             "DEFAULT": [],
             "ovr_config//cpu:arm64": [
-                "-DENABLE_QUANTIZED_MOE_FFN",
+                "-DEXECUTORCH_QUANTIZED_MOE_USE_TORCHAO",
                 "-DTORCHAO_BUILD_CPU_AARCH64=1",
                 "-DTORCHAO_ENABLE_ARM_NEON_DOT=1",
             ],

extension/llm/custom_ops/test_op_moe.cpp

@@ -44,8 +44,8 @@ TEST(OpQuantizedMoeFfnTest, RegistrationSmokeTest) {
   Tensor expert_bias = tff.zeros({0});
 
   // Use empty packed buffers; the kernel will fail loudly if it tries to
-  // dereference them. With ENABLE_QUANTIZED_MOE_FFN unset (CI x86 build
-  // without torchao linkage) the kernel ET_CHECK_MSGs out before doing
+  // dereference them. With EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO unset (CI x86
+  // build without torchao linkage) the kernel ET_CHECK_MSGs out before doing
   // any real work, which is what we want this test to verify.
   Tensor packed_w1 = tfb.zeros({E, 1});
   Tensor packed_w3 = tfb.zeros({E, 1});