CMake support for quantized MoE optimized build

[digantdesai]

For MoE custom op

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22547

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ You can merge normally! (1 Unrelated Failure), 1 Unclassified Failure

As of commit fa78d43 with merge base 1cfcdf8 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• pull / nxp-mcuxpresso-test / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command docker exec -t ad4ac4ab5f935018fe00a564262341fc4b24a1d2260265fe3d195087dc6dd046 /exec failed with exit code 127

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-moshi-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  E: Failed to fetch http://security.ubuntu.com/ubuntu/pool/main/libs/libssh/libssh-gcrypt-4_0.9.6-2ubuntu0.22.04.7_amd64.deb 404 Not Found [IP: 172.66.152.176 80]

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[digantdesai]

@claude review this PR.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[claude]

Claude finished @digantdesai's task in 4m 4s —— View job

---

PR Review: CMake support for quantized MoE optimized build

Reviewed the full diff (op_moe.cpp, CMakeLists.txt, targets.bzl, docs, test). Overall this is a clean improvement: the rename to EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED reads better as an optimization gate, the CMake now builds the torchao kernel sources standalone (torchao_moe_linear) when EXECUTORCH_BUILD_KERNELS_TORCHAO isn't set, and the architecture guard + shared _target_enable_quantized_moe_torchao helper remove the previous copy-paste. A few things worth addressing before merge.

Correctness / behavior

1. Runtime kernel selection changed from header-driven to hard-coded universal — but the surrounding comments still describe the old behavior.
   op_moe.cpp previously called select_ukernel_config<kWeightNbit>(header), which resolved universal or kleidi from the packed header. It now hard-rejects anything but universal (op_moe.cpp:254-265) and uses a fixed universal_ukernel_config<kWeightNbit>(). That's a legitimate simplification, but several comments now contradict the code:

   • op_moe.cpp:268-270 — "derive them from the chosen config rather than assuming a fixed layout" and the "universal or kleidi" framing. The config is now always universal/fixed.
   • op_moe.cpp:324-326 — "The torchao path validates against its own selected format (universal or kleidi)". There is no selection anymore.
   • targets.bzl:56-62 — the TODO claims "The runtime (op_moe.cpp) already selects the ukernel from the header, so kleidi headers resolve automatically once the kernels are compiled." This is now false: enabling KleidiAI would additionally require relaxing the hard universal check and adding a kleidi ukernel config in op_moe.cpp. The TODO understates the work.

   Please update these to reflect that the runtime is now universal-only and that kleidi support requires an op_moe.cpp change, not just build flags. Fix this →

2. TORCHAO_PARALLEL_EXECUTORCH is defined by CMake but not by Buck — the two builds compile different expert-parallelism branches.
   _target_enable_quantized_moe_torchao sets TORCHAO_PARALLEL_EXECUTORCH=1 on custom_ops (CMakeLists.txt:101), so the #if defined(...) && !defined(TORCHAO_PARALLEL_EXECUTORCH) block at op_moe.cpp:685-694 compiles parallelize_experts = false in CMake. The Buck flags (targets.bzl:66-68) do not define TORCHAO_PARALLEL_EXECUTORCH, so Buck compiles the linear_operator_num_threads() == 1 runtime guard instead. The runtime guard makes this safe (no nested-pool deadlock either way), but it's an intentional-looking divergence — worth confirming this is deliberate, and ideally aligning the two so the kernel's threading model is identical across build systems.

Minor

3. CMakeLists.txt:110-111 — the FATAL_ERROR message says "supported only on aarch64 and arm64" while the guard regex is ^(aarch64|arm64)$. Fine, but note _common_compile_options adds -march=armv8.2-a+dotprod for arm64|aarch64 (unanchored, line 23) whereas the optimized gate is anchored — consistent enough, just flagging the slightly different regexes in the same file.

4. quantized_moe_torchao_target is a plain (non-cache) variable set at file scope inside the OPTIMIZED block and reused in the AOT block (CMakeLists.txt:121/123 → 242). This works because the OPTIMIZED block runs first and both are top-level, but it's an implicit ordering dependency; a short comment or a cache/set(... PARENT_SCOPE)-free note would help future readers.

Nice touch validating the packed blob size against the derived uk.packed_weights_size(...) (op_moe.cpp:271-292) before dispatch — that keeps the loud-failure behavior the tests rely on.

• Read PR diff
• Review CMakeLists.txt changes
• Review op_moe.cpp changes
• Review docs / targets.bzl / test changes
• Post review feedback

[Copilot]

🟡 Changes recommended

There are newly misleading/incorrect comments in the optimized MoE build wiring and runtime dispatch documentation that should be updated to match the new universal-only behavior.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR adds/adjusts build-time wiring for the llama::quantized_moe_ffn custom op so that an optimized torchao-based implementation can be enabled in CMake (in addition to existing Buck wiring), while keeping a portable reference fallback as the default.

Changes:

1. Introduces a CMake option EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED that gates linking/compiling the optimized torchao linear operator (and optionally builds a minimal torchao linear static library when the full torchao kernel build isn’t enabled).
2. Renames the compile-time feature gate from ENABLE_QUANTIZED_MOE_FFN to EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO across code and build files, and tightens the optimized runtime path to require the universal packed-weights format.
3. Updates Llama export/runtime documentation strings to reference the new optimized-build control surface.

File summaries

File	Description
extension/llm/custom_ops/test_op_moe.cpp	Updates test comment to reflect the new optimized-path feature macro name.
extension/llm/custom_ops/targets.bzl	Renames the arm64 preprocessor flag used to enable the torchao optimized path in Buck builds.
extension/llm/custom_ops/op_moe.cpp	Switches to EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO, adds an explicit universal ukernel config path, and tightens packed-weights validation for the optimized implementation.
extension/llm/custom_ops/CMakeLists.txt	Adds EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED option and CMake-side linking/compile-def wiring for the optimized MoE kernel.
examples/models/llama/README.md	Updates runtime-build instructions to reference the new CMake option.
examples/models/llama/export_llama_lib.py	Updates CLI help text to describe optimized vs reference runtime behavior without referencing the old macro name.

Review details

Suppressed comments (1)

extension/llm/custom_ops/op_moe.cpp:326

• The comment here says the torchao path validates against a selected format “(universal or kleidi)”, but the new torchao path in this file explicitly requires the universal packed-weights type and rejects non-universal headers. This comment is now misleading and should be updated to match the actual behavior.

  // Reference path only: it unpacks the universal layout, so validate the blob
  // holds the header plus the universal packed weight-data bytes for the
  // claimed dims before any path dereferences it. The torchao path validates
  // against its own selected format (universal or kleidi) inside
  // torchao_linear.

• Files reviewed: 6/6 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

The reference (non-torchao) path still doesn’t validate the packed-weights header type, which can lead to misinterpreting non-universal blobs instead of failing loudly, and there is also misleading test commentary to correct.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 2
• Review effort level: Lite

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 2m 28s —— View job

---

PR Review: CMake support for quantized MoE optimized build

Re-reviewed the current tip (fa78d43) against merge base 1cfcdf8. The rename to EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO / EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED, the standalone torchao_moe_linear static lib, the shared _target_enable_quantized_moe_torchao helper, and the arch guard all look good. The stale comments flagged in my earlier pass (op_moe.cpp:266-267, 319-322 and the targets.bzl KleidiAI TODO) are now fixed — nice.

Two things from the latest Copilot pass are legitimate and still open, plus one behavioral change worth an explicit confirmation.

Correctness

1. Reference path never validates header.type. torchao_linear now hard-checks header.type == …_universal (op_moe.cpp:252-256), but the reference build (#ifndef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) only validates the blob size (op_moe.cpp:318-346) and reference_linear reads the header size without ever calling PackedWeightsHeader::read/checking the type (op_moe.cpp:182-205). If the AoT packer ever emits a non-universal format (e.g. kleidi), the reference path unpacks it as universal and silently returns wrong results instead of failing loudly. Add the same header.type == …_universal check to the reference path so both paths reject non-universal blobs identically. Fix this →

2. Expert-level parallelism was removed entirely — confirm intended. The old parallelize_experts logic (and the threadpool.h / thread_parallel_interface.h includes) is gone; quantized_moe_ffn_out now always calls run_experts(0, E) serially (op_moe.cpp:681). This resolves the earlier CMake-vs-Buck TORCHAO_PARALLEL_EXECUTORCH branch-divergence concern (op_moe.cpp no longer branches on it), so per-token/per-GEMM throughput now depends solely on the kernel's internal threading. But that internal threading still diverges: CMake defines TORCHAO_PARALLEL_EXECUTORCH=1 on the kernel via _target_enable_quantized_moe_torchao (CMakeLists.txt:96-102), while Buck's preproc flags (targets.bzl:59-66) do not. Worth confirming the Buck kernel isn't falling back to fully single-threaded execution now that experts run serially — otherwise this is a silent perf regression on the Buck path.

Minor

3. Misleading test comment (test_op_moe.cpp:46-49): the comment says "the kernel ET_CHECK_MSGs out before doing any real work, which is what we want this test to verify" — but the test explicitly does not call the kernel (see the comment at line 57-59 and the symbol-resolution-only EXPECT_NE(fn, nullptr)). Reword to say the packed buffers are placeholders documenting the expected schema, and the test only verifies the op symbol resolves/links. Fix this →

4. quantized_moe_torchao_target is a plain (non-cache) variable set in the OPTIMIZED block and reused in the AOT block (CMakeLists.txt:112/121 → 264). Works because both are top-level and the OPTIMIZED block runs first, but it's an implicit ordering dependency — a one-line comment noting that would help future readers.

5. CMakeLists.txt FATAL_ERROR message says "supported only on aarch64 and arm64" with an anchored guard ^(aarch64|arm64)$, while _common_compile_options uses the unanchored arm64|aarch64 (line 23). Consistent enough in practice; just flagging the two different regexes in one file.

Overall this is a clean, well-scoped change. The two items worth acting on before merge are #1 (loud-failure symmetry in the reference path) and #2 (confirm the Buck kernel-threading path is intentional); #3 is a quick doc fix.

• Read the diff (op_moe.cpp, CMakeLists.txt, targets.bzl, docs, test)
• Analyze correctness & build wiring
• Post review feedback
  · branch quantized-moe-fix

[meta-codesync]

@digantdesai has imported this pull request. If you are a Meta employee, you can view this in D118853081.