Skip to content

Commit 607c4c1

Browse files
authored
[ExecuTorch][llm] Add MoE quantized custom op documentation (pytorch#21539)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: pytorch#21125 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/75/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/75/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/74/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/75/orig Differential Revision: [D102382001](https://our.internmc.facebook.com/intern/diff/D102382001/) @diff-train-skip-merge Co-authored-by: Digant Desai <digantdesai@meta.com>
1 parent cb5e32f commit 607c4c1

1 file changed

Lines changed: 63 additions & 0 deletions

File tree

‎examples/models/llama/README.md‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)