From 767c5113c84f341035a0bb8f7f7831fb4e4aec0b Mon Sep 17 00:00:00 2001
From: Digant Desai <digantdesai@meta.com>
Date: Tue, 21 Jul 2026 20:38:20 -0700
Subject: [PATCH] Update

[ghstack-poisoned]
---
 examples/models/llama/README.md | 63 +++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/examples/models/llama/README.md b/examples/models/llama/README.md
index f674d454987..850c2997ec2 100644
--- a/examples/models/llama/README.md
+++ b/examples/models/llama/README.md
@@ -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.