[ExecuTorch][llm] Fuse w1+w3 into single GEMM in quantized_moe_ffn (pytorch#21538)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#21124 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/74/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/74/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/digantdesai/73/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/digantdesai/74/orig
Differential Revision:
[D102799854](https://our.internmc.facebook.com/intern/diff/D102799854/)
@diff-train-skip-merge

Co-authored-by: Digant Desai <digantdesai@meta.com>

Author: digantdesai

examples/models/llama/source_transformation/moe.py

@@ -104,17 +104,15 @@ class QuantizedMoEFFN(nn.Module):
     Buffers (registered, not parameters):
       gate_weight  [E, D]   fp32  — copied from `MOEFeedForward.gate.weight`
       expert_bias  [E] or [0]  fp32 — empty when `use_expert_bias=False`
-      packed_w1    [E, packed_bytes_w1]  uint8 — torchao opaque blobs
-      packed_w3    [E, packed_bytes_w3]  uint8
-      packed_w2    [E, packed_bytes_w2]  uint8
+      packed_w13   [E, packed_bytes_w13]  uint8 — fused w1+w3 torchao blobs
+      packed_w2    [E, packed_bytes_w2]   uint8
     """
 
     def __init__(
         self,
         gate_weight: torch.Tensor,
         expert_bias: torch.Tensor | None,
-        packed_w1: torch.Tensor,
-        packed_w3: torch.Tensor,
+        packed_w13: torch.Tensor,
         packed_w2: torch.Tensor,
         *,
         num_experts: int,
@@ -149,8 +147,7 @@ def __init__(
 
         # torchao packed blobs are int8 tensors; reinterpret the bytes as
         # uint8 (no value conversion) for the op schema.
-        self.register_buffer("packed_w1", packed_w1.view(torch.uint8))
-        self.register_buffer("packed_w3", packed_w3.view(torch.uint8))
+        self.register_buffer("packed_w13", packed_w13.view(torch.uint8))
         self.register_buffer("packed_w2", packed_w2.view(torch.uint8))
         self.shared_expert: nn.Module | None = None
 
@@ -178,8 +175,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x_flat,
             self.gate_weight,
             self.expert_bias,
-            self.packed_w1,
-            self.packed_w3,
+            self.packed_w13,
             self.packed_w2,
             self.num_activated_experts,
             self.num_experts,
@@ -261,18 +257,15 @@ def _build_quantized_moe_ffn_from_eager(
         _torchao_pack_int4_weight if weight_nbit == 4 else _torchao_pack_int8_weight
     )
 
-    packed_w1_list: list[torch.Tensor] = []
-    packed_w3_list: list[torch.Tensor] = []
+    packed_w13_list: list[torch.Tensor] = []
     packed_w2_list: list[torch.Tensor] = []
     for ei in range(e):
-        # w1, w3 share the [F, D] shape, group along K=D.
-        packed_w1_list.append(pack_fn(w1[ei], group_size))
-        packed_w3_list.append(pack_fn(w3[ei], group_size))
-        # w2 packed shape is [D, F], group along K=F.
+        # Fuse w1 and w3 into a single [2F, D] matrix before packing.
+        w13 = torch.cat([w1[ei], w3[ei]], dim=0)  # [2F, D]
+        packed_w13_list.append(pack_fn(w13, group_size))
         packed_w2_list.append(pack_fn(w2_packed_in[ei], group_size))
 
-    packed_w1 = _stack_per_expert_packed(packed_w1_list)
-    packed_w3 = _stack_per_expert_packed(packed_w3_list)
+    packed_w13 = _stack_per_expert_packed(packed_w13_list)
     packed_w2 = _stack_per_expert_packed(packed_w2_list)
 
     # `MOEFeedForward` registers `expert_bias` as a buffer that is None unless
@@ -284,8 +277,7 @@ def _build_quantized_moe_ffn_from_eager(
     replacement = QuantizedMoEFFN(
         gate_weight=moe.gate.weight.detach().clone(),
         expert_bias=expert_bias,
-        packed_w1=packed_w1,
-        packed_w3=packed_w3,
+        packed_w13=packed_w13,
         packed_w2=packed_w2,
         num_experts=e,
         num_activated_experts=moe.num_activated_experts,

extension/llm/custom_ops/custom_ops.py

@@ -168,8 +168,7 @@ def _validate_quantized_moe_ffn_params(
     x,
     gate_weight,
     expert_bias,
-    packed_w1,
-    packed_w3,
+    packed_w13,
     packed_w2,
     num_activated_experts,
     num_experts,
@@ -202,18 +201,14 @@ def _validate_quantized_moe_ffn_params(
         ), f"expert_bias must be float32, got {expert_bias.dtype}"
 
     for name, t in (
-        ("packed_w1", packed_w1),
-        ("packed_w3", packed_w3),
+        ("packed_w13", packed_w13),
         ("packed_w2", packed_w2),
     ):
         assert (
             t.dim() == 2 and t.size(0) == num_experts
         ), f"{name} must be [E={num_experts}, packed_bytes], got {list(t.size())}"
         assert t.dtype == torch.uint8, f"{name} must be uint8, got {t.dtype}"
 
-    assert packed_w1.size(1) == packed_w3.size(
-        1
-    ), "packed_w1 and packed_w3 per-expert blob sizes must match"
     assert (
         0 < num_activated_experts <= num_experts
     ), f"num_activated_experts ({num_activated_experts}) out of range [1, {num_experts}]"
@@ -241,8 +236,7 @@ def quantized_moe_ffn_meta(
     x,
     gate_weight,
     expert_bias,
-    packed_w1,
-    packed_w3,
+    packed_w13,
     packed_w2,
     num_activated_experts,
     num_experts,
@@ -257,8 +251,7 @@ def quantized_moe_ffn_meta(
         x,
         gate_weight,
         expert_bias,
-        packed_w1,
-        packed_w3,
+        packed_w13,
         packed_w2,
         num_activated_experts,
         num_experts,