Skip to content

[ExecuTorch][llm] Add quantized_moe_ffn custom op with C++ kernel #21119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
digantdesai wants to merge 2 commits into
base: gh/digantdesai/69/base
Choose a base branch
from
+1,024 −3
Open

[ExecuTorch][llm] Add quantized_moe_ffn custom op with C++ kernel #21119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f700b8a
Update
digantdesai Jul 22, 2026
4ec3a2d
Update
digantdesai Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 45 additions & 0 deletions extension/llm/custom_ops/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading