[ExecuTorch][llm] Add quantized_moe_ffn custom op with C++ kernel (pytorch#21533)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#21119 by
@digantdesai
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/digantdesai/69/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/69/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/69/orig
Differential Revision:
[D102382000](https://our.internmc.facebook.com/intern/diff/D102382000/)
@diff-train-skip-merge

Co-authored-by: Digant Desai <digantdesai@meta.com>

Author: digantdesai

extension/llm/custom_ops/CMakeLists.txt

@@ -83,6 +83,33 @@ target_include_directories(
 )
 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
+# dotprod).
+option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE
+       "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)
+    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()
+  # 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
+  )
+endif()
+
 target_compile_options(custom_ops PUBLIC ${_common_compile_options})
 
 install(
@@ -98,6 +125,7 @@ if(EXECUTORCH_BUILD_KERNELS_LLM_AOT)
     custom_ops_aot_lib SHARED
     ${_custom_ops__srcs}
     ${CMAKE_CURRENT_SOURCE_DIR}/op_sdpa_aot.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/op_moe_aot.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_fast_hadamard_transform_aten.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_tile_crop.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_tile_crop_aot.cpp
@@ -134,6 +162,23 @@ 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()
+    target_compile_definitions(
+      custom_ops_aot_lib PUBLIC ENABLE_QUANTIZED_MOE_FFN=1
+    )
+    target_link_libraries(
+      custom_ops_aot_lib
+      PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch
+    )
+  endif()
   if(WIN32)
     # There is no direct replacement for libpthread.so on Windows. For the
     # Windows build, link directly against pthreadpool and cpuinfo.