Feat/mps gemm 4bit bf16

benchmarks_wip/bench_gemm_baseline.py

@@ -7,6 +7,15 @@
 -> 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.
+
+bf16 question: bf16 used to have no native path at all (MPSMatrixMultiplication hard-asserts
+on it), so its "native" column WAS the fallback and the ratio was 1.00x by construction. It
+now runs the GEMM through MPSGraph, so the ratio is finally a real measurement. The shapes
+under `--cogkit` are the ones CogView4-6B QLoRA actually hits at 512x512 batch 1: hidden 4096,
+MLP 16384, M = 1024 image tokens (+ text).
+
+A `clone()` control is printed first. GPU timings on this machine are worthless under
+contention; if the control moves between runs, nothing else in the table is comparable.
 """
 
 import time
@@ -66,11 +75,67 @@ def linear_only():
     )
 
 
+def bench_bwd(M, N, K, dtype, quant_type="nf4", blocksize=64):
+    """gemm_4bit_backward: grad_A[M,K] = grad_output[M,N] . B_dq[N,K].
+
+    `fallback` is the composition MatMul4Bit.backward ran inline before this op existed --
+    native dequant (its own sync) then a torch matmul (torch's queue).
+    """
+    G = torch.randn(1, M, N, 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 native():
+        return torch.ops.bitsandbytes.gemm_4bit_backward(G, B_q, list(B.shape), qs.absmax, blocksize, quant_type)
+
+    def fallback():
+        B_dq = mps_ops._dequantize_4bit_impl(B_q, qs.absmax, blocksize, quant_type, list(B.shape), dtype)
+        return torch.matmul(G, B_dq)
+
+    t_nat, t_fb = timed(native), timed(fallback)
+    print(
+        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)"
+    )
+
+
+def control():
+    """Contention canary. A fixed 64 MB device clone, unrelated to any code under test: if
+    this number differs across runs, the machine was busy and the table is not comparable."""
+    x = torch.empty(16 * 1024 * 1024, dtype=torch.float32, device=DEV)
+    return timed(lambda: x.clone())
+
+
+# The shapes CogView4-6B QLoRA hits: (N, K) for qkv-proj, out-proj, mlp-in, mlp-out.
+COGKIT_SHAPES = ((12288, 4096), (4096, 4096), (16384, 4096), (4096, 16384))
+
+
 if __name__ == "__main__":
+    import sys
+
     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()
+    print(f"iters={ITERS} warmup={WARMUP} device={DEV} lib={native}")
+    print(f"control (64MB clone): {control():.3f}ms\n")
+
+    if "--cogkit" in sys.argv:
+        for dtype in (torch.bfloat16, torch.float16):
+            print(f"=== gemm_4bit, CogView4-6B shapes, {str(dtype).replace('torch.', '')} ===")
+            for M in (1024, 1280):
+                for N, K in COGKIT_SHAPES:
+                    bench(M, N, K, dtype)
+            print()
+        for dtype in (torch.bfloat16, torch.float16):
+            print(f"=== gemm_4bit_BACKWARD, CogView4-6B shapes, {str(dtype).replace('torch.', '')} ===")
+            for M in (1024, 1280):
+                for N, K in COGKIT_SHAPES:
+                    bench_bwd(M, N, K, dtype)
+            print()
+        # Second control read. If it has drifted from the first, the machine changed
+        # underneath the table and the table is not internally comparable.
+        print(f"control (64MB clone), after: {control():.3f}ms")
+    else:
+        for dtype in (torch.bfloat16, 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/_ops.py

@@ -236,6 +236,43 @@ def _(
     return out, absmax
 
 
+torch.library.define(
+    "bitsandbytes::gemm_4bit_backward",
+    "(Tensor grad_output, Tensor B, int[] shapeB, Tensor absmax, int blocksize, str quant_type, "
+    "Tensor? absmax_8bit=None, Tensor? absmax_code=None, Tensor? absmax_offset=None) -> Tensor",
+)
+
+
+@register_fake("bitsandbytes::gemm_4bit_backward")
+def _(
+    grad_output: torch.Tensor,
+    B: torch.Tensor,
+    shapeB: Sequence[int],
+    absmax: torch.Tensor,
+    blocksize: int,
+    quant_type: str,
+    absmax_8bit: Optional[torch.Tensor] = None,
+    absmax_code: Optional[torch.Tensor] = None,
+    absmax_offset: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # grad_A[..., K] = grad_output[..., N] @ B_dq[N, K]. Note the inner dim is N here, not K:
+    # this consumes B_dq in the orientation dequantize_4bit already produces, untransposed.
+    torch._check(len(shapeB) == 2, lambda: f"shapeB must be 2D [N, K], got {list(shapeB)}")
+    torch._check(
+        grad_output.shape[-1] == shapeB[0],
+        lambda: f"grad_output inner dim ({grad_output.shape[-1]}) must match shapeB[0] ({shapeB[0]})",
+    )
+    torch._check(
+        grad_output.dtype in (torch.float16, torch.bfloat16, torch.float32),
+        lambda: f"grad_output must be float16, bfloat16, or float32, got {grad_output.dtype}",
+    )
+    torch._check(blocksize in (32, 64, 128, 256, 512, 1024, 2048, 4096), lambda: f"invalid blocksize {blocksize}")
+    torch._check(quant_type in ("nf4", "fp4"), lambda: f"quant_type must be 'nf4' or 'fp4', got {quant_type!r}")
+    return torch.empty(
+        (*grad_output.shape[:-1], shapeB[1]), device=grad_output.device, dtype=grad_output.dtype
+    )
+
+
 torch.library.define(
     "bitsandbytes::gemm_4bit",
     "(Tensor A, Tensor B, int[] shapeB, Tensor absmax, int blocksize, str quant_type, "

bitsandbytes/autograd/_functions.py

@@ -380,8 +380,28 @@ def backward(ctx, grad_output):
         # if req_gradB: grad_B = torch.matmul(grad_output.t(), A)
         if req_gradA:
             # B in ctx.tensors is already in canonical [(N*K+1)//2, 1] form (normalized in forward).
-            # dequantize returns [N, K]; matmul(grad_output[M,N], [N,K]) = grad_A[M,K].
-            grad_A = torch.matmul(grad_output, F.dequantize_4bit(B, ctx.state).to(grad_output.dtype))
+            # dequantize returns [N, K]; matmul(grad_output[M,N], [N,K]) = grad_A[M,K]. The op's
+            # default kernel is exactly that composition; a backend may fuse it (MPS does, into
+            # one command buffer, which halves the cross-queue syncs this pays per layer).
+            state = ctx.state
+            if not state.nested:
+                grad_A = torch.ops.bitsandbytes.gemm_4bit_backward.default(
+                    grad_output, B, state.shape, state.absmax, state.blocksize, state.quant_type
+                )
+            elif state.state2.blocksize == 256:
+                grad_A = torch.ops.bitsandbytes.gemm_4bit_backward.default(
+                    grad_output,
+                    B,
+                    state.shape,
+                    state.state2.absmax,
+                    state.blocksize,
+                    state.quant_type,
+                    absmax_8bit=state.absmax,
+                    absmax_code=state.state2.code,
+                    absmax_offset=state.offset,
+                )
+            else:
+                raise NotImplementedError("nested quantization with state2.blocksize != 256 is not supported")
 
         return grad_A, grad_B, None, grad_bias, None
 

bitsandbytes/backends/default/ops.py

@@ -345,6 +345,36 @@ def _gemm_4bit_default_impl(
 register_kernel("bitsandbytes::gemm_4bit", "default")(_gemm_4bit_default_impl)
 
 
+def _gemm_4bit_backward_default_impl(
+    grad_output: torch.Tensor,
+    B: torch.Tensor,
+    shapeB: Sequence[int],
+    absmax: torch.Tensor,
+    blocksize: int,
+    quant_type: str,
+    absmax_8bit: Optional[torch.Tensor] = None,
+    absmax_code: Optional[torch.Tensor] = None,
+    absmax_offset: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """grad_A = grad_output @ dequantize_4bit(B) -- verbatim what MatMul4Bit.backward did inline.
+
+    Keeping it here means every device gets the op for free and the fused MPS kernel has an
+    oracle to be checked against.
+    """
+    if absmax_8bit is not None:
+        absmax = (
+            torch.ops.bitsandbytes.dequantize_blockwise.default(absmax_8bit, absmax, absmax_code, 256, torch.float32)
+            + absmax_offset
+        )
+    B_dq = torch.ops.bitsandbytes.dequantize_4bit.default(
+        B, absmax, blocksize, quant_type, shapeB, grad_output.dtype
+    )
+    return torch.matmul(grad_output, B_dq)
+
+
+register_kernel("bitsandbytes::gemm_4bit_backward", "default")(_gemm_4bit_backward_default_impl)
+
+
 MOMENTUM = 0
 RMSPROP = 1
 ADAGRAD = 2

bitsandbytes/backends/mps/ops.py

@@ -10,6 +10,7 @@
 
 from collections.abc import Sequence
 from math import prod
+import os
 import platform
 from typing import Optional
 
@@ -477,6 +478,28 @@ def _(
     out.copy_(result)
 
 
+def _bf16_gemm_enabled(backward: bool = False) -> bool:
+    """Whether the native bf16 gemm may run for this direction.
+
+    Two kinds of gate. The capability symbol is the correctness one: a dylib with
+    bnb_mps_gemm_4bit but built before bf16 support would take dtype_flag=2, fall through to
+    fp32 element size, and read a bf16 scratch as float -- silent garbage, no crash.
+
+    The env vars are the measurement ones, and they are per-direction on purpose. Forward and
+    backward are separate phases of a training step with separate costs, so a single switch
+    could only ever answer "both or neither"; these answer "what did the backward alone buy",
+    in one session against one binary:
+        BNB_MPS_DISABLE_BF16_GEMM=1      -- forces the fallback for BOTH directions
+        BNB_MPS_DISABLE_BF16_GEMM_BWD=1  -- forces it for the backward only
+    A speedup you cannot switch off is a speedup you cannot verify.
+    """
+    if os.environ.get("BNB_MPS_DISABLE_BF16_GEMM", "") not in ("", "0"):
+        return False
+    if backward and os.environ.get("BNB_MPS_DISABLE_BF16_GEMM_BWD", "") not in ("", "0"):
+        return False
+    return hasattr(_mps_native._lib, "bnb_mps_gemm_4bit_supports_bf16")
+
+
 def _gemm_4bit_native(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -496,10 +519,13 @@ def _gemm_4bit_native(
     per-call cross-queue sync is what dominates wall-clock at small/medium M -- see the
     Phase M2 finding in MPS_STATUS.md).
 
+    fp32/fp16 run the GEMM through MPSMatrixMultiplication; bf16 runs it through MPSGraph,
+    which has a bf16 matmul where MPSMatrixMultiplication hard-asserts on anything but
+    fp32/fp16/int8/int16. Same structure either way.
+
     `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.
+    A.dtype is fp32/fp16/bf16, K % 32 == 0, and power-of-two blocksize >= 32.
     """
     N, K = int(shapeB[0]), int(shapeB[-1])
     M = A.numel() // K
@@ -515,7 +541,7 @@ def _gemm_4bit_native(
         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]
+    dtype_flag = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2}[A.dtype]
     bs_shift = blocksize.bit_length() - 1
 
     torch.mps.synchronize()
@@ -536,6 +562,98 @@ def _gemm_4bit_native(
     return out.reshape(*A.shape[:-1], N)
 
 
+def _gemm_4bit_backward_native(
+    grad_output: torch.Tensor,
+    B: torch.Tensor,
+    shapeB: Sequence[int],
+    absmax: torch.Tensor,
+    blocksize: int,
+    quant_type: str,
+) -> torch.Tensor:
+    """grad_A[M, K] = grad_output[M, N] . B_dq[N, K], one command buffer / one commit / one wait.
+
+    Structurally identical to _gemm_4bit_native, and deliberately so: the same chunked dequant
+    fills the same private scratch, and only the matmul orientation differs (no transpose here --
+    dequantize_4bit already emits [N, K], which is the orientation grad_A wants).
+
+    The win is the sync, not the arithmetic. The composition this replaces runs a native dequant
+    on our queue (wait), hands the result to torch, and runs a matmul on torch's queue (wait) --
+    twice the cross-queue round trip, per Linear4bit, per step.
+
+    `absmax` must already be the plain per-block fp32 scale. Preconditions are the caller's.
+    """
+    N, K = int(shapeB[0]), int(shapeB[-1])
+    M = grad_output.numel() // N
+
+    B_flat = B if B.dtype == torch.uint8 else B.view(torch.uint8)
+    B_flat = _ensure_native_buffer(B_flat.reshape(-1))
+    G_flat = _ensure_native_buffer(grad_output.reshape(-1))
+    code_f = _ensure_native_buffer(_get_4bit_code(quant_type, grad_output.device).to(torch.float32))
+    absmax_f = _ensure_native_buffer(absmax.to(torch.float32))
+
+    out = torch.empty(M * K, dtype=grad_output.dtype, device=grad_output.device)
+    dtype_flag = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2}[grad_output.dtype]
+    bs_shift = blocksize.bit_length() - 1
+
+    torch.mps.synchronize()
+    _mps_native.bnb_mps_gemm_4bit_bwd(
+        code_f.data_ptr(),
+        B_flat.data_ptr(),
+        absmax_f.data_ptr(),
+        G_flat.data_ptr(),
+        out.data_ptr(),
+        M,
+        K,
+        N,
+        bs_shift,
+        dtype_flag,
+    )
+
+    return out.reshape(*grad_output.shape[:-1], K)
+
+
+@register_kernel("bitsandbytes::gemm_4bit_backward", "mps")
+def _(
+    grad_output: torch.Tensor,
+    B: torch.Tensor,
+    shapeB: Sequence[int],
+    absmax: torch.Tensor,
+    blocksize: int,
+    quant_type: str,
+    absmax_8bit: Optional[torch.Tensor] = None,
+    absmax_code: Optional[torch.Tensor] = None,
+    absmax_offset: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    N, K = int(shapeB[0]), int(shapeB[-1])
+
+    if absmax_8bit is not None:
+        absmax = (
+            torch.ops.bitsandbytes.dequantize_blockwise.default(absmax_8bit, absmax, absmax_code, 256, torch.float32)
+            + absmax_offset
+        )
+
+    # Same guards as the forward, minus bias (there is none) and with the inner dim being N.
+    # K % 32 == 0 is still what the chunked dequant needs; it fills B_dq[N, K] either way.
+    if (
+        _native_available()
+        and hasattr(_mps_native._lib, "bnb_mps_gemm_4bit_bwd")  # dylibs before the fused backward
+        and (
+            grad_output.dtype in (torch.float32, torch.float16)
+            or (grad_output.dtype == torch.bfloat16 and _bf16_gemm_enabled(backward=True))
+        )
+        and len(shapeB) == 2
+        and grad_output.shape[-1] == N
+        and K % 32 == 0
+        and blocksize >= 32
+        and (blocksize & (blocksize - 1)) == 0
+        and B.numel() * B.element_size() == (N * K) // 2
+    ):
+        return _gemm_4bit_backward_native(grad_output, B, shapeB, absmax, blocksize, quant_type)
+
+    B_dq = _dequantize_4bit_impl(B, absmax, blocksize, quant_type, shapeB, grad_output.dtype)
+    return torch.matmul(grad_output, B_dq)
+
+
 @register_kernel("bitsandbytes::gemm_4bit", "mps")
 def _(
     A: torch.Tensor,
@@ -561,16 +679,19 @@ 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.
+    # Native Metal path (dequant -> scratch -> MPSMatrixMultiplication (fp32/fp16) or MPSGraph
+    # (bf16) -> bias, one command buffer / one sync). Guards: 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. bf16 additionally requires the capability marker, since a dylib with
+    # bnb_mps_gemm_4bit but no bf16 support would misread the scratch as fp32.
     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 (
+            A.dtype in (torch.float32, torch.float16)
+            or (A.dtype == torch.bfloat16 and _bf16_gemm_enabled())
+        )
         and len(shapeB) == 2
         and shapeB[-1] == K
         and K % 32 == 0