diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index d09d54c6389..a29b9347373 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/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,