| """Dense Metal attention and explicit BF16 rounding boundaries.""" |
|
|
| from __future__ import annotations |
|
|
| import mlx.core as mx |
|
|
|
|
| def silu(x): |
| |
| y = x.astype(mx.float32) |
| return (y * mx.sigmoid(y)).astype(x.dtype) |
|
|
|
|
| def rms_norm(x, weight=None, eps=1e-5, *, rounded=False): |
| |
| |
| |
| if rounded: |
| normalized = mx.fast.rms_norm(x, None, eps) |
| return normalized if weight is None else normalized * weight |
| return mx.fast.rms_norm( |
| x.astype(mx.float32), None if weight is None else weight.astype(mx.float32), eps |
| ).astype(x.dtype) |
|
|
|
|
| def rotate(x, cos, sin): |
| n = cos.shape[-1] |
| a = x[..., :n] |
| half = n // 2 |
| rotated = mx.concatenate((-a[..., half:], a[..., :half]), axis=-1) |
| out = a * cos.astype(x.dtype) + rotated * sin.astype(x.dtype) |
| return mx.concatenate((out, x[..., n:]), axis=-1) |
|
|
|
|
| def attention(q, k, v, *, causal=False, query_chunk=0): |
| """Inputs [batch, sequence, heads, dim]; GQA without materializing repeated KV.""" |
| q, k, v = (a.transpose(0, 2, 1, 3) for a in (q, k, v)) |
| if not query_chunk or q.shape[2] <= query_chunk: |
| out = mx.fast.scaled_dot_product_attention( |
| q, k, v, scale=q.shape[-1] ** -0.5, mask="causal" if causal else None |
| ) |
| else: |
| rows = [] |
| for start in range(0, q.shape[2], query_chunk): |
| qq = q[:, :, start : start + query_chunk] |
| mask = None |
| if causal: |
| mask = ( |
| mx.arange(start, start + qq.shape[2])[:, None] |
| >= mx.arange(k.shape[2])[None] |
| ) |
| rows.append( |
| mx.fast.scaled_dot_product_attention( |
| qq, k, v, scale=q.shape[-1] ** -0.5, mask=mask |
| ) |
| ) |
| mx.eval(rows[-1]) |
| out = mx.concatenate(rows, axis=2) |
| return out.transpose(0, 2, 1, 3) |
|
|
|
|
| def modulate( |
| x, norm_weight, shift, scale, indices, eps=1e-5, *, kernel="native", threads=256 |
| ): |
| if ( |
| kernel == "metal-exact" |
| and x.ndim == 2 |
| and x.shape[1] == 5376 |
| and x.dtype == mx.bfloat16 |
| ): |
| from .tuning import norm_modulate_exact, supported_device |
|
|
| if supported_device() and 4096 <= x.shape[0] <= 9216: |
| return norm_modulate_exact(x, norm_weight, shift, scale, indices, eps) |
| if kernel in ("metal", "metal-qk"): |
| from .kernels import norm_modulate |
|
|
| return norm_modulate(x, norm_weight, shift, scale, indices, eps, threads) |
| h = rms_norm(x, norm_weight, eps) |
| return h * (mx.array(1, x.dtype) + scale[indices]) + shift[indices] |
|
|
|
|
| def gated_add(x, other, gate, indices): |
| return x + gate[indices] * other |
|
|
|
|
| def swiglu(x): |
| a, b = mx.split(x, 2, axis=-1) |
| return silu(a) * b |
|
|
|
|
| def layer_norm(x, weight, bias, eps=1e-5): |
| return mx.fast.layer_norm( |
| x.astype(mx.float32), weight.astype(mx.float32), bias.astype(mx.float32), eps |
| ).astype(x.dtype) |
|
|
