@@ -525,3 +525,66 @@ CMake Error at runner/CMakeLists.txt:72 (add_subdirectory):
525525
526526 does not contain a CMakeLists.txt file.
527527```
528+
529+ # Mixture-of-Experts (MoE) Quantized Custom Op
530+
531+ The ` --use_moe_quantized_op ` flag enables a portable-runtime custom op
532+ (` llama::quantized_moe_ffn ` ) that replaces the eager
533+ ` MOEFeedForward ` module with a fused INT4-weight, INT8-dyn-act MoE
534+ feed-forward kernel. On aarch64, the op uses torchao's optimized
535+ ` linear_8bit_act_xbit_weight::linear_operator ` (NEON i8mm/dotprod).
536+ On x86, a reference path unpacks torchao blobs, dequantizes, and uses
537+ cpublas::gemm — numerically equivalent, suitable for testing.
538+
539+ ## Supported configs
540+
541+ | Knob | Value |
542+ | ---| ---|
543+ | Weight quantization | symmetric INT4 (zero_point=0), ` group_size=32 ` |
544+ | Activation quantization | INT8 dynamic per-row (handled inside torchao) |
545+ | Routing | sigmoid+biased top-k+un-biased gather+renormalize×` route_scale ` OR softmax |
546+ | Shapes | works for any ` dim ` , ` hidden_dim ` divisible by ` group_size ` |
547+ | dtypes | fp32 activations only at the op boundary in v1 |
548+
549+ ## AOT export
550+
551+ Add ` --use_moe_quantized_op ` to the Llama export command described above.
552+
553+ The source transform (` source_transformation/moe.py ` ) walks the eager
554+ graph, INT4 group-quantizes each per-expert weight, packs each expert
555+ through ` torch.ops.torchao._pack_8bit_act_4bit_weight ` , and replaces
556+ each ` MOEFeedForward ` instance with ` QuantizedMoEFFN ` carrying
557+ ` [E, packed_bytes] ` opaque buffers. Tracing reaches the Meta kernel
558+ registered in ` executorch.extension.llm.custom_ops.custom_ops ` .
559+
560+ ## Runtime build
561+
562+ The runtime kernel ships in ` extension/llm/custom_ops/op_moe.cpp ` . It
563+ always compiles with a portable reference fallback (unpack + dequant +
564+ ` cpublas::gemm ` ) that works on any platform. ` ENABLE_QUANTIZED_MOE_FFN `
565+ is an ** optimization gate** , not a correctness requirement — when
566+ defined, the kernel uses torchao's fused ` linear_operator ` (NEON
567+ i8mm/dotprod on aarch64) instead of the reference path.
568+
569+ In CMake, opt in to the optimized path with:
570+
571+ ``` cmake
572+ -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON
573+ ```
574+
575+ In Buck, ` _get_quantized_moe_deps() ` in ` targets.bzl ` wires:
576+ - ** arm64** : optimized torchao ` linear_operator ` + weight-packing
577+ - ** other architectures** : weight-packing headers only (reference path)
578+
579+ ## Validation
580+
581+ - Python: ` buck2 test executorch/extension/llm/custom_ops:test_quantized_moe `
582+ runs the source-transform, Meta-kernel, validation, numerical-parity,
583+ export-wiring, and config tests. These are validated on x86, which uses
584+ the reference kernel path; the optimized torchao path is compiled and
585+ exercised by the same suite when built for aarch64.
586+ - C++: ` buck2 test executorch/extension/llm/custom_ops:op_moe_test `
587+ verifies the op is registered and links.
588+ - E2E (requires a real checkpoint): build a PTE with
589+ ` --use_moe_quantized_op ` and compare against ET eager;
590+ bar is SQNR > 20 dB and cosine > 0.99.
0 commit comments