Dequant-free Apple matmul; copy-free GQA decode attention

README.md

@@ -97,14 +97,14 @@ on this machine.
 
 | batch | decode | prefill |
 |---|---|---|
-| 1 | 28.6 tok/s | 166 tok/s |
-| 4 | 61.2 | 176 |
-| 8 | 74.4 | 175 |
-| 16 | 77.0 | 175 |
+| 1 | 31.3 tok/s | 169 tok/s |
+| 4 | 110.6 | 173 |
+| 8 | 144.1 | 168 |
+| 16 | 137.7 | 168 |
 
-Per matmul against dense fp16 of the same shape: 5.8–6.5x at batch 1, 1.3–1.7x at 16,
-0.7–1.0x from 32 up. At batch 1 the matmuls read 81–99 GB/s; a device-to-device copy
-reaches 98 GB/s.
+Per matmul against dense fp16 of the same shape: 4.7–5.6x at 1 token, 3.2–3.9x at 8,
+1.1–1.4x at 32, 0.7–0.9x from 64 up. From 1 to 4 tokens a matmul takes at most 1.25x
+the time to read its weights at the 98 GB/s a device-to-device copy reaches.
 
 ### Quality
 
@@ -161,17 +161,20 @@ While a forward runs at most 8 tokens under vLLM, a second stream loads each wei
 matrix into L2 two matmuls ahead of its use, paced by a progress flag the matmuls write.
 GPUs below compute capability 8.0 get a Triton kernel instead.
 
-**Apple** is three Metal kernels. Decode gives each output row several threads, each
-taking an interleaved share of the 256-column groups against a 4-entry value table; the
-partial sums meet in threadgroup memory. The 8-bit `lm_head` decodes the same way.
-Prefill feeds `simdgroup_half8x8` matrices from a weight slab rebuilt in threadgroup
-memory. None writes anything dense to device memory. The one-token decode step runs
-under `torch.compile`; the first step compiles for about 30 s.
+**Apple** multiplies without dequantising. Per 256-column group a row's output is
+`c0 * S0 + c1 * S1 + bias * sum(x)`, where `S0` and `S1` are the sums of the activations
+each bit-plane selects. For every 8 columns a threadgroup builds a 256-entry table of all
+subset sums of those 8 activations, 4 tokens per entry, and one byte of a bit-plane
+indexes it; the planes are the checkpoint's own, transposed at load. From 8 tokens on,
+the kernel is bound by random reads of those tables. The 8-bit `lm_head` has its own
+Metal kernel. One-token decode attention folds each KV head's query heads into one pass
+instead of repeating the KV cache, and the decode step runs under `torch.compile`; the
+first step compiles for about 30 s.
 
 ## Limits
 
 - **TP = 1.** No tensor parallelism.
-- On Apple, batches of 32 tokens and up run at 0.68–1.0x of a dense fp16 matmul. The win
+- On Apple, batches of 64 tokens and up run at 0.7–0.9x of a dense fp16 matmul. The win
   is at decode batch sizes.
 - At 2 bits the model invents confident detail about anything obscure, and past roughly
   250 tokens long open-ended answers drift. Use it for short exchanges and bounded tasks.

bpdq.py

@@ -1027,6 +1027,172 @@ def _matmul_mps8(x, codes, tbl, N, K):
     return y[:T]
 
 
+_METAL_LUT_SRC = r"""
+#include <metal_stdlib>
+using namespace metal;
+typedef {VT} vt;
+typedef {FT} ft;
+
+// 2-bit, dequant-free: y = sum_g c0 * S0 + c1 * S1 + bias * Sx, where S0 / S1 are the
+// activation sums selected by each bit-plane. Per 8 columns, a 256-entry table holds
+// every subset sum of the 8 activations for V tokens; one byte of a plane indexes it.
+// planes [K/32, 2, M]: bit i of a word is column 32w + i. coeffs [G, 3, M] = (c0, c1, bias).
+// A thread owns R rows; tables are double-buffered per KC columns; gridDim.y splits K.
+kernel void bpdq_lut(
+    device       float* part   [[buffer(0)]],    // [split, TT, M]
+    device const uint*  planes [[buffer(1)]],
+    device const float* coeffs [[buffer(2)]],
+    device const half*  x      [[buffer(3)]],    // [T, K]
+    constant     uint&  M      [[buffer(4)]],
+    constant     uint&  K      [[buffer(5)]],
+    constant     uint&  KS     [[buffer(6)]],    // columns per split
+    constant     uint&  T      [[buffer(7)]],    // real tokens; the rest read as zero
+    uint2 tid2 [[thread_position_in_threadgroup]],
+    uint2 gid  [[threadgroup_position_in_grid]])
+{{
+    const uint tid = tid2.x;
+    const uint TT = {TT}u, V = {V}u, NV = {TT}u / {V}u, TG = {TG}u, KC = {KC}u, R = {R}u;
+    const uint NTB = KC / 8u, TS = NV * NTB * 256u;
+    threadgroup vt tab2[2 * ({TT} / {V}) * ({KC} / 8) * 256];
+
+    uint mr[{R}];
+    for (uint r = 0; r < R; ++r) mr[r] = min(gid.x * TG * R + r * TG + tid, M - 1u);
+    const uint k0 = gid.y * KS, k1 = min(k0 + KS, K);
+    ft acc[{R}][{TT} / {V}];
+    vt s0[{R}][{TT} / {V}], s1[{R}][{TT} / {V}], sx[{TT} / {V}];
+    for (uint r = 0; r < R; ++r)
+        for (uint n = 0; n < NV; ++n) {{ acc[r][n] = ft(0); s0[r][n] = vt(0); s1[r][n] = vt(0); }}
+    for (uint n = 0; n < NV; ++n) sx[n] = vt(0);
+
+    // one thread per (token group, table, high nibble): 16 entries, one add each
+    auto build = [&](uint kc, threadgroup vt* tb) {{
+        for (uint e = tid; e < NV * NTB * 16u; e += TG) {{
+            const uint n = e / (NTB * 16u), b = (e / 16u) % NTB, hi = e % 16u;
+            vt xv[8];
+            for (uint u = 0; u < V; ++u) {{
+                const uint t = min(n * V + u, T - 1u);
+                const device half4* xp = (const device half4*)(x + t * K + kc + b * 8u);
+                const half4 a = n * V + u < T ? xp[0] : half4(0), c = n * V + u < T ? xp[1] : half4(0);
+                xv[0][u] = a[0]; xv[1][u] = a[1]; xv[2][u] = a[2]; xv[3][u] = a[3];
+                xv[4][u] = c[0]; xv[5][u] = c[1]; xv[6][u] = c[2]; xv[7][u] = c[3];
+            }}
+            vt cur = (hi & 1u ? xv[4] : vt(0)) + (hi & 2u ? xv[5] : vt(0))
+                   + (hi & 4u ? xv[6] : vt(0)) + (hi & 8u ? xv[7] : vt(0));
+            threadgroup vt* tp = tb + (n * NTB + b) * 256u + hi * 16u;
+            tp[0] = cur;
+            cur = cur + xv[0]; tp[1] = cur;
+            cur = cur + xv[1]; tp[3] = cur;
+            cur = cur - xv[0]; tp[2] = cur;
+            cur = cur + xv[2]; tp[6] = cur;
+            cur = cur + xv[0]; tp[7] = cur;
+            cur = cur - xv[1]; tp[5] = cur;
+            cur = cur - xv[0]; tp[4] = cur;
+            cur = cur + xv[3]; tp[12] = cur;
+            cur = cur + xv[0]; tp[13] = cur;
+            cur = cur + xv[1]; tp[15] = cur;
+            cur = cur - xv[0]; tp[14] = cur;
+            cur = cur - xv[2]; tp[10] = cur;
+            cur = cur + xv[0]; tp[11] = cur;
+            cur = cur - xv[1]; tp[9] = cur;
+            cur = cur - xv[0]; tp[8] = cur;
+        }}
+    }};
+
+    build(k0, tab2);
+    uint buf = 0;
+    for (uint kc = k0; kc < k1; kc += KC) {{
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        const threadgroup vt* tab = tab2 + buf * TS;
+        if (kc + KC < k1) build(kc + KC, tab2 + (buf ^ 1u) * TS);
+        buf ^= 1u;
+        #pragma clang loop unroll(full)
+        for (uint wq = 0; wq < KC / 32u; ++wq) {{
+            const uint w = kc / 32u + wq;
+            #pragma clang loop unroll(full)
+            for (uint j = 0; j < 4u; ++j)
+                for (uint n = 0; n < NV; ++n) sx[n] += tab[(n * NTB + wq * 4u + j) * 256u + 255u];
+            #pragma clang loop unroll(full)
+            for (uint r = 0; r < R; ++r) {{
+                const uchar4 p0 = as_type<uchar4>(planes[(w * 2u) * M + mr[r]]);
+                const uchar4 p1 = as_type<uchar4>(planes[(w * 2u + 1u) * M + mr[r]]);
+                #pragma clang loop unroll(full)
+                for (uint j = 0; j < 4u; ++j) {{
+                    #pragma clang loop unroll(full)
+                    for (uint n = 0; n < NV; ++n) {{
+                        const threadgroup vt* tp = tab + (n * NTB + wq * 4u + j) * 256u;
+                        s0[r][n] += tp[p0[j]];
+                        s1[r][n] += tp[p1[j]];
+                    }}
+                }}
+            }}
+        }}
+        if (((kc + KC) & 255u) == 0u || kc + KC >= k1) {{         // end of a 256-column group
+            const uint g = kc / 256u;
+            for (uint r = 0; r < R; ++r) {{
+                const uint m = mr[r];
+                const float c0 = coeffs[(g * 3u) * M + m], c1 = coeffs[(g * 3u + 1u) * M + m];
+                const float cb = coeffs[(g * 3u + 2u) * M + m];
+                for (uint n = 0; n < NV; ++n) {{
+                    acc[r][n] += c0 * ft(s0[r][n]) + c1 * ft(s1[r][n]) + cb * ft(sx[n]);
+                    s0[r][n] = vt(0); s1[r][n] = vt(0);
+                }}
+            }}
+            for (uint n = 0; n < NV; ++n) sx[n] = vt(0);
+        }}
+    }}
+    for (uint r = 0; r < R; ++r) {{
+        const uint row = gid.x * TG * R + r * TG + tid;
+        if (row < M)
+            for (uint t = 0; t < min(TT, T); ++t) part[(gid.y * T + t) * M + row] = acc[r][t / V][t % V];
+    }}
+}}
+"""
+
+
+@functools.lru_cache(maxsize=32)
+def _metal_lut_lib(tt, tg, kc, r):
+    v = min(tt, 4)
+    vt, ft = {2: ("half2", "float2"), 4: ("half4", "float4")}[v]
+    return torch.mps.compile_shader(_METAL_LUT_SRC.format(VT=vt, FT=ft, TT=tt, V=v, TG=tg, KC=kc, R=r))
+
+
+def mps_lut_supported(msbits, lut_bias, group_size, in_f, device):
+    return (torch.device(device).type == "mps" and msbits == 2 and lut_bias
+            and group_size == 256 and in_f % 256 == 0)
+
+
+def mps_lut_config(T, M, K):
+    """(padded tokens, threadgroup, KC, rows per thread, split) for T <= 8."""
+    tt = max(2, _next_pow2(T))
+    tg, r = (128, 4) if tt <= 4 else (256, 4)
+    groups, split = -(-M // (tg * r)), 1
+    while groups * split < 32 and split < 16 and K % (split * 512) == 0:
+        split *= 2
+    return tt, tg, 32, r, split
+
+
+def _matmul_mps_lut(x, planes, coeffs, M, K):
+    T = x.shape[0]
+    if T > 8:
+        return torch.cat([_matmul_mps_lut(c, planes, coeffs, M, K) for c in x.split(8)], 0)
+    tt, tg, kc, r, split = mps_lut_config(T, M, K)
+    if x.storage_offset() % 4:                           # half4 loads
+        x = x.clone()
+    part = torch.empty(split, T, M, device=x.device, dtype=torch.float32)
+    _metal_lut_lib(tt, tg, kc, r).bpdq_lut(
+        part, planes, coeffs, x, M, K, K // split, T,
+        threads=(-(-M // (tg * r)) * tg, split), group_size=(tg, 1))
+    return (part[0] if split == 1 else part.sum(0)).half()
+
+
+def _planes_dense(planes, coeffs, in_f):
+    """[K/32, 2, M] planes + [G, 3, M] coeffs -> dense fp32 [M, K]."""
+    sh = torch.arange(32, device=planes.device, dtype=torch.int32)
+    bits = ((planes[..., None] >> sh) & 1).permute(1, 2, 0, 3).reshape(2, planes.shape[2], -1)
+    c = coeffs.repeat_interleave(256, 0)[:in_f]               # [K, 3, M]
+    return (c[:, 2].T + bits[0] * c[:, 0].T + bits[1] * c[:, 1].T).float()
+
+
 def _select_expr(msbits, word, table):
     """Nested ternary over the register table: msbits selects, no conversions."""
     def rec(prefix, depth):
@@ -1138,6 +1304,10 @@ class PackedBPDQ:
         if mma_supported(msbits, group_size, in_f, out_f, device):
             codes, tbl = mma_repack(planes, coeffs, in_f, lut_bias, device)
             return cls(codes, tbl, dead, in_f, out_f, msbits, coeffs.shape[2], group_size)
+        if mps_lut_supported(msbits, lut_bias, group_size, in_f, device):
+            codes = planes.permute(2, 0, 1).contiguous().to(device)      # [K/32, 2, M]
+            rt = coeffs.to(device).float().permute(0, 2, 1).contiguous()
+            return cls(codes, rt, dead, in_f, out_f, msbits, coeffs.shape[2], group_size)
         cpw = 32 // msbits
         cwords = (in_f + cpw - 1) // cpw
         codes = torch.zeros(cwords, out_f, dtype=torch.int32, device=device)
@@ -1195,6 +1365,8 @@ class PackedBPDQ:
             return _mma_dense(self.codes, self.coeffs, self.in_f)
         if self.msbits == 8:
             return _gptq8_rows(self.codes, self.coeffs, 0, self.out_f).float()
+        if self.codes.dim() == 3:
+            return _planes_dense(self.codes, self.coeffs, self.in_f)
         return _dequant_rt(self.codes, self.coeffs, self.in_f, self.msbits,
                            self.n_coeff, self.group_size, self.cpw)
 
@@ -1210,8 +1382,8 @@ class PackedBPDQ:
 def packed_matmul(x, codes, coeffs, dead, in_f, msbits, n_coeff, group_size,
                   scratch=None):
     """x [T, in_f] -> [T, out_f] in x.dtype, straight off the packed weight.
-    4-D codes are the mma layout (CUDA); otherwise the row-major runtime layout."""
-    out_f = codes.shape[0] * 16 if codes.dim() == 4 else codes.shape[1]
+    4-D codes: mma layout (CUDA). 3-D: bit-planes (Apple, 2-bit). 2-D: row-major runtime layout."""
+    out_f = {4: codes.shape[0] * 16, 3: codes.shape[-1]}.get(codes.dim(), codes.shape[1])
     cpw = 32 // msbits
     wpg = group_size // cpw
     lut_bias = n_coeff > msbits
@@ -1224,6 +1396,8 @@ def packed_matmul(x, codes, coeffs, dead, in_f, msbits, n_coeff, group_size,
         y = torch.empty(x.shape[0], out_f, device=x.device, dtype=torch.float32)
         _matmul_cuda(x.contiguous(), codes, coeffs, y, out_f, in_f, msbits,
                      n_coeff, lut_bias, cpw, wpg)
+    elif dev == "mps" and codes.dim() == 3:
+        y = _matmul_mps_lut(x.contiguous().half(), codes, coeffs, out_f, in_f)
     elif dev == "mps" and msbits == 8:
         y = _matmul_mps8(x.contiguous().half(), codes, coeffs, out_f, in_f)
     elif dev == "mps":
@@ -1242,6 +1416,24 @@ _PROJ = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.
          "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj")
 
 
+def _gqa_decode_attention(module, query, key, value, attention_mask, dropout=0.0,
+                          scaling=None, is_causal=None, **kw):
+    """SDPA; one-token queries fold each KV head's query heads into the query axis
+    instead of repeating the KV cache."""
+    from transformers.integrations.sdpa_attention import sdpa_attention_forward
+    B, H, Tq, D = query.shape
+    Hkv = key.shape[1]
+    if Tq != 1 or H == Hkv:
+        return sdpa_attention_forward(module, query, key, value, attention_mask, dropout=dropout,
+                                      scaling=scaling, is_causal=is_causal, **kw)
+    mask = attention_mask
+    if mask is not None:
+        mask = mask[..., : key.shape[-2]].expand(B, 1, H // Hkv, key.shape[-2])
+    out = torch.nn.functional.scaled_dot_product_attention(
+        query.view(B, Hkv, H // Hkv, D), key, value, attn_mask=mask, scale=scaling)
+    return out.view(B, H, 1, D).transpose(1, 2).contiguous(), None
+
+
 @torch.library.custom_op("bpdq::matmul", mutates_args=())
 def _op_matmul(x: torch.Tensor, codes: torch.Tensor, coeffs: torch.Tensor,
                dead: Optional[torch.Tensor], in_f: int, msbits: int,
@@ -1252,7 +1444,7 @@ def _op_matmul(x: torch.Tensor, codes: torch.Tensor, coeffs: torch.Tensor,
 
 @_op_matmul.register_fake
 def _(x, codes, coeffs, dead, in_f, msbits, n_coeff, group_size):
-    out_f = codes.shape[0] * 16 if codes.dim() == 4 else codes.shape[1]
+    out_f = {4: codes.shape[0] * 16, 3: codes.shape[-1]}.get(codes.dim(), codes.shape[1])
     return x.new_empty((x.shape[0], out_f))
 
 
@@ -1372,6 +1564,11 @@ def load(path, device=None, dtype=torch.float16, verbose=False):
 
     model.to(device)
     if torch.device(device).type == "mps":
+        from transformers import AttentionInterface
+        from transformers.masking_utils import AttentionMaskInterface, sdpa_mask
+        AttentionInterface.register("bpdq_gqa", _gqa_decode_attention)
+        AttentionMaskInterface.register("bpdq_gqa", sdpa_mask)
+        model.set_attn_implementation("bpdq_gqa")
         model._bpdq_step = torch.compile(model.forward, dynamic=False)
     tok = AutoTokenizer.from_pretrained(path)
     return model, tok

demo.ipynb

@@ -13,7 +13,7 @@
     "| weights resident | **3.65 GiB** | 15.3 GiB |\n",
     "| 1 sequence, RTX 4090 | **322 tok/s** | 59 |\n",
     "| 16 sequences, RTX 4090 | **3452 tok/s** | 890 |\n",
-    "| Apple M4, 1 sequence | **29 tok/s** | *does not fit* |\n",
+    "| Apple M4, 1 sequence | **31 tok/s** | *does not fit* |\n",
     "\n",
     "The notebook works out which machine you are on and picks the kernels."
    ]