Skip to content
Open
Changes from all commits
Commits
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
63 changes: 63 additions & 0 deletions examples/models/llama/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -525,3 +525,66 @@ CMake Error at runner/CMakeLists.txt:72 (add_subdirectory):

does not contain a CMakeLists.txt file.
```

# Mixture-of-Experts (MoE) Quantized Custom Op

The `--use_moe_quantized_op` flag enables a portable-runtime custom op
(`llama::quantized_moe_ffn`) that replaces the eager
`MOEFeedForward` module with a fused INT4-weight, INT8-dyn-act MoE
feed-forward kernel. On aarch64, the op uses torchao's optimized
`linear_8bit_act_xbit_weight::linear_operator` (NEON i8mm/dotprod).
On x86, a reference path unpacks torchao blobs, dequantizes, and uses
cpublas::gemm — numerically equivalent, suitable for testing.

## Supported configs

| Knob | Value |
|---|---|
| Weight quantization | symmetric INT4 (zero_point=0), `group_size=32` |
| Activation quantization | INT8 dynamic per-row (handled inside torchao) |
| Routing | sigmoid+biased top-k+un-biased gather+renormalize×`route_scale` OR softmax |
| Shapes | works for any `dim`, `hidden_dim` divisible by `group_size` |
| dtypes | fp32 activations only at the op boundary in v1 |

## AOT export

Add `--use_moe_quantized_op` to the Llama export command described above.

The source transform (`source_transformation/moe.py`) walks the eager
graph, INT4 group-quantizes each per-expert weight, packs each expert
through `torch.ops.torchao._pack_8bit_act_4bit_weight`, and replaces
each `MOEFeedForward` instance with `QuantizedMoEFFN` carrying
`[E, packed_bytes]` opaque buffers. Tracing reaches the Meta kernel
registered in `executorch.extension.llm.custom_ops.custom_ops`.

## Runtime build

The runtime kernel ships in `extension/llm/custom_ops/op_moe.cpp`. It
always compiles with a portable reference fallback (unpack + dequant +
`cpublas::gemm`) that works on any platform. `ENABLE_QUANTIZED_MOE_FFN`
is an **optimization gate**, not a correctness requirement — when
defined, the kernel uses torchao's fused `linear_operator` (NEON
i8mm/dotprod on aarch64) instead of the reference path.

In CMake, opt in to the optimized path with:

```cmake
-DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON
```

In Buck, `_get_quantized_moe_deps()` in `targets.bzl` wires:
- **arm64**: optimized torchao `linear_operator` + weight-packing
- **other architectures**: weight-packing headers only (reference path)

## Validation

- Python: `buck2 test executorch/extension/llm/custom_ops:test_quantized_moe`
runs the source-transform, Meta-kernel, validation, numerical-parity,
export-wiring, and config tests. These are validated on x86, which uses
the reference kernel path; the optimized torchao path is compiled and
exercised by the same suite when built for aarch64.
- C++: `buck2 test executorch/extension/llm/custom_ops:op_moe_test`
verifies the op is registered and links.
- E2E (requires a real checkpoint): build a PTE with
`--use_moe_quantized_op` and compare against ET eager;
bar is SQNR > 20 dB and cosine > 0.99.
Loading