From b499b2672f5a811ed183dda939548131f8c93d38 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 2 Nov 2026 13:30:57 -0800 Subject: [PATCH 1/3] [MLX] Add int8 x int4 MPP GEMM custom op and Muse Glimmer benchmark Adds `mlx::int8_int4_gemm`: dynamic symmetric per-token int8 activations x symmetric per-channel int4 weights, INT32 accumulation via Metal Performance Primitives (Metal 4). Two lowerings: * full_k=False: one fused kernel that quantizes activations and unpacks int4 weights into threadgroup memory per BK tile, then runs MPP matmul2d. * full_k=True: per-token quantize kernel, then a full-K MPP `int8 x int4b_format` GEMM reading packed weights straight from device memory. Tile sizes (BM, BN, BK) are Metal template parameters. Adds bench_int8_int4_gemm.py, which sweeps tile configs for every Muse Glimmer GEMM shape, compares the best config against MLX 4-bit affine quantized_matmul (gs=32) for fp16 and bf16, and estimates total Muse Glimmer GEMM time per forward. --- .../mlx/custom_kernel_ops/int8_int4_gemm.py | 485 ++++++++++++++++++ .../test/bench_int8_int4_gemm.py | 375 ++++++++++++++ 2 files changed, 860 insertions(+) create mode 100644 backends/mlx/custom_kernel_ops/int8_int4_gemm.py create mode 100644 backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py 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..d320e2ffa3d --- /dev/null +++ b/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py @@ -0,0 +1,375 @@ +#!/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. +The fastest config is reported against MLX. The Muse Glimmer estimate is +``sum over shapes of (best 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 config / shape / M runs ``--iters`` times, built lazily and timed +as a single ``mx.eval``; the reported time is the average per call. Each kernel +runs once untimed first (the accuracy check doubles as this warmup). +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 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 = dict( + 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[[int], mx.array], iters: int) -> float: + """Average seconds per call over ``iters`` runs; ``fn(i)`` must build (not + evaluate) run ``i``.""" + outs = [fn(i) for i in range(iters)] + start = time.perf_counter() + mx.eval(outs) + return (time.perf_counter() - start) / iters + + +def tile_configs(K: int, args: argparse.Namespace) -> List[Tile]: + """Valid tiles; BK > 0 needs BK % 32 == 0, K % BK == 0, and tiles that fit + in threadgroup memory.""" + tiles: List[Tile] = [] + for bm in args.bm: + 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 best(entries: List[Entry]) -> Optional[Entry]: + """Fastest config that ran and passed the accuracy check.""" + ok = [e for e in entries if e[1] is not None and e[2] is not None and e[2] < 1e-2] + return min(ok, key=lambda e: e[1]) if ok else None + + +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]}}.""" + 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[[int], mx.array]) -> float: + return time_per_call(fn, args.iters) + + 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 i: int8_int4_gemm( + x, w_packed[i % copies], w_scale[i % copies], tile + ) + ) + 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(i: int) -> mx.array: + return mx.quantized_matmul( + x, + mlx_w[i % copies], + mlx_scales[i % copies], + mlx_biases[i % copies], + transpose=True, + group_size=gs, + bits=4, + mode="affine", + ) + + mx.eval(mlx_call(0)) + t_mlx = timed(mlx_call) + results[M] = {"mlx": t_mlx, "int8": entries} + + b = best(entries) + summary = ( + f"best {config_name(b[0])} {b[1] * 1e6:.1f} us, " + f"{t_mlx / b[1]:.2f}x vs mlx" + if b is not None + else "no int8_int4 config passed" + ) + print( + f" [{dtype_name}] {name:<8} M={M:<4} mlx {t_mlx * 1e6:>9.1f} us | {summary}", + flush=True, + ) + + del w_packed, w_scale, mlx_w, mlx_scales, mlx_biases + mx.clear_cache() + 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} ===") + 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 = best(shape_res[M]["int8"]) + mlx_cols = f"{t_mlx * 1e6:>10.1f} {ops / t_mlx / 1e12:>6.2f}" + if b is None: + int8_cols = f"{'failed':>12} {'':>6} {'':<20}" + speedup = f"{'-':>7}" + else: + 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: best(res[n][M]["int8"]) for n in res} + t_mlx = sum(SHAPES[n][2] * res[n][M]["mlx"] for n in res) + if any(b is None for b in bests.values()): + print(f"{M:>4} | {'n/a':>12} | {t_mlx * 1e3:>10.2f} | {'-':>7}") + continue + 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=[32, 64, 128]) + 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( + "--iters", type=int, default=10, help="Timed runs per config / shape / M." + ) + 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"iters: {args.iters}", + flush=True, + ) + + for dtype_name in args.dtypes: + res = {name: bench_shape(name, dtype_name, ms, args) for name in args.shapes} + print_dtype_report(dtype_name, ms, res, args.group_size) + print(flush=True) + + +if __name__ == "__main__": + main() From c0c33222947e396242768dc633899e5aff0885a4 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 2 Nov 2026 14:31:33 -0800 Subject: [PATCH 2/3] [MLX] int8_int4 bench: sweep BM=8/16, add --warmup and --dtype-pause - Default --bm now 8 16 32 64 128; BMs that are not a multiple of 8 are skipped (MPP matmul2d requires M % 8 == 0). - --warmup (default 1): untimed runs per config / shape / M before the timed --iters runs, so kernels are compiled first. - --dtype-pause (default 60s) between activation dtypes. --- .../test/bench_int8_int4_gemm.py | 51 ++++++++++++++----- 1 file changed, 37 insertions(+), 14 deletions(-) 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 index d320e2ffa3d..99ddeb040ce 100644 --- a/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py +++ b/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py @@ -28,9 +28,10 @@ and send back the output. ``--verbose`` prints every tile config. -Timing: each config / shape / M runs ``--iters`` times, built lazily and timed -as a single ``mx.eval``; the reported time is the average per call. Each kernel -runs once untimed first (the accuracy check doubles as this warmup). +Timing: each config / shape / M first runs ``--warmup`` untimed calls (so the +kernel is compiled), then ``--iters`` calls built lazily and timed as a single +``mx.eval``; the reported time is the average per call. The run pauses +``--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. @@ -147,9 +148,11 @@ def check_int8_int4( return err.item() -def time_per_call(fn: Callable[[int], mx.array], iters: int) -> float: - """Average seconds per call over ``iters`` runs; ``fn(i)`` must build (not - evaluate) run ``i``.""" +def time_per_call(fn: Callable[[int], mx.array], warmup: int, iters: int) -> float: + """Average seconds per call over ``iters`` timed runs after ``warmup`` + untimed runs; ``fn(i)`` must build (not evaluate) run ``i``.""" + for i in range(warmup): + mx.eval(fn(i)) outs = [fn(i) for i in range(iters)] start = time.perf_counter() mx.eval(outs) @@ -157,10 +160,10 @@ def time_per_call(fn: Callable[[int], mx.array], iters: int) -> float: def tile_configs(K: int, args: argparse.Namespace) -> List[Tile]: - """Valid tiles; BK > 0 needs BK % 32 == 0, K % BK == 0, and tiles that fit - in threadgroup memory.""" + """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 args.bm: + 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: @@ -203,7 +206,7 @@ def bench_shape( mx.eval(w_packed, w_scale, mlx_w, mlx_scales, mlx_biases) def timed(fn: Callable[[int], mx.array]) -> float: - return time_per_call(fn, args.iters) + return time_per_call(fn, args.warmup, args.iters) results: Dict[int, dict] = {} for M in ms: @@ -242,7 +245,6 @@ def mlx_call(i: int) -> mx.array: mode="affine", ) - mx.eval(mlx_call(0)) t_mlx = timed(mlx_call) results[M] = {"mlx": t_mlx, "int8": entries} @@ -330,7 +332,13 @@ def main() -> None: parser.add_argument( "--dtypes", nargs="+", choices=list(DTYPES), default=list(DTYPES) ) - parser.add_argument("--bm", type=int, nargs="+", default=[32, 64, 128]) + 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", @@ -342,9 +350,21 @@ def main() -> None: 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 config / shape / M." ) + parser.add_argument( + "--dtype-pause", + type=float, + default=60, + help="Seconds to pause between activation dtypes.", + ) parser.add_argument( "--rotate-mb", type=float, @@ -361,11 +381,14 @@ def main() -> None: print( f"device: {info.get('device_name', '?')} mlx: {mx.__version__} " f"M: {ms} dtypes: {args.dtypes} mlx group_size: {args.group_size} " - f"iters: {args.iters}", + f"warmup: {args.warmup} iters: {args.iters}", flush=True, ) - for dtype_name in args.dtypes: + 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 = {name: bench_shape(name, dtype_name, ms, args) for name in args.shapes} print_dtype_report(dtype_name, ms, res, args.group_size) print(flush=True) From 6c85770b46f01994292f50d77c1baedc78b7dd17 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 3 Nov 2026 17:00:25 -0800 Subject: [PATCH 3/3] [MLX] int8_int4 bench: serialized timing, two-round tile selection, official re-run - Time each config as a dependent chain (each call consumes the previous output) and subtract the chain-only cost, so calls cannot overlap on the GPU as they cannot between model layers. Independent calls hid per-call latency and overstated small-M speedups. - Pick the tile in two rounds: every config with --iters (10), then the --refine-top (3) fastest per lowering (staged, full-K) with --refine-iters (20). - Re-time the refined winner against MLX alternately --final-repeats (5) times with --final-iters (100) and report the medians; MLX is no longer timed after the sweep. --- .../test/bench_int8_int4_gemm.py | 220 +++++++++++++----- 1 file changed, 160 insertions(+), 60 deletions(-) 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 index 99ddeb040ce..e4f2f9ad79a 100644 --- a/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py +++ b/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py @@ -18,8 +18,13 @@ multiple of 32 dividing K). * ``BK = -1``: quantize kernel + full-K MPP ``int8 x int4b`` GEMM; MPP handles the whole K itself. -The fastest config is reported against MLX. The Muse Glimmer estimate is -``sum over shapes of (best time x times the shape occurs in the model)``. +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:: @@ -28,10 +33,16 @@ and send back the output. ``--verbose`` prints every tile config. -Timing: each config / shape / M first runs ``--warmup`` untimed calls (so the -kernel is compiled), then ``--iters`` calls built lazily and timed as a single -``mx.eval``; the reported time is the average per call. The run pauses -``--dtype-pause`` seconds between activation dtypes. +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. @@ -40,6 +51,7 @@ import argparse import math +import statistics import time from typing import Callable, Dict, List, Optional, Tuple @@ -108,12 +120,12 @@ def int8_int4_gemm( M, K = x.shape N = w_packed.shape[0] bm, bn, bk = tile - launch = dict( - 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], - ) + 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( @@ -148,15 +160,25 @@ def check_int8_int4( return err.item() -def time_per_call(fn: Callable[[int], mx.array], warmup: int, iters: int) -> float: - """Average seconds per call over ``iters`` timed runs after ``warmup`` - untimed runs; ``fn(i)`` must build (not evaluate) run ``i``.""" +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(i)) - outs = [fn(i) for i in range(iters)] - start = time.perf_counter() - mx.eval(outs) - return (time.perf_counter() - start) / iters + 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]: @@ -175,10 +197,17 @@ def tile_configs(K: int, args: argparse.Namespace) -> List[Tile]: return tiles -def best(entries: List[Entry]) -> Optional[Entry]: - """Fastest config that ran and passed the accuracy check.""" - ok = [e for e in entries if e[1] is not None and e[2] is not None and e[2] < 1e-2] - return min(ok, key=lambda e: e[1]) if ok else None +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: @@ -188,7 +217,8 @@ def _random_bytes(shape: Tuple[int, ...]) -> mx.array: def bench_shape( name: str, dtype_name: str, ms: List[int], args: argparse.Namespace ) -> Dict[int, dict]: - """Returns {M: {"mlx": seconds, "int8": [Entry]}}.""" + """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 @@ -205,8 +235,13 @@ def bench_shape( ] mx.eval(w_packed, w_scale, mlx_w, mlx_scales, mlx_biases) - def timed(fn: Callable[[int], mx.array]) -> float: - return time_per_call(fn, args.warmup, args.iters) + 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: @@ -217,9 +252,7 @@ def timed(fn: Callable[[int], mx.array]) -> float: try: err = check_int8_int4(x, w_packed[0], w_scale[0], tile) t = timed( - lambda i: int8_int4_gemm( - x, w_packed[i % copies], w_scale[i % copies], tile - ) + 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] @@ -233,9 +266,9 @@ def timed(fn: Callable[[int], mx.array]) -> float: flush=True, ) - def mlx_call(i: int) -> mx.array: + def mlx_call(xi: mx.array, i: int) -> mx.array: return mx.quantized_matmul( - x, + xi, mlx_w[i % copies], mlx_scales[i % copies], mlx_biases[i % copies], @@ -245,23 +278,50 @@ def mlx_call(i: int) -> mx.array: mode="affine", ) - t_mlx = timed(mlx_call) - results[M] = {"mlx": t_mlx, "int8": entries} + 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 - b = best(entries) - summary = ( - f"best {config_name(b[0])} {b[1] * 1e6:.1f} us, " - f"{t_mlx / b[1]:.2f}x vs mlx" - if b is not None - else "no int8_int4 config passed" - ) + 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} mlx {t_mlx * 1e6:>9.1f} us | {summary}", + 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, ) - del w_packed, w_scale, mlx_w, mlx_scales, mlx_biases - mx.clear_cache() return results @@ -270,6 +330,7 @@ def print_dtype_report( ) -> 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}" @@ -281,17 +342,16 @@ def print_dtype_report( for M in ms: ops = 2 * M * N * K t_mlx = shape_res[M]["mlx"] - b = best(shape_res[M]["int8"]) - mlx_cols = f"{t_mlx * 1e6:>10.1f} {ops / t_mlx / 1e12:>6.2f}" + b = shape_res[M]["best"] if b is None: - int8_cols = f"{'failed':>12} {'':>6} {'':<20}" - speedup = f"{'-':>7}" - else: - 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} | {'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}" ) @@ -305,11 +365,11 @@ def print_dtype_report( print(header) print("-" * len(header)) for M in ms: - bests = {n: best(res[n][M]["int8"]) for n in res} - t_mlx = sum(SHAPES[n][2] * res[n][M]["mlx"] for n in res) + 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} | {t_mlx * 1e3:>10.2f} | {'-':>7}") + 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} | " @@ -357,7 +417,33 @@ def main() -> None: help="Untimed runs per config / shape / M before timing (compiles kernels).", ) parser.add_argument( - "--iters", type=int, default=10, help="Timed runs per config / shape / M." + "--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", @@ -365,6 +451,12 @@ def main() -> None: 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, @@ -381,7 +473,9 @@ def main() -> None: 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} iters: {args.iters}", + 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, ) @@ -389,7 +483,13 @@ def main() -> None: 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 = {name: bench_shape(name, dtype_name, ms, args) for name in args.shapes} + 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)