| """ |
| TITAN kernel dispatch layer. |
| |
| Every primitive has: |
| * a FAST path (chosen by measurement — see reports/01_pmlf_audit.md) |
| * a REFERENCE path (simple, obviously-correct) |
| Toggle with TITAN_KERNELS=ref|fast (env) or kernels.set_mode(). |
| |
| Measured decisions (Apple M5 Max, MLX 0.32.0, bf16): |
| rmsnorm : mx.fast.rms_norm 0.276ms vs PMLF fuse 0.670ms -> NATIVE (2.43x) |
| swiglu : PMLF fuse.swiglu 2.212ms vs native 2.110ms -> PMLF (tied, modular) |
| attention : mx.fast.sdpa 1.364ms vs manual 2.420ms -> NATIVE (PMLF has none) |
| rope : mx.fast.rope 0.278ms vs manual 0.507ms -> NATIVE (PMLF has none) |
| packing : PMLF pack 100% vs 26.5% efficiency -> PMLF (3.77x FLOP win) |
| """ |
| import os, sys |
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
| _PMLF_PATH = '/Users/jonah/dataset_prep/kernels' |
| if _PMLF_PATH not in sys.path: |
| sys.path.insert(0, _PMLF_PATH) |
|
|
| |
| PMLF_AVAILABLE = False |
| PMLF_VERSION = 'unavailable' |
| try: |
| from pmlf.fuse import swiglu as _pmlf_swiglu |
| from pmlf.pack import plan_packing as _pmlf_plan_packing |
| PMLF_AVAILABLE = True |
| PMLF_VERSION = 'local-2026-08-06+sha:3597f8a2' |
| except Exception as _e: |
| _pmlf_swiglu = None |
| _pmlf_plan_packing = None |
|
|
| MODE = os.environ.get('TITAN_KERNELS', 'fast') |
|
|
|
|
| def set_mode(m): |
| global MODE |
| assert m in ('fast', 'ref') |
| MODE = m |
|
|
|
|
| def used_kernels(): |
| return { |
| 'mode': MODE, |
| 'pmlf_available': PMLF_AVAILABLE, |
| 'pmlf_version': PMLF_VERSION, |
| 'rmsnorm': 'mx.fast.rms_norm' if MODE == 'fast' else 'reference', |
| 'swiglu': ('pmlf.fuse.swiglu' if (MODE == 'fast' and PMLF_AVAILABLE) |
| else 'reference'), |
| 'attention': 'mx.fast.scaled_dot_product_attention' if MODE == 'fast' else 'reference', |
| 'rope': 'mx.fast.rope' if MODE == 'fast' else 'reference', |
| 'packing': 'pmlf.pack.plan_packing (FFD)' if PMLF_AVAILABLE else 'reference FFD', |
| } |
|
|
|
|
| |
| def rmsnorm(x, weight, eps=1e-6): |
| if MODE == 'fast': |
| return mx.fast.rms_norm(x, weight, eps) |
| v = mx.mean(x.astype(mx.float32) ** 2, axis=-1, keepdims=True) |
| return (x.astype(mx.float32) * mx.rsqrt(v + eps)).astype(x.dtype) * weight |
|
|
|
|
| def rmsnorm_ref(x, weight, eps=1e-6): |
| v = mx.mean(x.astype(mx.float32) ** 2, axis=-1, keepdims=True) |
| return (x.astype(mx.float32) * mx.rsqrt(v + eps)).astype(x.dtype) * weight |
|
|
|
|
| |
| def swiglu(x, w_gate, w_up, w_down): |
| if MODE == 'fast' and PMLF_AVAILABLE: |
| return _pmlf_swiglu(x, w_gate, w_up, w_down) |
| g = x @ w_gate |
| return ((g * mx.sigmoid(g)) * (x @ w_up)) @ w_down |
|
|
|
|
| def swiglu_ref(x, w_gate, w_up, w_down): |
| g = x @ w_gate |
| return ((g * mx.sigmoid(g)) * (x @ w_up)) @ w_down |
|
|
|
|
| |
| def rope(x, dims, base=10000.0, offset=0): |
| """x: (B, H, T, Dh). offset may be an int or an int array (packed restarts).""" |
| if MODE == 'fast': |
| return mx.fast.rope(x, dims, traditional=False, base=base, scale=1.0, |
| offset=offset) |
| return rope_ref(x, dims, base=base, offset=offset) |
|
|
|
|
| def rope_ref(x, dims, base=10000.0, offset=0): |
| import numpy as np |
| T, d = x.shape[-2], dims |
| inv = mx.array((1.0 / (base ** (np.arange(0, d, 2) / d))).astype(np.float32)) |
| pos = mx.arange(T, dtype=mx.float32) + (offset if isinstance(offset, int) else 0) |
| ang = pos[:, None] * inv[None, :] |
| cos, sin = mx.cos(ang).astype(x.dtype), mx.sin(ang).astype(x.dtype) |
| x1, x2 = x[..., 0::2], x[..., 1::2] |
| o1 = x1 * cos - x2 * sin |
| o2 = x1 * sin + x2 * cos |
| return mx.stack([o1, o2], axis=-1).reshape(x.shape) |
|
|
|
|
| |
| def sdpa(q, k, v, scale, mask=None): |
| """q,k,v: (B, H, T, Dh). mask: None | 'causal' | bool array.""" |
| if MODE == 'fast': |
| return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask) |
| return sdpa_ref(q, k, v, scale, mask) |
|
|
|
|
| def sdpa_ref(q, k, v, scale, mask=None): |
| s = (q * scale) @ mx.swapaxes(k, -1, -2) |
| if isinstance(mask, str) and mask == 'causal': |
| T = q.shape[-2] |
| import numpy as np |
| cm = mx.array(np.tril(np.ones((T, T), dtype=bool))) |
| s = mx.where(cm, s, mx.array(-1e9, dtype=s.dtype)) |
| elif mask is not None and not isinstance(mask, str): |
| s = mx.where(mask, s, mx.array(-1e9, dtype=s.dtype)) |
| return mx.softmax(s.astype(mx.float32), axis=-1).astype(q.dtype) @ v |
|
|
|
|
| |
| def repeat_kv(x, n_rep): |
| """(B, n_kv, T, Dh) -> (B, n_kv*n_rep, T, Dh)""" |
| if n_rep == 1: |
| return x |
| B, H, T, D = x.shape |
| return mx.repeat(x, n_rep, axis=1) |
|
|
|
|
| |
| def cross_entropy_chunked(hidden, w_out, targets, mask=None, chunks=4, |
| label_smoothing=0.0, z_loss=0.0): |
| """Memory-bounded CE. hidden (N,d) bf16, w_out (d,V), targets (N,) int32. |
| mask (N,) bool/float: 1 where the token counts. Returns (loss, n_tokens).""" |
| N = hidden.shape[0] |
| cs = max(1, (N + chunks - 1) // chunks) |
| tot = mx.zeros((), dtype=mx.float32) |
| cnt = mx.zeros((), dtype=mx.float32) |
| for i in range(0, N, cs): |
| h = hidden[i:i + cs] |
| y = targets[i:i + cs] |
| lg = (h @ w_out).astype(mx.float32) |
| lse = mx.logsumexp(lg, axis=-1) |
| tgt = mx.take_along_axis(lg, y[:, None], axis=-1).squeeze(-1) |
| nll = lse - tgt |
| if label_smoothing > 0.0: |
| V = lg.shape[-1] |
| smooth = lse - mx.mean(lg, axis=-1) |
| nll = (1 - label_smoothing) * nll + label_smoothing * smooth |
| if z_loss > 0.0: |
| nll = nll + z_loss * (lse ** 2) |
| if mask is not None: |
| m = mask[i:i + cs].astype(mx.float32) |
| tot = tot + mx.sum(nll * m) |
| cnt = cnt + mx.sum(m) |
| else: |
| tot = tot + mx.sum(nll) |
| cnt = cnt + nll.shape[0] |
| return tot / mx.maximum(cnt, 1.0), cnt |
|
|
