diff --git a/backends/mlx/custom_kernel_ops/fp_int4_gemm.py b/backends/mlx/custom_kernel_ops/fp_int4_gemm.py new file mode 100644 index 00000000000..7259011299d --- /dev/null +++ b/backends/mlx/custom_kernel_ops/fp_int4_gemm.py @@ -0,0 +1,518 @@ +# +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# + +""" +``mlx::fp_int4_gemm``: fp16 / bf16 activation x 4-bit per-channel weight linear, +lowered to Metal Performance Primitives (Metal 4) kernels that multiply +``half`` (or ``bfloat``) by ``int4b_format`` with FP32 accumulation. + +Weights are symmetric per output channel (no zero points, no groups), the same +format as ``mlx::int8_intx_gemm`` with ``bits=4``: + + W[n, k] = w_scale[n] * wq[n, k], wq in [-8, 7] + Y[m, n] = w_scale[n] * sum_k x[m, k] * wq[n, k] + +``w_packed`` is ``uint8 [N, K // 2]`` from +:func:`executorch.backends.mlx.custom_kernel_ops.int8_intx_gemm.pack_int4`. + +The kernel is picked at runtime from M (see :func:`full_k_plans`): + * matvec for M <= 2: streams packed weights like MLX's qmv. + * otherwise MPP over the full K straight from device memory, optionally + split over K into FP32 partials plus a reduce kernel when that adds + enough threadgroups. + +Constraints for MLX lowering: + * ``K`` must be statically known and a multiple of 64. + * ``N`` must be statically known. + * Activations must be fp16 or bf16. + * Requires an MPP-capable GPU (Metal 4). + +Usage:: + + import executorch.backends.mlx.custom_kernel_ops.fp_int4_gemm # noqa: F401 + + y = torch.ops.mlx.fp_int4_gemm(x, w_packed, w_scale) + # x: (..., K) fp16 / bf16 + # w_packed: (N, K // 2) uint8 + # w_scale: (N,) float + # y: (..., N) x.dtype +""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch +from executorch.backends.mlx.custom_kernel_ops.int8_intx_gemm import unpack_int4 +from torch import nn, Tensor +from torch.fx.node import Node + + +class FpInt4Linear(nn.Module): + """Bias-free linear computed by ``mlx::fp_int4_gemm``.""" + + def __init__(self, w_packed: Tensor, w_scale: Tensor): + super().__init__() + self.in_features = 2 * w_packed.shape[1] + self.out_features = w_packed.shape[0] + self.register_buffer("w_packed", w_packed) + self.register_buffer("w_scale", w_scale.float()) + + def _apply(self, fn, recurse=True): + # Keep fp32 scales through model.to(activation_dtype); move device only. + w_scale = self.w_scale + result = super()._apply(fn, recurse) + self.w_scale = w_scale.to(self.w_packed.device) + return result + + def forward(self, x: Tensor) -> Tensor: + return torch.ops.mlx.fp_int4_gemm(x, self.w_packed, self.w_scale) + + def extra_repr(self) -> str: + return f"in_features={self.in_features}, out_features={self.out_features}" + + +@torch.library.custom_op("mlx::fp_int4_gemm", mutates_args=()) +def fp_int4_gemm(x: Tensor, w_packed: Tensor, w_scale: Tensor) -> Tensor: + wq = unpack_int4(w_packed).float() + y = (x.float() @ wq.T) * w_scale.float() + return y.to(x.dtype) + + +@torch.library.register_fake("mlx::fp_int4_gemm") +def fp_int4_gemm_fake(x: Tensor, w_packed: Tensor, w_scale: Tensor) -> Tensor: + return x.new_empty((*x.shape[:-1], w_packed.shape[0])) + + +# --------------------------------------------------------------------------- +# MLX handler +# --------------------------------------------------------------------------- + +from executorch.backends.mlx.builder.op_helpers import ( + emit_ceil_div, + emit_floordiv, + emit_if_else, + emit_product, + emit_shape, + emit_sub_int, + torch_dtype_to_scalar_type, +) +from executorch.backends.mlx.builder.op_registry import REGISTRY +from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder +from executorch.backends.mlx.builder.slot_manager import Slot +from executorch.backends.mlx.serialization.mlx_graph_schema import ( + IntOrVid, + MetalKernelNode, +) + +_K_ALIGN = 64 +_THREADS = 128 +_MATVEC_SIMDGROUPS = 2 +_MATVEC_ROWS_PER_SIMDGROUP = 4 +_MATVEC_BLOCK = 512 +_MAX_M = 2**31 - 1 + + +class FpInt4Plan(NamedTuple): + """Kernel choice for ``M <= max_m``. ``bn == 0`` selects the matvec kernel + with ``bm`` activation rows per threadgroup; otherwise a full-K MPP GEMM + with a ``bm x bn`` tile, split over K when ``splits > 1``.""" + + max_m: int + bm: int + bn: int + splits: int = 1 + + +def full_k_plans(N: int, K: int) -> list[FpInt4Plan]: + """Kernel plans, smallest ``max_m`` first; the last covers every larger M. + Untuned: the int8_intx fallback pattern (more splits when N gives few + output tiles) until ``test/bench_muse_glimmer_gemm.py --kernels fp_int4`` picks + per-shape tiles.""" + plans = [] + if K % _MATVEC_BLOCK == 0: + plans += [FpInt4Plan(1, 1, 0), FpInt4Plan(2, 2, 0)] + small_bn, small_splits = (64, 8) if N <= 16384 else (128, 2) + tiles = [(8, small_bn, small_splits), (16, small_bn, small_splits), (32, 128, 1)] + for max_m, (bm, bn, splits) in zip((8, 16, _MAX_M), tiles): + while K % (_K_ALIGN * splits): + splits //= 2 + plans.append(FpInt4Plan(max_m, bm, bn, splits)) + return plans + + +_HEADER = """ +#include +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; +""" + +# Template args T, N, K, BM, BN. MPP multiplies a T [K, M] tensor by an int4b +# [K, N] tensor (W transposed) over the whole K for each BM x BN output tile. +_FULL_K_SOURCE = """ + int M = 1; + for (int d = 0; d < x_ndim - 1; ++d) { + M *= x_shape[d]; + } + const int n0 = int(threadgroup_position_in_grid.x) * BN; + const int m0 = int(threadgroup_position_in_grid.y) * BM; + + tensor, tensor_inline> X( + (device T*)x, dextents(K, M)); + tensor, tensor_inline> W( + (device uchar*)w_packed, dextents(K, N), array{1, K}); + constexpr auto desc = matmul2d_descriptor( + BM, BN, static_cast(dynamic_extent), false, true, false, + matmul2d_descriptor::mode::multiply); + matmul2d> mm; + auto mA = X.slice(0, m0); + auto mB = W.slice(0, n0); + auto acc = mm.template get_destination_cooperative_tensor(); + mm.run(mA, mB, acc); + + #pragma clang loop unroll(full) + for (ushort i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) { + const auto coord = acc.get_multidimensional_index(i); + const int n = n0 + int(coord[0]); + const int m = m0 + int(coord[1]); + if (m < M && n < N) { + out[size_t(m) * N + n] = T(acc[i] * float(w_scale[n])); + } + } + } +""" + +# Template args T, N, K, BM, BN, SPLITS (K % (64 * SPLITS) == 0): grid z picks +# one of SPLITS contiguous K ranges; each threadgroup writes FP32 partial sums +# to partial[s, m, n]. +_SPLIT_K_SOURCE = """ + constexpr int KS = K / SPLITS; + int M = 1; + for (int d = 0; d < x_ndim - 1; ++d) { + M *= x_shape[d]; + } + const int n0 = int(threadgroup_position_in_grid.x) * BN; + const int m0 = int(threadgroup_position_in_grid.y) * BM; + const int s = int(threadgroup_position_in_grid.z); + const int k0 = s * KS; + + tensor, tensor_inline> X( + (device T*)x + k0, dextents(KS, M), array{1, K}); + tensor, tensor_inline> W( + (device uchar*)w_packed + k0 / 2, dextents(KS, N), + array{1, K}); + constexpr auto desc = matmul2d_descriptor( + BM, BN, static_cast(dynamic_extent), false, true, false, + matmul2d_descriptor::mode::multiply); + matmul2d> mm; + auto mA = X.slice(0, m0); + auto mB = W.slice(0, n0); + auto acc = mm.template get_destination_cooperative_tensor(); + mm.run(mA, mB, acc); + + #pragma clang loop unroll(full) + for (ushort i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) { + const auto coord = acc.get_multidimensional_index(i); + const int n = n0 + int(coord[0]); + const int m = m0 + int(coord[1]); + if (m < M && n < N) { + partial[(size_t(s) * M + m) * N + n] = acc[i]; + } + } + } +""" + +# Template args T, N, SPLITS: one thread per output sums the partials and +# applies the channel scale. +_SPLIT_K_REDUCE_SOURCE = """ + const int M = partial_shape[1]; + const size_t idx = thread_position_in_grid.x; + if (idx >= size_t(M) * N) { + return; + } + float sum = 0.0f; + for (int s = 0; s < SPLITS; ++s) { + sum += partial[size_t(s) * M * N + idx]; + } + out[idx] = T(sum * float(w_scale[idx % N])); +""" + +# Template args T, N, K, MR (K % 512 == 0) for small M. Each SIMD group +# computes 4 output channels for up to MR activation rows; each lane covers 16 +# K values per 512-wide step. XOR 0x88 makes nibbles unsigned (q = w + 8) so +# masked ushort words convert directly to float against activations pre-scaled +# by 1, 1/16, 1/256, 1/4096 (exact powers of two); subtracting 8 * sum(x) +# restores the signed dot product. +_MATVEC_SOURCE = """ + constexpr int SG = 2; + constexpr int RPS = 4; + constexpr int VPT = 16; + constexpr int BLOCK = 32 * VPT; + constexpr int KP = K / 2; + + int M = 1; + for (int d = 0; d < x_ndim - 1; ++d) { + M *= x_shape[d]; + } + const int lane = int(thread_index_in_simdgroup); + const int sg = int(simdgroup_index_in_threadgroup); + const int n0 = (int(threadgroup_position_in_grid.x) * SG + sg) * RPS; + const int m0 = int(threadgroup_position_in_grid.y) * MR; + if (n0 >= N) { + return; + } + const int mrows = min(MR, M - m0); + + device const uchar* w_rows[RPS]; + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + w_rows[r] = w_packed + size_t(min(n0 + r, N - 1)) * KP + lane * (VPT / 2); + } + device const T* x_base = x + size_t(m0) * K + lane * VPT; + const float4 prescale = + float4(1.0f, 1.0f / 16.0f, 1.0f / 256.0f, 1.0f / 4096.0f); + + float acc[MR][RPS]; + float xsum[MR]; + #pragma clang loop unroll(full) + for (int m = 0; m < MR; ++m) { + xsum[m] = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + acc[m][r] = 0.0f; + } + } + + for (int k = 0; k < K; k += BLOCK) { + // Converted once per block and reused for every activation row. + float4 wf[RPS][4]; + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + const ushort4 q = as_type( + *(device const uint2*)(w_rows[r] + k / 2) ^ uint2(0x88888888u)); + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + wf[r][i] = float4( + q[i] & 0x000f, q[i] & 0x00f0, q[i] & 0x0f00, q[i] & 0xf000); + } + } + #pragma clang loop unroll(full) + for (int m = 0; m < MR; ++m) { + if (m < mrows) { + device const vec* xv = + (device const vec*)(x_base + size_t(m) * K + k); + float4 xs[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const float4 v = float4(xv[i]); + xsum[m] += (v.x + v.y) + (v.z + v.w); + xs[i] = v * prescale; + } + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + float a = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + a += dot(xs[i], wf[r][i]); + } + acc[m][r] += a; + } + } + } + } + + #pragma clang loop unroll(full) + for (int m = 0; m < MR; ++m) { + if (m < mrows) { + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + const float total = simd_sum(acc[m][r] - 8.0f * xsum[m]); + const int n = n0 + r; + if (lane == 0 && n < N) { + out[size_t(m0 + m) * N + n] = T(total * float(w_scale[n])); + } + } + } + } +""" + + +def _lit(*values: int) -> list[IntOrVid]: + return [IntOrVid.from_literal(v) for v in values] + + +@REGISTRY.register(target=[torch.ops.mlx.fp_int4_gemm.default]) +def _fp_int4_gemm_handler(P: MLXProgramBuilder, n: Node) -> Slot: + x_slot, w_packed_slot, w_scale_slot = P.args(n) + x_node = n.args[0] + x_meta = x_node.meta["val"] + N, KP = n.args[1].meta["val"].shape + if not isinstance(N, int) or not isinstance(KP, int): + raise NotImplementedError( + "mlx::fp_int4_gemm: w_packed shape must be statically known" + ) + K = 2 * KP + if K % _K_ALIGN != 0: + raise NotImplementedError( + f"mlx::fp_int4_gemm: K must be a multiple of {_K_ALIGN}; got K={K}" + ) + if x_meta.shape[-1] != K: + raise ValueError( + f"mlx::fp_int4_gemm: x last dim {x_meta.shape[-1]} does not match " + f"K={K} from w_packed" + ) + if x_meta.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError( + f"mlx::fp_int4_gemm: activations must be fp16 or bf16; got {x_meta.dtype}" + ) + + dtype_int = torch_dtype_to_scalar_type(x_meta.dtype) + out = P.make_or_get_slot(n) + leading = emit_shape(P, x_node, x_slot, end_dim=-1) + out_shape_flat = leading + [IntOrVid.from_literal(N)] + m_iov = emit_product(P, leading) + _, partial_slot = P.make_tmp_slot() + n_lit = IntOrVid.from_literal(N) + + def emit_matvec(mr: int) -> None: + rows_per_tg = _MATVEC_SIMDGROUPS * _MATVEC_ROWS_PER_SIMDGROUP + P.emit( + MetalKernelNode( + name="fp_int4_gemm_matvec", + source=_MATVEC_SOURCE, + header=_HEADER, + inputs=[ + P.slot_to_tid(s) for s in (x_slot, w_packed_slot, w_scale_slot) + ], + outputs=[P.slot_to_tid(out)], + grid=[ + IntOrVid.from_literal( + (N + rows_per_tg - 1) // rows_per_tg * _MATVEC_SIMDGROUPS * 32 + ), + emit_ceil_div(P, m_iov, mr), + IntOrVid.from_literal(1), + ], + threadgroup=_lit(_MATVEC_SIMDGROUPS * 32, 1, 1), + input_names=["x", "w_packed", "w_scale"], + output_names=["out"], + output_shapes_flat=out_shape_flat, + output_shape_lengths=[len(out_shape_flat)], + output_dtypes=[dtype_int], + template_arg_names=["T", "N", "K", "MR"], + template_arg_kinds=[2, 0, 0, 0], + template_arg_values=[dtype_int, N, K, mr], + ) + ) + + def emit_mpp(bm: int, bn: int, splits: int) -> None: + grid = [ + IntOrVid.from_literal((N + bn - 1) // bn * _THREADS), + emit_ceil_div(P, m_iov, bm), + IntOrVid.from_literal(1), + ] + if splits == 1: + P.emit( + MetalKernelNode( + name="fp_int4_gemm_full_k", + source=_FULL_K_SOURCE, + header=_HEADER, + inputs=[ + P.slot_to_tid(s) for s in (x_slot, w_packed_slot, w_scale_slot) + ], + outputs=[P.slot_to_tid(out)], + grid=grid, + threadgroup=_lit(_THREADS, 1, 1), + input_names=["x", "w_packed", "w_scale"], + output_names=["out"], + output_shapes_flat=out_shape_flat, + output_shape_lengths=[len(out_shape_flat)], + output_dtypes=[dtype_int], + template_arg_names=["T", "N", "K", "BM", "BN"], + template_arg_kinds=[2, 0, 0, 0, 0], + template_arg_values=[dtype_int, N, K, bm, bn], + ) + ) + return + P.emit( + MetalKernelNode( + name="fp_int4_gemm_split_k", + source=_SPLIT_K_SOURCE, + header=_HEADER, + inputs=[P.slot_to_tid(s) for s in (x_slot, w_packed_slot)], + outputs=[P.slot_to_tid(partial_slot)], + grid=grid[:2] + [IntOrVid.from_literal(splits)], + threadgroup=_lit(_THREADS, 1, 1), + input_names=["x", "w_packed"], + output_names=["partial"], + output_shapes_flat=[IntOrVid.from_literal(splits), m_iov, n_lit], + output_shape_lengths=[3], + output_dtypes=[torch_dtype_to_scalar_type(torch.float32)], + template_arg_names=["T", "N", "K", "BM", "BN", "SPLITS"], + template_arg_kinds=[2, 0, 0, 0, 0, 0], + template_arg_values=[dtype_int, N, K, bm, bn, splits], + ) + ) + reduce_threads = emit_product( + P, + [ + emit_ceil_div(P, emit_product(P, [m_iov, n_lit]), 256), + IntOrVid.from_literal(256), + ], + ) + P.emit( + MetalKernelNode( + name="fp_int4_gemm_split_k_reduce", + source=_SPLIT_K_REDUCE_SOURCE, + header=_HEADER, + inputs=[P.slot_to_tid(s) for s in (partial_slot, w_scale_slot)], + outputs=[P.slot_to_tid(out)], + grid=[ + reduce_threads, + IntOrVid.from_literal(1), + IntOrVid.from_literal(1), + ], + threadgroup=_lit(256, 1, 1), + input_names=["partial", "w_scale"], + output_names=["out"], + output_shapes_flat=out_shape_flat, + output_shape_lengths=[len(out_shape_flat)], + output_dtypes=[dtype_int], + template_arg_names=["T", "N", "SPLITS"], + template_arg_kinds=[2, 0, 0], + template_arg_values=[dtype_int, N, splits], + ) + ) + + def emit_plan(plan: FpInt4Plan) -> None: + if plan.bn == 0: + emit_matvec(plan.bm) + else: + emit_mpp(plan.bm, plan.bn, plan.splits) + + # (M - 1) // max_m == 0 exactly when M <= max_m, so each IfNode's else + # branch takes the first plan that covers M. + m_minus_1 = emit_sub_int(P, m_iov, IntOrVid.from_literal(1)) + + def emit_plans(plans: list[FpInt4Plan]) -> None: + plan, rest = plans[0], plans[1:] + if not rest: + emit_plan(plan) + return + emit_if_else( + P, + emit_floordiv(P, m_minus_1, IntOrVid.from_literal(plan.max_m)), + emit_then=lambda: emit_plans(rest), + emit_else=lambda: emit_plan(plan), + ) + + emit_plans(full_k_plans(N, K)) + return out diff --git a/backends/mlx/custom_kernel_ops/int8_intx_gemm.py b/backends/mlx/custom_kernel_ops/int8_intx_gemm.py new file mode 100644 index 00000000000..48f37dc8dfb --- /dev/null +++ b/backends/mlx/custom_kernel_ops/int8_intx_gemm.py @@ -0,0 +1,807 @@ +# +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# + +""" +``mlx::int8_intx_gemm``: 8-bit dynamic activation x 4- or 2-bit per-channel +weight linear, lowered to Metal Performance Primitives (Metal 4) kernels that +do INT8 x INT{4,2} -> INT32 accumulation over the full K dimension. + +Quantization scheme (symmetric, no zero points), with ``bits`` in {4, 2} and +``qmin, qmax = -2**(bits - 1), 2**(bits - 1) - 1``: + + x_scale[m] = max(max_k |x[m, k]|, eps) / 127 per token + xq[m, k] = clamp(round(x[m, k] * (1 / x_scale[m])), -128, 127) + W[n, k] = w_scale[n] * wq[n, k], wq in [qmin, qmax] per channel + Y[m, n] = x_scale[m] * w_scale[n] * sum_k xq[m, k] * wq[n, k] + +``round`` is round-half-to-even and the scale math is fp32, so the kernels and +the eager reference quantize identically. + +Weight packing (see :func:`pack_intx`): ``w_packed`` is +``uint8 [N, K * bits // 8]``; element ``k`` of row ``n`` sits in byte +``k * bits // 8`` at bit offset ``(k * bits) % 8`` (low bits first), as +two's-complement ``bits``-bit integers (the layout of Metal's ``int4b_format`` +and ``int2b_format``). ``bits`` is inferred from ``w_packed``'s width. + +INT32 accumulation cannot overflow while ``K * 128 * 8 < 2**31``. + +Lowering: a per-token kernel writes int8 activations and scales; then the GEMM +is picked at runtime from M (see :func:`full_k_plans`): a matvec for M <= 2, +otherwise MPP ``int8 x int{4,2}b_format`` over the full K straight from device +memory, on activations zero-padded to whole BM tiles, split over K into INT32 +partials plus a reduce kernel when that adds enough threadgroups. All paths +produce bit-identical results. + +Constraints for MLX lowering: + * ``K`` must be statically known and a multiple of 64. + * ``N`` must be statically known. + * Requires an MPP-capable GPU (Metal 4). + +Usage:: + + import executorch.backends.mlx.custom_kernel_ops.int8_intx_gemm # noqa: F401 + + y = torch.ops.mlx.int8_intx_gemm(x, w_packed, w_scale) + # x: (..., K) fp16 / bf16 / fp32 + # w_packed: (N, K * bits // 8) uint8, bits in {4, 2} + # w_scale: (N,) float + # y: (..., N) x.dtype +""" + +from __future__ import annotations + +from typing import Callable, NamedTuple + +import torch +from torch import nn, Tensor +from torch.fx.node import Node + +_EPS = 1e-10 +BITS = (4, 2) + + +# --------------------------------------------------------------------------- +# Packing / quantization helpers +# --------------------------------------------------------------------------- + + +def _check_bits(bits: int) -> None: + if bits not in BITS: + raise ValueError(f"bits must be one of {BITS}; got {bits}") + + +def pack_intx(wq: Tensor, bits: int) -> Tensor: + """Pack ``bits``-bit values of shape ``[N, K]`` into + ``uint8 [N, K * bits // 8]``, low bits first.""" + _check_bits(bits) + per_byte = 8 // bits + if wq.shape[-1] % per_byte != 0: + raise ValueError( + f"pack_intx: K must be a multiple of {per_byte}; got {wq.shape[-1]}" + ) + fields = wq.to(torch.int16) & ((1 << bits) - 1) + out = fields[..., 0::per_byte] + for i in range(1, per_byte): + out = out | (fields[..., i::per_byte] << (i * bits)) + return out.to(torch.uint8) + + +def unpack_intx(w_packed: Tensor, bits: int) -> Tensor: + """Inverse of :func:`pack_intx`; returns ``int8 [N, K]``.""" + _check_bits(bits) + b = w_packed.to(torch.int16) + mask = (1 << bits) - 1 + wq = torch.stack([(b >> s) & mask for s in range(0, 8, bits)], dim=-1).flatten(-2) + return torch.where(wq >= 1 << (bits - 1), wq - (1 << bits), wq).to(torch.int8) + + +def pack_int4(wq: Tensor) -> Tensor: + return pack_intx(wq, 4) + + +def unpack_int4(w_packed: Tensor) -> Tensor: + return unpack_intx(w_packed, 4) + + +def infer_bits(w_packed: Tensor, K: int) -> int: + """Weight bit width from ``w_packed [N, K * bits // 8]`` and ``K``.""" + bits = 8 * w_packed.shape[-1] // K + if bits not in BITS or K * bits != 8 * w_packed.shape[-1]: + raise ValueError( + f"w_packed width {w_packed.shape[-1]} is not K * bits // 8 for " + f"K={K} and bits in {BITS}" + ) + return bits + + +def quantize_per_token_int8(x: Tensor) -> tuple[Tensor, Tensor]: + """Kernel-matching symmetric per-token quantization of ``x [M, K]``. + + Returns ``(xq int8 [M, K], x_scale fp32 [M, 1])``. + """ + xf = x.float() + x_scale = (xf.abs().amax(dim=-1, keepdim=True) / 127).clamp(min=_EPS) + xq = torch.round(xf * torch.reciprocal(x_scale)).clamp(-128, 127) + return xq.to(torch.int8), x_scale + + +def quantize_intx_per_channel( + rows: Callable[[int, int], Tensor], N: int, bits: int, chunk_rows: int = 4096 +) -> tuple[Tensor, Tensor]: + """Symmetric per-output-channel ``bits``-bit quantization of an ``[N, K]`` + weight. The scale maps max |w| to ``2**(bits - 1) - 0.5`` (7.5 for int4, + 1.5 for int2) so the range is used symmetrically before clamping. + + ``rows(start, end)`` returns rows ``[start, end)`` as a float tensor, so a + large weight is only ever materialized ``chunk_rows`` rows at a time. + Returns ``(w_packed uint8 [N, K * bits // 8], w_scale fp32 [N])``. + """ + _check_bits(bits) + qmax = (1 << (bits - 1)) - 1 + packed, scales = [], [] + for start in range(0, N, chunk_rows): + w = rows(start, min(start + chunk_rows, N)).float() + scale = (w.abs().amax(dim=-1) / (qmax + 0.5)).clamp(min=_EPS) + wq = torch.round(w / scale[:, None]).clamp(-qmax - 1, qmax) + packed.append(pack_intx(wq, bits)) + scales.append(scale) + del w, wq + return torch.cat(packed), torch.cat(scales) + + +# --------------------------------------------------------------------------- +# Custom op + eager fallback +# --------------------------------------------------------------------------- + + +class Int8IntxLinear(nn.Module): + """Bias-free linear computed by ``mlx::int8_intx_gemm``.""" + + def __init__(self, w_packed: Tensor, w_scale: Tensor, bits: int): + super().__init__() + _check_bits(bits) + self.bits = bits + self.in_features = w_packed.shape[1] * 8 // bits + self.out_features = w_packed.shape[0] + self.register_buffer("w_packed", w_packed) + self.register_buffer("w_scale", w_scale.float()) + + def _apply(self, fn, recurse=True): + # Keep fp32 scales through model.to(activation_dtype); move device only. + w_scale = self.w_scale + result = super()._apply(fn, recurse) + self.w_scale = w_scale.to(self.w_packed.device) + return result + + def forward(self, x: Tensor) -> Tensor: + return torch.ops.mlx.int8_intx_gemm(x, self.w_packed, self.w_scale) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"bits={self.bits}" + ) + + +@torch.library.custom_op("mlx::int8_intx_gemm", mutates_args=()) +def int8_intx_gemm(x: Tensor, w_packed: Tensor, w_scale: Tensor) -> Tensor: + K = x.shape[-1] + xq, x_scale = quantize_per_token_int8(x.reshape(-1, K)) + wq = unpack_intx(w_packed, infer_bits(w_packed, K)) + acc = xq.to(torch.int64) @ wq.to(torch.int64).T + y = (x_scale * w_scale.float()) * acc.float() + return y.to(x.dtype).reshape(*x.shape[:-1], w_packed.shape[0]) + + +@torch.library.register_fake("mlx::int8_intx_gemm") +def int8_intx_gemm_fake(x: Tensor, w_packed: Tensor, w_scale: Tensor) -> Tensor: + return x.new_empty((*x.shape[:-1], w_packed.shape[0])) + + +# --------------------------------------------------------------------------- +# MLX handler +# --------------------------------------------------------------------------- + +from executorch.backends.mlx.builder.op_helpers import ( + emit_ceil_div, + emit_floordiv, + emit_if_else, + emit_product, + emit_shape, + emit_sub_int, + torch_dtype_to_scalar_type, +) +from executorch.backends.mlx.builder.op_registry import REGISTRY +from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder +from executorch.backends.mlx.builder.slot_manager import Slot +from executorch.backends.mlx.serialization.mlx_graph_schema import ( + IntOrVid, + MetalKernelNode, +) + +_K_ALIGN = 64 +_THREADS = 128 +_QUANT_THREADS = 256 +_MATVEC_SIMDGROUPS = 2 +_MATVEC_ROWS_PER_SIMDGROUP = 4 +_MATVEC_BLOCK = 512 +_MAX_M = 2**31 - 1 + + +class Int8IntxPlan(NamedTuple): + """Kernel choice for ``M <= max_m``. ``bn == 0`` selects the matvec kernel + with ``bm`` activation rows per threadgroup; otherwise a padded full-K MPP + GEMM with a ``bm x bn`` tile, split over K when ``splits > 1``.""" + + max_m: int + bm: int + bn: int + splits: int = 1 + + +# Muse Glimmer (N, K) -> (bm, bn, splits) for M <= 8, M <= 16, and larger M, +# from the int4 M4 Max sweep (test/bench_muse_glimmer_gemm.py --m 1 2 4 8 16 256, +# fp16). lm_head uses full-K past M = 16: split-K was within noise +# there and would need a [2, M, 202048] int32 partial buffer. int2 has not been +# swept and uses the int4 tiles. +_MG_TILES: dict[tuple[int, int], list[tuple[int, int, int]]] = { + (8704, 6656): [(8, 128, 8), (16, 128, 8), (64, 64, 4)], # qkv + (6656, 4096): [(8, 64, 8), (16, 64, 8), (128, 32, 2)], # o + (39936, 6656): [(8, 128, 4), (16, 128, 8), (32, 128, 1)], # gate_up + (6656, 19968): [(8, 64, 8), (16, 64, 8), (32, 128, 4)], # down + (202048, 6656): [(8, 128, 2), (16, 128, 2), (64, 64, 1)], # lm_head +} + + +def full_k_plans(N: int, K: int) -> list[Int8IntxPlan]: + """Kernel plans, smallest ``max_m`` first; the last covers every larger M. + Matvec for M <= 2, then padded (split-K) MPP tiles. Other shapes use the + Muse Glimmer pattern: more splits when N gives few output tiles.""" + plans = [] + if K % _MATVEC_BLOCK == 0: + plans += [Int8IntxPlan(1, 1, 0), Int8IntxPlan(2, 2, 0)] + tiles = _MG_TILES.get((N, K)) + if tiles is None: + bn, splits = (64, 8) if N <= 16384 else (128, 2) + tiles = [(8, bn, splits), (16, bn, splits), (32, 128, 1)] + for max_m, (bm, bn, splits) in zip((8, 16, _MAX_M), tiles): + while K % (_K_ALIGN * splits): + splits //= 2 + plans.append(Int8IntxPlan(max_m, bm, bn, splits)) + return plans + + +_HEADER = f""" +#include +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; + +constant float INT8_GEMM_EPS = {_EPS}f; + +template +inline uint quantize_int8x4(vec v, float inv_scale) {{ + const float4 q = clamp(rint(float4(v) * inv_scale), -128.0f, 127.0f); + return as_type(char4(q)); +}} + +inline float int8_scale_from_amax(float amax) {{ + return max(precise::divide(amax, 127.0f), INT8_GEMM_EPS); +}} +""" + +_W_FORMAT = {4: "int4b_format", 2: "int2b_format"} + + +# Template args T, K, TPR: one threadgroup of TPR threads per xq row. Rows +# below M (from x) write xq and x_scale; padding rows past M write zeros to xq +# so the GEMM sees whole BM tiles. +_QUANT_SOURCE = """ + constexpr int SG = TPR / 32; + threadgroup float sg_amax[SG]; + + int M = 1; + for (int d = 0; d < x_ndim - 1; ++d) { + M *= x_shape[d]; + } + const int m = int(threadgroup_position_in_grid.y); + const int tid = int(thread_index_in_threadgroup); + device uint* xq_row = (device uint*)(xq + size_t(m) * K); + if (m >= M) { + for (int i = tid; i < K / 4; i += TPR) { + xq_row[i] = 0; + } + return; + } + device const vec* x_row = (device const vec*)(x + size_t(m) * K); + + float amax = 0.0f; + for (int i = tid; i < K / 4; i += TPR) { + const float4 v = abs(float4(x_row[i])); + amax = max(amax, max(max(v.x, v.y), max(v.z, v.w))); + } + amax = simd_max(amax); + if (SG > 1) { + if (thread_index_in_simdgroup == 0) { + sg_amax[simdgroup_index_in_threadgroup] = amax; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + amax = sg_amax[0]; + for (int s = 1; s < SG; ++s) { + amax = max(amax, sg_amax[s]); + } + } + const float scale = int8_scale_from_amax(amax); + const float inv_scale = precise::divide(1.0f, scale); + + for (int i = tid; i < K / 4; i += TPR) { + xq_row[i] = quantize_int8x4(x_row[i], inv_scale); + } + if (tid == 0) { + x_scale[m] = scale; + } +""" + + +# Template args T, N, K, BM, BN; {W_FORMAT} is substituted per bit width. MPP +# multiplies an int8 [K, rows] tensor by a packed [K, N] weight tensor (W +# transposed) over the whole K for each BM x BN output tile. xq may have zero +# padding rows past M = len(x_scale). +def _full_k_source(bits: int) -> str: + return """ + const int rows = xq_shape[0]; + const int M = x_scale_shape[0]; + const int n0 = int(threadgroup_position_in_grid.x) * BN; + const int m0 = int(threadgroup_position_in_grid.y) * BM; + + tensor, tensor_inline> X( + (device int8_t*)xq, dextents(K, rows)); + tensor, tensor_inline> W( + (device uchar*)w_packed, dextents(K, N), array{1, K}); + constexpr auto desc = matmul2d_descriptor( + BM, BN, static_cast(dynamic_extent), false, true, false, + matmul2d_descriptor::mode::multiply); + matmul2d> mm; + auto mA = X.slice(0, m0); + auto mB = W.slice(0, n0); + auto acc = mm.template get_destination_cooperative_tensor(); + mm.run(mA, mB, acc); + + #pragma clang loop unroll(full) + for (ushort i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) { + const auto coord = acc.get_multidimensional_index(i); + const int n = n0 + int(coord[0]); + const int m = m0 + int(coord[1]); + if (m < M && n < N) { + const float s = x_scale[m] * float(w_scale[n]); + out[size_t(m) * N + n] = T(s * float(acc[i])); + } + } + } +""".replace( + "W_FORMAT", _W_FORMAT[bits] + ) + + +# Template args T, N, K, BM, BN, SPLITS (K % (64 * SPLITS) == 0): grid z picks +# one of SPLITS contiguous K ranges. Each threadgroup runs MPP over its range +# and writes exact INT32 partial sums to partial[s, m, n]. +def _split_k_source(bits: int) -> str: + return """ + constexpr int KS = K / SPLITS; + const int rows = xq_shape[0]; + const int M = x_scale_shape[0]; + const int n0 = int(threadgroup_position_in_grid.x) * BN; + const int m0 = int(threadgroup_position_in_grid.y) * BM; + const int s = int(threadgroup_position_in_grid.z); + const int k0 = s * KS; + + tensor, tensor_inline> X( + (device int8_t*)xq + k0, dextents(KS, rows), array{1, K}); + tensor, tensor_inline> W( + (device uchar*)w_packed + k0 / (8 / BITS), dextents(KS, N), + array{1, K}); + constexpr auto desc = matmul2d_descriptor( + BM, BN, static_cast(dynamic_extent), false, true, false, + matmul2d_descriptor::mode::multiply); + matmul2d> mm; + auto mA = X.slice(0, m0); + auto mB = W.slice(0, n0); + auto acc = mm.template get_destination_cooperative_tensor(); + mm.run(mA, mB, acc); + + #pragma clang loop unroll(full) + for (ushort i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) { + const auto coord = acc.get_multidimensional_index(i); + const int n = n0 + int(coord[0]); + const int m = m0 + int(coord[1]); + if (m < M && n < N) { + partial[(size_t(s) * M + m) * N + n] = acc[i]; + } + } + } +""".replace( + "W_FORMAT", _W_FORMAT[bits] + ).replace( + "BITS", str(bits) + ) + + +# Template args T, N, SPLITS: one thread per output sums the INT32 partials and +# applies the scales in the same order as the full-K epilogue, so results match +# the unsplit kernel exactly. +_SPLIT_K_REDUCE_SOURCE = """ + const int M = x_scale_shape[0]; + const size_t idx = thread_position_in_grid.x; + if (idx >= size_t(M) * N) { + return; + } + const int m = int(idx / N); + const int n = int(idx % N); + int sum = 0; + for (int s = 0; s < SPLITS; ++s) { + sum += partial[size_t(s) * M * N + idx]; + } + const float scale = x_scale[m] * float(w_scale[n]); + out[idx] = T(scale * float(sum)); +""" + + +# Matvec (template args T, N, K, MR; K % 512 == 0) for small M: reads int8 +# activations from the quantize kernel (unpadded) and streams packed weights +# like MLX's qmv. Each SIMD group computes 4 output channels for up to MR +# activation rows; each lane covers 16 K values per 512-wide step. XORing every +# field's sign bit makes fields unsigned (q = w + 2**(bits - 1)) so masked +# ushort words convert directly to float against activations pre-scaled by the +# matching power of two (exact in fp32); subtracting 2**(bits - 1) * sum(x) +# restores the signed dot product. Per-lane sums stay below 2^24, so they are +# exact integers and the INT32 simd_sum matches the MPP kernels bit for bit. +_MATVEC_PROLOGUE = """ + constexpr int SG = 2; + constexpr int RPS = 4; + constexpr int VPT = 16; + constexpr int BLOCK = 32 * VPT; + constexpr int KP = K * BITS / 8; + + const int M = x_scale_shape[0]; + const int lane = int(thread_index_in_simdgroup); + const int sg = int(simdgroup_index_in_threadgroup); + const int n0 = (int(threadgroup_position_in_grid.x) * SG + sg) * RPS; + const int m0 = int(threadgroup_position_in_grid.y) * MR; + if (n0 >= N) { + return; + } + const int mrows = min(MR, M - m0); + + device const uchar* w_rows[RPS]; + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + w_rows[r] = w_packed + size_t(min(n0 + r, N - 1)) * KP + lane * (VPT * BITS / 8); + } + device const char* x_base = (device const char*)xq + size_t(m0) * K + lane * VPT; + + float acc[MR][RPS]; + float xsum[MR]; + #pragma clang loop unroll(full) + for (int m = 0; m < MR; ++m) { + xsum[m] = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + acc[m][r] = 0.0f; + } + } + + for (int k = 0; k < K; k += BLOCK) { + // Converted once per block and reused for every activation row. + float4 wf[RPS][4]; + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { +""" + +_MATVEC_LOAD = { + 4: """ + const ushort4 q = as_type( + *(device const uint2*)(w_rows[r] + k / 2) ^ uint2(0x88888888u)); + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + wf[r][i] = float4( + q[i] & 0x000f, q[i] & 0x00f0, q[i] & 0x0f00, q[i] & 0xf000); + } +""", + 2: """ + const ushort2 q = as_type( + *(device const uint*)(w_rows[r] + k / 4) ^ 0xAAAAAAAAu); + #pragma clang loop unroll(full) + for (int h = 0; h < 2; ++h) { + wf[r][2 * h] = float4( + q[h] & 0x0003, q[h] & 0x000c, q[h] & 0x0030, q[h] & 0x00c0); + wf[r][2 * h + 1] = float4( + q[h] & 0x0300, q[h] & 0x0c00, q[h] & 0x3000, q[h] & 0xc000); + } +""", +} + +_MATVEC_BODY = """ + } + #pragma clang loop unroll(full) + for (int m = 0; m < MR; ++m) { + if (m < mrows) { + const uint4 xb = *(device const uint4*)(x_base + size_t(m) * K + k); + float4 xs[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const float4 v = float4(as_type(xb[i])); + xsum[m] += (v.x + v.y) + (v.z + v.w); + xs[i] = v * PRESCALE; + } + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + float a = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + a += dot(xs[i], wf[r][i]); + } + acc[m][r] += a; + } + } + } + } + + #pragma clang loop unroll(full) + for (int m = 0; m < MR; ++m) { + if (m < mrows) { + #pragma clang loop unroll(full) + for (int r = 0; r < RPS; ++r) { + const int total = simd_sum(int(acc[m][r] - OFFSET * xsum[m])); + const int n = n0 + r; + if (lane == 0 && n < N) { + const float s = x_scale[m0 + m] * float(w_scale[n]); + out[size_t(m0 + m) * N + n] = T(s * float(total)); + } + } + } + } +""" + + +# Activation pre-scales that undo each weight field's bit position within its +# ushort word; int2 groups of 4 alternate between the low and high byte. +_MATVEC_PRESCALE = { + 4: "float4(1.0f, 1.0f / 16.0f, 1.0f / 256.0f, 1.0f / 4096.0f)", + 2: ( + "((i % 2) ? float4(1.0f / 256.0f, 1.0f / 1024.0f, 1.0f / 4096.0f, " + "1.0f / 16384.0f) : float4(1.0f, 1.0f / 4.0f, 1.0f / 16.0f, 1.0f / 64.0f))" + ), +} + + +def _matvec_source(bits: int) -> str: + body = _MATVEC_BODY.replace("PRESCALE", _MATVEC_PRESCALE[bits]).replace( + "OFFSET", f"{float(1 << (bits - 1))}f" + ) + return (_MATVEC_PROLOGUE + _MATVEC_LOAD[bits] + body).replace("BITS", str(bits)) + + +_FULL_K_SOURCES = {b: _full_k_source(b) for b in BITS} +_SPLIT_K_SOURCES = {b: _split_k_source(b) for b in BITS} +_MATVEC_SOURCES = {b: _matvec_source(b) for b in BITS} + + +def _lit(*values: int) -> list[IntOrVid]: + return [IntOrVid.from_literal(v) for v in values] + + +@REGISTRY.register(target=[torch.ops.mlx.int8_intx_gemm.default]) +def _int8_intx_gemm_handler(P: MLXProgramBuilder, n: Node) -> Slot: # noqa: C901 + x_slot, w_packed_slot, w_scale_slot = P.args(n) + x_node = n.args[0] + x_meta = x_node.meta["val"] + N, KP = n.args[1].meta["val"].shape + K = x_meta.shape[-1] + if not all(isinstance(v, int) for v in (N, KP, K)): + raise NotImplementedError( + "mlx::int8_intx_gemm: w_packed shape and K must be statically known" + ) + bits = 8 * KP // K + if bits not in BITS or K * bits != 8 * KP: + raise ValueError( + f"mlx::int8_intx_gemm: w_packed width {KP} is not K * bits // 8 for " + f"K={K} and bits in {BITS}" + ) + if K % _K_ALIGN != 0: + raise NotImplementedError( + f"mlx::int8_intx_gemm: K must be a multiple of {_K_ALIGN}; got K={K}" + ) + + dtype_int = torch_dtype_to_scalar_type(x_meta.dtype) + out = P.make_or_get_slot(n) + leading = emit_shape(P, x_node, x_slot, end_dim=-1) + out_shape_flat = leading + [IntOrVid.from_literal(N)] + m_iov = emit_product(P, leading) + _, xq_slot = P.make_tmp_slot() + _, x_scale_slot = P.make_tmp_slot() + _, partial_slot = P.make_tmp_slot() + int8 = torch_dtype_to_scalar_type(torch.int8) + int32 = torch_dtype_to_scalar_type(torch.int32) + fp32 = torch_dtype_to_scalar_type(torch.float32) + k_lit = IntOrVid.from_literal(K) + n_lit = IntOrVid.from_literal(N) + tag = f"int8_int{bits}_gemm" + + def emit_quantize(rows: IntOrVid) -> None: + # rows > M zero-pads xq to whole BM tiles; x_scale stays [M]. + P.emit( + MetalKernelNode( + name="int8_intx_gemm_quantize", + source=_QUANT_SOURCE, + header=_HEADER, + inputs=[P.slot_to_tid(x_slot)], + outputs=[P.slot_to_tid(xq_slot), P.slot_to_tid(x_scale_slot)], + grid=[ + IntOrVid.from_literal(_QUANT_THREADS), + rows, + IntOrVid.from_literal(1), + ], + threadgroup=_lit(_QUANT_THREADS, 1, 1), + input_names=["x"], + output_names=["xq", "x_scale"], + output_shapes_flat=[rows, k_lit, m_iov], + output_shape_lengths=[2, 1], + output_dtypes=[int8, fp32], + template_arg_names=["T", "K", "TPR"], + template_arg_kinds=[2, 0, 0], + template_arg_values=[dtype_int, K, _QUANT_THREADS], + ) + ) + + def emit_matvec(mr: int) -> None: + emit_quantize(m_iov) + rows_per_tg = _MATVEC_SIMDGROUPS * _MATVEC_ROWS_PER_SIMDGROUP + P.emit( + MetalKernelNode( + name=f"{tag}_matvec", + source=_MATVEC_SOURCES[bits], + header=_HEADER, + inputs=[ + P.slot_to_tid(s) + for s in (xq_slot, x_scale_slot, w_packed_slot, w_scale_slot) + ], + outputs=[P.slot_to_tid(out)], + grid=[ + IntOrVid.from_literal( + (N + rows_per_tg - 1) // rows_per_tg * _MATVEC_SIMDGROUPS * 32 + ), + emit_ceil_div(P, m_iov, mr), + IntOrVid.from_literal(1), + ], + threadgroup=_lit(_MATVEC_SIMDGROUPS * 32, 1, 1), + input_names=["xq", "x_scale", "w_packed", "w_scale"], + output_names=["out"], + output_shapes_flat=out_shape_flat, + output_shape_lengths=[len(out_shape_flat)], + output_dtypes=[dtype_int], + template_arg_names=["T", "N", "K", "MR"], + template_arg_kinds=[2, 0, 0, 0], + template_arg_values=[dtype_int, N, K, mr], + ) + ) + + def emit_mpp(bm: int, bn: int, splits: int) -> None: + emit_quantize( + emit_product(P, [emit_ceil_div(P, m_iov, bm), IntOrVid.from_literal(bm)]) + ) + grid = [ + IntOrVid.from_literal((N + bn - 1) // bn * _THREADS), + emit_ceil_div(P, m_iov, bm), + IntOrVid.from_literal(1), + ] + if splits == 1: + P.emit( + MetalKernelNode( + name=f"{tag}_full_k", + source=_FULL_K_SOURCES[bits], + header=_HEADER, + inputs=[ + P.slot_to_tid(s) + for s in (xq_slot, x_scale_slot, w_packed_slot, w_scale_slot) + ], + outputs=[P.slot_to_tid(out)], + grid=grid, + threadgroup=_lit(_THREADS, 1, 1), + input_names=["xq", "x_scale", "w_packed", "w_scale"], + output_names=["out"], + output_shapes_flat=out_shape_flat, + output_shape_lengths=[len(out_shape_flat)], + output_dtypes=[dtype_int], + template_arg_names=["T", "N", "K", "BM", "BN"], + template_arg_kinds=[2, 0, 0, 0, 0], + template_arg_values=[dtype_int, N, K, bm, bn], + ) + ) + return + P.emit( + MetalKernelNode( + name=f"{tag}_split_k", + source=_SPLIT_K_SOURCES[bits], + header=_HEADER, + inputs=[ + P.slot_to_tid(s) for s in (xq_slot, x_scale_slot, w_packed_slot) + ], + outputs=[P.slot_to_tid(partial_slot)], + grid=grid[:2] + [IntOrVid.from_literal(splits)], + threadgroup=_lit(_THREADS, 1, 1), + input_names=["xq", "x_scale", "w_packed"], + output_names=["partial"], + output_shapes_flat=[IntOrVid.from_literal(splits), m_iov, n_lit], + output_shape_lengths=[3], + output_dtypes=[int32], + template_arg_names=["T", "N", "K", "BM", "BN", "SPLITS"], + template_arg_kinds=[2, 0, 0, 0, 0, 0], + template_arg_values=[dtype_int, N, K, bm, bn, splits], + ) + ) + reduce_threads = emit_product( + P, + [ + emit_ceil_div(P, emit_product(P, [m_iov, n_lit]), 256), + IntOrVid.from_literal(256), + ], + ) + P.emit( + MetalKernelNode( + name="int8_intx_gemm_split_k_reduce", + source=_SPLIT_K_REDUCE_SOURCE, + header=_HEADER, + inputs=[ + P.slot_to_tid(s) for s in (partial_slot, x_scale_slot, w_scale_slot) + ], + outputs=[P.slot_to_tid(out)], + grid=[ + reduce_threads, + IntOrVid.from_literal(1), + IntOrVid.from_literal(1), + ], + threadgroup=_lit(256, 1, 1), + input_names=["partial", "x_scale", "w_scale"], + output_names=["out"], + output_shapes_flat=out_shape_flat, + output_shape_lengths=[len(out_shape_flat)], + output_dtypes=[dtype_int], + template_arg_names=["T", "N", "SPLITS"], + template_arg_kinds=[2, 0, 0], + template_arg_values=[dtype_int, N, splits], + ) + ) + + def emit_plan(plan: Int8IntxPlan) -> None: + if plan.bn == 0: + emit_matvec(plan.bm) + else: + emit_mpp(plan.bm, plan.bn, plan.splits) + + # (M - 1) // max_m == 0 exactly when M <= max_m, so each IfNode's else + # branch takes the first plan that covers M. + m_minus_1 = emit_sub_int(P, m_iov, IntOrVid.from_literal(1)) + + def emit_plans(plans: list[Int8IntxPlan]) -> None: + plan, rest = plans[0], plans[1:] + if not rest: + emit_plan(plan) + return + emit_if_else( + P, + emit_floordiv(P, m_minus_1, IntOrVid.from_literal(plan.max_m)), + emit_then=lambda: emit_plans(rest), + emit_else=lambda: emit_plan(plan), + ) + + emit_plans(full_k_plans(N, K)) + return out diff --git a/backends/mlx/custom_kernel_ops/test/bench_muse_glimmer_gemm.py b/backends/mlx/custom_kernel_ops/test/bench_muse_glimmer_gemm.py new file mode 100644 index 00000000000..28f227fd6e0 --- /dev/null +++ b/backends/mlx/custom_kernel_ops/test/bench_muse_glimmer_gemm.py @@ -0,0 +1,810 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Benchmark ``mlx::int8_intx_gemm`` (dynamic per-token int8 activations x +per-channel int4 or int2 weights, INT32 accumulation; kernels ``int8_int4`` and +``int8_int2``) and ``mlx::fp_int4_gemm`` (fp16 / bf16 activations x per-channel +int4 weights, FP32 accumulation) against MLX ``quantized_matmul`` (4-bit affine, +group size 32) on every Muse Glimmer GEMM shape, and estimate total Muse +Glimmer GEMM time per forward pass. + +For each activation dtype, shape, M (token batch size), and kernel +(``--kernels``), every tile config is checked against a reference and timed. +The sweep matches what the op lowers to: + * M <= ``--matvec-max-m`` (2): only the matvec kernel, which streams packed + weights like MLX's qmv. Its rows per threadgroup (MR) is M. + * larger M: only MPP tiles (BM x BN, BM capped at max(8, M) since larger + tiles only add padding rows, and BM >= 32 for M >= 128). Full-K + (``--splits`` 1) runs MPP over all + of K; split-K (S > 1) gives each of S threadgroups per output tile K / S + and a reduce kernel adds the partials and applies the scales. +int8_int4 / int8_int2 quantize activations first (included in their times). +Tile selection has two rounds: every config is timed with ``--iters``, then +the ``--refine-top`` fastest configs of each lowering (full-K, split-K) are +re-timed with ``--refine-iters`` and the fastest is kept. Each kernel's winner +is then re-timed against MLX: they alternate ``--final-repeats`` times and the +median of each is the reported (official) number, so all are measured under +the same GPU conditions. The Muse Glimmer estimate is ``sum over shapes of +(official time x times the shape occurs in the model)``. + +Just run:: + + python -m executorch.backends.mlx.custom_kernel_ops.test.bench_muse_glimmer_gemm + +and send back the output. ``--verbose`` prints every tile config. + +Timing: each timing first runs ``--warmup`` untimed calls (so the kernel is +compiled), then N calls built lazily as a dependent chain (each call's input +is ``x0 + 0 * previous_output[:, :1]``) and timed as a single ``mx.eval``. +The chain serializes calls like consecutive layers in a model; independent +calls would let the GPU overlap them and hide per-call latency. The same chain +without the GEMM is timed and subtracted, and the result is the average per +call. N targets ``--time-budget-ms`` of GPU time per timing (estimated from one +call), between ``--min-iters`` and a per-round cap: ``--iters`` for the tile +sweep, ``--refine-iters`` for the refine round, and ``--final-iters`` for the +official re-run. ``--shape-pause`` seconds between shapes and +``--dtype-pause`` seconds between activation dtypes. +Calls rotate across enough weight copies to exceed ``--rotate-mb`` so weights +are not served from cache. + +Requires MLX and an MPP-capable GPU (Metal 4). +""" + +import argparse +import math +import statistics +import time +from typing import Callable, Dict, List, Optional, Tuple + +import mlx.core as mx +from executorch.backends.mlx.custom_kernel_ops import ( + fp_int4_gemm as fp4, + int8_intx_gemm as op, +) + +# Muse Glimmer (MuseGlimmerConfig defaults): 52 identical decoder layers with +# fused Q/K/V/output-gate and fused gate/up projections, plus lm_head. +_DIM = 6656 +_N_LAYERS = 52 +_Q_DIM = 32 * 128 +_KV_DIM = 2 * 128 +_FFN_DIM = 3 * _DIM +_VOCAB = 202048 + +# name -> (N, K, occurrences per forward) for y[M, N] = x[M, K] @ W[N, K].T +SHAPES: Dict[str, Tuple[int, int, int]] = { + "qkv": (_Q_DIM + 2 * _KV_DIM + _Q_DIM, _DIM, _N_LAYERS), + "o": (_DIM, _Q_DIM, _N_LAYERS), + "gate_up": (2 * _FFN_DIM, _DIM, _N_LAYERS), + "down": (_DIM, _FFN_DIM, _N_LAYERS), + "lm_head": (_VOCAB, _DIM, 1), +} +DTYPES = {"fp16": mx.float16, "bf16": mx.bfloat16} +KERNELS = ("int8_int4", "int8_int2", "fp_int4") +# Weight bits per kernel. +KERNEL_BITS = {"int8_int4": 4, "int8_int2": 2, "fp_int4": 4} +MATVEC = 0 + +# (BM, BN, splits): an MPP tile, split over K when splits > 1. BN == MATVEC +# selects the matvec kernel with BM activation rows per threadgroup. +Tile = Tuple[int, int, int] +# (tile, seconds or None on failure, check error or None) +Entry = Tuple[Tile, Optional[float], Optional[float]] + +_quant_kernel = mx.fast.metal_kernel( + name="int8_intx_gemm_quantize", + input_names=["x"], + output_names=["xq", "x_scale"], + source=op._QUANT_SOURCE, + header=op._HEADER, +) +_split_k_reduce_kernel = mx.fast.metal_kernel( + name="int8_intx_gemm_split_k_reduce", + input_names=["partial", "x_scale", "w_scale"], + output_names=["out"], + source=op._SPLIT_K_REDUCE_SOURCE, + header=op._HEADER, +) +_full_k_kernels = { + b: mx.fast.metal_kernel( + name=f"int8_int{b}_gemm_full_k", + input_names=["xq", "x_scale", "w_packed", "w_scale"], + output_names=["out"], + source=op._FULL_K_SOURCES[b], + header=op._HEADER, + ) + for b in op.BITS +} +_split_k_kernels = { + b: mx.fast.metal_kernel( + name=f"int8_int{b}_gemm_split_k", + input_names=["xq", "x_scale", "w_packed"], + output_names=["partial"], + source=op._SPLIT_K_SOURCES[b], + header=op._HEADER, + ) + for b in op.BITS +} +_matvec_kernels = { + b: mx.fast.metal_kernel( + name=f"int8_int{b}_gemm_matvec", + input_names=["xq", "x_scale", "w_packed", "w_scale"], + output_names=["out"], + source=op._MATVEC_SOURCES[b], + header=op._HEADER, + ) + for b in op.BITS +} +_fp4_full_k_kernel = mx.fast.metal_kernel( + name="fp_int4_gemm_full_k", + input_names=["x", "w_packed", "w_scale"], + output_names=["out"], + source=fp4._FULL_K_SOURCE, + header=fp4._HEADER, +) +_fp4_split_k_kernel = mx.fast.metal_kernel( + name="fp_int4_gemm_split_k", + input_names=["x", "w_packed"], + output_names=["partial"], + source=fp4._SPLIT_K_SOURCE, + header=fp4._HEADER, +) +_fp4_split_k_reduce_kernel = mx.fast.metal_kernel( + name="fp_int4_gemm_split_k_reduce", + input_names=["partial", "w_scale"], + output_names=["out"], + source=fp4._SPLIT_K_REDUCE_SOURCE, + header=fp4._HEADER, +) +_fp4_matvec_kernel = mx.fast.metal_kernel( + name="fp_int4_gemm_matvec", + input_names=["x", "w_packed", "w_scale"], + output_names=["out"], + source=fp4._MATVEC_SOURCE, + header=fp4._HEADER, +) + + +def config_name(tile: Tile) -> str: + bm, bn, splits = tile + if bn == MATVEC: + return f"matvec mr{bm}" + if splits > 1: + return f"split_k{splits} {bm}x{bn}" + return f"full_k {bm}x{bn}" + + +def lowering(tile: Tile) -> str: + return config_name(tile).split()[0].rstrip("0123456789") + + +def int8_intx_gemm( + x: mx.array, + w_packed: mx.array, + w_scale: mx.array, + tile: Tile, + bits: int, + quant_threads: int = op._QUANT_THREADS, +) -> mx.array: + M, K = x.shape + N = w_packed.shape[0] + bm, bn, splits = tile + if bn == MATVEC: + xq, x_scale = _quant_kernel( + inputs=[x], + template=[("T", x.dtype), ("K", K), ("TPR", quant_threads)], + grid=(quant_threads, M, 1), + threadgroup=(quant_threads, 1, 1), + output_shapes=[(M, K), (M,)], + output_dtypes=[mx.int8, mx.float32], + ) + rows_per_tg = op._MATVEC_SIMDGROUPS * op._MATVEC_ROWS_PER_SIMDGROUP + return _matvec_kernels[bits]( + inputs=[xq, x_scale, w_packed, w_scale], + template=[("T", x.dtype), ("N", N), ("K", K), ("MR", bm)], + grid=( + math.ceil(N / rows_per_tg) * op._MATVEC_SIMDGROUPS * 32, + math.ceil(M / bm), + 1, + ), + threadgroup=(op._MATVEC_SIMDGROUPS * 32, 1, 1), + output_shapes=[(M, N)], + output_dtypes=[x.dtype], + )[0] + launch = { + "grid": (math.ceil(N / bn) * op._THREADS, math.ceil(M / bm), 1), + "threadgroup": (op._THREADS, 1, 1), + "output_shapes": [(M, N)], + "output_dtypes": [x.dtype], + } + template = [("T", x.dtype), ("N", N), ("K", K), ("BM", bm), ("BN", bn)] + # Zero-pad xq to whole BM tiles, as the op always does on the MPP path. + rows = math.ceil(M / bm) * bm + xq, x_scale = _quant_kernel( + inputs=[x], + template=[("T", x.dtype), ("K", K), ("TPR", quant_threads)], + grid=(quant_threads, rows, 1), + threadgroup=(quant_threads, 1, 1), + output_shapes=[(rows, K), (M,)], + output_dtypes=[mx.int8, mx.float32], + ) + if splits > 1: + partial = _split_k_kernels[bits]( + inputs=[xq, x_scale, w_packed], + template=template + [("SPLITS", splits)], + grid=(math.ceil(N / bn) * op._THREADS, math.ceil(M / bm), splits), + threadgroup=(op._THREADS, 1, 1), + output_shapes=[(splits, M, N)], + output_dtypes=[mx.int32], + )[0] + return _split_k_reduce_kernel( + inputs=[partial, x_scale, w_scale], + template=[("T", x.dtype), ("N", N), ("SPLITS", splits)], + grid=(math.ceil(M * N / 256) * 256, 1, 1), + threadgroup=(256, 1, 1), + output_shapes=[(M, N)], + output_dtypes=[x.dtype], + )[0] + return _full_k_kernels[bits]( + inputs=[xq, x_scale, w_packed, w_scale], template=template, **launch + )[0] + + +def fp_int4_gemm( + x: mx.array, w_packed: mx.array, w_scale: mx.array, tile: Tile +) -> mx.array: + M, K = x.shape + N = w_packed.shape[0] + bm, bn, splits = tile + if bn == MATVEC: + rows_per_tg = fp4._MATVEC_SIMDGROUPS * fp4._MATVEC_ROWS_PER_SIMDGROUP + return _fp4_matvec_kernel( + inputs=[x, w_packed, w_scale], + template=[("T", x.dtype), ("N", N), ("K", K), ("MR", bm)], + grid=( + math.ceil(N / rows_per_tg) * fp4._MATVEC_SIMDGROUPS * 32, + math.ceil(M / bm), + 1, + ), + threadgroup=(fp4._MATVEC_SIMDGROUPS * 32, 1, 1), + output_shapes=[(M, N)], + output_dtypes=[x.dtype], + )[0] + template = [("T", x.dtype), ("N", N), ("K", K), ("BM", bm), ("BN", bn)] + grid = (math.ceil(N / bn) * fp4._THREADS, math.ceil(M / bm), 1) + if splits > 1: + partial = _fp4_split_k_kernel( + inputs=[x, w_packed], + template=template + [("SPLITS", splits)], + grid=grid[:2] + (splits,), + threadgroup=(fp4._THREADS, 1, 1), + output_shapes=[(splits, M, N)], + output_dtypes=[mx.float32], + )[0] + return _fp4_split_k_reduce_kernel( + inputs=[partial, w_scale], + template=[("T", x.dtype), ("N", N), ("SPLITS", splits)], + grid=(math.ceil(M * N / 256) * 256, 1, 1), + threadgroup=(256, 1, 1), + output_shapes=[(M, N)], + output_dtypes=[x.dtype], + )[0] + return _fp4_full_k_kernel( + inputs=[x, w_packed, w_scale], + template=template, + grid=grid, + threadgroup=(fp4._THREADS, 1, 1), + output_shapes=[(M, N)], + output_dtypes=[x.dtype], + )[0] + + +def _unpack_wq(w_packed: mx.array, cols: int, bits: int) -> mx.array: + w = w_packed[:cols].astype(mx.int32) + mask = (1 << bits) - 1 + fields = mx.stack([(w >> s) & mask for s in range(0, 8, bits)], axis=-1) + fields = fields.reshape(cols, -1) + half = 1 << (bits - 1) + return mx.where(fields >= half, fields - 2 * half, fields).astype(mx.float32) + + +def _rel_err(out: mx.array, ref: mx.array) -> float: + err = mx.max(mx.abs(out.astype(mx.float32) - ref)) / mx.max(mx.abs(ref)) + return err.item() + + +def check_int8_intx( + x: mx.array, + w_packed: mx.array, + w_scale: mx.array, + tile: Tile, + bits: int, + cols: int = 256, + **kwargs, +) -> float: + """Max error of the first ``cols`` outputs, relative to max |reference|.""" + out = int8_intx_gemm(x, w_packed, w_scale, tile, bits, **kwargs)[:, :cols] + xf = x.astype(mx.float32) + x_scale = mx.maximum(mx.max(mx.abs(xf), axis=-1, keepdims=True) / 127, op._EPS) + xq = mx.clip(mx.round(xf * mx.reciprocal(x_scale)), -128, 127) + wq = _unpack_wq(w_packed, cols, bits) + ref = (xq @ wq.T) * x_scale * w_scale[:cols].astype(mx.float32) + return _rel_err(out, ref) + + +def check_fp_int4( + x: mx.array, w_packed: mx.array, w_scale: mx.array, tile: Tile, cols: int = 256 +) -> float: + """Max error of the first ``cols`` outputs, relative to max |reference|.""" + out = fp_int4_gemm(x, w_packed, w_scale, tile)[:, :cols] + wq = _unpack_wq(w_packed, cols, 4) + ref = (x.astype(mx.float32) @ wq.T) * w_scale[:cols].astype(mx.float32) + return _rel_err(out, ref) + + +def time_per_call( + fn: Callable[[mx.array, int], mx.array], + x0: mx.array, + warmup: int, + max_iters: int, + budget_s: float, + min_iters: int, +) -> float: + """Average serialized seconds per call; ``fn(x, i)`` must build (not + evaluate) run ``i`` on input ``x``. Runs about ``budget_s`` worth of calls + (estimated from one evaluated call), clamped to [min_iters, max_iters].""" + + def chain(call: Optional[Callable[[mx.array, int], mx.array]], iters: int) -> float: + x = x0 + for i in range(iters): + y = call(x, i) if call is not None else x + x = x0 + 0 * y[:, :1] + start = time.perf_counter() + mx.eval(x) + return time.perf_counter() - start + + for i in range(warmup): + mx.eval(fn(x0, i)) + estimate = chain(fn, 1) + iters = max(min_iters, min(max_iters, math.ceil(budget_s / max(estimate, 1e-9)))) + chain(None, iters) + return max(chain(fn, iters) - chain(None, iters), 0.0) / iters + + +def tile_configs(K: int, M: int, args: argparse.Namespace) -> List[Tile]: + """Matvec only for M <= --matvec-max-m (needs K % 512 == 0); otherwise MPP + tiles only. MPP needs BM % 8 == 0, BM > max(8, M) only adds padding rows, + BM < 32 never won at M >= 128, and split-K needs K % (64 * splits) == 0.""" + if M <= args.matvec_max_m and K % op._MATVEC_BLOCK == 0: + return [(M, MATVEC, 1)] + min_bm = 32 if M >= 128 else 8 + return [ + (bm, bn, splits) + for bm in args.bm + if bm % 8 == 0 and min_bm <= bm <= max(8, M) + for bn in args.bn + for splits in args.splits + if splits == 1 or K % (64 * splits) == 0 + ] + + +def passed(entries: List[Entry]) -> List[Entry]: + """Configs that ran and passed the accuracy check.""" + return [e for e in entries if e[1] is not None and e[2] is not None and e[2] < 1e-2] + + +def refine_candidates(entries: List[Entry], top: int) -> List[Tile]: + """The ``top`` fastest passing tiles of each lowering (full-K, split-K, + matvec).""" + ok = sorted(passed(entries), key=lambda e: e[1]) + by_lowering: Dict[str, List[Tile]] = {} + for tile, _, _ in ok: + group = by_lowering.setdefault(lowering(tile), []) + if len(group) < top: + group.append(tile) + return [tile for group in by_lowering.values() for tile in group] + + +def _random_bytes(shape: Tuple[int, ...]) -> mx.array: + return mx.random.randint(0, 256, shape).astype(mx.uint8) + + +def bench_shape( # noqa: C901 + name: str, dtype_name: str, ms: List[int], args: argparse.Namespace +) -> Dict[int, dict]: + """Returns {M: {"mlx": seconds or None, "best": {kernel: (Tile, seconds) + or None}}}; ``mlx`` and ``best`` are the official medians.""" + N, K, _ = SHAPES[name] + dtype = DTYPES[dtype_name] + gs = args.group_size + copies = max(1, math.ceil(args.rotate_mb * 2**20 / (N * K // 2))) + bits_used = sorted({KERNEL_BITS[k] for k in args.kernels}) + # Same copy count for every bit width keeps the rotation simple; int2 + # copies are half the size, still well past cache. + w_packed = { + b: [_random_bytes((N, K * b // 8)) for _ in range(copies)] for b in bits_used + } + w_scale = [mx.random.uniform(0.005, 0.02, (N,)) for _ in range(copies)] + mlx_w = [_random_bytes((N, K // 2)).view(mx.uint32) for _ in range(copies)] + mlx_scales = [ + mx.random.uniform(0.005, 0.02, (N, K // gs)).astype(dtype) + for _ in range(copies) + ] + mlx_biases = [ + mx.random.uniform(-0.1, 0.0, (N, K // gs)).astype(dtype) for _ in range(copies) + ] + mx.eval(list(w_packed.values()), w_scale, mlx_w, mlx_scales, mlx_biases) + quant_opts = {"quant_threads": args.quant_threads} + + def timed( + fn: Callable[[mx.array, int], mx.array], x0: mx.array, iters: int + ) -> float: + return time_per_call( + fn, x0, args.warmup, iters, args.time_budget_ms / 1e3, args.min_iters + ) + + def call(kernel: str, xi: mx.array, i: int, tile: Tile) -> mx.array: + bits = KERNEL_BITS[kernel] + wp, ws = w_packed[bits][i % copies], w_scale[i % copies] + if kernel == "fp_int4": + return fp_int4_gemm(xi, wp, ws, tile) + return int8_intx_gemm(xi, wp, ws, tile, bits, **quant_opts) + + def check(kernel: str, x: mx.array, tile: Tile) -> float: + bits = KERNEL_BITS[kernel] + wp = w_packed[bits][0] + if kernel == "fp_int4": + return check_fp_int4(x, wp, w_scale[0], tile) + return check_int8_intx(x, wp, w_scale[0], tile, bits, **quant_opts) + + def mlx_call(xi: mx.array, i: int) -> mx.array: + return mx.quantized_matmul( + xi, + mlx_w[i % copies], + mlx_scales[i % copies], + mlx_biases[i % copies], + transpose=True, + group_size=gs, + bits=4, + mode="affine", + ) + + def select(kernel: str, x: mx.array, M: int) -> Optional[Tuple[Tile, float]]: + """Sweep + refine; returns (tile, sweep or refine seconds) or None. A + single config (matvec) is only checked; the official re-run times it.""" + tiles = tile_configs(K, M, args) + if len(tiles) == 1: + err = check(kernel, x, tiles[0]) + if err >= 1e-2: + print( + f" [{dtype_name}] {name:<8} M={M:<4} {kernel} " + f"{config_name(tiles[0])} failed check ({err:.1e})", + flush=True, + ) + return None + return tiles[0], 0.0 + entries: List[Entry] = [] + for tile in tiles: + try: + err = check(kernel, x, tile) + t = timed( + lambda xi, i, tile=tile: call(kernel, xi, i, tile), x, args.iters + ) + except Exception as e: # compile or launch failure for this tile + reason = str(e).strip().splitlines()[0][:100] + print(f" skip {kernel} {config_name(tile)}: {reason}", flush=True) + err, t = None, None + entries.append((tile, t, err)) + if args.verbose and t is not None: + print( + f" {kernel:<9} {config_name(tile):<20} {t * 1e6:>10.1f} us " + f"check {err:.1e}", + flush=True, + ) + candidates = refine_candidates(entries, args.refine_top) + if len(candidates) == 1: + return candidates[0], next(e[1] for e in entries if e[0] == candidates[0]) + if not candidates: + print( + f" [{dtype_name}] {name:<8} M={M:<4} no {kernel} config passed", + flush=True, + ) + return None + refined = { + tile: timed( + lambda xi, i, tile=tile: call(kernel, xi, i, tile), + x, + args.refine_iters, + ) + for tile in candidates + } + if args.verbose: + ranked = ", ".join( + f"{config_name(c)} {t * 1e6:.1f}" + for c, t in sorted(refined.items(), key=lambda kv: kv[1]) + ) + print(f" {kernel} refine (us): {ranked}", flush=True) + tile = min(refined, key=refined.get) + return tile, refined[tile] + + results: Dict[int, dict] = {} + for M in ms: + x = mx.random.normal((M, K)).astype(dtype) + mx.eval(x) + winners = {k: select(k, x, M) for k in args.kernels} + live = {k: w[0] for k, w in winners.items() if w is not None} + + times: Dict[str, List[float]] = {k: [] for k in live} + t_mlxs: List[float] = [] + for _ in range(args.final_repeats): + for kernel, tile in live.items(): + times[kernel].append( + timed( + lambda xi, i, kernel=kernel, tile=tile: call( + kernel, xi, i, tile + ), + x, + args.final_iters, + ) + ) + t_mlxs.append(timed(mlx_call, x, args.final_iters)) + t_mlx = statistics.median(t_mlxs) + best = { + k: (live[k], statistics.median(times[k])) if k in live else None + for k in args.kernels + } + results[M] = {"mlx": t_mlx, "best": best} + parts = [ + f"{k} {config_name(b[0])} {b[1] * 1e6:.1f} us ({t_mlx / b[1]:.2f}x)" + for k, b in best.items() + if b is not None + ] + print( + f" [{dtype_name}] {name:<8} M={M:<4} mlx {t_mlx * 1e6:.1f} us | " + + " | ".join(parts), + flush=True, + ) + + return results + + +def print_dtype_report( + dtype_name: str, + ms: List[int], + res: Dict[str, Dict[int, dict]], + gs: int, + kernels: List[str], +) -> None: + print() + print(f"=== {dtype_name}: best config per kernel vs MLX 4-bit affine gs={gs} ===") + print("(official: refined winners and MLX re-timed alternately, medians)") + header = f"{'shape':<8} {'N':>7} {'K':>6} {'M':>4} | {'mlx us':>9} {'TOPS':>5}" + for k in kernels: + header += f" | {k + ' us':>12} {'TOPS':>5} {'config':<18} {'vs mlx':>6}" + print(header) + print("-" * len(header)) + for name, shape_res in res.items(): + N, K, _ = SHAPES[name] + for M in ms: + ops = 2 * M * N * K + t_mlx = shape_res[M]["mlx"] + row = ( + f"{name:<8} {N:>7} {K:>6} {M:>4} | " + f"{t_mlx * 1e6:>9.1f} {ops / t_mlx / 1e12:>5.2f}" + ) + for k in kernels: + b = shape_res[M]["best"][k] + if b is None: + row += f" | {'failed':>12} {'':>5} {'':<18} {'':>6}" + continue + tile, t = b + row += ( + f" | {t * 1e6:>12.1f} {ops / t / 1e12:>5.2f} " + f"{config_name(tile):<18} {t_mlx / t:>5.2f}x" + ) + print(row) + + print() + counts = ", ".join(f"{n} x{SHAPES[n][2]}" for n in res) + print( + f"=== {dtype_name}: estimated Muse Glimmer GEMM time per forward ({counts}) ===" + ) + header = f"{'M':>4} | {'mlx ms':>8}" + for k in kernels: + header += f" | {k + ' ms':>12} {'vs mlx':>6}" + print(header) + print("-" * len(header)) + for M in ms: + t_mlx = sum(SHAPES[n][2] * res[n][M]["mlx"] for n in res) + row = f"{M:>4} | {t_mlx * 1e3:>8.2f}" + for k in kernels: + bests = [res[n][M]["best"][k] for n in res] + if any(b is None for b in bests): + row += f" | {'n/a':>12} {'-':>6}" + continue + t = sum(SHAPES[n][2] * b[1] for n, b in zip(res, bests)) + row += f" | {t * 1e3:>12.2f} {t_mlx / t:>5.2f}x" + print(row) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument( + "--m", + type=int, + nargs="+", + default=[1, 2, 4, 8, 16, 256], + help="Token batch sizes (GEMM M). 1-2 use matvec (decode); 256 is the " + "Muse Glimmer prefill chunk; larger M is GPU-saturated and scales " + "linearly.", + ) + parser.add_argument( + "--shapes", nargs="+", choices=list(SHAPES), default=list(SHAPES) + ) + parser.add_argument( + "--dtypes", + nargs="+", + choices=list(DTYPES), + default=["fp16"], + help="Activation dtypes (default fp16; add bf16 to also run it).", + ) + parser.add_argument( + "--kernels", + nargs="+", + choices=KERNELS, + default=list(KERNELS), + help="int8_int4 / int8_int2: int8 dynamic activations x int4 / int2; " + "fp_int4: fp16/bf16 activations x int4. Each is swept and compared " + "against MLX 4-bit.", + ) + parser.add_argument( + "--bm", + type=int, + nargs="+", + default=[8, 16, 32, 64, 128], + help="MPP output tile rows (multiple of 8; capped at max(8, M)).", + ) + parser.add_argument( + "--bn", + type=int, + nargs="+", + default=[32, 64, 128], + help="MPP output tile columns. BN = 16 never won a full M4 sweep.", + ) + parser.add_argument( + "--splits", + type=int, + nargs="+", + default=[1, 2, 4, 8], + help="Full-K split-K factors (1 = no split); needs K %% (64 * splits) == 0.", + ) + parser.add_argument( + "--group-size", type=int, default=32, help="MLX baseline group size." + ) + parser.add_argument( + "--warmup", + type=int, + default=1, + help="Untimed runs per config / shape / M before timing (compiles kernels).", + ) + parser.add_argument( + "--iters", + type=int, + default=10, + help="Max timed runs per tile config in the sweep.", + ) + parser.add_argument( + "--time-budget-ms", + type=float, + default=50, + help="Target GPU time per timing; the iteration count is budget / " + "(one-call estimate), clamped to [--min-iters, round cap]. Keeps short " + "kernels precise and long ones (large M) cheap, limiting heat.", + ) + parser.add_argument( + "--min-iters", + type=int, + default=3, + help="Minimum timed runs per timing.", + ) + parser.add_argument( + "--refine-top", + type=int, + default=3, + help="Fastest sweep configs per lowering (full-K, split-K) re-timed in " + "the refine round.", + ) + parser.add_argument( + "--refine-iters", + type=int, + default=20, + help="Max timed runs per config in the refine round.", + ) + parser.add_argument( + "--final-iters", + type=int, + default=100, + help="Max timed runs per official re-timing of the refined winner and " "MLX.", + ) + parser.add_argument( + "--final-repeats", + type=int, + default=5, + help="Alternating re-timings of the refined winner and MLX; the medians " + "are the official numbers.", + ) + parser.add_argument( + "--dtype-pause", + type=float, + default=60, + help="Seconds to pause between activation dtypes.", + ) + parser.add_argument( + "--shape-pause", + type=float, + default=10, + help="Seconds to pause between shapes within a dtype.", + ) + parser.add_argument( + "--rotate-mb", + type=float, + default=512, + help="Rotate weight copies until their total exceeds this size.", + ) + parser.add_argument( + "--matvec-max-m", + type=int, + default=2, + help="M up to this value sweeps only the matvec kernel (the op's M <= 2 " + "path); larger M sweeps only MPP tiles. 0 disables matvec.", + ) + parser.add_argument( + "--quant-threads", + type=int, + default=256, + help="int8_int4: threads per row in the activation quantize kernel " + "(multiple of 32, at most 1024).", + ) + parser.add_argument( + "--verbose", action="store_true", help="Print every tile config." + ) + args = parser.parse_args() + ms = sorted(set(args.m)) + + info = mx.device_info() + print( + f"device: {info.get('device_name', '?')} mlx: {mx.__version__} " + f"M: {ms} dtypes: {args.dtypes} kernels: {args.kernels} " + f"mlx group_size: {args.group_size} " + f"warmup: {args.warmup} time budget: {args.time_budget_ms:g} ms " + f"(min iters {args.min_iters}) sweep iters: {args.iters} " + f"refine: top {args.refine_top}/lowering x {args.refine_iters} iters " + f"final iters: {args.final_iters} final repeats: {args.final_repeats} " + f"quant threads: {args.quant_threads} " + f"bm: {args.bm} bn: {args.bn} splits: {args.splits} " + f"matvec max M: {args.matvec_max_m}", + flush=True, + ) + + for i, dtype_name in enumerate(args.dtypes): + if i > 0 and args.dtype_pause > 0: + print(f"pausing {args.dtype_pause:g}s before {dtype_name}...", flush=True) + time.sleep(args.dtype_pause) + res = {} + for j, name in enumerate(args.shapes): + if j > 0 and args.shape_pause > 0: + print(f"pausing {args.shape_pause:g}s before {name}...", flush=True) + time.sleep(args.shape_pause) + res[name] = bench_shape(name, dtype_name, ms, args) + mx.clear_cache() + print_dtype_report(dtype_name, ms, res, args.group_size, args.kernels) + print(flush=True) + + +if __name__ == "__main__": + main()