perf(metal): x-tiled BF16 prefill matmul — 5.6x TTFT (was 4.1x)

The plain batched BF16 matmul was bound by strided token-major x reads
(per 2-byte weight it streamed T×4 bytes of activation, ~16× the weight
traffic). `matmul_bf16_tiled` stages the T token columns into threadgroup
memory once per in_dim chunk (CHUNK=256), shared across all nsg=8
simdgroups in the threadgroup, so x global traffic drops ~8× and the
matmul becomes weight-bound. All threads cooperate on the load + barriers;
only o < out_dim writes.

508-token prompt (prefill-check --gpu), exact (max|Δlogit| = 0.000000):
  Qwen3-4B BF16  batched 4216 -> 3055 ms  (4.07x -> 5.59x vs 17.1 s)

forward.rs routes BF16 to the tiled kernel; the plain `matmul_bf16`
remains for the ullm-metal unit test. Next: x-tile the k-quant kernels
(same shared-x win) and tile weights too (full mul_mm).

Author: nobottomline

crates/ullm-metal/src/forward.rs

@@ -179,7 +179,7 @@ pub struct GpuForward {
     p_gelu_mul: ComputePipelineState,
     p_add: ComputePipelineState,
     p_matvec_mlx4: ComputePipelineState,
-    p_matmul_bf16: ComputePipelineState,
+    p_matmul_bf16_tiled: ComputePipelineState,
     p_matmul_mlx4: ComputePipelineState,
     p_matmul_q4k: ComputePipelineState,
     p_matmul_q6k: ComputePipelineState,
@@ -323,7 +323,7 @@ impl GpuForward {
             p_gelu_mul: pso("gelu_mul")?,
             p_add: pso("add_inplace")?,
             p_matvec_mlx4: pso("matvec_mlx4")?,
-            p_matmul_bf16: pso("matmul_bf16")?,
+            p_matmul_bf16_tiled: pso("matmul_bf16_tiled")?,
             p_matmul_mlx4: pso("matmul_mlx4")?,
             p_matmul_q4k: pso("matmul_q4k")?,
             p_matmul_q6k: pso("matmul_q6k")?,
@@ -700,7 +700,7 @@ impl GpuForward {
         }
         // BF16 and the k-quants share the (w, x, y, in, out, n_cols) signature.
         let pso = match w.dtype {
-            DType::BF16 => &self.p_matmul_bf16,
+            DType::BF16 => &self.p_matmul_bf16_tiled,
             DType::Q4K => &self.p_matmul_q4k,
             DType::Q6K => &self.p_matmul_q6k,
             _ => return false,

crates/ullm-metal/src/shader.metal

@@ -510,6 +510,57 @@ kernel void matmul_bf16(
     }
 }
 
+// Batched BF16 matmul, x-tiled. Same result as `matmul_bf16` but the T token
+// columns of the activation are staged into threadgroup memory ONCE per in_dim
+// chunk and shared by all `nsg` simdgroups (each computing a different output
+// row), so the strided token-major x reads — the bottleneck of the plain kernel
+// — drop ~nsg-fold and the matmul becomes weight-bound. All threads cooperate on
+// the load + barriers (no early-out), only o < out_dim writes.
+kernel void matmul_bf16_tiled(
+    device const ushort* w       [[buffer(0)]],
+    device const float*  x       [[buffer(1)]],
+    device float*        y       [[buffer(2)]],
+    constant uint&       in_dim  [[buffer(3)]],
+    constant uint&       out_dim [[buffer(4)]],
+    constant uint&       n_cols  [[buffer(5)]],
+    uint2  tgpig [[threadgroup_position_in_grid]],
+    ushort tiisg [[thread_index_in_simdgroup]],
+    ushort sgitg [[simdgroup_index_in_threadgroup]],
+    ushort nsg   [[simdgroups_per_threadgroup]])
+{
+    const uint T = 8;
+    const uint CHUNK = 256;
+    threadgroup float xt[T * CHUNK]; // T token columns x CHUNK in-dim
+    uint o = (uint)tgpig.x * nsg + sgitg;
+    uint col0 = (uint)tgpig.y * T;
+    uint tid = (uint)sgitg * 32u + tiisg;
+    uint nthreads = (uint)nsg * 32u;
+    device const ushort* wrow = w + (uint)o * in_dim;
+    float acc[T];
+    for (uint t = 0; t < T; ++t) acc[t] = 0.f;
+
+    for (uint c0 = 0; c0 < in_dim; c0 += CHUNK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup); // prev chunk done reading xt
+        for (uint idx = tid; idx < T * CHUNK; idx += nthreads) {
+            uint t = idx / CHUNK, j = idx % CHUNK, i = c0 + j, s = col0 + t;
+            xt[idx] = (i < in_dim && s < n_cols) ? x[s * in_dim + i] : 0.f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup); // xt filled
+        if (o < out_dim) {
+            uint jmax = min(CHUNK, in_dim - c0);
+            for (uint j = tiisg; j < jmax; j += 32u) {
+                float wv = as_type<float>((uint)wrow[c0 + j] << 16);
+                for (uint t = 0; t < T; ++t) acc[t] += wv * xt[t * CHUNK + j];
+            }
+        }
+    }
+    for (uint t = 0; t < T; ++t) {
+        float r = simd_sum(acc[t]);
+        uint s = col0 + t;
+        if (tiisg == 0 && o < out_dim && s < n_cols) y[s * out_dim + o] = r;
+    }
+}
+
 // Batched MLX 4-bit matmul (prompt prefill): packed-u32 weights dequantized in
 // the kernel (q*scale+bias, 8 nibbles/word LSB-first), W[out,in] x X[S,in] ->
 // Y[S,out]. One simdgroup per output row computes a tile of T columns, reading