Speed up Maple decode with and without the native extension

maple.py

@@ -182,6 +182,7 @@ def _add_rms_norm_ok(dim, dtype, w, eps):
     )
 
 
+@mx.compile
 def _aggregate_add_rms_norm(h, r, scores, w, eps):
     return _make_add_rms_norm_kernel(eps, aggregate=True)(
         inputs=[h.reshape(-1), r.reshape(-1), scores.reshape(-1), w],
@@ -314,7 +315,7 @@ def _make_qk_norm_rope_kernel():
             }
             return;
         }
-        const device T_* wh = w + head * HEAD_DIM;
+        const device T_* wh = head < (uint)NQ ? qw : kw;
 
         float ss = 0.0f;
         for (int i = 0; i < per_lane; ++i) {
@@ -332,7 +333,8 @@ def _make_qk_norm_rope_kernel():
             if (ROPE_DIM > 0 && j < ROPE_DIM) {
                 constexpr int rhalf = ROPE_DIM > 0 ? ROPE_DIM / 2 : 1;
                 int p = j < rhalf ? j : j - rhalf;
-                float theta = pos * inv_freq[p];
+                float inv = metal::exp2(-(float(p) / float(rhalf)) * pos_eps[2]);
+                float theta = pos * inv;
                 float c = metal::fast::cos(theta);
                 float s = metal::fast::sin(theta);
                 int j2 = j < rhalf ? j + rhalf : j - rhalf;
@@ -344,7 +346,7 @@ def _make_qk_norm_rope_kernel():
     """
     return mx.fast.metal_kernel(
         name="maple_qk_norm_rope",
-        input_names=["x", "w", "inv_freq", "pos_eps"],
+        input_names=["x", "qw", "kw", "pos_eps"],
         output_names=["out"],
         source=source,
     )
@@ -378,8 +380,8 @@ _row_qmv_kernel = mx.fast.metal_kernel(
     uint expert=GATHER?ids[slot]:0;
     uint row0=threadgroup_position_in_grid.x*8+sg*4;
     const device uchar* w=(const device uchar*)weight+expert*N*(K/4);
-    const device bfloat16_t* sc=scales+expert*N;
-    const device bfloat16_t* bi=biases+expert*N;
+    const device bfloat16_t* sc=scales+expert*N*(GROUPED?K/128:1);
+    const device bfloat16_t* bi=biases+expert*N*(GROUPED?K/128:1);
     const device bfloat16_t* xv=x+(SPLIT?slot*K:0);
     float result[4]={0};
     for(int k=0;k<K;k+=512) {
@@ -392,7 +394,8 @@ _row_qmv_kernel = mx.fast.metal_kernel(
         }
         for(int r=0;r<4;++r) {
             uint row=row0+r;
-            result[r]+=maple_qdot(w+row*(K/4)+k/4+lid*4,xt,(float)sc[row],(float)bi[row],sum);
+            uint si=GROUPED?row*(K/128)+(k+lid*16)/128:row;
+            result[r]+=maple_qdot(w+row*(K/4)+k/4+lid*4,xt,(float)sc[si],(float)bi[si],sum);
         }
     }
     for(int r=0;r<4;++r) {
@@ -415,8 +418,8 @@ _row_up_kernel = mx.fast.metal_kernel(
     uint expert=ids[slot];
     uint first=threadgroup_position_in_grid.x*(SGS*R)+sg*R;
     const device uchar* w=(const device uchar*)weight+expert*1024*512;
-    const device bfloat16_t* sc=scales+expert*1024;
-    const device bfloat16_t* bi=biases+expert*1024;
+    const device bfloat16_t* sc=scales+expert*1024*(GROUPED?16:1);
+    const device bfloat16_t* bi=biases+expert*1024*(GROUPED?16:1);
     float up[R]={0},gate[R]={0};
     for(int k=0;k<2048;k+=512) {
         const device bfloat16_t* xx=x+k+lid*16;
@@ -430,9 +433,9 @@ _row_up_kernel = mx.fast.metal_kernel(
         for(int r=0;r<R;++r) {
             uint row=first+r;
             uint wi=row*512+k/4+lid*4;
-            uint si=row;
+            uint si=GROUPED?row*16+(k+lid*16)/128:row;
             up[r]+=maple_qdot(w+wi,xt,(float)sc[si],(float)bi[si],sum);
-            gate[r]+=maple_qdot(w+wi+512*512,xt,(float)sc[si+512],(float)bi[si+512],sum);
+            gate[r]+=maple_qdot(w+wi+512*512,xt,(float)sc[si+512*(GROUPED?16:1)],(float)bi[si+512*(GROUPED?16:1)],sum);
         }
     }
     for(int r=0;r<R;++r) {
@@ -453,10 +456,9 @@ _row_up_kernel = mx.fast.metal_kernel(
 
 
 def _row_quantized_metadata(p):
-    """Compact repeated affine metadata, without changing quantized weights."""
+    """Use compact row metadata with snapshots, live group metadata without."""
     if (
         mx.__version__ != "0.32.0"
-        or not hasattr(_maple_native, "ArraySnapshot")
         or type(p) not in (nn.QuantizedLinear, QuantizedSwitchLinear)
         or p.bits != 2
         or p.group_size != 128
@@ -467,16 +469,16 @@ def _row_quantized_metadata(p):
         or p.biases.dtype != mx.bfloat16
     ):
         return None
-    sources = (p.weight, p.scales, p.biases)
+    native = hasattr(_maple_native, "ArraySnapshot")
+    sources = (p.weight, p.scales, p.biases) if native else None
     state = p.get("_maple_row_state")
-    if state is not None and state["sources"].matches(sources):
+    if (
+        native
+        and state is not None
+        and state.get("sources") is not None
+        and state["sources"].matches(sources)
+    ):
         return state if state["supported"] else None
-    state = {
-        "sources": _maple_native.ArraySnapshot(sources),
-        "supported": False,
-        "qmv_ok": None,
-    }
-    p["_maple_row_state"] = state
     k, n = p.scales.shape[-1] * 128, p.weight.shape[-2]
     if (
         k not in (512, 2048)
@@ -486,6 +488,23 @@ def _row_quantized_metadata(p):
         or p.biases.shape != p.scales.shape
     ):
         return None
+    if not native:
+        if (
+            state is None
+            or state.get("sources") is not None
+            or (state.get("k"), state.get("n")) != (k, n)
+        ):
+            state = dict(k=k, n=n, qmv_ok=None, zero=mx.array([0], mx.uint32), sources=None)
+            p["_maple_row_state"] = state
+        # Live arrays need no snapshot: no derived weight values are cached.
+        state.update(scales=p.scales, biases=p.biases)
+        return state
+    state = {
+        "sources": _maple_native.ArraySnapshot(sources),
+        "supported": False,
+        "qmv_ok": None,
+    }
+    p["_maple_row_state"] = state
     sc, bi = p.scales[..., :1], p.biases[..., :1]
     same = mx.all(p.scales.view(mx.uint16) == sc.view(mx.uint16)) & mx.all(
         p.biases.view(mx.uint16) == bi.view(mx.uint16)
@@ -507,7 +526,10 @@ def _row_qmv_arrays(x, weight, scales, biases, ids, gather, split):
     shape = (8, 1, n) if gather else (*x.shape[:-1], n)
     return _row_qmv_kernel(
         inputs=[x, weight, scales, biases, ids],
-        template=[("K", k), ("N", n), ("GATHER", gather), ("SPLIT", split)],
+        template=[
+            ("K", k), ("N", n), ("GATHER", gather), ("SPLIT", split),
+            ("GROUPED", scales.ndim == weight.ndim),
+        ],
         grid=(n // 4 * 32, 8 if gather else 1, 1),
         threadgroup=(64, 1, 1),
         output_shapes=[shape],
@@ -520,6 +542,7 @@ def _row_experts_arrays(x, indices, uw, us, ub, dw, ds, db):
     ids = indices.reshape(-1).astype(mx.uint32)
     y = _row_up_kernel(
         inputs=[x, uw, us, ub, ids],
+        template=[("GROUPED", us.ndim == uw.ndim)],
         grid=(4096, 8, 1),
         threadgroup=(64, 1, 1),
         output_shapes=[(8, 1, 512)],
@@ -630,17 +653,18 @@ class MapleAttention(nn.Module):
 
     def _qk_fused(self, qkv, offset):
         """Normalize/rotate Q and K; copy any trailing V heads unchanged."""
-        self._ensure_qk_state()
-
         # cache.offset is a Python int for a plain cache but an mx.array for
-        # the batched caches; coerce so the pos/eps pair is always uniform.
-        pos_eps = mx.array([float(offset), self._eps], dtype=mx.float32)
+        # the batched caches; coerce before constructing the kernel input.
+        pos_eps = mx.array(
+            [float(offset), self._eps, math.log2(self._rope_base)], dtype=mx.float32
+        )
         return _qk_norm_rope_kernel(
-            inputs=[qkv, self._qk_w, self._inv_freq, pos_eps],
+            inputs=[qkv, self.q_norm.weight, self.k_norm.weight, pos_eps],
             template=[
                 ("T_", qkv.dtype),
                 ("HEAD_DIM", self.head_dim),
                 ("ROPE_DIM", self.rope.dims if self.use_rope else 0),
+                ("NQ", self.num_attention_heads),
                 ("NQK", self.num_attention_heads + self.num_key_value_heads),
             ],
             grid=(32, qkv.shape[0], 1),
@@ -1099,15 +1123,18 @@ _router_select_kernel = mx.fast.metal_kernel(
 
 
 @mx.compile
-def _norm_router_arrays(h, r, norm_weight, router_weight, eps):
+def _norm_router_arrays(h, r, norm_weight, router_weight, eps, portable=False):
     h, hn = _add_rms_norm(h, r, norm_weight, eps)
-    logits = _router_gemv_kernel(
-        inputs=[hn, router_weight],
-        grid=(2048, 1, 1),
-        threadgroup=(128, 1, 1),
-        output_shapes=[(1, 1, 256)],
-        output_dtypes=[mx.float32],
-    )[0]
+    if portable:
+        logits = hn.astype(mx.float32) @ router_weight.astype(mx.float32).T
+    else:
+        logits = _router_gemv_kernel(
+            inputs=[hn, router_weight],
+            grid=(2048, 1, 1),
+            threadgroup=(128, 1, 1),
+            output_shapes=[(1, 1, 256)],
+            output_dtypes=[mx.float32],
+        )[0]
     indices, scores = _router_select_kernel(
         inputs=[logits],
         grid=(128, 1, 1),
@@ -1125,7 +1152,7 @@ def _decode_norm_router(h, r, norm, gate):
         return hh, hn, indices, scores
 
     if (
-        not hasattr(_maple_native, "ArraySnapshot")
+        mx.__version__ != "0.32.0"
         or type(norm) is not MapleRMSNorm
         or type(gate) is not MapleGate
         or h.shape != (1, 1, 2048)
@@ -1141,15 +1168,17 @@ def _decode_norm_router(h, r, norm, gate):
         or norm.eps <= 0
     ):
         return fallback()
+    native = hasattr(_maple_native, "ArraySnapshot")
     sources = (norm.weight, gate.weight)
     state = gate.get("_maple_norm_router_state")
     if (
         state is None
         or state["eps"] != norm.eps
-        or not state["sources"].matches(sources)
+        or native != (state["sources"] is not None)
+        or (native and not state["sources"].matches(sources))
     ):
         state = {
-            "sources": _maple_native.ArraySnapshot(sources),
+            "sources": _maple_native.ArraySnapshot(sources) if native else None,
             "eps": norm.eps,
             "ok": None,
         }
@@ -1158,7 +1187,7 @@ def _decode_norm_router(h, r, norm, gate):
         return fallback()
 
     def fast():
-        return _norm_router_arrays(h, r, norm.weight, gate.weight, norm.eps)
+        return _norm_router_arrays(h, r, norm.weight, gate.weight, norm.eps, not native)
 
     if state["ok"] is None:
         state["ok"] = False
@@ -1349,6 +1378,102 @@ def _flash_scatter(top, token_map, logits, force_ids, force_logits, vocab_size):
     )[0]
 
 
+# Exact bf16 top-k set: two radix histograms, with MLX's stable tie order.
+# One threadgroup owns selection and compaction; no cross-group synchronization.
+_head_select_kernel = mx.fast.metal_kernel(
+    name="maple_head_select",
+    input_names=["x"],
+    output_names=["out"],
+    header=r"""
+    inline uint head_key(bfloat16_t x) {
+        uint raw=as_type<ushort>(x), magnitude=raw&0x7fffu;
+        if(magnitude>0x7f80u) return 0xffffu;
+        if(magnitude<0x80u) return 0x8000u;
+        return (raw&0x8000u)?((~raw)&0xffffu):(raw^0x8000u);
+    }
+    inline uint head_prefix(uint value,uint tid,threadgroup uint* groups) {
+        uint lane=tid%32u,sg=tid/32u;
+        uint prefix=simd_prefix_exclusive_sum(value);
+        if(lane==31u) groups[sg]=prefix+value;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for(uint w=0;w<sg;++w) prefix+=groups[w];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        return prefix;
+    }
+    """,
+    source=r"""
+    uint tid=thread_position_in_threadgroup.x;
+    constexpr uint VALUES=(N+1023u)/1024u;
+    uint keys[VALUES];
+    threadgroup atomic_uint hist[256];
+    threadgroup uint groups[32],cut[5];
+    if(tid<256u) atomic_store_explicit(hist+tid,0u,memory_order_relaxed);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for(uint j=0;j<VALUES;++j) {
+        uint i=tid*VALUES+j;
+        keys[j]=i<N?head_key(x[i]):0u;
+        uint digit=keys[j]>>8;
+        uint matches=uint((simd_vote::vote_t)simd_ballot(i<N));
+        for(uint b=0;b<8;++b) {
+            bool bit=(digit&(1u<<b))!=0u;
+            uint votes=uint((simd_vote::vote_t)simd_ballot(bit));
+            matches &= bit?votes:~votes;
+        }
+        if(i<N && tid%32u==ctz(matches))
+            atomic_fetch_add_explicit(hist+digit,popcount(matches),memory_order_relaxed);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    uint count=tid<256u?atomic_load_explicit(hist+tid,memory_order_relaxed):0u;
+    uint prefix=head_prefix(count,tid,groups);
+    if(tid<256u && prefix<=N-K && prefix+count>N-K) {cut[0]=tid;cut[1]=prefix;}
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if(tid<256u) atomic_store_explicit(hist+tid,0u,memory_order_relaxed);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for(uint j=0;j<VALUES;++j) {
+        if(tid*VALUES+j<N && (keys[j]>>8)==cut[0])
+            atomic_fetch_add_explicit(hist+(keys[j]&255u),1u,memory_order_relaxed);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    count=tid<256u?atomic_load_explicit(hist+tid,memory_order_relaxed):0u;
+    prefix=head_prefix(count,tid,groups);
+    uint rank=N-K-cut[1];
+    if(tid<256u && prefix<=rank && prefix+count>rank) {
+        cut[2]=(cut[0]<<8)|tid;
+        cut[3]=N-(cut[1]+prefix+count);
+        cut[4]=count;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    uint greater=0u,equal=0u;
+    for(uint j=0;j<VALUES;++j) if(tid*VALUES+j<N) {
+        greater+=keys[j]>cut[2]; equal+=keys[j]==cut[2];
+    }
+    uint greater_prefix=head_prefix(greater,tid,groups);
+    uint equal_prefix=head_prefix(equal,tid,groups);
+    uint skip=cut[4]-(K-cut[3]);
+    for(uint j=0;j<VALUES;++j) {
+        uint i=tid*VALUES+j;
+        if(i>=N) break;
+        if(keys[j]>cut[2]) out[greater_prefix++]=i;
+        else if(keys[j]==cut[2]) {
+            if(equal_prefix>=skip) out[cut[3]+equal_prefix-skip]=i;
+            equal_prefix++;
+        }
+    }
+    """,
+)
+
+
+def _head_select(x, k):
+    return _head_select_kernel(
+        inputs=[x],
+        template=[("N", x.size), ("K", k)],
+        grid=(1024, 1, 1),
+        threadgroup=(1024, 1, 1),
+        output_shapes=[(1, k)],
+        output_dtypes=[mx.uint32],
+    )[0]
+
+
 class FlashHead(nn.Module):
     """Two-phase approximate lm_head for single-stream decode.
 
@@ -1406,9 +1531,8 @@ class FlashHead(nn.Module):
             ),
         }
         self._force_ids = mx.array(meta.get("force_tokens", []), dtype=mx.int32)
-        self._force_rows = None
-        self._force_sources = None
         self._scatter = None
+        self._select = None
 
     def _scatter_reference(self, top, logits, force_logits, vocab_size):
         oids = self.token_map[top[0]].reshape(-1)
@@ -1421,9 +1545,26 @@ class FlashHead(nn.Module):
 
     def __call__(self, h: mx.array, lm_head: nn.Module) -> mx.array:
         hv = h[:, -1, :]
-        top = mx.argpartition(self.centroids(hv), kth=-self.n_probes, axis=-1)[
+        scores = self.centroids(hv)
+        reference = lambda: mx.argpartition(scores, kth=-self.n_probes, axis=-1)[
             ..., -self.n_probes :
-        ]  # [1, n_probes]
+        ]
+        if (
+            mx.__version__ == "0.32.0"
+            and mx.default_device() == mx.gpu
+            and scores.dtype == mx.bfloat16
+            and scores.ndim == 2
+            and scores.shape[0] == 1
+            and 0 < self.n_probes <= scores.size <= 8192
+        ):
+            if self._select is None:
+                self._select = _matches(
+                    lambda: (mx.sort(_head_select(scores, self.n_probes)),),
+                    lambda: (mx.sort(reference()),),
+                )
+            top = _head_select(scores, self.n_probes) if self._select else reference()
+        else:
+            top = reference()
 
         logits = mx.gather_qmm(
             hv.reshape(1, 1, 1, 1, -1),
@@ -1438,27 +1579,17 @@ class FlashHead(nn.Module):
 
         force_logits = logits[:0]
         if self._force_ids.size:
-            sources = (lm_head.weight, lm_head.scales, lm_head.biases, self._force_ids)
-            if self._force_sources is None or not self._force_sources.matches(sources):
-                self._force_rows = (
-                    lm_head.weight[self._force_ids],
-                    lm_head.scales[self._force_ids],
-                    lm_head.biases[self._force_ids],
-                )
-                if hasattr(_maple_native, "ArraySnapshot"):
-                    mx.eval(*self._force_rows)
-                    self._force_sources = _maple_native.ArraySnapshot(sources)
-            fw, fs, fb = self._force_rows
-            force_logits = mx.quantized_matmul(
-                hv,
-                fw,
-                scales=fs,
-                biases=fb,
+            force_logits = mx.gather_qmm(
+                hv.reshape(1, 1, -1),
+                lm_head.weight.reshape(lm_head.weight.shape[0], 1, -1),
+                lm_head.scales.reshape(lm_head.weight.shape[0], 1, -1),
+                lm_head.biases.reshape(lm_head.weight.shape[0], 1, -1),
+                rhs_indices=self._force_ids,
                 transpose=True,
                 group_size=lm_head.group_size,
                 bits=lm_head.bits,
                 mode=getattr(lm_head, "mode", "affine"),
-            )[0]
+            ).reshape(-1)
         vocab_size = lm_head.weight.shape[0]
         if mx.__version__ == "0.32.0" and self._force_ids.size <= 8:
             fast = lambda: _flash_scatter(