[ExecuTorch][llm] Add Meta kernel and registration probe for quantized_moe_ffn

extension/llm/custom_ops/custom_ops.py

@@ -27,6 +27,8 @@
     assert op is not None
     op2 = torch.ops.llama.fast_hadamard_transform.default
     assert op2 is not None
+    op3 = torch.ops.llama.quantized_moe_ffn.default
+    assert op3 is not None
 except:
     # This is needed to ensure that custom ops are registered
     from executorch.extension.pybindings import portable_lib  # noqa # usort: skip
@@ -48,6 +50,8 @@
     assert op is not None
     op2 = torch.ops.llama.fast_hadamard_transform.default
     assert op2 is not None
+    op3 = torch.ops.llama.quantized_moe_ffn.default
+    assert op3 is not None
 
 custom_ops_lib = torch.library.Library("llama", "IMPL")
 
@@ -160,6 +164,114 @@ def fast_hadamard_transform_meta(mat):
     return torch.empty_like(mat)
 
 
+def _validate_quantized_moe_ffn_params(
+    x,
+    gate_weight,
+    expert_bias,
+    packed_w1,
+    packed_w3,
+    packed_w2,
+    num_activated_experts,
+    num_experts,
+    hidden_dim,
+    dim,
+    group_size,
+    weight_nbit,
+    score_func,
+    route_scale,
+):
+    assert x.dim() == 2, f"Expected x to be 2D [T, D], got {x.dim()}D"
+    assert x.size(1) == dim, f"x last dim ({x.size(1)}) must equal dim ({dim})"
+    assert x.dtype == torch.float32, f"Expected x to be float32, got {x.dtype}"
+
+    assert (
+        gate_weight.dim() == 2
+        and gate_weight.size(0) == num_experts
+        and gate_weight.size(1) == dim
+    ), f"gate_weight must be [E={num_experts}, D={dim}], got {list(gate_weight.size())}"
+    assert (
+        gate_weight.dtype == torch.float32
+    ), f"gate_weight must be float32, got {gate_weight.dtype}"
+
+    if expert_bias.numel() > 0:
+        assert (
+            expert_bias.dim() == 1 and expert_bias.size(0) == num_experts
+        ), f"expert_bias must be [E={num_experts}] when non-empty, got {list(expert_bias.size())}"
+        assert (
+            expert_bias.dtype == torch.float32
+        ), f"expert_bias must be float32, got {expert_bias.dtype}"
+
+    for name, t in (
+        ("packed_w1", packed_w1),
+        ("packed_w3", packed_w3),
+        ("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}]"
+    assert score_func in (
+        "sigmoid",
+        "softmax",
+    ), f"score_func must be 'sigmoid' or 'softmax', got {score_func!r}"
+    assert weight_nbit in (
+        4,
+        8,
+    ), f"weight_nbit must be 4 or 8 (v1), got {weight_nbit}"
+    assert group_size > 0, f"group_size must be positive, got {group_size}"
+    assert dim > 0, f"dim must be positive, got {dim}"
+    assert hidden_dim > 0, f"hidden_dim must be positive, got {hidden_dim}"
+    assert (
+        dim % group_size == 0
+    ), f"dim ({dim}) must be divisible by group_size ({group_size})"
+    assert (
+        hidden_dim % group_size == 0
+    ), f"hidden_dim ({hidden_dim}) must be divisible by group_size ({group_size})"
+
+
+@impl(custom_ops_lib, "quantized_moe_ffn", "Meta")
+def quantized_moe_ffn_meta(
+    x,
+    gate_weight,
+    expert_bias,
+    packed_w1,
+    packed_w3,
+    packed_w2,
+    num_activated_experts,
+    num_experts,
+    hidden_dim,
+    dim,
+    group_size,
+    weight_nbit,
+    score_func,
+    route_scale,
+):
+    _validate_quantized_moe_ffn_params(
+        x,
+        gate_weight,
+        expert_bias,
+        packed_w1,
+        packed_w3,
+        packed_w2,
+        num_activated_experts,
+        num_experts,
+        hidden_dim,
+        dim,
+        group_size,
+        weight_nbit,
+        score_func,
+        route_scale,
+    )
+    return torch.empty((x.size(0), dim), dtype=torch.float32, device=x.device)
+
+
 @impl(custom_ops_lib, "custom_sdpa", "Meta")
 def custom_sdpa(
     query,