Phase M2: fused native gemv_4bit Metal kernel (dequant + dot product, no B_dq)

One SIMD-group (32 threads) per output element n: threads stride the packed
row in uint4 units (16 B = 32 nibbles), dequantize nf4/fp4 in registers
(code[nib] * absmax, rounded to the activation dtype to match the oracle's
B_dq.to(dtype)), and accumulate the dot product in fp32 with split
accumulators + explicit fma, reduced via simd_sum. Per-dtype kernel variants
(fp32/fp16/bf16) bind A and out in torch's own dtype, so the steady-state
call launches zero torch cast kernels. The dequantized B is never written to
device memory.

- csrc/mps_kernels.metal: gemv_4bit_body<T> + fp32/fp16/bf16 kernels
- csrc/mps_ops.mm: bnb_mps_gemv_4bit dispatch (own queue, blocking wait,
  per-dtype pipeline; BNB_MPS_PROFILE=1 logs kernel-only GPU time)
- cextension.py: argtypes (guarded by hasattr so stale dylibs keep working)
- backends/mps/ops.py: route _gemv_4bit_impl to native when available;
  guards (M==1, K % 32 == 0, pow2 blocksize, 16-entry code, packed-size
  check) fall back to the unchanged dequant + F.linear path
- tests: fused-native spy test (asserts native is actually hit), unaligned-K
  fallback test, gemv added to the graceful-fallback matrix

Parity: tests/test_mps_parity.py -k gemv all green under
BNB_MPS_REQUIRE_NATIVE=1 (nf4+fp4, bs 64/256, fp32/fp16/bf16).
Speedup vs dequant+F.linear baseline (nf4/bs64, M1 shapes): 3.4-6.2x
wall-clock (fp16 4096x4096: 0.31ms vs 1.64ms). Kernel-only ~0.11ms on the
~25 MB shapes when clocked up (~230 GB/s read vs the standalone dequant
kernel's ~54 GB/s); per-call cross-queue sync now dominates wall-clock.

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

Author: eaglstun

benchmarks_wip/bench_gemm_baseline.py

@@ -0,0 +1,64 @@
+"""Phase M1 baseline, gemm side: as M grows, does the GEMM cost overtake dequant?
+
+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.
+"""
+
+import time
+
+import torch
+
+import bitsandbytes.functional as F
+
+DEV = "mps"
+ITERS = 30
+WARMUP = 8
+
+
+def sync():
+    torch.mps.synchronize()
+
+
+def timed(fn):
+    for _ in range(WARMUP):
+        fn()
+    sync()
+    t0 = time.perf_counter()
+    for _ in range(ITERS):
+        fn()
+    sync()
+    return (time.perf_counter() - t0) / ITERS * 1e3
+
+
+def bench(M, N, K, dtype, quant_type="nf4", blocksize=64):
+    A = torch.randn(1, M, K, dtype=dtype, device=DEV)
+    B = torch.randn(N, K, dtype=dtype, device=DEV)
+    B_q, qs = F.quantize_4bit(B, blocksize=blocksize, quant_type=quant_type)
+
+    def full():
+        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)
+
+    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)
+    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)"
+    )
+
+
+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.','')} ===")
+        for M in (8, 64, 512, 2048):
+            bench(M, 4096, 4096, dtype)
+        print()

benchmarks_wip/bench_gemv_fused.py

@@ -0,0 +1,79 @@
+"""Phase M2: fused native gemv_4bit vs the Phase-M1 baseline (dequant -> F.linear).
+
+Per shape, times:
+  - fused    torch.ops.bitsandbytes.gemv_4bit (routes to the fused Metal kernel)
+  - baseline native dequant of B + F.linear (what gemv_4bit did before Phase M2)
+and reports ms/iter plus the fused kernel's effective read bandwidth
+(packed B + absmax + A, i.e. the memory the fused kernel actually touches).
+"""
+
+import time
+
+import torch
+
+from bitsandbytes.backends.mps import ops as mps_ops
+import bitsandbytes.functional as F
+
+DEV = "mps"
+ITERS = 50
+WARMUP = 10
+
+assert mps_ops._native_available(), "native MPS library required for this benchmark"
+assert hasattr(mps_ops._mps_native._lib, "bnb_mps_gemv_4bit"), "fused gemv kernel missing"
+
+
+def sync():
+    torch.mps.synchronize()
+
+
+def timed(fn):
+    for _ in range(WARMUP):
+        fn()
+    sync()
+    t0 = time.perf_counter()
+    for _ in range(ITERS):
+        fn()
+    sync()
+    return (time.perf_counter() - t0) / ITERS * 1e3  # ms/iter
+
+
+def bench(N, K, dtype, quant_type="nf4", blocksize=64):
+    A = torch.randn(1, 1, K, dtype=dtype, device=DEV)
+    B = torch.randn(N, K, dtype=dtype, device=DEV)
+    B_q, absmax = torch.ops.bitsandbytes.quantize_4bit(B, blocksize, quant_type, torch.uint8)
+    code = F.get_4bit_type(quant_type, device=DEV, blocksize=blocksize)
+
+    def fused():
+        return torch.ops.bitsandbytes.gemv_4bit(A, B_q, B.shape, absmax, code, blocksize)
+
+    def baseline():
+        B_dq = torch.ops.bitsandbytes.dequantize_4bit(B_q, absmax, blocksize, quant_type, list(B.shape), dtype)
+        return torch.nn.functional.linear(A, B_dq)
+
+    # Sanity: fused output matches the baseline it replaces.
+    ref = baseline()
+    got = fused()
+    max_err = (got.float() - ref.float()).abs().max().item()
+
+    t_fused = timed(fused)
+    t_base = timed(baseline)
+
+    # Memory the fused kernel reads: packed B (N*K/2 bytes) + absmax (N*K/blocksize fp32)
+    # + A (K fp32); writes out (N fp32).
+    bytes_moved = N * K // 2 + (N * K // blocksize) * 4 + K * 4 + N * 4
+    gbps = bytes_moved / (t_fused * 1e-3) / 1e9
+
+    print(
+        f"  N={N:>6} K={K:>6} {str(dtype).replace('torch.', ''):>8}  "
+        f"fused={t_fused:7.3f}ms  baseline={t_base:7.3f}ms  "
+        f"speedup={t_base / t_fused:5.1f}x  fused-read={gbps:6.1f} GB/s  max|err|={max_err:.2e}"
+    )
+
+
+if __name__ == "__main__":
+    print(f"iters={ITERS} warmup={WARMUP} device={DEV}\n")
+    for dtype in (torch.float16, torch.bfloat16, torch.float32):
+        print(f"=== gemv_4bit (M=1), {str(dtype).replace('torch.', '')} ===")
+        for N, K in [(4096, 4096), (11008, 4096), (4096, 11008)]:
+            bench(N, K, dtype)
+        print()

benchmarks_wip/bench_matmul_baseline.py

@@ -0,0 +1,70 @@
+"""Phase M1 baseline: how slow is today's unfused gemv_4bit (dequant -> F.linear) on MPS?
+
+Times, per shape:
+  - total  gemv_4bit (native dequant of B + torch F.linear)
+  - dequant-only (native _dequantize_4bit_impl of B)  -> the isolable cost
+  - F.linear on an already-materialized B_dq          -> the GEMM cost
+so we can see where the wall-clock actually goes and whether fusing dequant into
+the matmul (Option B) or just moving the GEMM on-device (Option A) is the win.
+"""
+
+import time
+
+import torch
+
+import bitsandbytes.functional as F
+
+DEV = "mps"
+ITERS = 50
+WARMUP = 10
+
+
+def sync():
+    torch.mps.synchronize()
+
+
+def timed(fn):
+    for _ in range(WARMUP):
+        fn()
+    sync()
+    t0 = time.perf_counter()
+    for _ in range(ITERS):
+        fn()
+    sync()
+    return (time.perf_counter() - t0) / ITERS * 1e3  # ms/iter
+
+
+def bench(N, K, dtype, quant_type="nf4", blocksize=64):
+    A = torch.randn(1, 1, K, dtype=dtype, device=DEV)
+    B = torch.randn(N, K, dtype=dtype, device=DEV)
+    B_q, absmax = torch.ops.bitsandbytes.quantize_4bit(B, blocksize, quant_type, torch.uint8)
+    code = F.get_4bit_type(quant_type, device=DEV, blocksize=blocksize)
+
+    def full():
+        return torch.ops.bitsandbytes.gemv_4bit(A, B_q, B.shape, absmax, code, blocksize)
+
+    def dequant_only():
+        return torch.ops.bitsandbytes.dequantize_4bit(B_q, absmax, blocksize, quant_type, list(B.shape), dtype)
+
+    B_dq = dequant_only()
+
+    def linear_only():
+        return torch.nn.functional.linear(A, B_dq)
+
+    t_full = timed(full)
+    t_deq = timed(dequant_only)
+    t_lin = timed(linear_only)
+    print(
+        f"  N={N:>6} K={K:>6} {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"(dequant is {100 * t_deq / t_full:4.1f}% of total)"
+    )
+
+
+if __name__ == "__main__":
+    print(f"iters={ITERS} warmup={WARMUP} device={DEV}\n")
+    for dtype in (torch.float16, torch.bfloat16):
+        print(f"=== gemv_4bit (M=1), {str(dtype).replace('torch.', '')} ===")
+        for N, K in [(4096, 4096), (11008, 4096), (4096, 11008)]:
+            bench(N, K, dtype)
+        print()

bitsandbytes/backends/mps/ops.py

@@ -365,6 +365,52 @@ def _(
     out.copy_(result)
 
 
+def _gemv_4bit_native(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    shapeB: Sequence[int],
+    absmax: torch.Tensor,
+    code: torch.Tensor,
+    blocksize: int,
+) -> torch.Tensor:
+    """Route gemv_4bit (M == 1) through the fused hand-written Metal kernel.
+
+    The kernel reads packed 4-bit B + per-block absmax + the 16-entry code table and
+    computes out[n] = sum_k A[k] * dequant(B[n, k]) directly -- the dequantized B is never
+    materialized. Dequantized weights are rounded to A's dtype in-kernel (reproducing the
+    reference's B_dq.to(dtype)); accumulation is fp32, so only accumulation order differs
+    from the oracle. A and out bind in A's own dtype (per-dtype kernel variants), so the
+    steady-state call launches no torch cast kernels. Preconditions (checked by the
+    caller): K % 32 == 0 and power-of-two blocksize.
+    """
+    N, K = int(shapeB[0]), int(shapeB[-1])
+
+    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(code.to(torch.float32))
+    absmax_f = _ensure_native_buffer(absmax.to(torch.float32))
+
+    out = torch.empty(N, dtype=A.dtype, device=A.device)
+    dtype_flag = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2}[A.dtype]
+    bs_shift = blocksize.bit_length() - 1
+
+    torch.mps.synchronize()
+    _mps_native.bnb_mps_gemv_4bit(
+        code_f.data_ptr(),
+        B_flat.data_ptr(),
+        absmax_f.data_ptr(),
+        A_flat.data_ptr(),
+        out.data_ptr(),
+        K,
+        N,
+        bs_shift,
+        dtype_flag,
+    )
+
+    return out.reshape(*A.shape[:-1], N)
+
+
 def _gemv_4bit_impl(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -373,6 +419,24 @@ def _gemv_4bit_impl(
     code: torch.Tensor,
     blocksize: int,
 ) -> torch.Tensor:
+    # Fused native Metal kernel when available. Guards: true gemv (M == 1), 2-D shapeB
+    # matching A's K, K % 32 == 0 (uint4 row loads need 16-byte-aligned rows), power-of-two
+    # blocksize (the kernel indexes absmax with a shift), and a plain 16-entry code table
+    # whose packed B has the expected size.
+    if (
+        _native_available()
+        and hasattr(_mps_native._lib, "bnb_mps_gemv_4bit")  # stale dylibs predate the fused kernel
+        and A.numel() == A.shape[-1]
+        and len(shapeB) == 2
+        and shapeB[-1] == A.shape[-1]
+        and shapeB[-1] % 32 == 0
+        and blocksize >= 32
+        and (blocksize & (blocksize - 1)) == 0
+        and code.numel() == 16
+        and B.numel() * B.element_size() == (shapeB[0] * shapeB[1]) // 2
+    ):
+        return _gemv_4bit_native(A, B, shapeB, absmax, code, blocksize)
+
     if blocksize in (64, 128, 256) and (k := _get_kernel()) is not None:
         if B.dtype != torch.uint8:
             B = B.view(torch.uint8)

bitsandbytes/cextension.py

@@ -181,6 +181,22 @@ def __init__(self, lib: ct.CDLL, metallib_path: Path):
             ct.c_int64,  # blocksize
         ]
 
+        # Older native builds predate the fused matmul kernels; guard so a stale dylib
+        # keeps the quant/dequant native path without breaking library load.
+        if hasattr(lib, "bnb_mps_gemv_4bit"):
+            lib.bnb_mps_gemv_4bit.restype = None
+            lib.bnb_mps_gemv_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 [K])
+                ct.c_void_p,  # out (activation dtype [N])
+                ct.c_int64,  # K
+                ct.c_int64,  # N
+                ct.c_int64,  # bs_shift = log2(blocksize)
+                ct.c_int64,  # dtype_flag (0=fp32, 1=fp16, 2=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