Update

[ghstack-poisoned]

Author: digantdesai

examples/models/llama/export_llama_lib.py

@@ -565,6 +565,19 @@ def build_args_parser() -> argparse.ArgumentParser:
         help="Use SpinQuant for better quantization performance. Only support cuda and native.",
     )
 
+    parser.add_argument(
+        "--use_moe_quantized_op",
+        action="store_true",
+        default=False,
+        help=(
+            "Replace eager MoE feed-forward modules with the "
+            "`llama::quantized_moe_ffn` portable-runtime custom op (INT4 "
+            "weights, INT8 dyn-quant activations via torchao). On aarch64 "
+            "with ENABLE_QUANTIZED_MOE_FFN the optimized torchao kernel is "
+            "used; otherwise a portable reference fallback runs."
+        ),
+    )
+
     parser.add_argument(
         "-qat",
         "--use_qat",
@@ -809,6 +822,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager:
             use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear,
             use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding,
             quantize_with_hqq=llm_config.quantization.use_hqq,
+            use_moe_quantized_op=llm_config.model.use_moe_quantized_op,
         )
     )
 
@@ -1712,6 +1726,7 @@ def _get_source_transforms(  # noqa
     use_torchao_kernels_linear: bool = False,
     use_torchao_kernels_tied_embedding: bool = False,
     quantize_with_hqq: bool = True,
+    use_moe_quantized_op: bool = False,
 ) -> List[Callable[[torch.nn.Module], torch.nn.Module]]:
     """
     Return a list of functions that transform a graph.
@@ -1769,6 +1784,17 @@ def _get_source_transforms(  # noqa
 
             transforms.append(inject_fast_hadamard_transform_native_for_spin_quant)
 
+    if use_moe_quantized_op:
+        from .source_transformation.moe import replace_moe_with_quantized_op
+
+        transforms.append(
+            partial(
+                replace_moe_with_quantized_op,
+                group_size=group_size or 32,
+                weight_nbit=4,
+            )
+        )
+
     if embedding_quantize:
         """
         When this option is selected, it finds all embedding layers and transforms

extension/llm/custom_ops/BUCK

@@ -100,6 +100,9 @@ fbcode_target(_kind = runtime.python_test,
     preload_deps = [
         ":custom_ops_aot_lib_mkl_noomp",
         ":custom_ops_aot_py",
+        # The Python source transform calls
+        # `torch.ops.torchao._pack_8bit_act_4bit_weight` to pack each
+        # expert's INT4 weights, so preload the torchao AOT library too.
         "//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten",
     ],
     deps = [
@@ -108,5 +111,7 @@ fbcode_target(_kind = runtime.python_test,
         "//executorch/examples/models/llama:llama_transformer",
         "//executorch/examples/models/llama:transformer_modules",
         "//executorch/examples/models/llama:source_transformation",
+        "//executorch/examples/models/llama:export_library",
+        "//executorch/extension/llm/export/config:llm_config",
     ],
 )

extension/llm/custom_ops/test_quantized_moe.py

@@ -585,3 +585,110 @@ def test_no_shared_expert_is_none(self) -> None:
         wrapper.m = moe
         replace_moe_with_quantized_op(wrapper, group_size=32, weight_nbit=4)
         self.assertIsNone(wrapper.m.shared_expert)
+
+
+class TestExportPipelineWiring(unittest.TestCase):
+    """The export pipeline correctly includes the MoE transform."""
+
+    def test_get_source_transforms_includes_moe_when_enabled(self) -> None:
+        from functools import partial
+
+        from executorch.examples.models.llama.export_llama_lib import (
+            _get_source_transforms,
+        )
+
+        transforms = _get_source_transforms(
+            dtype_override=torch.float32, use_moe_quantized_op=True
+        )
+        moe_transforms = [
+            t
+            for t in transforms
+            if isinstance(t, partial)
+            and t.func.__name__ == "replace_moe_with_quantized_op"
+        ]
+        self.assertEqual(len(moe_transforms), 1)
+        self.assertEqual(moe_transforms[0].keywords["group_size"], 32)
+        self.assertEqual(moe_transforms[0].keywords["weight_nbit"], 4)
+
+    def test_get_source_transforms_excludes_moe_when_disabled(self) -> None:
+        from functools import partial
+
+        from executorch.examples.models.llama.export_llama_lib import (
+            _get_source_transforms,
+        )
+
+        transforms = _get_source_transforms(
+            dtype_override=torch.float32, use_moe_quantized_op=False
+        )
+        moe_transforms = [
+            t
+            for t in transforms
+            if isinstance(t, partial)
+            and hasattr(t.func, "__name__")
+            and t.func.__name__ == "replace_moe_with_quantized_op"
+        ]
+        self.assertEqual(len(moe_transforms), 0)
+
+    def test_get_source_transforms_passes_custom_group_size(self) -> None:
+        from functools import partial
+
+        from executorch.examples.models.llama.export_llama_lib import (
+            _get_source_transforms,
+        )
+
+        transforms = _get_source_transforms(
+            dtype_override=torch.float32,
+            use_moe_quantized_op=True,
+            group_size=64,
+        )
+        moe_transforms = [
+            t
+            for t in transforms
+            if isinstance(t, partial)
+            and t.func.__name__ == "replace_moe_with_quantized_op"
+        ]
+        self.assertEqual(len(moe_transforms), 1)
+        self.assertEqual(moe_transforms[0].keywords["group_size"], 64)
+
+    def test_sentinel_op_is_registered(self) -> None:
+        self.assertTrue(hasattr(torch.ops.llama, "_quantized_moe_ffn_active"))
+        self.assertTrue(torch.ops.llama._quantized_moe_ffn_active())
+
+
+class TestLlmConfigMoeFlag(unittest.TestCase):
+    """llm_config wires the --use_moe_quantized_op flag correctly."""
+
+    def test_from_args_sets_flag(self) -> None:
+        import argparse
+
+        from executorch.extension.llm.export.config.llm_config import LlmConfig
+
+        args = argparse.Namespace(use_moe_quantized_op=True)
+        config = LlmConfig.from_args(args)
+        self.assertTrue(config.model.use_moe_quantized_op)
+
+    def test_default_is_false(self) -> None:
+        from executorch.extension.llm.export.config.llm_config import ModelConfig
+
+        self.assertFalse(ModelConfig().use_moe_quantized_op)
+
+    def test_from_args_missing_field_defaults_false(self) -> None:
+        import argparse
+
+        from executorch.extension.llm.export.config.llm_config import LlmConfig
+
+        args = argparse.Namespace()
+        config = LlmConfig.from_args(args)
+        self.assertFalse(config.model.use_moe_quantized_op)
+
+    def test_argparser_flag_true(self) -> None:
+        from executorch.examples.models.llama.export_llama_lib import build_args_parser
+
+        args = build_args_parser().parse_args(["--use_moe_quantized_op"])
+        self.assertTrue(args.use_moe_quantized_op)
+
+    def test_argparser_flag_default_false(self) -> None:
+        from executorch.examples.models.llama.export_llama_lib import build_args_parser
+
+        args = build_args_parser().parse_args([])
+        self.assertFalse(args.use_moe_quantized_op)

extension/llm/export/config/llm_config.py

@@ -205,6 +205,9 @@ class ModelConfig:
     use_kv_cache: bool = False
     quantize_kv_cache: bool = False
     local_global_attention: Optional[List[int]] = None
+    # Replace eager MOEFeedForward modules with the
+    # `llama::quantized_moe_ffn` portable-runtime custom op.
+    use_moe_quantized_op: bool = False
 
     def __post_init__(self):
         self._validate_attention_sink()
@@ -728,6 +731,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig":  # noqa: C901
             llm_config.model.quantize_kv_cache = args.quantize_kv_cache
         if hasattr(args, "local_global_attention"):
             llm_config.model.local_global_attention = args.local_global_attention
+        if hasattr(args, "use_moe_quantized_op"):
+            llm_config.model.use_moe_quantized_op = args.use_moe_quantized_op
 
         # ExportConfig
         if hasattr(args, "max_seq_length"):