Skip to content

Commit 1a6b1c7

Browse files
authored
[ExecuTorch][llm] Add Meta kernel and registration probe for quantized_moe_ffn (pytorch#21534)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: pytorch#21120 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/70/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/70/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/digantdesai/69/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/digantdesai/70/orig Differential Revision: [D102381995](https://our.internmc.facebook.com/intern/diff/D102381995/) @diff-train-skip-merge Co-authored-by: Digant Desai <digantdesai@meta.com>
1 parent b623687 commit 1a6b1c7

1 file changed

Lines changed: 112 additions & 0 deletions

File tree

‎extension/llm/custom_ops/custom_ops.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
assert op is not None
2828
op2 = torch.ops.llama.fast_hadamard_transform.default
2929
assert op2 is not None
30+
op3 = torch.ops.llama.quantized_moe_ffn.default
31+
assert op3 is not None
3032
except:
3133
# This is needed to ensure that custom ops are registered
3234
from executorch.extension.pybindings import portable_lib # noqa # usort: skip
@@ -48,6 +50,8 @@
4850
assert op is not None
4951
op2 = torch.ops.llama.fast_hadamard_transform.default
5052
assert op2 is not None
53+
op3 = torch.ops.llama.quantized_moe_ffn.default
54+
assert op3 is not None
5155

5256
custom_ops_lib = torch.library.Library("llama", "IMPL")
5357

@@ -160,6 +164,114 @@ def fast_hadamard_transform_meta(mat):
160164
return torch.empty_like(mat)
161165

162166

167+
def _validate_quantized_moe_ffn_params(
168+
x,
169+
gate_weight,
170+
expert_bias,
171+
packed_w1,
172+
packed_w3,
173+
packed_w2,
174+
num_activated_experts,
175+
num_experts,
176+
hidden_dim,
177+
dim,
178+
group_size,
179+
weight_nbit,
180+
score_func,
181+
route_scale,
182+
):
183+
assert x.dim() == 2, f"Expected x to be 2D [T, D], got {x.dim()}D"
184+
assert x.size(1) == dim, f"x last dim ({x.size(1)}) must equal dim ({dim})"
185+
assert x.dtype == torch.float32, f"Expected x to be float32, got {x.dtype}"
186+
187+
assert (
188+
gate_weight.dim() == 2
189+
and gate_weight.size(0) == num_experts
190+
and gate_weight.size(1) == dim
191+
), f"gate_weight must be [E={num_experts}, D={dim}], got {list(gate_weight.size())}"
192+
assert (
193+
gate_weight.dtype == torch.float32
194+
), f"gate_weight must be float32, got {gate_weight.dtype}"
195+
196+
if expert_bias.numel() > 0:
197+
assert (
198+
expert_bias.dim() == 1 and expert_bias.size(0) == num_experts
199+
), f"expert_bias must be [E={num_experts}] when non-empty, got {list(expert_bias.size())}"
200+
assert (
201+
expert_bias.dtype == torch.float32
202+
), f"expert_bias must be float32, got {expert_bias.dtype}"
203+
204+
for name, t in (
205+
("packed_w1", packed_w1),
206+
("packed_w3", packed_w3),
207+
("packed_w2", packed_w2),
208+
):
209+
assert (
210+
t.dim() == 2 and t.size(0) == num_experts
211+
), f"{name} must be [E={num_experts}, packed_bytes], got {list(t.size())}"
212+
assert t.dtype == torch.uint8, f"{name} must be uint8, got {t.dtype}"
213+
214+
assert packed_w1.size(1) == packed_w3.size(
215+
1
216+
), "packed_w1 and packed_w3 per-expert blob sizes must match"
217+
assert (
218+
0 < num_activated_experts <= num_experts
219+
), f"num_activated_experts ({num_activated_experts}) out of range [1, {num_experts}]"
220+
assert score_func in (
221+
"sigmoid",
222+
"softmax",
223+
), f"score_func must be 'sigmoid' or 'softmax', got {score_func!r}"
224+
assert weight_nbit in (
225+
4,
226+
8,
227+
), f"weight_nbit must be 4 or 8 (v1), got {weight_nbit}"
228+
assert group_size > 0, f"group_size must be positive, got {group_size}"
229+
assert dim > 0, f"dim must be positive, got {dim}"
230+
assert hidden_dim > 0, f"hidden_dim must be positive, got {hidden_dim}"
231+
assert (
232+
dim % group_size == 0
233+
), f"dim ({dim}) must be divisible by group_size ({group_size})"
234+
assert (
235+
hidden_dim % group_size == 0
236+
), f"hidden_dim ({hidden_dim}) must be divisible by group_size ({group_size})"
237+
238+
239+
@impl(custom_ops_lib, "quantized_moe_ffn", "Meta")
240+
def quantized_moe_ffn_meta(
241+
x,
242+
gate_weight,
243+
expert_bias,
244+
packed_w1,
245+
packed_w3,
246+
packed_w2,
247+
num_activated_experts,
248+
num_experts,
249+
hidden_dim,
250+
dim,
251+
group_size,
252+
weight_nbit,
253+
score_func,
254+
route_scale,
255+
):
256+
_validate_quantized_moe_ffn_params(
257+
x,
258+
gate_weight,
259+
expert_bias,
260+
packed_w1,
261+
packed_w3,
262+
packed_w2,
263+
num_activated_experts,
264+
num_experts,
265+
hidden_dim,
266+
dim,
267+
group_size,
268+
weight_nbit,
269+
score_func,
270+
route_scale,
271+
)
272+
return torch.empty((x.size(0), dim), dtype=torch.float32, device=x.device)
273+
274+
163275
@impl(custom_ops_lib, "custom_sdpa", "Meta")
164276
def custom_sdpa(
165277
query,

0 commit comments

Comments
 (0)