diff --git a/backends/mlx/custom_kernel_ops/int8_int4_gemm.py b/backends/mlx/custom_kernel_ops/int8_int4_gemm.py new file mode 100644 index 00000000000..14b40353646 --- /dev/null +++ b/backends/mlx/custom_kernel_ops/int8_int4_gemm.py @@ -0,0 +1,485 @@ +# +# 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_int4_gemm``: 8-bit dynamic activation x 4-bit per-channel weight +linear, lowered to Metal Performance Primitives (Metal 4) kernels that do +INT8 x INT4 -> INT32 accumulation over the full K dimension. + +Quantization scheme (symmetric, no zero points): + + 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] 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_int4`): ``w_packed`` is ``uint8 [N, K // 2]``; +byte ``j`` of row ``n`` holds ``wq[n, 2j]`` in its low nibble and +``wq[n, 2j + 1]`` in its high nibble, both as two's-complement int4 (the +layout of Metal's ``int4b_format``). + +INT32 accumulation cannot overflow while ``K * 128 * 8 < 2**31`` +(``K < 2**21``). + +Lowering variants (``full_k``): + * ``False``: one kernel. Each threadgroup quantizes its activation rows and + unpacks int4 weights into threadgroup memory one K tile at a time, and + MPP multiplies the int8 tiles. + * ``True``: two kernels. A per-token kernel writes int8 activations and + scales; then MPP runs ``int8 x int4b_format`` over the full K straight + from device memory, choosing its own K tiling and buffering. + +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_int4_gemm # noqa: F401 + + y = torch.ops.mlx.int8_int4_gemm(x, w_packed, w_scale, full_k=True) + # x: (..., K) fp16 / bf16 / fp32 + # w_packed: (N, K // 2) uint8 + # w_scale: (N,) float + # y: (..., N) x.dtype +""" + +from __future__ import annotations + +import torch +from torch import Tensor +from torch.fx.node import Node + +_EPS = 1e-10 + + +# --------------------------------------------------------------------------- +# Packing / quantization helpers +# --------------------------------------------------------------------------- + + +def pack_int4(wq: Tensor) -> Tensor: + """Pack int4 values in ``[-8, 7]`` of shape ``[N, K]`` into ``uint8 [N, K // 2]``.""" + if wq.shape[-1] % 2 != 0: + raise ValueError(f"pack_int4: K must be even; got {wq.shape[-1]}") + nibbles = wq.to(torch.int16) & 0xF + return (nibbles[..., 0::2] | (nibbles[..., 1::2] << 4)).to(torch.uint8) + + +def unpack_int4(w_packed: Tensor) -> Tensor: + """Inverse of :func:`pack_int4`; returns ``int8 [N, K]``.""" + b = w_packed.to(torch.int16) + lo = b & 0xF + hi = (b >> 4) & 0xF + wq = torch.stack([lo, hi], dim=-1).flatten(-2) + return torch.where(wq >= 8, wq - 16, wq).to(torch.int8) + + +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 + + +# --------------------------------------------------------------------------- +# Custom op + eager fallback +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("mlx::int8_int4_gemm", mutates_args=()) +def int8_int4_gemm( + x: Tensor, w_packed: Tensor, w_scale: Tensor, full_k: bool = False +) -> Tensor: + K = x.shape[-1] + xq, x_scale = quantize_per_token_int8(x.reshape(-1, K)) + wq = unpack_int4(w_packed) + 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_int4_gemm") +def int8_int4_gemm_fake( + x: Tensor, w_packed: Tensor, w_scale: Tensor, full_k: bool = False +) -> 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_product, + emit_shape, + 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, +) + +_BM = 32 +_BN = 32 +_BK = 64 +_THREADS = 128 + + +_HEADER = f""" +#include +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; + +constant float INT8_GEMM_EPS = {_EPS}f; + +// 8 packed bytes -> 16 int8 values in K order (low nibble first). +inline uint4 unpack_int4x16(uint2 packed) {{ + const uchar4 b0 = as_type(packed.x); + const uchar4 b1 = as_type(packed.y); + const char4 lo0 = as_type(uchar4(b0 << 4)) >> 4; + const char4 hi0 = as_type(b0) >> 4; + const char4 lo1 = as_type(uchar4(b1 << 4)) >> 4; + const char4 hi1 = as_type(b1) >> 4; + return uint4( + as_type(char4(lo0.x, hi0.x, lo0.y, hi0.y)), + as_type(char4(lo0.z, hi0.z, lo0.w, hi0.w)), + as_type(char4(lo1.x, hi1.x, lo1.y, hi1.y)), + as_type(char4(lo1.z, hi1.z, lo1.w, hi1.w))); +}} + +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); +}} +""" + + +# full_k=False. Template args BM, BN, BK (BK a multiple of 32 dividing K). One +# threadgroup of 128 threads (4 SIMD groups) computes a BM x BN output tile. +# 1. Pre-pass: one SIMD group per tile row reduces |x| over K into the row's +# scale and inverse scale. +# 2. Main loop per BK step: threads stride over 16-byte chunks, quantizing +# activations and unpacking int4 weights into threadgroup memory (global +# loads are issued before the barrier); MPP matmul2d accumulates INT32. +# 3. Epilogue: scale each INT32 sum by x_scale[m] * w_scale[n]. +_SOURCE = """ + constexpr int THREADS = 128; + constexpr int KP = K / 2; + constexpr int CPR = BK / 16; + constexpr int X_CHUNKS = BM * CPR; + constexpr int W_CHUNKS = BN * CPR; + constexpr int X_PER_THREAD = (X_CHUNKS + THREADS - 1) / THREADS; + constexpr int W_PER_THREAD = (W_CHUNKS + THREADS - 1) / THREADS; + + 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 tid = int(thread_index_in_threadgroup); + const int lane = int(thread_index_in_simdgroup); + const int sg = int(simdgroup_index_in_threadgroup); + + threadgroup uint4 x_tile_v[X_CHUNKS]; + threadgroup uint4 w_tile_v[W_CHUNKS]; + threadgroup float row_scale[BM]; + threadgroup float row_inv_scale[BM]; + + for (int r = sg; r < BM; r += THREADS / 32) { + float amax = 0.0f; + if (m0 + r < M) { + device const vec* xr = + (device const vec*)(x + size_t(m0 + r) * K); + for (int i = lane; i < K / 4; i += 32) { + const float4 v = abs(float4(xr[i])); + amax = max(amax, max(max(v.x, v.y), max(v.z, v.w))); + } + } + amax = simd_max(amax); + if (lane == 0) { + const float s = int8_scale_from_amax(amax); + row_scale[r] = s; + row_inv_scale[r] = precise::divide(1.0f, s); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Tensor dimension 0 is contiguous: the X tile is [K, M] and the W tile is + // [K, N], so W enters transposed. + using TileX = tensor, tensor_inline>; + using TileW = tensor, tensor_inline>; + TileX a((threadgroup int8_t*)x_tile_v, extents{}); + TileW b((threadgroup int8_t*)w_tile_v, extents{}); + constexpr auto desc = matmul2d_descriptor( + BM, BN, BK, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> mm; + auto acc = mm.template get_destination_cooperative_tensor(); + #pragma clang loop unroll(full) + for (ushort i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) { + acc[i] = 0; + } + } + + for (int k0 = 0; k0 < K; k0 += BK) { + uint4 xv[X_PER_THREAD]; + uint2 wv[W_PER_THREAD]; + #pragma clang loop unroll(full) + for (int j = 0; j < X_PER_THREAD; ++j) { + const int c = tid + j * THREADS; + const int row = c / CPR; + xv[j] = uint4(0); + if (c < X_CHUNKS && m0 + row < M) { + device const vec* src = (device const vec*)( + x + size_t(m0 + row) * K + k0 + (c % CPR) * 16); + const float inv = row_inv_scale[row]; + #pragma clang loop unroll(full) + for (int q = 0; q < 4; ++q) { + xv[j][q] = quantize_int8x4(src[q], inv); + } + } + } + #pragma clang loop unroll(full) + for (int j = 0; j < W_PER_THREAD; ++j) { + const int c = tid + j * THREADS; + const int row = c / CPR; + wv[j] = (c < W_CHUNKS && n0 + row < N) + ? *(device const uint2*)( + w_packed + size_t(n0 + row) * KP + (k0 + (c % CPR) * 16) / 2) + : uint2(0); + } + // Also orders the previous step's MMA reads before these writes. + threadgroup_barrier(mem_flags::mem_threadgroup); + #pragma clang loop unroll(full) + for (int j = 0; j < X_PER_THREAD; ++j) { + const int c = tid + j * THREADS; + if (c < X_CHUNKS) { + x_tile_v[c] = xv[j]; + } + } + #pragma clang loop unroll(full) + for (int j = 0; j < W_PER_THREAD; ++j) { + const int c = tid + j * THREADS; + if (c < W_CHUNKS) { + w_tile_v[c] = unpack_int4x16(wv[j]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + mm.run(a, b, 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 = row_scale[coord[1]] * float(w_scale[n]); + out[size_t(m) * N + n] = T(s * float(acc[i])); + } + } + } +""" + + +# full_k=True, kernel 1: one SIMD group per token row writes xq and x_scale. +_QUANT_SOURCE = """ + const size_t m = threadgroup_position_in_grid.y; + const uint lane = thread_index_in_simdgroup; + device const vec* x_row = (device const vec*)(x + m * K); + + float amax = 0.0f; + for (int i = lane; i < K / 4; i += 32) { + const float4 v = abs(float4(x_row[i])); + amax = max(amax, max(max(v.x, v.y), max(v.z, v.w))); + } + const float scale = int8_scale_from_amax(simd_max(amax)); + const float inv_scale = precise::divide(1.0f, scale); + + device uint* xq_row = (device uint*)(xq + m * K); + for (int i = lane; i < K / 4; i += 32) { + xq_row[i] = quantize_int8x4(x_row[i], inv_scale); + } + if (lane == 0) { + x_scale[m] = scale; + } +""" + + +# full_k=True, kernel 2 (template args BM, BN): MPP multiplies an int8 [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 = """ + const int M = xq_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, 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) { + const float s = x_scale[m] * float(w_scale[n]); + out[size_t(m) * N + n] = T(s * float(acc[i])); + } + } + } +""" + + +def _lit(*values: int) -> list[IntOrVid]: + return [IntOrVid.from_literal(v) for v in values] + + +@REGISTRY.register(target=[torch.ops.mlx.int8_int4_gemm.default]) +def _int8_int4_gemm_handler(P: MLXProgramBuilder, n: Node) -> Slot: + args = P.args(n) + if len(args) not in (3, 4): + raise ValueError(f"mlx::int8_int4_gemm: expected 3-4 args, got {len(args)}") + x_slot, w_packed_slot, w_scale_slot = args[:3] + full_k = bool(args[3]) if len(args) == 4 else bool(n.kwargs.get("full_k", False)) + x_node = n.args[0] + x_meta = x_node.meta["val"] + w_packed_meta = n.args[1].meta["val"] + + N, KP = w_packed_meta.shape + if not isinstance(N, int) or not isinstance(KP, int): + raise NotImplementedError( + "mlx::int8_int4_gemm: w_packed shape must be statically known" + ) + K = 2 * KP + if K % _BK != 0: + raise NotImplementedError( + f"mlx::int8_int4_gemm: K must be a multiple of {_BK}; got K={K}" + ) + if x_meta.shape[-1] != K: + raise ValueError( + f"mlx::int8_int4_gemm: x last dim {x_meta.shape[-1]} does not match " + f"K={K} from w_packed" + ) + + 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) + gemm_grid = [ + IntOrVid.from_literal((N + _BN - 1) // _BN * _THREADS), + emit_ceil_div(P, m_iov, _BM), + IntOrVid.from_literal(1), + ] + + if not full_k: + P.emit( + MetalKernelNode( + name="int8_int4_gemm", + source=_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=gemm_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", "BK"], + template_arg_kinds=[2, 0, 0, 0, 0, 0], + template_arg_values=[dtype_int, N, K, _BM, _BN, _BK], + ) + ) + return out + + _, xq_slot = P.make_tmp_slot() + _, x_scale_slot = P.make_tmp_slot() + P.emit( + MetalKernelNode( + name="int8_int4_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(32), m_iov, IntOrVid.from_literal(1)], + threadgroup=_lit(32, 1, 1), + input_names=["x"], + output_names=["xq", "x_scale"], + output_shapes_flat=[m_iov, IntOrVid.from_literal(K), m_iov], + output_shape_lengths=[2, 1], + output_dtypes=[ + torch_dtype_to_scalar_type(torch.int8), + torch_dtype_to_scalar_type(torch.float32), + ], + template_arg_names=["T", "K"], + template_arg_kinds=[2, 0], + template_arg_values=[dtype_int, K], + ) + ) + P.emit( + MetalKernelNode( + name="int8_int4_gemm_full_k", + source=_FULL_K_SOURCE, + 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=gemm_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 out diff --git a/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py b/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py new file mode 100644 index 00000000000..e4f2f9ad79a --- /dev/null +++ b/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py @@ -0,0 +1,498 @@ +#!/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_int4_gemm`` (dynamic per-token int8 activations x +per-channel int4 weights, MPP INT32 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, and M (token batch size), every int8_int4 +tile config (BM x BN x BK) is checked against a reference and timed. BK picks +the lowering: + * ``BK > 0``: one fused kernel that steps over K in BK tiles (BK a + multiple of 32 dividing K). + * ``BK = -1``: quantize kernel + full-K MPP ``int8 x int4b`` GEMM; MPP + handles the whole K itself. +Tile selection has two rounds: every config is timed with ``--iters``, then +the ``--refine-top`` fastest configs of each lowering (staged, full-K) are +re-timed with ``--refine-iters`` and the fastest is kept. That config is then +re-timed against MLX: the two alternate ``--final-repeats`` times and the +median of each is the reported (official) number, so both 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)``. + +Defaults cover speculative decoding / prefill (M = 4..512), fp16 and bf16, all +shapes, and the full tile sweep. Just run:: + + python -m executorch.backends.mlx.custom_kernel_ops.test.bench_int8_int4_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 is ``--iters`` for the tile sweep 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. int8_int4 times include activation quantization. + +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 int8_int4_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} +FULL_K = -1 +_THREADGROUP_MEMORY = 32768 + +# (BM, BN, BK); BK == FULL_K selects the full-K lowering. +Tile = Tuple[int, int, int] +# (tile, seconds or None on failure, check error or None) +Entry = Tuple[Tile, Optional[float], Optional[float]] + +_staged_kernel = mx.fast.metal_kernel( + name="int8_int4_gemm", + input_names=["x", "w_packed", "w_scale"], + output_names=["out"], + source=op._SOURCE, + header=op._HEADER, +) +_quant_kernel = mx.fast.metal_kernel( + name="int8_int4_gemm_quantize", + input_names=["x"], + output_names=["xq", "x_scale"], + source=op._QUANT_SOURCE, + header=op._HEADER, +) +_full_k_kernel = mx.fast.metal_kernel( + name="int8_int4_gemm_full_k", + input_names=["xq", "x_scale", "w_packed", "w_scale"], + output_names=["out"], + source=op._FULL_K_SOURCE, + header=op._HEADER, +) + + +def config_name(tile: Tile) -> str: + bm, bn, bk = tile + if bk == FULL_K: + return f"full_k {bm}x{bn}" + return f"staged {bm}x{bn}x{bk}" + + +def int8_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, bk = tile + 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)] + if bk != FULL_K: + return _staged_kernel( + inputs=[x, w_packed, w_scale], template=template + [("BK", bk)], **launch + )[0] + xq, x_scale = _quant_kernel( + inputs=[x], + template=[("T", x.dtype), ("K", K)], + grid=(32, M, 1), + threadgroup=(32, 1, 1), + output_shapes=[(M, K), (M,)], + output_dtypes=[mx.int8, mx.float32], + ) + return _full_k_kernel( + inputs=[xq, x_scale, w_packed, w_scale], template=template, **launch + )[0] + + +def check_int8_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 = int8_int4_gemm(x, w_packed, w_scale, tile)[:, :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) + w = w_packed[:cols].astype(mx.int32) + nibbles = mx.stack([w & 0xF, w >> 4], axis=-1).reshape(cols, -1) + wq = mx.where(nibbles >= 8, nibbles - 16, nibbles).astype(mx.float32) + ref = (xq @ wq.T) * x_scale * w_scale[:cols].astype(mx.float32) + err = mx.max(mx.abs(out.astype(mx.float32) - ref)) / mx.max(mx.abs(ref)) + return err.item() + + +def time_per_call( + fn: Callable[[mx.array, int], mx.array], x0: mx.array, warmup: int, iters: int +) -> float: + """Average serialized seconds per call; ``fn(x, i)`` must build (not + evaluate) run ``i`` on input ``x``.""" + + def chain(call: Optional[Callable[[mx.array, int], mx.array]]) -> 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)) + chain(None) + return max(chain(fn) - chain(None), 0.0) / iters + + +def tile_configs(K: int, args: argparse.Namespace) -> List[Tile]: + """Valid tiles; MPP needs BM % 8 == 0; BK > 0 needs BK % 32 == 0, + K % BK == 0, and tiles that fit in threadgroup memory.""" + tiles: List[Tile] = [] + for bm in (bm for bm in args.bm if bm % 8 == 0): + for bn in args.bn: + for bk in args.bk: + if bk == FULL_K: + tiles.append((bm, bn, bk)) + continue + tg_bytes = (bm + bn) * bk + 8 * bm + if bk % 32 == 0 and K % bk == 0 and tg_bytes <= _THREADGROUP_MEMORY: + tiles.append((bm, bn, bk)) + return tiles + + +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 (staged, full-K).""" + ok = sorted(passed(entries), key=lambda e: e[1]) + full_k = [e[0] for e in ok if e[0][2] == FULL_K][:top] + staged = [e[0] for e in ok if e[0][2] != FULL_K][:top] + return full_k + staged + + +def _random_bytes(shape: Tuple[int, ...]) -> mx.array: + return mx.random.randint(0, 256, shape).astype(mx.uint8) + + +def bench_shape( + name: str, dtype_name: str, ms: List[int], args: argparse.Namespace +) -> Dict[int, dict]: + """Returns {M: {"mlx": seconds, "int8": [Entry], "best": (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))) + w_packed = [_random_bytes((N, K // 2)) for _ in range(copies)] + 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(w_packed, w_scale, mlx_w, mlx_scales, mlx_biases) + + def timed( + fn: Callable[[mx.array, int], mx.array], x0: mx.array, iters: int + ) -> float: + return time_per_call(fn, x0, args.warmup, iters) + + def int8_call(xi: mx.array, i: int, tile: Tile) -> mx.array: + return int8_int4_gemm(xi, w_packed[i % copies], w_scale[i % copies], tile) + + results: Dict[int, dict] = {} + for M in ms: + x = mx.random.normal((M, K)).astype(dtype) + mx.eval(x) + entries: List[Entry] = [] + for tile in tile_configs(K, args): + try: + err = check_int8_int4(x, w_packed[0], w_scale[0], tile) + t = timed( + lambda xi, i, tile=tile: int8_call(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 {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" {config_name(tile):<20} {t * 1e6:>10.1f} us " + f"check {err:.1e}", + flush=True, + ) + + 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", + ) + + candidates = refine_candidates(entries, args.refine_top) + if not candidates: + results[M] = {"mlx": None, "int8": entries, "best": None} + print( + f" [{dtype_name}] {name:<8} M={M:<4} no int8_int4 config passed", + flush=True, + ) + continue + + refined = { + tile: timed( + lambda xi, i, tile=tile: int8_call(xi, i, tile), x, args.refine_iters + ) + for tile in candidates + } + tile = min(refined, key=refined.get) + 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" refine (us): {ranked}", flush=True) + + t_int8s, t_mlxs = [], [] + for _ in range(args.final_repeats): + t_int8s.append( + timed( + lambda xi, i, tile=tile: int8_call(xi, i, tile), + x, + args.final_iters, + ) + ) + t_mlxs.append(timed(mlx_call, x, args.final_iters)) + t_int8 = statistics.median(t_int8s) + t_mlx = statistics.median(t_mlxs) + results[M] = {"mlx": t_mlx, "int8": entries, "best": (tile, t_int8)} + print( + f" [{dtype_name}] {name:<8} M={M:<4} {config_name(tile):<20} " + f"int8_int4 {t_int8 * 1e6:>9.1f} us | mlx {t_mlx * 1e6:>9.1f} us | " + f"{t_mlx / t_int8:.2f}x (refine {refined[tile] * 1e6:.1f} us, " + f"{len(candidates)} candidates)", + flush=True, + ) + + return results + + +def print_dtype_report( + dtype_name: str, ms: List[int], res: Dict[str, Dict[int, dict]], gs: int +) -> None: + print() + print(f"=== {dtype_name}: best int8_int4 config vs MLX 4-bit affine gs={gs} ===") + print("(official: refined winner and MLX re-timed alternately, medians)") + header = ( + f"{'shape':<8} {'N':>7} {'K':>6} {'M':>4} | {'int8_int4 us':>12} " + f"{'TOPS':>6} {'config':<20} | {'mlx us':>10} {'TOPS':>6} | {'speedup':>7}" + ) + 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"] + b = shape_res[M]["best"] + if b is None: + print(f"{name:<8} {N:>7} {K:>6} {M:>4} | {'failed':>12}") + continue + mlx_cols = f"{t_mlx * 1e6:>10.1f} {ops / t_mlx / 1e12:>6.2f}" + tile, t = b + int8_cols = ( + f"{t * 1e6:>12.1f} {ops / t / 1e12:>6.2f} {config_name(tile):<20}" + ) + speedup = f"{t_mlx / t:>6.2f}x" + print( + f"{name:<8} {N:>7} {K:>6} {M:>4} | {int8_cols} | {mlx_cols} | {speedup}" + ) + + 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} | {'int8_int4 ms':>12} | {'mlx ms':>10} | {'speedup':>7}" + print(header) + print("-" * len(header)) + for M in ms: + bests = {n: res[n][M]["best"] for n in res} + if any(b is None for b in bests.values()): + print(f"{M:>4} | {'n/a':>12} | {'n/a':>10} | {'-':>7}") + continue + t_mlx = sum(SHAPES[n][2] * res[n][M]["mlx"] for n in res) + t_int8 = sum(SHAPES[n][2] * bests[n][1] for n in res) + print( + f"{M:>4} | {t_int8 * 1e3:>12.2f} | {t_mlx * 1e3:>10.2f} | " + f"{t_mlx / t_int8:>6.2f}x" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument( + "--m", + type=int, + nargs="+", + default=[4, 8, 16, 256, 512], + help="Token batch sizes (GEMM M).", + ) + parser.add_argument( + "--shapes", nargs="+", choices=list(SHAPES), default=list(SHAPES) + ) + parser.add_argument( + "--dtypes", nargs="+", choices=list(DTYPES), default=list(DTYPES) + ) + parser.add_argument( + "--bm", + type=int, + nargs="+", + default=[8, 16, 32, 64, 128], + help="Output tile rows (MPP requires a multiple of 8).", + ) + parser.add_argument("--bn", type=int, nargs="+", default=[32, 64, 128]) + parser.add_argument( + "--bk", + type=int, + nargs="+", + default=[32, 64, 128, FULL_K], + help=f"K step (multiple of 32 dividing K), or {FULL_K} for full K.", + ) + 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="Timed runs per tile config in the sweep." + ) + parser.add_argument( + "--refine-top", + type=int, + default=3, + help="Fastest sweep configs per lowering (staged, full-K) re-timed in " + "the refine round.", + ) + parser.add_argument( + "--refine-iters", + type=int, + default=20, + help="Timed runs per config in the refine round.", + ) + parser.add_argument( + "--final-iters", + type=int, + default=100, + help="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( + "--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} mlx group_size: {args.group_size} " + f"warmup: {args.warmup} 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}", + 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) + print(flush=True) + + +if __name__ == "__main__": + main()