Phase M3: native gemm_4bit -- chunked dequant + MPSMatrixMultiplication, one command buffer

Option A from the M1 decision: bnb_mps_gemm_4bit (csrc/mps_ops.mm) encodes, on ONE
command buffer with ONE commit + ONE blocking wait, (1) a chunked 4-bit dequant kernel
(dequantize_4bit_chunked_fp32/fp16: one thread per 32-element uint4 chunk, writing
(T)(code[nib]*absmax) into a growable private scratch MTLBuffer -- the same rounding as
the oracle's B_dq.to(dtype)), (2) a shape-cached MPSMatrixMultiplication computing
A[M,K] . B_dq[N,K]^T (row-major, transposeRight), and (3) an optional bias epilogue
kernel out[m,n] += bias[n]. The single sync is the structural win: the previous
dequant + F.linear tail paid the ~0.15-0.25ms cross-queue sync twice per call.

bf16 is excluded by the Python router: MPSMatrixMultiplication hard-asserts on anything
but fp32/fp16/int8/int16 (probed on macOS 26.4.1), so bf16 keeps the existing
dequant + F.linear fallback verbatim. Other guards mirror the fused gemv (K % 32 == 0,
power-of-two blocksize >= 32, packed-size check, hasattr guard for stale dylibs); nested
absmax is still unpacked to plain fp32 absmax before routing, unchanged.

Parity: tests/test_mps_parity.py -k "gemm_4bit or gemv" = 69 passed under
BNB_MPS_REQUIRE_NATIVE=1. Native path asserted via spy (fp32/fp16 x nf4/fp4 x +/-bias x
+/-nested absmax); bf16 and K%32!=0 fallbacks asserted; graceful fallback with the
native handle off covered. fp32 vs MPSMatMul accumulation stays within the documented
1e-5 atol at the calibrated K <= 256 (the one trip found was in the PURE-TORCH fallback
composition at K=256/M=4; the fallback test pins K=64 like the main gemm test -- the
documented tolerance is unchanged).

Speedup vs the dequant+F.linear fallback (nf4/bs64, N=K=4096, 30 iters,
benchmarks_wip/bench_gemm_baseline.py): fp16 2.5x (M=8), 1.5x (M=64/512), 1.08x
(M=2048); fp32 1.6x/1.5x (M=8/64), 1.1x (M=512), ~1.0x (M=2048). Win = one sync + a
much faster chunked dequant at small/medium M; flat at M=2048 where the GEMM dominates
and MPSMatMul ~= F.linear.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: eaglstun

benchmarks_wip/bench_gemm_baseline.py

@@ -1,14 +1,19 @@
-"""Phase M1 baseline, gemm side: as M grows, does the GEMM cost overtake dequant?
+"""Phase M1 baseline + Phase M3 comparison for gemm_4bit.
 
-If dequant stays ~fixed (it's B-only, independent of M) while linear grows with M,
-there's a crossover M above which Option A (fast on-device GEMM) starts to matter.
-Below it, gemm is just as dequant-bound as gemv and Option B wins there too.
+M1 question: as M grows, does the GEMM cost overtake dequant? (Answer: crossover near
+M~512; below it gemm is dequant-bound.)
+
+M3 question: what does the native path (chunked dequant -> scratch -> MPSMatrixMultiplication
+-> bias, ONE command buffer / ONE sync) buy over the dequant + F.linear fallback (native
+dequant wait + torch GEMM + second sync)? `gemm_4bit` routes native automatically when built,
+so `native` here is just the op; `fallback` reproduces the old tail verbatim.
 """
 
 import time
 
 import torch
 
+import bitsandbytes.backends.mps.ops as mps_ops
 import bitsandbytes.functional as F
 
 DEV = "mps"
@@ -36,29 +41,36 @@ def bench(M, N, K, dtype, quant_type="nf4", blocksize=64):
     B = torch.randn(N, K, dtype=dtype, device=DEV)
     B_q, qs = F.quantize_4bit(B, blocksize=blocksize, quant_type=quant_type)
 
-    def full():
+    def native():
+        # Routes through bnb_mps_gemm_4bit when the native library is built (fp32/fp16).
         return torch.ops.bitsandbytes.gemm_4bit(A, B_q, list(B.shape), qs.absmax, blocksize, quant_type)
 
     def dequant_only():
         return torch.ops.bitsandbytes.dequantize_4bit(B_q, qs.absmax, blocksize, quant_type, list(B.shape), dtype)
 
+    def fallback():
+        # The pre-M3 tail: native dequant (its own sync) + torch F.linear (torch's queue).
+        B_dq = mps_ops._dequantize_4bit_impl(B_q, qs.absmax, blocksize, quant_type, list(B.shape), dtype)
+        return torch.nn.functional.linear(A, B_dq)
+
     B_dq = dequant_only()
 
     def linear_only():
         return torch.nn.functional.linear(A, B_dq)
 
-    t_full, t_deq, t_lin = timed(full), timed(dequant_only), timed(linear_only)
+    t_nat, t_fb, t_deq, t_lin = timed(native), timed(fallback), timed(dequant_only), timed(linear_only)
     print(
-        f"  M={M:>4} N={N:>5} K={K:>5} {str(dtype).replace('torch.',''):>8}  "
-        f"total={t_full:7.3f}ms  dequant={t_deq:7.3f}ms  linear={t_lin:7.3f}ms  "
-        f"(linear is {100*t_lin/t_full:4.1f}% of total)"
+        f"  M={M:>4} N={N:>5} K={K:>5} {str(dtype).replace('torch.', ''):>8}  "
+        f"native={t_nat:7.3f}ms  fallback={t_fb:7.3f}ms  ({t_fb / t_nat:4.2f}x)  "
+        f"[fallback = dequant {t_deq:6.3f} + linear {t_lin:6.3f}]"
     )
 
 
 if __name__ == "__main__":
-    print(f"iters={ITERS} warmup={WARMUP} device={DEV}\n")
-    for dtype in (torch.float16,):
-        print(f"=== gemm_4bit, N=K=4096, {str(dtype).replace('torch.','')} ===")
+    native = "native" if mps_ops._native_available() else "FALLBACK-ONLY (no native build)"
+    print(f"iters={ITERS} warmup={WARMUP} device={DEV} lib={native}\n")
+    for dtype in (torch.float16, torch.float32):
+        print(f"=== gemm_4bit, N=K=4096, {str(dtype).replace('torch.', '')} ===")
         for M in (8, 64, 512, 2048):
             bench(M, 4096, 4096, dtype)
         print()

bitsandbytes/backends/mps/ops.py

@@ -477,6 +477,65 @@ def _(
     out.copy_(result)
 
 
+def _gemm_4bit_native(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    shapeB: Sequence[int],
+    absmax: torch.Tensor,
+    blocksize: int,
+    quant_type: str,
+    bias: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Route gemm_4bit (general M) through the native Metal entry point.
+
+    One command buffer, one commit, one blocking wait: a chunked Metal kernel dequantizes
+    packed B into a private scratch MTLBuffer in A's dtype (reproducing the reference's
+    B_dq.to(dtype) rounding), MPSMatrixMultiplication computes A[M,K] . B_dq[N,K]^T, and an
+    optional bias epilogue adds bias[N] -- all on the same command buffer. Compared with the
+    dequant + F.linear fallback this removes the torch round-trip and its second sync (the
+    per-call cross-queue sync is what dominates wall-clock at small/medium M -- see the
+    Phase M2 finding in MPS_STATUS.md).
+
+    `absmax` must already be the plain per-block fp32 scale: nested/compressed absmax is
+    unpacked by the caller BEFORE this function. Preconditions (checked by the caller):
+    A.dtype is fp32/fp16 (MPSMatrixMultiplication asserts on bf16), K % 32 == 0, and
+    power-of-two blocksize >= 32.
+    """
+    N, K = int(shapeB[0]), int(shapeB[-1])
+    M = A.numel() // K
+
+    B_flat = B if B.dtype == torch.uint8 else B.view(torch.uint8)
+    B_flat = _ensure_native_buffer(B_flat.reshape(-1))
+    A_flat = _ensure_native_buffer(A.reshape(-1))
+    code_f = _ensure_native_buffer(_get_4bit_code(quant_type, A.device).to(torch.float32))
+    absmax_f = _ensure_native_buffer(absmax.to(torch.float32))
+    bias_ptr = None
+    if bias is not None:
+        bias_f = _ensure_native_buffer(bias.reshape(-1))
+        bias_ptr = bias_f.data_ptr()
+
+    out = torch.empty(M * N, dtype=A.dtype, device=A.device)
+    dtype_flag = {torch.float32: 0, torch.float16: 1}[A.dtype]
+    bs_shift = blocksize.bit_length() - 1
+
+    torch.mps.synchronize()
+    _mps_native.bnb_mps_gemm_4bit(
+        code_f.data_ptr(),
+        B_flat.data_ptr(),
+        absmax_f.data_ptr(),
+        A_flat.data_ptr(),
+        bias_ptr,
+        out.data_ptr(),
+        M,
+        K,
+        N,
+        bs_shift,
+        dtype_flag,
+    )
+
+    return out.reshape(*A.shape[:-1], N)
+
+
 @register_kernel("bitsandbytes::gemm_4bit", "mps")
 def _(
     A: torch.Tensor,
@@ -502,6 +561,26 @@ def _(
             + absmax_offset
         )
 
+    # Native Metal path (dequant -> scratch -> MPSMatrixMultiplication -> bias, one command
+    # buffer / one sync). Guards: fp32/fp16 only (MPSMatrixMultiplication hard-asserts on
+    # bf16 -- verified on macOS 26.4.1 -- so bf16 keeps the dequant + F.linear fallback),
+    # 2-D shapeB matching A's K, K % 32 == 0 (uint4 loads in the chunked dequant kernel),
+    # power-of-two blocksize (absmax indexed with a shift), a bias matching out's dtype and
+    # width, and packed B of the expected size.
+    if (
+        _native_available()
+        and hasattr(_mps_native._lib, "bnb_mps_gemm_4bit")  # stale dylibs predate the native gemm
+        and A.dtype in (torch.float32, torch.float16)
+        and len(shapeB) == 2
+        and shapeB[-1] == K
+        and K % 32 == 0
+        and blocksize >= 32
+        and (blocksize & (blocksize - 1)) == 0
+        and (bias is None or (bias.dtype == A.dtype and bias.numel() == N))
+        and B.numel() * B.element_size() == (N * K) // 2
+    ):
+        return _gemm_4bit_native(A, B, shapeB, absmax, blocksize, quant_type, bias)
+
     # Use HF Hub kernel when supported for GEMV.
     if M == 1 and blocksize in (64, 128, 256) and (k := _get_kernel()) is not None:
         if B.dtype != torch.uint8:

bitsandbytes/cextension.py

@@ -197,6 +197,22 @@ def __init__(self, lib: ct.CDLL, metallib_path: Path):
                 ct.c_int64,  # dtype_flag (0=fp32, 1=fp16, 2=bf16)
             ]
 
+        if hasattr(lib, "bnb_mps_gemm_4bit"):
+            lib.bnb_mps_gemm_4bit.restype = None
+            lib.bnb_mps_gemm_4bit.argtypes = [
+                ct.c_void_p,  # code (float32[16])
+                ct.c_void_p,  # B (uint8 packed, N*K/2 bytes)
+                ct.c_void_p,  # absmax (float32[blocks])
+                ct.c_void_p,  # A (activation dtype [M*K])
+                ct.c_void_p,  # bias (activation dtype [N]; None when absent)
+                ct.c_void_p,  # out (activation dtype [M*N])
+                ct.c_int64,  # M
+                ct.c_int64,  # K
+                ct.c_int64,  # N
+                ct.c_int64,  # bs_shift = log2(blocksize)
+                ct.c_int64,  # dtype_flag (0=fp32, 1=fp16; MPSMatrixMultiplication has no bf16)
+            ]
+
     def verify_buffer_contract(self) -> None:
         """Verify the undocumented torch contract that an MPS tensor's data_ptr() is its
         id<MTLBuffer>. Raises RuntimeError if a future torch has broken it, so callers can

csrc/mps_kernels.metal

@@ -234,6 +234,72 @@ BNB_GEMV_4BIT_KERNEL(gemv_4bit_fp32, float)
 BNB_GEMV_4BIT_KERNEL(gemv_4bit_fp16, half)
 BNB_GEMV_4BIT_KERNEL(gemv_4bit_bf16, bfloat)
 
+// ---- Phase M3: chunked 4-bit dequant into the ACTIVATION dtype (gemm_4bit scratch) ----
+// Fills the scratch B_dq consumed by MPSMatrixMultiplication in bnb_mps_gemm_4bit. One
+// thread per 32-element chunk (16 packed bytes, one uint4 load), writing
+// (T)(code[nib] * absmax) -- the same rounding as the reference's B_dq.to(dtype), so the
+// GEMM multiplies exactly the weights the oracle multiplies. Preconditions match the fused
+// gemv kernel (enforced by the Python router): K % 32 == 0 so rows are 16-byte aligned and
+// total elements are a multiple of 32; blocksize is a power of two >= 32 (bs_shift =
+// log2(blocksize)), so a chunk never straddles an absmax block.
+template <typename T>
+static inline void dequantize_4bit_chunked_body(
+    device const float* code,
+    device const uchar* B,
+    device const float* absmax,
+    device T* out,
+    uint bs_shift,
+    uint chunk) {
+    const ulong base = (ulong)chunk << 5;  // first element index of this chunk
+    device const uint4* p = (device const uint4*)(B + (base >> 1));
+    const uint4 packed = *p;
+    const float am = absmax[base >> bs_shift];
+
+#pragma unroll
+    for (uint w = 0; w < 4; ++w) {
+        const uint word = packed[w];
+        // Little-endian: byte b of `word` is packed byte (base/2 + w*4 + b), holding
+        // elements base + w*8 + 2b (high nibble) and base + w*8 + 2b + 1 (low nibble).
+#pragma unroll
+        for (uint b = 0; b < 4; ++b) {
+            const uint byte = (word >> (b << 3)) & 0xFFu;
+            const ulong j = base + (ulong)((w << 3) | (b << 1));
+            out[j] = (T)(code[byte >> 4] * am);
+            out[j + 1] = (T)(code[byte & 0x0Fu] * am);
+        }
+    }
+}
+
+#define BNB_DEQUANT_4BIT_CHUNKED_KERNEL(NAME, T)                                                                     \
+    kernel void NAME(                                                                                                \
+        device const float* code [[buffer(0)]],   /* 16-entry 4-bit code (NF4 or FP4) */                            \
+        device const uchar* B [[buffer(1)]],      /* packed nibbles, row-major [N, K], N*K/2 bytes */               \
+        device const float* absmax [[buffer(2)]], /* per-block scales over the flattened [N*K] index */             \
+        device T* out [[buffer(3)]],              /* N*K elements of T (the GEMM scratch) */                        \
+        constant uint& bs_shift [[buffer(4)]],    /* log2(blocksize) */                                             \
+        uint chunk [[thread_position_in_grid]]) {                                                                   \
+        dequantize_4bit_chunked_body<T>(code, B, absmax, out, bs_shift, chunk);                                      \
+    }
+
+BNB_DEQUANT_4BIT_CHUNKED_KERNEL(dequantize_4bit_chunked_fp32, float)
+BNB_DEQUANT_4BIT_CHUNKED_KERNEL(dequantize_4bit_chunked_fp16, half)
+
+// ---- Phase M3: bias epilogue for gemm_4bit ----
+// out[m, n] += bias[n], broadcast over rows, in the activation dtype (reproducing
+// F.linear's bias add on the T-typed matmul result). 2-D grid: x = n (column), y = m (row).
+#define BNB_GEMM_BIAS_ADD_KERNEL(NAME, T)                                                                            \
+    kernel void NAME(                                                                                                \
+        device T* out [[buffer(0)]],           /* [M, N] row-major */                                               \
+        device const T* bias [[buffer(1)]],    /* [N] */                                                            \
+        constant uint& N [[buffer(2)]],                                                                              \
+        uint2 gid [[thread_position_in_grid]]) {                                                                     \
+        const ulong idx = (ulong)gid.y * (ulong)N + (ulong)gid.x;                                                    \
+        out[idx] = (T)(out[idx] + bias[gid.x]);                                                                      \
+    }
+
+BNB_GEMM_BIAS_ADD_KERNEL(gemm_bias_add_fp32, float)
+BNB_GEMM_BIAS_ADD_KERNEL(gemm_bias_add_fp16, half)
+
 // ---- 4-bit blockwise quantize (NF4/FP4): A (float32) -> packed out + absmax ----
 // `bounds` are the 15 midpoints of the SORTED 16-entry code; `order` maps the searchsorted
 // index back to the stored 4-bit index (identity for NF4, the argsort remap for FP4).