From f700b8a1ebb029fb9b62a5b99121b003752eff88 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 21 Jul 2026 20:37:49 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- extension/llm/custom_ops/CMakeLists.txt | 45 ++ extension/llm/custom_ops/op_moe.cpp | 616 ++++++++++++++++++ extension/llm/custom_ops/op_moe.h | 80 +++ extension/llm/custom_ops/op_moe_aot.cpp | 121 ++++ extension/llm/custom_ops/targets.bzl | 86 ++- extension/llm/custom_ops/test_op_moe.cpp | 78 +++ .../executorch/build/build_variables.bzl | 1 + 7 files changed, 1024 insertions(+), 3 deletions(-) create mode 100644 extension/llm/custom_ops/op_moe.cpp create mode 100644 extension/llm/custom_ops/op_moe.h create mode 100644 extension/llm/custom_ops/op_moe_aot.cpp create mode 100644 extension/llm/custom_ops/test_op_moe.cpp diff --git a/extension/llm/custom_ops/CMakeLists.txt b/extension/llm/custom_ops/CMakeLists.txt index 2cdfe547430..8a43a5ddf5c 100644 --- a/extension/llm/custom_ops/CMakeLists.txt +++ b/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. diff --git a/extension/llm/custom_ops/op_moe.cpp b/extension/llm/custom_ops/op_moe.cpp new file mode 100644 index 00000000000..6b3233b8390 --- /dev/null +++ b/extension/llm/custom_ops/op_moe.cpp @@ -0,0 +1,616 @@ +/* + * 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. + */ + +#include + +#include +#include +#include + +#include +#include + +#ifdef ENABLE_QUANTIZED_MOE_FFN +#include +#include +#include // std::nullopt, used only by the optimized aarch64 path +#endif // ENABLE_QUANTIZED_MOE_FFN + +#include +#include +#include +#include + +namespace torch { +namespace executor { +namespace native { + +namespace { + +using executorch::aten::string_view; + +// SwiGLU epilogue: out = silu(h1) * h3 (in place over h1). +// silu(x) = x * sigmoid(x) = x / (1 + exp(-x)). +inline void swiglu_inplace(float* h1, const float* h3, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const float v = h1[i]; + const float s = v / (1.0f + std::exp(-v)); + h1[i] = s * h3[i]; + } +} + +// Stable softmax over a row of length E (in place). +inline void softmax_row(float* row, int64_t e) { + float maxv = row[0]; + for (int64_t i = 1; i < e; ++i) { + if (row[i] > maxv) { + maxv = row[i]; + } + } + float sum = 0.0f; + for (int64_t i = 0; i < e; ++i) { + row[i] = std::exp(row[i] - maxv); + sum += row[i]; + } + const float inv_sum = (sum > 0.0f) ? (1.0f / sum) : 0.0f; + for (int64_t i = 0; i < e; ++i) { + row[i] *= inv_sum; + } +} + +// In-place sigmoid over a contiguous buffer. +inline void sigmoid_inplace(float* p, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + p[i] = 1.0f / (1.0f + std::exp(-p[i])); + } +} + +// Simple O(E*K) top-k selection. E and K are both small in MoE models +// (E <= 128, K <= 8), so this is cheaper than partial_sort and needs no +// scratch buffer. Stable across ties via index ordering. +inline void +topk_indices(const float* scores, int64_t e, int64_t k, int32_t* out_indices) { + for (int64_t ki = 0; ki < k; ++ki) { + int32_t best = -1; + float best_score = -std::numeric_limits::infinity(); + for (int64_t ei = 0; ei < e; ++ei) { + bool already_selected = false; + for (int64_t j = 0; j < ki; ++j) { + if (out_indices[j] == static_cast(ei)) { + already_selected = true; + break; + } + } + if (already_selected) { + continue; + } + if (scores[ei] > best_score || + (scores[ei] == best_score && (best < 0 || ei < best))) { + best_score = scores[ei]; + best = static_cast(ei); + } + } + // Non-finite (NaN) scores compare false against everything, so `best` can + // remain -1 when fewer than `k` finite candidates are left. Writing -1 + // would later index cursor[-1]/expert_offsets out of bounds, so fall back + // to the lowest-indexed unselected expert (num_activated <= num_experts + // guarantees one exists) to keep the selected id valid. + if (best < 0) { + for (int64_t ei = 0; ei < e; ++ei) { + bool already_selected = false; + for (int64_t j = 0; j < ki; ++j) { + if (out_indices[j] == static_cast(ei)) { + already_selected = true; + break; + } + } + if (!already_selected) { + best = static_cast(ei); + break; + } + } + } + out_indices[ki] = best; + } +} + +// Reference linear: unpack torchao blob → dequantize → cpublas::gemm. +// TODO: move to torchao? +template +inline void reference_linear( + const uint8_t* packed_w_blob, + int64_t packed_blob_bytes, + const float* x, + int64_t m, + int64_t n, + int64_t k, + int64_t group_size, + float* out) { + constexpr int kNr = 8, kKr = 16, kSr = 2; + const auto header_size = + static_cast(torchao::ops::PackedWeightsHeader::size()); + ET_CHECK_MSG( + packed_blob_bytes >= header_size, + "torchao packed blob too small to contain header"); + + const void* packed_data = packed_w_blob + header_size; + const auto n_int = static_cast(n); + const auto k_int = static_cast(k); + const auto gs_int = static_cast(group_size); + + std::vector qvals(static_cast(n * k)); + std::vector scales(static_cast(n * (k / group_size))); + torchao::weight_packing::unpack_weights( + qvals.data(), + scales.data(), + /*weight_zeros=*/nullptr, + /*bias=*/nullptr, + n_int, + k_int, + gs_int, + /*has_weight_zeros=*/false, + /*has_bias=*/false, + packed_data); + + std::vector w_fp32(static_cast(n * k)); + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t ki = 0; ki < k; ++ki) { + const int64_t group_idx = ni * (k / group_size) + ki / group_size; + w_fp32[static_cast(ni * k + ki)] = + static_cast(qvals[static_cast(ni * k + ki)]) * + scales[static_cast(group_idx)]; + } + } + + // w_fp32 is [N, K] row-major. We need out = x [M, K] @ w^T [K, N]. + ::executorch::cpublas::gemm( + ::executorch::cpublas::TransposeType::Transpose, + ::executorch::cpublas::TransposeType::NoTranspose, + /*m=*/n, + /*n=*/m, + /*k=*/k, + /*alpha=*/1.0f, + /*a=*/w_fp32.data(), + /*lda=*/k, + /*b=*/x, + /*ldb=*/k, + /*beta=*/0.0f, + /*c=*/out, + /*ldc=*/n); +} + +// Dispatch a single per-expert grouped GEMM through torchao's +// linear_operator (optimized, aarch64) or reference unpack+dequant+gemm. +#ifdef ENABLE_QUANTIZED_MOE_FFN +template +inline void torchao_linear( + const uint8_t* packed_w_blob, + int64_t packed_blob_bytes, + const float* x, + int64_t m, + int64_t n, + int64_t k, + int64_t group_size, + float* out) { + ET_CHECK_MSG( + packed_blob_bytes >= + static_cast(torchao::ops::PackedWeightsHeader::size()), + "torchao packed blob too small to contain header"); + auto header = torchao::ops::PackedWeightsHeader::read(packed_w_blob); + auto uk = torchao::ops::linear_8bit_act_xbit_weight::select_ukernel_config< + kWeightNbit>(header); + + torchao::ops::linear_8bit_act_xbit_weight::linear_operator( + uk, + /*tiling_params=*/std::nullopt, + /*output=*/out, + /*m=*/static_cast(m), + /*n=*/static_cast(n), + /*k=*/static_cast(k), + /*group_size=*/static_cast(group_size), + /*packed_weights=*/packed_w_blob + + torchao::ops::PackedWeightsHeader::size(), + /*activations=*/x, + /*has_clamp=*/false, + /*clamp_min=*/0.0f, + /*clamp_max=*/0.0f); +} +#endif // ENABLE_QUANTIZED_MOE_FFN + +inline void expert_linear_dispatch( + int64_t weight_nbit, + const uint8_t* packed_w_blob, + int64_t packed_blob_bytes, + const float* x, + int64_t m, + int64_t n, + int64_t k, + int64_t group_size, + float* out) { + // Validate the blob holds at least the header plus the torchao packed + // weight-data bytes for the claimed dims before any path dereferences it. + constexpr int kNr = 8, kKr = 16, kSr = 2; + const int64_t required_bytes = + static_cast(torchao::ops::PackedWeightsHeader::size()) + + static_cast(torchao::weight_packing::packed_weights_size( + static_cast(n), + static_cast(k), + static_cast(group_size), + static_cast(weight_nbit), + /*has_weight_zeros=*/false, + /*has_bias=*/false, + kNr, + kKr, + kSr)); + ET_CHECK_MSG( + packed_blob_bytes >= required_bytes, + "torchao packed blob too small: have %lld bytes, need >= %lld for " + "(n=%lld, k=%lld, group_size=%lld, weight_nbit=%lld)", + static_cast(packed_blob_bytes), + static_cast(required_bytes), + static_cast(n), + static_cast(k), + static_cast(group_size), + static_cast(weight_nbit)); + switch (weight_nbit) { + case 4: +#ifdef ENABLE_QUANTIZED_MOE_FFN + torchao_linear<4>( + packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); +#else + reference_linear<4>( + packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); +#endif + return; + case 8: +#ifdef ENABLE_QUANTIZED_MOE_FFN + torchao_linear<8>( + packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); +#else + reference_linear<8>( + packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); +#endif + return; + default: + ET_CHECK_MSG( + false, + "quantized_moe_ffn: unsupported weight_nbit=%lld", + static_cast(weight_nbit)); + } +} + +} // namespace + +Tensor& quantized_moe_ffn_out( + KernelRuntimeContext& ctx, + const Tensor& x, + const Tensor& gate_weight, + const Tensor& expert_bias, + const Tensor& packed_w1, + const Tensor& packed_w3, + const Tensor& packed_w2, + int64_t num_activated_experts, + int64_t num_experts, + int64_t hidden_dim, + int64_t dim, + int64_t group_size, + int64_t weight_nbit, + string_view score_func, + double route_scale, + Tensor& out) { + (void)ctx; + + // ----- Shape & dtype checks ----- + ET_CHECK_MSG(x.dim() == 2, "x must be 2D [T, D]"); + ET_CHECK_MSG( + x.size(1) == dim, + "x last dim must equal dim (got %lld vs %lld)", + static_cast(x.size(1)), + static_cast(dim)); + ET_CHECK_MSG(x.scalar_type() == ScalarType::Float, "x must be fp32"); + + ET_CHECK_MSG( + gate_weight.dim() == 2 && gate_weight.size(0) == num_experts && + gate_weight.size(1) == dim, + "gate_weight must be [E, D]"); + ET_CHECK_MSG( + gate_weight.scalar_type() == ScalarType::Float, + "gate_weight must be fp32"); + + const bool use_expert_bias = expert_bias.numel() > 0; + if (use_expert_bias) { + ET_CHECK_MSG( + expert_bias.dim() == 1 && expert_bias.size(0) == num_experts, + "expert_bias must be [E] when provided"); + ET_CHECK_MSG( + expert_bias.scalar_type() == ScalarType::Float, + "expert_bias must be fp32"); + } + + ET_CHECK_MSG( + packed_w1.dim() == 2 && packed_w1.size(0) == num_experts, + "packed_w1 must be [E, packed_bytes]"); + ET_CHECK_MSG( + packed_w3.dim() == 2 && packed_w3.size(0) == num_experts, + "packed_w3 must be [E, packed_bytes]"); + ET_CHECK_MSG( + packed_w2.dim() == 2 && packed_w2.size(0) == num_experts, + "packed_w2 must be [E, packed_bytes]"); + ET_CHECK_MSG( + packed_w1.scalar_type() == ScalarType::Byte, "packed_w1 must be uint8"); + ET_CHECK_MSG( + packed_w3.scalar_type() == ScalarType::Byte, "packed_w3 must be uint8"); + ET_CHECK_MSG( + packed_w2.scalar_type() == ScalarType::Byte, "packed_w2 must be uint8"); + ET_CHECK_MSG( + packed_w1.size(1) == packed_w3.size(1), + "packed_w1 and packed_w3 per-expert blob sizes must match"); + + ET_CHECK_MSG( + num_activated_experts > 0 && num_activated_experts <= num_experts, + "num_activated_experts out of range"); + + ET_CHECK_MSG( + score_func == "sigmoid" || score_func == "softmax", + "score_func must be \"sigmoid\" or \"softmax\""); + // expert_bias only shifts sigmoid top-k selection; the softmax path never + // reads it, so reject the combination instead of silently ignoring it. + ET_CHECK_MSG( + !(use_expert_bias && score_func == "softmax"), + "expert_bias is only supported with score_func=\"sigmoid\""); + + ET_CHECK_MSG(group_size > 0, "group_size must be positive"); + ET_CHECK_MSG(dim > 0, "dim must be positive"); + ET_CHECK_MSG(hidden_dim > 0, "hidden_dim must be positive"); + ET_CHECK_MSG( + dim % group_size == 0, + "dim (%lld) must be divisible by group_size (%lld)", + static_cast(dim), + static_cast(group_size)); + ET_CHECK_MSG( + hidden_dim % group_size == 0, + "hidden_dim (%lld) must be divisible by group_size (%lld)", + static_cast(hidden_dim), + static_cast(group_size)); + + const int64_t T = x.size(0); + const int64_t D = dim; + const int64_t F = hidden_dim; + const int64_t E = num_experts; + const int64_t K = num_activated_experts; + const int64_t pw1_bytes = packed_w1.size(1); + const int64_t pw3_bytes = packed_w3.size(1); + const int64_t pw2_bytes = packed_w2.size(1); + + // Resize the output to [T, D]. + ET_KERNEL_CHECK_MSG( + ctx, + out.scalar_type() == ScalarType::Float, + InvalidArgument, + out, + "output must be fp32"); + Tensor::SizesType expected_out_sizes[2] = { + static_cast(T), static_cast(D)}; + ET_KERNEL_CHECK( + ctx, + resize_tensor(out, {expected_out_sizes, 2}) == Error::Ok, + InvalidArgument, + out); + + const float* x_ptr = x.const_data_ptr(); + const float* gate_w_ptr = gate_weight.const_data_ptr(); + const float* expert_bias_ptr = + use_expert_bias ? expert_bias.const_data_ptr() : nullptr; + const uint8_t* pw1_ptr = packed_w1.const_data_ptr(); + const uint8_t* pw3_ptr = packed_w3.const_data_ptr(); + const uint8_t* pw2_ptr = packed_w2.const_data_ptr(); + float* out_ptr = out.mutable_data_ptr(); + + // Zero the output: we'll do a weighted scatter-add into it at the end. + std::memset(out_ptr, 0, sizeof(float) * static_cast(T * D)); + + // ----- 1. Router GEMM: scores [T, E] = x [T, D] @ gate_weight^T ----- + // + // Row-major equivalent of `cpublas::gemm` arguments: see the + // op_sdpa_impl.h comment about column-major BLAS interpretation. With + // `transa=Transpose, transb=NoTranspose, m=E, n=T, k=D`, the call + // produces `scores [T, E]` row-major. + std::vector scores(static_cast(T * E), 0.0f); + ::executorch::cpublas::gemm( + ::executorch::cpublas::TransposeType::Transpose, // transa: gate_weight + ::executorch::cpublas::TransposeType::NoTranspose, // transb: x + /*m=*/E, + /*n=*/T, + /*k=*/D, + /*alpha=*/1.0f, + /*a=*/gate_w_ptr, + /*lda=*/D, + /*b=*/x_ptr, + /*ldb=*/D, + /*beta=*/0.0f, + /*c=*/scores.data(), + /*ldc=*/E); + + // ----- 2/3. Score gating + top-k ----- + // + // We allocate per-token expert indices [T, K] and per-token unbiased + // routing weights [T, K]. For sigmoid: scores are sigmoid'd in place, + // top-k selection uses (scores + bias) but the gathered weights come + // from the un-biased scores, then are renormalized by row-sum and + // multiplied by route_scale. For softmax: top-k of raw scores, then + // softmax over the k gathered values. + std::vector expert_indices(static_cast(T * K), 0); + std::vector expert_weights(static_cast(T * K), 0.0f); + + const bool is_sigmoid = (score_func == "sigmoid"); + if (is_sigmoid) { + sigmoid_inplace(scores.data(), T * E); + } + std::vector scores_for_topk; + if (is_sigmoid && use_expert_bias) { + scores_for_topk.assign(scores.begin(), scores.end()); + for (int64_t t = 0; t < T; ++t) { + float* row = scores_for_topk.data() + t * E; + for (int64_t e = 0; e < E; ++e) { + row[e] += expert_bias_ptr[e]; + } + } + } + const float* topk_src = + scores_for_topk.empty() ? scores.data() : scores_for_topk.data(); + + for (int64_t t = 0; t < T; ++t) { + int32_t* idx_row = expert_indices.data() + t * K; + float* w_row = expert_weights.data() + t * K; + + topk_indices(topk_src + t * E, E, K, idx_row); + + if (is_sigmoid) { + // Gather UN-biased sigmoid scores, then renormalize by row-sum and + // scale by route_scale. + float row_sum = 0.0f; + for (int64_t k = 0; k < K; ++k) { + const float v = scores[t * E + idx_row[k]]; + w_row[k] = v; + row_sum += v; + } + const float scale = static_cast(route_scale) / (row_sum + 1e-20f); + for (int64_t k = 0; k < K; ++k) { + w_row[k] *= scale; + } + } else { + // Softmax path: gather pre-softmax scores at the top-k indices, + // then softmax the K-vector. + for (int64_t k = 0; k < K; ++k) { + w_row[k] = scores[t * E + idx_row[k]]; + } + softmax_row(w_row, K); + } + } + + // ----- 4. Permute (counting sort over flattened [T*K] expert IDs) ----- + std::vector expert_offsets(static_cast(E + 1), 0); + const int64_t total_pairs = T * K; + for (int64_t i = 0; i < total_pairs; ++i) { + ++expert_offsets[expert_indices[i] + 1]; + } + for (int64_t e = 0; e < E; ++e) { + expert_offsets[e + 1] += expert_offsets[e]; + } + // permuted_token_idx[r] = original token index that produced gathered row r. + // permuted_gate[r] = routing weight to apply on the way out. + std::vector permuted_token_idx(static_cast(total_pairs), 0); + std::vector permuted_gate(static_cast(total_pairs), 0.0f); + std::vector cursor(static_cast(E), 0); + for (int64_t e = 0; e < E; ++e) { + cursor[e] = expert_offsets[e]; + } + for (int64_t t = 0; t < T; ++t) { + for (int64_t k = 0; k < K; ++k) { + const int32_t e = expert_indices[t * K + k]; + const int64_t pos = cursor[e]++; + permuted_token_idx[pos] = static_cast(t); + permuted_gate[pos] = expert_weights[t * K + k]; + } + } + + // ----- 5. Gather x_perm [total_pairs, D] ----- + // Temporary buffers scale with T*K*D + max_m_e*(2F+D) floats. + // For typical MoE decode (T=1, K=4, D=768, F=384): ~15 KB total. + std::vector x_perm(static_cast(total_pairs * D), 0.0f); + for (int64_t r = 0; r < total_pairs; ++r) { + const int32_t t = permuted_token_idx[r]; + std::memcpy( + x_perm.data() + r * D, + x_ptr + static_cast(t) * D, + sizeof(float) * static_cast(D)); + } + + // ----- 6. Per-expert grouped GEMM via torchao ----- + // Reusable per-expert scratch (sized for the worst case max_m_e). + int64_t max_m_e = 0; + for (int64_t e = 0; e < E; ++e) { + const int64_t m_e = expert_offsets[e + 1] - expert_offsets[e]; + if (m_e > max_m_e) { + max_m_e = m_e; + } + } + std::vector h1_buf(static_cast(max_m_e * F), 0.0f); + std::vector h3_buf(static_cast(max_m_e * F), 0.0f); + std::vector out_e_buf(static_cast(max_m_e * D), 0.0f); + + for (int64_t e = 0; e < E; ++e) { + const int64_t m_e = expert_offsets[e + 1] - expert_offsets[e]; + if (m_e == 0) { + continue; + } + const float* x_e = x_perm.data() + expert_offsets[e] * D; + const uint8_t* w1_e = pw1_ptr + e * pw1_bytes; + const uint8_t* w3_e = pw3_ptr + e * pw3_bytes; + const uint8_t* w2_e = pw2_ptr + e * pw2_bytes; + + // Up-projection (D -> F): H1 = x_e @ w1[e]^T + expert_linear_dispatch( + weight_nbit, + w1_e, + pw1_bytes, + x_e, + m_e, + F, + D, + group_size, + h1_buf.data()); + // Gate-projection (D -> F): H3 = x_e @ w3[e]^T + expert_linear_dispatch( + weight_nbit, + w3_e, + pw3_bytes, + x_e, + m_e, + F, + D, + group_size, + h3_buf.data()); + // SwiGLU epilogue (in place over h1_buf). + swiglu_inplace(h1_buf.data(), h3_buf.data(), m_e * F); + // Down-projection (F -> D): OUT_e = MID @ w2[e]^T + expert_linear_dispatch( + weight_nbit, + w2_e, + pw2_bytes, + h1_buf.data(), + m_e, + D, + F, + group_size, + out_e_buf.data()); + + // ----- 7. Weighted scatter-add unpermute for this expert's block ----- + for (int64_t r_local = 0; r_local < m_e; ++r_local) { + const int64_t r = expert_offsets[e] + r_local; + const int32_t t = permuted_token_idx[r]; + const float gate = permuted_gate[r]; + float* dst = out_ptr + static_cast(t) * D; + const float* src = out_e_buf.data() + r_local * D; + for (int64_t d = 0; d < D; ++d) { + dst[d] += gate * src[d]; + } + } + } + + return out; +} + +} // namespace native +} // namespace executor +} // namespace torch + +EXECUTORCH_LIBRARY( + llama, + "quantized_moe_ffn.out", + torch::executor::native::quantized_moe_ffn_out); diff --git a/extension/llm/custom_ops/op_moe.h b/extension/llm/custom_ops/op_moe.h new file mode 100644 index 00000000000..64c57c34124 --- /dev/null +++ b/extension/llm/custom_ops/op_moe.h @@ -0,0 +1,80 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace torch { +namespace executor { +namespace native { + +// llama::quantized_moe_ffn.out +// +// Single fused MoE feed-forward op for token-choice routing with +// torchao-packed group-quantized INT4 (or up to INT8) experts and INT8 +// dynamic-quantized activations. +// +// Implements the same math as +// executorch.examples.models.llama.llama_transformer.MOEFeedForward.forward: +// - router GEMM on `gate_weight`, +// - "sigmoid" or "softmax" scoring (with optional `expert_bias` shifting +// top-k selection but not the gathered routing weights), +// - top-k expert selection, +// - permute-then-grouped-GEMM expert evaluation +// (3 GEMMs per active expert via +// torchao::ops::linear_8bit_act_xbit_weight), +// - SwiGLU activation, +// - weighted scatter-add unpermute. +// +// On aarch64, the torchao kernel selector picks the best available +// ukernel (NEON i8mm, dotprod, or scalar) based on runtime CPU +// feature detection. On x86, a reference path unpacks the torchao +// blob, dequantizes, and calls cpublas::gemm. +// +// Inputs: +// x [T, D] fp32 — flattened activations +// gate_weight [E, D] fp32 — router weight (block_sparse_moe.gate.weight) +// expert_bias [E] fp32 — additive bias for top-k selection. May be empty +// (numel == 0) if `moe_gate_bias=False`. +// packed_w1, packed_w3 [E, packed_bytes] uint8 — torchao opaque packed +// weight blobs for the up/gate projections (D -> F per expert). +// packed_w2 [E, packed_bytes] uint8 — torchao packed blob for the down +// projection (F -> D per expert). +// num_activated_experts — top-k. +// num_experts E, hidden_dim F, dim D — shapes. +// group_size — INT4 group size (e.g. 32). +// weight_nbit — quantized weight bit width. INT4 (weight_nbit==4) and +// INT8 (weight_nbit==8) are supported; other values trip a runtime check. +// score_func — "sigmoid" or "softmax". +// route_scale — multiplier applied to the gathered (unbiased) routing +// weights for the sigmoid path. Ignored for softmax. +// +// Output: +// out [T, D] fp32. +Tensor& quantized_moe_ffn_out( + KernelRuntimeContext& ctx, + const Tensor& x, + const Tensor& gate_weight, + const Tensor& expert_bias, + const Tensor& packed_w1, + const Tensor& packed_w3, + const Tensor& packed_w2, + int64_t num_activated_experts, + int64_t num_experts, + int64_t hidden_dim, + int64_t dim, + int64_t group_size, + int64_t weight_nbit, + executorch::aten::string_view score_func, + double route_scale, + Tensor& out); + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/extension/llm/custom_ops/op_moe_aot.cpp b/extension/llm/custom_ops/op_moe_aot.cpp new file mode 100644 index 00000000000..85c4a846a5a --- /dev/null +++ b/extension/llm/custom_ops/op_moe_aot.cpp @@ -0,0 +1,121 @@ +/* + * 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. + */ + +#include +#include +#include + +#include + +namespace torch { +namespace executor { +namespace native { + +// AOT (libtorch) shim. The ExecuTorch runtime kernel is registered in +// op_moe.cpp via EXECUTORCH_LIBRARY; here we expose the same op to the +// torch.library system so it is discoverable from torch.export traces. + +Tensor& quantized_moe_ffn_out_no_context( + const Tensor& x, + const Tensor& gate_weight, + const Tensor& expert_bias, + const Tensor& packed_w1, + const Tensor& packed_w3, + const Tensor& packed_w2, + const int64_t num_activated_experts, + const int64_t num_experts, + const int64_t hidden_dim, + const int64_t dim, + const int64_t group_size, + const int64_t weight_nbit, + executorch::aten::string_view score_func, + const double route_scale, + Tensor& out) { + executorch::runtime::KernelRuntimeContext context{}; + return torch::executor::native::quantized_moe_ffn_out( + context, + x, + gate_weight, + expert_bias, + packed_w1, + packed_w3, + packed_w2, + num_activated_experts, + num_experts, + hidden_dim, + dim, + group_size, + weight_nbit, + score_func, + route_scale, + out); +} + +at::Tensor quantized_moe_ffn_aten( + const at::Tensor& x, + const at::Tensor& gate_weight, + const at::Tensor& expert_bias, + const at::Tensor& packed_w1, + const at::Tensor& packed_w3, + const at::Tensor& packed_w2, + const int64_t num_activated_experts, + const int64_t num_experts, + const int64_t hidden_dim, + const int64_t dim, + const int64_t group_size, + const int64_t weight_nbit, + c10::string_view score_func, + const double route_scale) { + auto output = at::empty({x.size(0), dim}, x.options().dtype(at::kFloat)); + WRAP_TO_ATEN(quantized_moe_ffn_out_no_context, 14) + (x, + gate_weight, + expert_bias, + packed_w1, + packed_w3, + packed_w2, + num_activated_experts, + num_experts, + hidden_dim, + dim, + group_size, + weight_nbit, + score_func, + route_scale, + output); + return output; +} + +} // namespace native +} // namespace executor +} // namespace torch + +TORCH_LIBRARY_FRAGMENT(llama, m) { + m.def( + "quantized_moe_ffn(Tensor x, Tensor gate_weight, Tensor expert_bias, " + "Tensor packed_w1, Tensor packed_w3, Tensor packed_w2, " + "int num_activated_experts, int num_experts, int hidden_dim, int dim, " + "int group_size, int weight_nbit, str score_func, float route_scale) " + "-> Tensor"); + m.def( + "quantized_moe_ffn.out(Tensor x, Tensor gate_weight, Tensor expert_bias, " + "Tensor packed_w1, Tensor packed_w3, Tensor packed_w2, " + "int num_activated_experts, int num_experts, int hidden_dim, int dim, " + "int group_size, int weight_nbit, str score_func, float route_scale, " + "*, Tensor(a!) out) -> Tensor(a!)"); + m.def("_quantized_moe_ffn_active() -> bool"); +} + +TORCH_LIBRARY_IMPL(llama, CompositeExplicitAutograd, m) { + m.impl("quantized_moe_ffn", torch::executor::native::quantized_moe_ffn_aten); + m.impl( + "quantized_moe_ffn.out", + WRAP_TO_ATEN( + torch::executor::native::quantized_moe_ffn_out_no_context, 14)); + m.impl("_quantized_moe_ffn_active", []() { return true; }); +} diff --git a/extension/llm/custom_ops/targets.bzl b/extension/llm/custom_ops/targets.bzl index 1c69d082029..add6660f159 100644 --- a/extension/llm/custom_ops/targets.bzl +++ b/extension/llm/custom_ops/targets.bzl @@ -1,5 +1,5 @@ load("@fbsource//xplat/executorch/build:build_variables.bzl", "EXTENSION_LLM_CUSTOM_OPS_BUCK_SRCS") -load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "is_xplat", "runtime") load( "@fbsource//xplat/executorch/kernels/optimized:lib_defs.bzl", "get_vec_preprocessor_flags", @@ -22,6 +22,61 @@ def _get_quantized_preproc_flags(): else: return ["-DENABLE_CUSTOM_QUANTIZED_SDPA"] +def _get_quantized_moe_deps(): + """torchao dependencies for `llama::quantized_moe_ffn`. + + xplat: full optimized torchao linear kernel (aarch64 NEON + portable + scalar), which also transitively brings in weight_packing headers. + + fbcode: weight-packing headers only (for the reference + unpack+dequant+cpublas::gemm fallback). The optimized torchao + linear_operator is not linked on fbcode. + """ + if runtime.is_oss: + return [] + if is_xplat(): + return [ + "//xplat/pytorch/ao/torchao/csrc/cpu/shared_kernels:op_library_common", + "//xplat/pytorch/ao/torchao/csrc/cpu/torch_free_kernels:weight_packing", + ] + select({ + "DEFAULT": [], + "ovr_config//cpu:arm64": [ + "//xplat/pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_torch_free", + ], + }) + return [ + "fbcode//pytorch/ao/torchao/csrc/cpu/shared_kernels:op_library_common", + "fbcode//pytorch/ao/torchao/csrc/cpu/torch_free_kernels:weight_packing", + ] + +def _get_quantized_moe_preproc_flags(): + if runtime.is_oss: + return [] + if is_xplat(): + return select({ + "DEFAULT": [], + "ovr_config//cpu:arm64": [ + "-DENABLE_QUANTIZED_MOE_FFN", + "-DTORCHAO_BUILD_CPU_AARCH64=1", + "-DTORCHAO_ENABLE_ARM_NEON_DOT=1", + ], + }) + return [] + +def _get_quantized_moe_aot_packer_deps(): + """The Python source transform calls `torch.ops.torchao._pack_8bit_act_4bit_weight`, + which is registered by torchao's AOT (aten) library. Pull this in + only on fbcode so that Python tests can pack experts; xplat builds + do not need the AOT packer. + """ + if runtime.is_oss: + return [] + if is_xplat(): + return [] + return [ + "fbcode//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten", + ] + def define_common_targets(): """Defines targets that should be shared between fbcode and xplat. @@ -35,6 +90,7 @@ def define_common_targets(): exported_headers = [ "op_fallback.h", "op_fast_hadamard_transform.h", + "op_moe.h", "op_sdpa.h", "op_update_cache.h", ], @@ -42,7 +98,8 @@ def define_common_targets(): "op_sdpa_impl.h", ], exported_preprocessor_flags = get_vec_preprocessor_flags() + - _get_quantized_preproc_flags(), + _get_quantized_preproc_flags() + + _get_quantized_moe_preproc_flags(), exported_deps = [ "//executorch/runtime/kernel:kernel_includes", "//executorch/kernels/portable/cpu:scalar_utils", @@ -54,8 +111,9 @@ def define_common_targets(): deps = [ "//executorch/kernels/portable/cpu/util:reduce_util", "//executorch/extension/llm/custom_ops/spinquant:fast_hadamard_transform", - ] + get_vec_deps() + _get_quantized_sdpa_deps(), + ] + get_vec_deps() + _get_quantized_sdpa_deps() + _get_quantized_moe_deps(), compiler_flags = ["-Wno-missing-prototypes", "-Wno-global-constructors", "-Wno-error=pass-failed"] + get_compiler_optimization_flags() + + (["-isystem", "xplat/pytorch/ao"] if is_xplat() else []) + select({ "DEFAULT": [], "ovr_config//cpu:arm64": ["-march=armv8.2-a+dotprod"], @@ -70,6 +128,7 @@ def define_common_targets(): name = "custom_ops_aot_lib" + mkl_dep, srcs = [ "op_fast_hadamard_transform_aten.cpp", + "op_moe_aot.cpp", "op_sdpa_aot.cpp", "op_tile_crop.cpp", "op_tile_crop_aot.cpp", @@ -125,6 +184,27 @@ def define_common_targets(): ], ) + runtime.cxx_test( + name = "op_moe_test", + srcs = [ + "test_op_moe.cpp", + ], + visibility = ["//executorch/..."], + deps = [ + "//executorch/runtime/core/exec_aten:lib", + "//executorch/runtime/core/exec_aten/testing_util:tensor_util", + "//executorch/kernels/test:test_util", + ":custom_ops", + ], + ) + + # NOTE: the Python test for `llama::quantized_moe_ffn` lives in the + # fbcode-only BUCK file because it needs fbcode-specific + # `preload_deps` (`:custom_ops_aot_lib_mkl_noomp`, + # `:custom_ops_aot_py`) and the `portable_lib` python binding to + # bootstrap the AOT-side custom op registration. See the + # `test_quantized_moe` target in `BUCK`. + ## For preprocess runtime.python_library( name = "preprocess_custom_ops_py", diff --git a/extension/llm/custom_ops/test_op_moe.cpp b/extension/llm/custom_ops/test_op_moe.cpp new file mode 100644 index 00000000000..30e05734bf0 --- /dev/null +++ b/extension/llm/custom_ops/test_op_moe.cpp @@ -0,0 +1,78 @@ +/* + * 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. + */ + +#include + +#include +#include +#include + +#include + +using namespace ::testing; +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::runtime::testing::TensorFactory; + +namespace { + +// Smoke test: schema is registered and the op can be called via its +// out-variant entry point. Numerical correctness is exercised by the +// Python tests in test_quantized_moe.py against a Python q-dq reference, +// because constructing a meaningful torchao packed-weight blob from C++ +// here would require duplicating the AOT packer logic. +TEST(OpQuantizedMoeFfnTest, RegistrationSmokeTest) { + // Tiny shapes + constexpr int T = 2; + constexpr int D = 32; + constexpr int F = 32; + constexpr int E = 4; + constexpr int K = 2; + constexpr int kGroupSize = 32; + constexpr int kWeightNbit = 4; + + TensorFactory tff; + TensorFactory tfb; + + Tensor x = tff.zeros({T, D}); + Tensor gate = tff.zeros({E, D}); + 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 + // 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}); + Tensor packed_w2 = tfb.zeros({E, 1}); + + Tensor out = tff.zeros({T, D}); + + executorch::runtime::KernelRuntimeContext ctx{}; + // We don't actually call the kernel here in the registration smoke test + // because the empty packed buffers would not be valid torchao blobs. + // Just verify the op symbol resolves at link time. + auto fn = &torch::executor::native::quantized_moe_ffn_out; + EXPECT_NE(fn, nullptr); + // Silence unused-variable warnings on the input tensors above; they + // demonstrate the expected schema. + (void)x; + (void)gate; + (void)expert_bias; + (void)packed_w1; + (void)packed_w3; + (void)packed_w2; + (void)out; + (void)ctx; + (void)kGroupSize; + (void)kWeightNbit; + (void)K; + (void)F; +} + +} // namespace diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index 2f29203b53b..4c4abe0b24a 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -496,6 +496,7 @@ VULKAN_SCHEMA_SRCS = [ EXTENSION_LLM_CUSTOM_OPS_BUCK_SRCS = [ "op_fallback.cpp", "op_fast_hadamard_transform.cpp", + "op_moe.cpp", "op_sdpa.cpp", "op_update_cache.cpp", ]