|
27 | 27 | assert op is not None |
28 | 28 | op2 = torch.ops.llama.fast_hadamard_transform.default |
29 | 29 | assert op2 is not None |
| 30 | + op3 = torch.ops.llama.quantized_moe_ffn.default |
| 31 | + assert op3 is not None |
30 | 32 | except: |
31 | 33 | # This is needed to ensure that custom ops are registered |
32 | 34 | from executorch.extension.pybindings import portable_lib # noqa # usort: skip |
|
48 | 50 | assert op is not None |
49 | 51 | op2 = torch.ops.llama.fast_hadamard_transform.default |
50 | 52 | assert op2 is not None |
| 53 | + op3 = torch.ops.llama.quantized_moe_ffn.default |
| 54 | + assert op3 is not None |
51 | 55 |
|
52 | 56 | custom_ops_lib = torch.library.Library("llama", "IMPL") |
53 | 57 |
|
@@ -160,6 +164,114 @@ def fast_hadamard_transform_meta(mat): |
160 | 164 | return torch.empty_like(mat) |
161 | 165 |
|
162 | 166 |
|
| 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 | + |
163 | 275 | @impl(custom_ops_lib, "custom_sdpa", "Meta") |
164 | 276 | def custom_sdpa( |
165 | 277 | query, |
|
0 commit comments