From b499b2672f5a811ed183dda939548131f8c93d38 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 2 Nov 2026 13:30:57 -0800 Subject: [PATCH 1/5] [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/5] [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/5] [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) From d92fee358227abff4a9f68f03c806feb3dad915f Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 4 Nov 2026 11:36:08 -0800 Subject: [PATCH 4/5] [MLX] int8_int4: split-K and matvec kernels, parallel padded quantize; prune bench sweep Kernels (int8_int4_gemm.py): - Quantize kernel: one threadgroup of TPR (256) threads per row instead of one SIMD group; can zero-pad xq to whole BM tiles. Full-K GEMM reads M from x_scale so padded xq works. - Split-K full-K GEMM: SPLITS threadgroups per output tile each cover K/SPLITS and write INT32 partials; a reduce kernel sums them and applies the scales. - Matvec for small M: streams packed weights like MLX's qmv with exact integer accumulation. - Both new paths are bit-identical to the full-K kernel. The handler does not emit them yet. Benchmark (bench_int8_int4_gemm.py): - Sweeps split-K (--splits 1 2 4 8) and matvec (--matvec-max-m 2); refine keeps the top configs of each lowering. - --pad-m and --quant-threads. - Smaller default sweep: staged (BK > 0) is off (never won a serialized run), BM <= max(8, M), M = 4 8 16 256. - Time-based iteration counts: each timing targets --time-budget-ms (50) of GPU time, clamped to [--min-iters, round cap], so large-M timings stay short and the GPU stays cooler. --- .../mlx/custom_kernel_ops/int8_int4_gemm.py | 237 +++++++++++++-- .../test/bench_int8_int4_gemm.py | 272 ++++++++++++++---- 2 files changed, 439 insertions(+), 70 deletions(-) diff --git a/backends/mlx/custom_kernel_ops/int8_int4_gemm.py b/backends/mlx/custom_kernel_ops/int8_int4_gemm.py index 14b40353646..1cbf7a04d6a 100644 --- a/backends/mlx/custom_kernel_ops/int8_int4_gemm.py +++ b/backends/mlx/custom_kernel_ops/int8_int4_gemm.py @@ -37,6 +37,11 @@ scales; then MPP runs ``int8 x int4b_format`` over the full K straight from device memory, choosing its own K tiling and buffering. +``_SPLIT_K_SOURCE`` / ``_SPLIT_K_REDUCE_SOURCE`` (split-K full-K GEMM) and +``_MATVEC_SOURCE`` (small-M matvec) produce bit-identical results to the full-K +kernel. The handler does not emit them yet; ``test/bench_int8_int4_gemm.py`` +times them. + Constraints for MLX lowering: * ``K`` must be statically known and a multiple of 64. * ``N`` must be statically known. @@ -141,6 +146,7 @@ def int8_int4_gemm_fake( _BN = 32 _BK = 64 _THREADS = 128 +_QUANT_THREADS = 256 _HEADER = f""" @@ -310,40 +316,67 @@ def int8_int4_gemm_fake( """ -# full_k=True, kernel 1: one SIMD group per token row writes xq and x_scale. +# full_k=True, kernel 1 (template args 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 = """ - 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); + 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 = lane; i < K / 4; i += 32) { + 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))); } - const float scale = int8_scale_from_amax(simd_max(amax)); + 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); - device uint* xq_row = (device uint*)(xq + m * K); - for (int i = lane; i < K / 4; i += 32) { + for (int i = tid; i < K / 4; i += TPR) { xq_row[i] = quantize_int8x4(x_row[i], inv_scale); } - if (lane == 0) { + if (tid == 0) { x_scale[m] = scale; } """ -# full_k=True, kernel 2 (template args BM, BN): MPP multiplies an int8 [K, M] +# full_k=True, kernel 2 (template args BM, BN): MPP multiplies an int8 [K, rows] # tensor by an int4b [K, N] tensor (W transposed) over the whole K for each -# BM x BN output tile. +# BM x BN output tile. xq may have zero padding rows past M = len(x_scale). _FULL_K_SOURCE = """ - const int M = xq_shape[0]; + 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, M)); + (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( @@ -370,6 +403,170 @@ def int8_int4_gemm_fake( """ +# Split-K full-K, kernel 2 (template args 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]. +_SPLIT_K_SOURCE = """ + 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 / 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]; + } + } + } +""" + + +# Split-K, kernel 3 (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_SIMDGROUPS = 2 +_MATVEC_ROWS_PER_SIMDGROUP = 4 +_MATVEC_BLOCK = 512 + +# 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. 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 in +# fp32); subtracting 8 * 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_SOURCE = """ + constexpr int SG = 2; + constexpr int RPS = 4; + constexpr int VPT = 16; + constexpr int BLOCK = 32 * VPT; + constexpr int KP = K / 2; + + 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 / 2); + } + device const char* x_base = (device const char*)xq + 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) { + 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] - 8.0f * 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)); + } + } + } + } +""" + + def _lit(*values: int) -> list[IntOrVid]: return [IntOrVid.from_literal(v) for v in values] @@ -445,8 +642,12 @@ def _int8_int4_gemm_handler(P: MLXProgramBuilder, n: Node) -> Slot: 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), + grid=[ + IntOrVid.from_literal(_QUANT_THREADS), + m_iov, + IntOrVid.from_literal(1), + ], + threadgroup=_lit(_QUANT_THREADS, 1, 1), input_names=["x"], output_names=["xq", "x_scale"], output_shapes_flat=[m_iov, IntOrVid.from_literal(K), m_iov], @@ -455,9 +656,9 @@ def _int8_int4_gemm_handler(P: MLXProgramBuilder, n: Node) -> Slot: 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], + template_arg_names=["T", "K", "TPR"], + template_arg_kinds=[2, 0, 0], + template_arg_values=[dtype_int, K, _QUANT_THREADS], ) ) P.emit( 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 e4f2f9ad79a..ada2f3f2c4a 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 @@ -14,20 +14,27 @@ 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 > 0`` (staged, off by default): 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. + handles the whole K itself. With ``--splits`` S > 1 (split-K), S + threadgroups per output tile each cover K / S and write INT32 partial + sums, and a reduce kernel adds them and applies the scales. + * matvec (M <= ``--matvec-max-m``): quantize kernel + a SIMD matvec that + streams packed weights like MLX's qmv, with exact integer accumulation + (same results as the MPP kernels). +BM is capped at max(8, M), since larger tiles only add padding rows. Tile selection has two rounds: every config is timed with ``--iters``, then -the ``--refine-top`` fastest configs of each lowering (staged, full-K) are +the ``--refine-top`` fastest configs of each lowering (staged, full-K, split-K, +matvec) 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:: +Defaults cover speculative decoding / prefill (M = 4..256), fp16 and bf16, all +shapes, and the full-K / split-K / matvec tile sweep. Just run:: python -m executorch.backends.mlx.custom_kernel_ops.test.bench_int8_int4_gemm @@ -39,8 +46,10 @@ 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. +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. Short kernels get the cap; long ones (large M) stay cheap. ``--shape-pause`` seconds between shapes and ``--dtype-pause`` seconds between activation dtypes. Calls rotate across enough weight copies to exceed ``--rotate-mb`` so weights @@ -77,10 +86,13 @@ } DTYPES = {"fp16": mx.float16, "bf16": mx.bfloat16} FULL_K = -1 +MATVEC = -2 _THREADGROUP_MEMORY = 32768 -# (BM, BN, BK); BK == FULL_K selects the full-K lowering. -Tile = Tuple[int, int, int] +# (BM, BN, BK, splits); BK == FULL_K selects the full-K lowering, and splits > 1 +# splits K across that many threadgroups per output tile (split-K). BK == MATVEC +# selects the matvec kernel with BM activation rows per threadgroup. +Tile = Tuple[int, int, int, int] # (tile, seconds or None on failure, check error or None) Entry = Tuple[Tile, Optional[float], Optional[float]] @@ -105,21 +117,77 @@ source=op._FULL_K_SOURCE, header=op._HEADER, ) +_split_k_kernel = mx.fast.metal_kernel( + name="int8_int4_gemm_split_k", + input_names=["xq", "x_scale", "w_packed"], + output_names=["partial"], + source=op._SPLIT_K_SOURCE, + header=op._HEADER, +) +_split_k_reduce_kernel = mx.fast.metal_kernel( + name="int8_int4_gemm_split_k_reduce", + input_names=["partial", "x_scale", "w_scale"], + output_names=["out"], + source=op._SPLIT_K_REDUCE_SOURCE, + header=op._HEADER, +) +_matvec_kernel = mx.fast.metal_kernel( + name="int8_int4_gemm_matvec", + input_names=["xq", "x_scale", "w_packed", "w_scale"], + output_names=["out"], + source=op._MATVEC_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}" + bm, bn, bk, splits = tile + if bk == MATVEC: + return f"matvec mr{bm}" + if bk != FULL_K: + return f"staged {bm}x{bn}x{bk}" + 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_int4_gemm( - x: mx.array, w_packed: mx.array, w_scale: mx.array, tile: Tile + x: mx.array, + w_packed: mx.array, + w_scale: mx.array, + tile: Tile, + pad_m: bool = False, + quant_threads: int = op._QUANT_THREADS, ) -> mx.array: M, K = x.shape N = w_packed.shape[0] - bm, bn, bk = tile + bm, bn, bk, splits = tile + if bk == 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_kernel( + 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), @@ -131,24 +199,47 @@ def int8_int4_gemm( return _staged_kernel( inputs=[x, w_packed, w_scale], template=template + [("BK", bk)], **launch )[0] + rows = math.ceil(M / bm) * bm if pad_m else M 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,)], + 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_kernel( + 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_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 + x: mx.array, + w_packed: mx.array, + w_scale: mx.array, + tile: Tile, + cols: int = 256, + **kwargs, ) -> float: """Max error of the first ``cols`` outputs, relative to max |reference|.""" - out = int8_int4_gemm(x, w_packed, w_scale, tile)[:, :cols] + out = int8_int4_gemm(x, w_packed, w_scale, tile, **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) @@ -161,12 +252,18 @@ def check_int8_int4( def time_per_call( - fn: Callable[[mx.array, int], mx.array], x0: mx.array, warmup: int, iters: int + 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``.""" + 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]]) -> float: + 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 @@ -177,23 +274,34 @@ def chain(call: Optional[Callable[[mx.array, int], mx.array]]) -> float: for i in range(warmup): mx.eval(fn(x0, i)) - chain(None) - return max(chain(fn) - chain(None), 0.0) / iters + 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, 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.""" +def tile_configs(K: int, M: int, args: argparse.Namespace) -> List[Tile]: + """Valid tiles; MPP needs BM % 8 == 0, and BM > max(8, M) only adds + padding rows; BK > 0 needs BK % 32 == 0, K % BK == 0, and tiles that fit in + threadgroup memory; split-K needs K % (64 * splits) == 0; matvec needs + K % 512 == 0 and M <= --matvec-max-m.""" tiles: List[Tile] = [] - for bm in (bm for bm in args.bm if bm % 8 == 0): + if M <= args.matvec_max_m and K % op._MATVEC_BLOCK == 0: + mr = 1 + while mr < min(M, 8): + mr *= 2 + tiles.append((mr, 0, MATVEC, 1)) + for bm in (bm for bm in args.bm if bm % 8 == 0 and bm <= max(8, M)): for bn in args.bn: for bk in args.bk: if bk == FULL_K: - tiles.append((bm, bn, bk)) + for splits in args.splits: + if splits == 1 or K % (64 * splits) == 0: + tiles.append((bm, bn, bk, splits)) 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)) + tiles.append((bm, bn, bk, 1)) return tiles @@ -203,11 +311,15 @@ def passed(entries: List[Entry]) -> List[Entry]: def refine_candidates(entries: List[Entry], top: int) -> List[Tile]: - """The ``top`` fastest passing tiles of each lowering (staged, full-K).""" + """The ``top`` fastest passing tiles of each lowering (staged, full-K, + split-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 + 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: @@ -234,23 +346,28 @@ def bench_shape( 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) + quant_opts = {"pad_m": args.pad_m, "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) + return time_per_call( + fn, x0, args.warmup, iters, args.time_budget_ms / 1e3, args.min_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) + return int8_int4_gemm( + xi, w_packed[i % copies], w_scale[i % copies], tile, **quant_opts + ) 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): + for tile in tile_configs(K, M, args): try: - err = check_int8_int4(x, w_packed[0], w_scale[0], tile) + err = check_int8_int4(x, w_packed[0], w_scale[0], tile, **quant_opts) t = timed( lambda xi, i, tile=tile: int8_call(xi, i, tile), x, args.iters ) @@ -383,8 +500,9 @@ def main() -> None: "--m", type=int, nargs="+", - default=[4, 8, 16, 256, 512], - help="Token batch sizes (GEMM M).", + default=[4, 8, 16, 256], + help="Token batch sizes (GEMM M). 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) @@ -399,13 +517,22 @@ def main() -> None: 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("--bn", type=int, nargs="+", default=[16, 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.", + default=[FULL_K], + help=f"K step (multiple of 32 dividing K) for the staged kernel, or " + f"{FULL_K} for full K. Staged never won a serialized M4 run, so only " + f"full K is swept by default; add e.g. 32 64 128 to include staged.", + ) + 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." @@ -417,26 +544,44 @@ 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 tile config in the sweep." + "--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 (staged, full-K) re-timed in " + help="Fastest sweep configs per lowering (staged, full-K, split-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.", + help="Max 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.", + help="Max timed runs per official re-timing of the refined winner and " "MLX.", ) parser.add_argument( "--final-repeats", @@ -463,6 +608,26 @@ def main() -> None: default=512, help="Rotate weight copies until their total exceeds this size.", ) + parser.add_argument( + "--matvec-max-m", + type=int, + default=2, + help="Include the matvec kernel in the sweep for M up to this value " + "(0 disables it).", + ) + parser.add_argument( + "--pad-m", + action="store_true", + help="Full-K: the quantize kernel zero-pads activations to a multiple of " + "BM rows so MPP always sees whole tiles.", + ) + parser.add_argument( + "--quant-threads", + type=int, + default=256, + help="Full-K: 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." ) @@ -473,9 +638,12 @@ 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} sweep iters: {args.iters} " + 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"final iters: {args.final_iters} final repeats: {args.final_repeats} " + f"pad_m: {args.pad_m} quant threads: {args.quant_threads} " + f"splits: {args.splits} matvec max M: {args.matvec_max_m}", flush=True, ) From c0c4172cf7d280f67eebbe727c000a9cbd67131a Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 6 Nov 2026 15:57:51 -0800 Subject: [PATCH 5/5] [MLX] int8_intx (int4/int2) and fp_int4 MPP GEMMs; rename bench to bench_muse_glimmer_gemm int8_intx_gemm (replaces int8_int4_gemm): - mlx::int8_intx_gemm(x, w_packed, w_scale): int8 dynamic activations x symmetric per-channel int4 or int2 weights, INT32 accumulation. Bits are inferred from w_packed's width (K * bits / 8). - Full-K and split-K MPP kernels are templated on the weight format (int4b_format / int2b_format); matvec has a decode per bit width. Quantize and split-K reduce kernels are shared. All paths are bit-identical. - Drops the staged (full_k=False) lowering and the full_k argument. - Helpers: pack_intx / unpack_intx / quantize_intx_per_channel(bits), Int8IntxLinear(w_packed, w_scale, bits). int2 reuses the int4 tiles. fp_int4_gemm (new): - mlx::fp_int4_gemm(x, w_packed, w_scale): fp16 / bf16 activations x the same per-channel int4 weights. MPP half/bfloat x int4b_format with FP32 accumulation (full-K, split-K with FP32 partials) and a matvec for M <= 2. Untuned tile plans. bench_muse_glimmer_gemm (renamed from bench_int8_int4_gemm): - --kernels int8_int4 int8_int2 fp_int4, one sweep and report column each. - Sweep matches the op: matvec only for M <= 2, MPP only for larger M; no staged. BM >= 32 for M >= 128, BN 16 dropped. Activations are always padded to whole BM tiles (--pad-m removed). - Defaults: --m 1 2 4 8 16 256, fp16 only. --- .../mlx/custom_kernel_ops/fp_int4_gemm.py | 518 +++++++++++ .../mlx/custom_kernel_ops/int8_int4_gemm.py | 686 --------------- .../mlx/custom_kernel_ops/int8_intx_gemm.py | 807 ++++++++++++++++++ ...nt4_gemm.py => bench_muse_glimmer_gemm.py} | 576 ++++++++----- 4 files changed, 1685 insertions(+), 902 deletions(-) create mode 100644 backends/mlx/custom_kernel_ops/fp_int4_gemm.py delete mode 100644 backends/mlx/custom_kernel_ops/int8_int4_gemm.py create mode 100644 backends/mlx/custom_kernel_ops/int8_intx_gemm.py rename backends/mlx/custom_kernel_ops/test/{bench_int8_int4_gemm.py => bench_muse_glimmer_gemm.py} (51%) 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_int4_gemm.py b/backends/mlx/custom_kernel_ops/int8_int4_gemm.py deleted file mode 100644 index 1cbf7a04d6a..00000000000 --- a/backends/mlx/custom_kernel_ops/int8_int4_gemm.py +++ /dev/null @@ -1,686 +0,0 @@ -# -# 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. - -``_SPLIT_K_SOURCE`` / ``_SPLIT_K_REDUCE_SOURCE`` (split-K full-K GEMM) and -``_MATVEC_SOURCE`` (small-M matvec) produce bit-identical results to the full-K -kernel. The handler does not emit them yet; ``test/bench_int8_int4_gemm.py`` -times them. - -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 -_QUANT_THREADS = 256 - - -_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 (template args 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; - } -""" - - -# full_k=True, kernel 2 (template args BM, BN): MPP multiplies an int8 [K, rows] -# tensor by an int4b [K, N] 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). -_FULL_K_SOURCE = """ - 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])); - } - } - } -""" - - -# Split-K full-K, kernel 2 (template args 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]. -_SPLIT_K_SOURCE = """ - 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 / 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]; - } - } - } -""" - - -# Split-K, kernel 3 (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_SIMDGROUPS = 2 -_MATVEC_ROWS_PER_SIMDGROUP = 4 -_MATVEC_BLOCK = 512 - -# 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. 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 in -# fp32); subtracting 8 * 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_SOURCE = """ - constexpr int SG = 2; - constexpr int RPS = 4; - constexpr int VPT = 16; - constexpr int BLOCK = 32 * VPT; - constexpr int KP = K / 2; - - 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 / 2); - } - device const char* x_base = (device const char*)xq + 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) { - 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] - 8.0f * 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)); - } - } - } - } -""" - - -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(_QUANT_THREADS), - m_iov, - IntOrVid.from_literal(1), - ], - threadgroup=_lit(_QUANT_THREADS, 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", "TPR"], - template_arg_kinds=[2, 0, 0], - template_arg_values=[dtype_int, K, _QUANT_THREADS], - ) - ) - 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/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_int8_int4_gemm.py b/backends/mlx/custom_kernel_ops/test/bench_muse_glimmer_gemm.py similarity index 51% rename from backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py rename to backends/mlx/custom_kernel_ops/test/bench_muse_glimmer_gemm.py index ada2f3f2c4a..28f227fd6e0 100644 --- a/backends/mlx/custom_kernel_ops/test/bench_int8_int4_gemm.py +++ b/backends/mlx/custom_kernel_ops/test/bench_muse_glimmer_gemm.py @@ -6,37 +6,35 @@ # 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`` (staged, off by default): 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. With ``--splits`` S > 1 (split-K), S - threadgroups per output tile each cover K / S and write INT32 partial - sums, and a reduce kernel adds them and applies the scales. - * matvec (M <= ``--matvec-max-m``): quantize kernel + a SIMD matvec that - streams packed weights like MLX's qmv, with exact integer accumulation - (same results as the MPP kernels). -BM is capped at max(8, M), since larger tiles only add padding rows. +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 (staged, full-K, split-K, -matvec) 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 ``--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)``. -Defaults cover speculative decoding / prefill (M = 4..256), fp16 and bf16, all -shapes, and the full-K / split-K / matvec tile sweep. Just run:: +Just run:: - python -m executorch.backends.mlx.custom_kernel_ops.test.bench_int8_int4_gemm + python -m executorch.backends.mlx.custom_kernel_ops.test.bench_muse_glimmer_gemm and send back the output. ``--verbose`` prints every tile config. @@ -49,11 +47,10 @@ 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. Short kernels get the cap; long ones (large M) stay cheap. -``--shape-pause`` seconds between shapes and ``--dtype-pause`` seconds between -activation dtypes. +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. +are not served from cache. Requires MLX and an MPP-capable GPU (Metal 4). """ @@ -65,7 +62,10 @@ 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 +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. @@ -85,67 +85,95 @@ "lm_head": (_VOCAB, _DIM, 1), } DTYPES = {"fp16": mx.float16, "bf16": mx.bfloat16} -FULL_K = -1 -MATVEC = -2 -_THREADGROUP_MEMORY = 32768 +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, BK, splits); BK == FULL_K selects the full-K lowering, and splits > 1 -# splits K across that many threadgroups per output tile (split-K). BK == MATVEC +# (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, int] +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", + name="int8_intx_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"], +_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._FULL_K_SOURCE, + source=op._SPLIT_K_REDUCE_SOURCE, header=op._HEADER, ) -_split_k_kernel = mx.fast.metal_kernel( - name="int8_int4_gemm_split_k", - input_names=["xq", "x_scale", "w_packed"], +_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=op._SPLIT_K_SOURCE, - header=op._HEADER, + source=fp4._SPLIT_K_SOURCE, + header=fp4._HEADER, ) -_split_k_reduce_kernel = mx.fast.metal_kernel( - name="int8_int4_gemm_split_k_reduce", - input_names=["partial", "x_scale", "w_scale"], +_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=op._SPLIT_K_REDUCE_SOURCE, - header=op._HEADER, + source=fp4._SPLIT_K_REDUCE_SOURCE, + header=fp4._HEADER, ) -_matvec_kernel = mx.fast.metal_kernel( - name="int8_int4_gemm_matvec", - input_names=["xq", "x_scale", "w_packed", "w_scale"], +_fp4_matvec_kernel = mx.fast.metal_kernel( + name="fp_int4_gemm_matvec", + input_names=["x", "w_packed", "w_scale"], output_names=["out"], - source=op._MATVEC_SOURCE, - header=op._HEADER, + source=fp4._MATVEC_SOURCE, + header=fp4._HEADER, ) def config_name(tile: Tile) -> str: - bm, bn, bk, splits = tile - if bk == MATVEC: + bm, bn, splits = tile + if bn == MATVEC: return f"matvec mr{bm}" - if bk != FULL_K: - return f"staged {bm}x{bn}x{bk}" if splits > 1: return f"split_k{splits} {bm}x{bn}" return f"full_k {bm}x{bn}" @@ -155,18 +183,18 @@ def lowering(tile: Tile) -> str: return config_name(tile).split()[0].rstrip("0123456789") -def int8_int4_gemm( +def int8_intx_gemm( x: mx.array, w_packed: mx.array, w_scale: mx.array, tile: Tile, - pad_m: bool = False, + bits: int, quant_threads: int = op._QUANT_THREADS, ) -> mx.array: M, K = x.shape N = w_packed.shape[0] - bm, bn, bk, splits = tile - if bk == MATVEC: + bm, bn, splits = tile + if bn == MATVEC: xq, x_scale = _quant_kernel( inputs=[x], template=[("T", x.dtype), ("K", K), ("TPR", quant_threads)], @@ -176,7 +204,7 @@ def int8_int4_gemm( output_dtypes=[mx.int8, mx.float32], ) rows_per_tg = op._MATVEC_SIMDGROUPS * op._MATVEC_ROWS_PER_SIMDGROUP - return _matvec_kernel( + return _matvec_kernels[bits]( inputs=[xq, x_scale, w_packed, w_scale], template=[("T", x.dtype), ("N", N), ("K", K), ("MR", bm)], grid=( @@ -195,11 +223,8 @@ def int8_int4_gemm( "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] - rows = math.ceil(M / bm) * bm if pad_m else M + # 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)], @@ -209,7 +234,7 @@ def int8_int4_gemm( output_dtypes=[mx.int8, mx.float32], ) if splits > 1: - partial = _split_k_kernel( + 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), @@ -225,30 +250,101 @@ def int8_int4_gemm( output_shapes=[(M, N)], output_dtypes=[x.dtype], )[0] - return _full_k_kernel( + return _full_k_kernels[bits]( inputs=[xq, x_scale, w_packed, w_scale], template=template, **launch )[0] -def check_int8_int4( +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_int4_gemm(x, w_packed, w_scale, tile, **kwargs)[:, :cols] + 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) - 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) + wq = _unpack_wq(w_packed, cols, bits) 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() + 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( @@ -281,28 +377,20 @@ def chain(call: Optional[Callable[[mx.array, int], mx.array]], iters: int) -> fl def tile_configs(K: int, M: int, args: argparse.Namespace) -> List[Tile]: - """Valid tiles; MPP needs BM % 8 == 0, and BM > max(8, M) only adds - padding rows; BK > 0 needs BK % 32 == 0, K % BK == 0, and tiles that fit in - threadgroup memory; split-K needs K % (64 * splits) == 0; matvec needs - K % 512 == 0 and M <= --matvec-max-m.""" - tiles: 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: - mr = 1 - while mr < min(M, 8): - mr *= 2 - tiles.append((mr, 0, MATVEC, 1)) - for bm in (bm for bm in args.bm if bm % 8 == 0 and bm <= max(8, M)): - for bn in args.bn: - for bk in args.bk: - if bk == FULL_K: - for splits in args.splits: - if splits == 1 or K % (64 * splits) == 0: - tiles.append((bm, bn, bk, splits)) - 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, 1)) - return tiles + 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]: @@ -311,8 +399,8 @@ def passed(entries: List[Entry]) -> List[Entry]: def refine_candidates(entries: List[Entry], top: int) -> List[Tile]: - """The ``top`` fastest passing tiles of each lowering (staged, full-K, - split-K).""" + """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: @@ -326,16 +414,21 @@ def _random_bytes(shape: Tuple[int, ...]) -> mx.array: return mx.random.randint(0, 256, shape).astype(mx.uint8) -def bench_shape( +def bench_shape( # noqa: C901 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.""" + """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))) - w_packed = [_random_bytes((N, K // 2)) for _ in range(copies)] + 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 = [ @@ -345,8 +438,8 @@ def bench_shape( 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) - quant_opts = {"pad_m": args.pad_m, "quant_threads": args.quant_threads} + 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 @@ -355,87 +448,125 @@ def timed( fn, x0, args.warmup, iters, args.time_budget_ms / 1e3, args.min_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, **quant_opts + 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", ) - results: Dict[int, dict] = {} - for M in ms: - x = mx.random.normal((M, K)).astype(dtype) - mx.eval(x) + 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 tile_configs(K, M, args): + for tile in tiles: try: - err = check_int8_int4(x, w_packed[0], w_scale[0], tile, **quant_opts) + err = check(kernel, x, tile) t = timed( - lambda xi, i, tile=tile: int8_call(xi, i, tile), x, args.iters + 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 {config_name(tile)}: {reason}", flush=True) + 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" {config_name(tile):<20} {t * 1e6:>10.1f} us " + f" {kernel:<9} {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 len(candidates) == 1: + return candidates[0], next(e[1] for e in entries if e[0] == candidates[0]) 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", + f" [{dtype_name}] {name:<8} M={M:<4} no {kernel} config passed", flush=True, ) - continue - + return None refined = { tile: timed( - lambda xi, i, tile=tile: int8_call(xi, i, tile), x, args.refine_iters + lambda xi, i, tile=tile: call(kernel, 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) + print(f" {kernel} refine (us): {ranked}", flush=True) + tile = min(refined, key=refined.get) + return tile, refined[tile] - t_int8s, t_mlxs = [], [] + 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): - t_int8s.append( - timed( - lambda xi, i, tile=tile: int8_call(xi, i, tile), - x, - args.final_iters, + 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_int8 = statistics.median(t_int8s) t_mlx = statistics.median(t_mlxs) - results[M] = {"mlx": t_mlx, "int8": entries, "best": (tile, t_int8)} + 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} {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)", + f" [{dtype_name}] {name:<8} M={M:<4} mlx {t_mlx * 1e6:.1f} us | " + + " | ".join(parts), flush=True, ) @@ -443,15 +574,18 @@ def mlx_call(xi: mx.array, i: int) -> mx.array: def print_dtype_report( - dtype_name: str, ms: List[int], res: Dict[str, Dict[int, dict]], gs: int + dtype_name: str, + ms: List[int], + res: Dict[str, Dict[int, dict]], + gs: int, + kernels: List[str], ) -> 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(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(): @@ -459,39 +593,43 @@ def print_dtype_report( 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}" + 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} | {'int8_int4 ms':>12} | {'mlx ms':>10} | {'speedup':>7}" + 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: - 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" - ) + 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: @@ -500,32 +638,43 @@ def main() -> None: "--m", type=int, nargs="+", - default=[4, 8, 16, 256], - help="Token batch sizes (GEMM M). 256 is the Muse Glimmer prefill chunk; " - "larger M is GPU-saturated and scales linearly.", + 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=list(DTYPES) + "--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="Output tile rows (MPP requires a multiple of 8).", + help="MPP output tile rows (multiple of 8; capped at max(8, M)).", ) - parser.add_argument("--bn", type=int, nargs="+", default=[16, 32, 64, 128]) parser.add_argument( - "--bk", + "--bn", type=int, nargs="+", - default=[FULL_K], - help=f"K step (multiple of 32 dividing K) for the staged kernel, or " - f"{FULL_K} for full K. Staged never won a serialized M4 run, so only " - f"full K is swept by default; add e.g. 32 64 128 to include staged.", + default=[32, 64, 128], + help="MPP output tile columns. BN = 16 never won a full M4 sweep.", ) parser.add_argument( "--splits", @@ -567,8 +716,7 @@ def main() -> None: "--refine-top", type=int, default=3, - help="Fastest sweep configs per lowering (staged, full-K, split-K) " - "re-timed in " + help="Fastest sweep configs per lowering (full-K, split-K) re-timed in " "the refine round.", ) parser.add_argument( @@ -612,20 +760,14 @@ def main() -> None: "--matvec-max-m", type=int, default=2, - help="Include the matvec kernel in the sweep for M up to this value " - "(0 disables it).", - ) - parser.add_argument( - "--pad-m", - action="store_true", - help="Full-K: the quantize kernel zero-pads activations to a multiple of " - "BM rows so MPP always sees whole tiles.", + 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="Full-K: threads per row in the activation quantize kernel " + help="int8_int4: threads per row in the activation quantize kernel " "(multiple of 32, at most 1024).", ) parser.add_argument( @@ -637,13 +779,15 @@ def main() -> None: 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"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"pad_m: {args.pad_m} quant threads: {args.quant_threads} " - f"splits: {args.splits} matvec max M: {args.matvec_max_m}", + 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, ) @@ -658,7 +802,7 @@ def main() -> None: 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_dtype_report(dtype_name, ms, res, args.group_size, args.kernels) print(flush=True)