From 6b193c4cae89ca79de74da41eed9910d7b131d51 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 21 Jul 2026 20:38:14 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../models/llama/source_transformation/moe.py | 30 +++---- extension/llm/custom_ops/custom_ops.py | 15 +--- extension/llm/custom_ops/op_moe.cpp | 78 +++++++------------ extension/llm/custom_ops/op_moe.h | 3 +- extension/llm/custom_ops/op_moe_aot.cpp | 24 ++---- .../llm/custom_ops/test_quantized_moe.py | 37 ++------- 6 files changed, 60 insertions(+), 127 deletions(-) diff --git a/examples/models/llama/source_transformation/moe.py b/examples/models/llama/source_transformation/moe.py index 8770ca6c2ba..64dad66a458 100644 --- a/examples/models/llama/source_transformation/moe.py +++ b/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, diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index a29b9347373..874f2b4860d 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/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,8 +201,7 @@ 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 ( @@ -211,9 +209,6 @@ def _validate_quantized_moe_ffn_params( ), 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, diff --git a/extension/llm/custom_ops/op_moe.cpp b/extension/llm/custom_ops/op_moe.cpp index 6b3233b8390..f8fb0801000 100644 --- a/extension/llm/custom_ops/op_moe.cpp +++ b/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(); const float* expert_bias_ptr = use_expert_bias ? expert_bias.const_data_ptr() : nullptr; - const uint8_t* pw1_ptr = packed_w1.const_data_ptr(); - const uint8_t* pw3_ptr = packed_w3.const_data_ptr(); + const uint8_t* pw13_ptr = packed_w13.const_data_ptr(); const uint8_t* pw2_ptr = packed_w2.const_data_ptr(); float* out_ptr = out.mutable_data_ptr(); @@ -540,8 +519,10 @@ Tensor& quantized_moe_ffn_out( max_m_e = m_e; } } - std::vector h1_buf(static_cast(max_m_e * F), 0.0f); - std::vector h3_buf(static_cast(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 h13_buf(static_cast(max_m_e * 2 * F), 0.0f); + std::vector mid_buf(static_cast(max_m_e * F), 0.0f); std::vector out_e_buf(static_cast(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, diff --git a/extension/llm/custom_ops/op_moe.h b/extension/llm/custom_ops/op_moe.h index 64c57c34124..0e6127930d3 100644 --- a/extension/llm/custom_ops/op_moe.h +++ b/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, diff --git a/extension/llm/custom_ops/op_moe_aot.cpp b/extension/llm/custom_ops/op_moe_aot.cpp index 85c4a846a5a..b22c52e0107 100644 --- a/extension/llm/custom_ops/op_moe_aot.cpp +++ b/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; }); } diff --git a/extension/llm/custom_ops/test_quantized_moe.py b/extension/llm/custom_ops/test_quantized_moe.py index 0b4299e8f12..f529dfb9578 100644 --- a/extension/llm/custom_ops/test_quantized_moe.py +++ b/extension/llm/custom_ops/test_quantized_moe.py @@ -243,10 +243,8 @@ def test_replaces_module_in_place_with_correct_buffers(self) -> None: self.assertEqual(replaced.hidden_dim, 32) self.assertEqual(tuple(replaced.gate_weight.shape), (4, 32)) self.assertEqual(replaced.expert_bias.numel(), 4) - self.assertEqual(replaced.packed_w1.shape[0], 4) - self.assertEqual(replaced.packed_w3.shape[0], 4) + self.assertEqual(replaced.packed_w13.shape[0], 4) self.assertEqual(replaced.packed_w2.shape[0], 4) - self.assertEqual(replaced.packed_w1.shape[1], replaced.packed_w3.shape[1]) def test_torchao_quantized_gate_exports_as_fp32_buffer(self) -> None: moe = _build_moe_eager( @@ -298,7 +296,7 @@ def test_buffers_stay_fp32_after_to_half(self) -> None: replaced = wrapper.block_sparse_moe self.assertEqual(replaced.gate_weight.dtype, torch.float32) self.assertEqual(replaced.expert_bias.dtype, torch.float32) - self.assertEqual(replaced.packed_w1.dtype, torch.uint8) + self.assertEqual(replaced.packed_w13.dtype, torch.uint8) def test_buffers_stay_fp32_after_to_bfloat16(self) -> None: moe = _build_moe_eager( @@ -325,8 +323,7 @@ def test_3d_input_shape_preserved_through_forward(self) -> None: x_meta.view(-1, 32), replaced.gate_weight.to("meta"), replaced.expert_bias.to("meta"), - replaced.packed_w1.to("meta"), - replaced.packed_w3.to("meta"), + replaced.packed_w13.to("meta"), replaced.packed_w2.to("meta"), replaced.num_activated_experts, replaced.num_experts, @@ -375,8 +372,7 @@ def test_meta_kernel_returns_correct_output_shape(self) -> None: x_meta, replaced.gate_weight.to("meta"), replaced.expert_bias.to("meta"), - replaced.packed_w1.to("meta"), - replaced.packed_w3.to("meta"), + replaced.packed_w13.to("meta"), replaced.packed_w2.to("meta"), replaced.num_activated_experts, replaced.num_experts, @@ -401,8 +397,7 @@ def test_meta_kernel_single_token(self) -> None: x_meta, replaced.gate_weight.to("meta"), replaced.expert_bias.to("meta"), - replaced.packed_w1.to("meta"), - replaced.packed_w3.to("meta"), + replaced.packed_w13.to("meta"), replaced.packed_w2.to("meta"), replaced.num_activated_experts, replaced.num_experts, @@ -465,8 +460,7 @@ def _make_valid_inputs(self) -> dict: x=torch.empty((4, r.dim), dtype=torch.float32, device="meta"), gate_weight=r.gate_weight.to("meta"), expert_bias=r.expert_bias.to("meta"), - packed_w1=r.packed_w1.to("meta"), - packed_w3=r.packed_w3.to("meta"), + packed_w13=r.packed_w13.to("meta"), packed_w2=r.packed_w2.to("meta"), num_activated_experts=r.num_activated_experts, num_experts=r.num_experts, @@ -502,9 +496,9 @@ def test_rejects_num_activated_experts_zero(self) -> None: with self.assertRaises(AssertionError): torch.ops.llama.quantized_moe_ffn(**kw) - def test_rejects_non_uint8_packed_w1(self) -> None: + def test_rejects_non_uint8_packed_w13(self) -> None: kw = self._make_valid_inputs() - kw["packed_w1"] = kw["packed_w1"].to(torch.float32) + kw["packed_w13"] = kw["packed_w13"].to(torch.float32) with self.assertRaises(AssertionError): torch.ops.llama.quantized_moe_ffn(**kw) @@ -514,21 +508,6 @@ def test_rejects_non_uint8_packed_w2(self) -> None: with self.assertRaises(AssertionError): torch.ops.llama.quantized_moe_ffn(**kw) - def test_rejects_non_uint8_packed_w3(self) -> None: - kw = self._make_valid_inputs() - kw["packed_w3"] = kw["packed_w3"].to(torch.float32) - with self.assertRaises(AssertionError): - torch.ops.llama.quantized_moe_ffn(**kw) - - def test_rejects_mismatched_packed_w1_w3_sizes(self) -> None: - kw = self._make_valid_inputs() - e = kw["num_experts"] - kw["packed_w3"] = torch.empty( - (e, kw["packed_w3"].size(1) + 1), dtype=torch.uint8, device="meta" - ) - with self.assertRaises(AssertionError): - torch.ops.llama.quantized_moe_ffn(**kw) - def test_rejects_non_fp32_gate_weight(self) -> None: kw = self._make_valid_inputs() kw["gate_weight"] = kw["gate_weight"].to(torch.float16)