[ExecuTorch][llm] Fuse w1+w3 into single GEMM in quantized_moe_ffn

Pull Request resolved: #21124

Fuse the up-projection (w1) and gate-projection (w3) into a single [2F, D] GEMM per expert. This halves the number of torchao activation quantizations per expert (from 2 to 1) and reduces total GEMM calls from 3 to 2 per active expert.

At AOT time, w1 and w3 are concatenated before packing: pack_fn(cat([w1, w3], dim=0)). At runtime, a single expert_linear_dispatch produces [m_e, 2F], then a fused swiglu_and_compact pass reads the interleaved h1/h3 and writes [m_e, F] for the w2 down-projection.

Schema changes from (packed_w1, packed_w3, packed_w2) to (packed_w13, packed_w2) — one fewer tensor arg (14 -> 13).
ghstack-source-id: 409183086
@exported-using-ghexport

Differential Revision: [D102799854](https://our.internmc.facebook.com/intern/diff/D102799854/)

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,