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

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,
@@ -253,18 +249,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)
 
     expert_bias = (
@@ -276,8 +269,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,

extension/llm/custom_ops/op_moe.cpp

@@ -34,16 +34,6 @@ namespace {
 
 using executorch::aten::string_view;
 
-// SwiGLU epilogue: out = silu(h1) * h3 (in place over h1).
-// silu(x) = x * sigmoid(x) = x / (1 + exp(-x)).
-inline void swiglu_inplace(float* h1, const float* h3, int64_t n) {
-  for (int64_t i = 0; i < n; ++i) {
-    const float v = h1[i];
-    const float s = v / (1.0f + std::exp(-v));
-    h1[i] = s * h3[i];
-  }
-}
-
 // Stable softmax over a row of length E (in place).
 inline void softmax_row(float* row, int64_t e) {
   float maxv = row[0];
@@ -291,8 +281,7 @@ Tensor& quantized_moe_ffn_out(
     const Tensor& x,
     const Tensor& gate_weight,
     const Tensor& expert_bias,
-    const Tensor& packed_w1,
-    const Tensor& packed_w3,
+    const Tensor& packed_w13,
     const Tensor& packed_w2,
     int64_t num_activated_experts,
     int64_t num_experts,
@@ -333,23 +322,15 @@ Tensor& quantized_moe_ffn_out(
   }
 
   ET_CHECK_MSG(
-      packed_w1.dim() == 2 && packed_w1.size(0) == num_experts,
-      "packed_w1 must be [E, packed_bytes]");
-  ET_CHECK_MSG(
-      packed_w3.dim() == 2 && packed_w3.size(0) == num_experts,
-      "packed_w3 must be [E, packed_bytes]");
+      packed_w13.dim() == 2 && packed_w13.size(0) == num_experts,
+      "packed_w13 must be [E, packed_bytes]");
   ET_CHECK_MSG(
       packed_w2.dim() == 2 && packed_w2.size(0) == num_experts,
       "packed_w2 must be [E, packed_bytes]");
   ET_CHECK_MSG(
-      packed_w1.scalar_type() == ScalarType::Byte, "packed_w1 must be uint8");
-  ET_CHECK_MSG(
-      packed_w3.scalar_type() == ScalarType::Byte, "packed_w3 must be uint8");
+      packed_w13.scalar_type() == ScalarType::Byte, "packed_w13 must be uint8");
   ET_CHECK_MSG(
       packed_w2.scalar_type() == ScalarType::Byte, "packed_w2 must be uint8");
-  ET_CHECK_MSG(
-      packed_w1.size(1) == packed_w3.size(1),
-      "packed_w1 and packed_w3 per-expert blob sizes must match");
 
   ET_CHECK_MSG(
       num_activated_experts > 0 && num_activated_experts <= num_experts,
@@ -383,8 +364,7 @@ Tensor& quantized_moe_ffn_out(
   const int64_t F = hidden_dim;
   const int64_t E = num_experts;
   const int64_t K = num_activated_experts;
-  const int64_t pw1_bytes = packed_w1.size(1);
-  const int64_t pw3_bytes = packed_w3.size(1);
+  const int64_t pw13_bytes = packed_w13.size(1);
   const int64_t pw2_bytes = packed_w2.size(1);
 
   // Resize the output to [T, D].
@@ -406,8 +386,7 @@ Tensor& quantized_moe_ffn_out(
   const float* gate_w_ptr = gate_weight.const_data_ptr<float>();
   const float* expert_bias_ptr =
       use_expert_bias ? expert_bias.const_data_ptr<float>() : nullptr;
-  const uint8_t* pw1_ptr = packed_w1.const_data_ptr<uint8_t>();
-  const uint8_t* pw3_ptr = packed_w3.const_data_ptr<uint8_t>();
+  const uint8_t* pw13_ptr = packed_w13.const_data_ptr<uint8_t>();
   const uint8_t* pw2_ptr = packed_w2.const_data_ptr<uint8_t>();
   float* out_ptr = out.mutable_data_ptr<float>();
 
@@ -540,8 +519,10 @@ Tensor& quantized_moe_ffn_out(
       max_m_e = m_e;
     }
   }
-  std::vector<float> h1_buf(static_cast<size_t>(max_m_e * F), 0.0f);
-  std::vector<float> h3_buf(static_cast<size_t>(max_m_e * F), 0.0f);
+  // h13_buf holds the fused [m_e, 2F] output of the w1+w3 GEMM.
+  // mid_buf holds the [m_e, F] SwiGLU result compacted for the w2 GEMM.
+  std::vector<float> h13_buf(static_cast<size_t>(max_m_e * 2 * F), 0.0f);
+  std::vector<float> mid_buf(static_cast<size_t>(max_m_e * F), 0.0f);
   std::vector<float> out_e_buf(static_cast<size_t>(max_m_e * D), 0.0f);
 
   for (int64_t e = 0; e < E; ++e) {
@@ -550,40 +531,37 @@ Tensor& quantized_moe_ffn_out(
       continue;
     }
     const float* x_e = x_perm.data() + expert_offsets[e] * D;
-    const uint8_t* w1_e = pw1_ptr + e * pw1_bytes;
-    const uint8_t* w3_e = pw3_ptr + e * pw3_bytes;
+    const uint8_t* w13_e = pw13_ptr + e * pw13_bytes;
     const uint8_t* w2_e = pw2_ptr + e * pw2_bytes;
 
-    // Up-projection (D -> F): H1 = x_e @ w1[e]^T
+    // Fused up+gate projection (D -> 2F): H13 = x_e @ w13[e]^T
+    // First F columns are h1 (up), last F are h3 (gate).
     expert_linear_dispatch(
         weight_nbit,
-        w1_e,
-        pw1_bytes,
+        w13_e,
+        pw13_bytes,
         x_e,
         m_e,
-        F,
+        2 * F,
         D,
         group_size,
-        h1_buf.data());
-    // Gate-projection (D -> F): H3 = x_e @ w3[e]^T
-    expert_linear_dispatch(
-        weight_nbit,
-        w3_e,
-        pw3_bytes,
-        x_e,
-        m_e,
-        F,
-        D,
-        group_size,
-        h3_buf.data());
-    // SwiGLU epilogue (in place over h1_buf).
-    swiglu_inplace(h1_buf.data(), h3_buf.data(), m_e * F);
+        h13_buf.data());
+    // SwiGLU + compact: read [m_e, 2F] interleaved, write [m_e, F]
+    for (int64_t i = 0; i < m_e; ++i) {
+      const float* src = h13_buf.data() + i * 2 * F;
+      float* dst = mid_buf.data() + i * F;
+      for (int64_t j = 0; j < F; ++j) {
+        const float v = src[j];
+        const float s = v / (1.0f + std::exp(-v));
+        dst[j] = s * src[F + j];
+      }
+    }
     // Down-projection (F -> D): OUT_e = MID @ w2[e]^T
     expert_linear_dispatch(
         weight_nbit,
         w2_e,
         pw2_bytes,
-        h1_buf.data(),
+        mid_buf.data(),
         m_e,
         D,
         F,

extension/llm/custom_ops/op_moe.h

@@ -62,8 +62,7 @@ Tensor& quantized_moe_ffn_out(
     const Tensor& x,
     const Tensor& gate_weight,
     const Tensor& expert_bias,
-    const Tensor& packed_w1,
-    const Tensor& packed_w3,
+    const Tensor& packed_w13,
     const Tensor& packed_w2,
     int64_t num_activated_experts,
     int64_t num_experts,

extension/llm/custom_ops/op_moe_aot.cpp

@@ -16,16 +16,11 @@ namespace torch {
 namespace executor {
 namespace native {
 
-// AOT (libtorch) shim. The ExecuTorch runtime kernel is registered in
-// op_moe.cpp via EXECUTORCH_LIBRARY; here we expose the same op to the
-// torch.library system so it is discoverable from torch.export traces.
-
 Tensor& quantized_moe_ffn_out_no_context(
     const Tensor& x,
     const Tensor& gate_weight,
     const Tensor& expert_bias,
-    const Tensor& packed_w1,
-    const Tensor& packed_w3,
+    const Tensor& packed_w13,
     const Tensor& packed_w2,
     const int64_t num_activated_experts,
     const int64_t num_experts,
@@ -42,8 +37,7 @@ Tensor& quantized_moe_ffn_out_no_context(
       x,
       gate_weight,
       expert_bias,
-      packed_w1,
-      packed_w3,
+      packed_w13,
       packed_w2,
       num_activated_experts,
       num_experts,
@@ -60,8 +54,7 @@ at::Tensor quantized_moe_ffn_aten(
     const at::Tensor& x,
     const at::Tensor& gate_weight,
     const at::Tensor& expert_bias,
-    const at::Tensor& packed_w1,
-    const at::Tensor& packed_w3,
+    const at::Tensor& packed_w13,
     const at::Tensor& packed_w2,
     const int64_t num_activated_experts,
     const int64_t num_experts,
@@ -72,12 +65,11 @@ at::Tensor quantized_moe_ffn_aten(
     c10::string_view score_func,
     const double route_scale) {
   auto output = at::empty({x.size(0), dim}, x.options().dtype(at::kFloat));
-  WRAP_TO_ATEN(quantized_moe_ffn_out_no_context, 14)
+  WRAP_TO_ATEN(quantized_moe_ffn_out_no_context, 13)
   (x,
    gate_weight,
    expert_bias,
-   packed_w1,
-   packed_w3,
+   packed_w13,
    packed_w2,
    num_activated_experts,
    num_experts,
@@ -98,13 +90,13 @@ at::Tensor quantized_moe_ffn_aten(
 TORCH_LIBRARY_FRAGMENT(llama, m) {
   m.def(
       "quantized_moe_ffn(Tensor x, Tensor gate_weight, Tensor expert_bias, "
-      "Tensor packed_w1, Tensor packed_w3, Tensor packed_w2, "
+      "Tensor packed_w13, Tensor packed_w2, "
       "int num_activated_experts, int num_experts, int hidden_dim, int dim, "
       "int group_size, int weight_nbit, str score_func, float route_scale) "
       "-> Tensor");
   m.def(
       "quantized_moe_ffn.out(Tensor x, Tensor gate_weight, Tensor expert_bias, "
-      "Tensor packed_w1, Tensor packed_w3, Tensor packed_w2, "
+      "Tensor packed_w13, Tensor packed_w2, "
       "int num_activated_experts, int num_experts, int hidden_dim, int dim, "
       "int group_size, int weight_nbit, str score_func, float route_scale, "
       "*, Tensor(a!) out) -> Tensor(a!)");
@@ -116,6 +108,6 @@ TORCH_LIBRARY_IMPL(llama, CompositeExplicitAutograd, m) {
   m.impl(
       "quantized_moe_ffn.out",
       WRAP_TO_ATEN(
-          torch::executor::native::quantized_moe_ffn_out_no_context, 14));
+          torch::executor::native::quantized_moe_ffn_out_no_context, 13));
   m.impl("_quantized_moe_ffn_active", []() { return true; });
 }