File size: 4,919 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """Experimental Metal RMSNorm + indexed AdaLN, with BF16 operation boundaries.
One threadgroup owns one row. The tuning harness checks output error as well as
latency for 128/256/512 threads; this is opt-in because reduction order differs.
"""
from functools import lru_cache
import mlx.core as mx
@lru_cache(maxsize=8)
def _kernel(threads):
if threads not in (128, 256, 512):
raise ValueError("norm threads must be 128, 256 or 512")
return mx.fast.metal_kernel(
name=f"h3_norm_mod_{threads}",
input_names=["x", "weight", "shift", "scale", "indices", "epsilon"],
output_names=["out"],
source=r"""
uint tid = thread_position_in_threadgroup.x;
uint row = threadgroup_position_in_grid.y;
uint width = x_shape[1];
float sum = 0;
for (uint c = tid; c < width; c += TG) {
float a = float(x[row * width + c]);
sum += a * a;
}
sum = simd_sum(sum);
threadgroup float partial[32];
uint lane = tid % 32, warp = tid / 32;
if (lane == 0) partial[warp] = sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (warp == 0) {
float v = lane < TG / 32 ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0) partial[0] = metal::rsqrt(v / float(width) + epsilon[0]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float inv = partial[0];
uint index = indices[row] * width;
for (uint c = tid; c < width; c += TG) {
T affine = T(float(x[row * width + c]) * inv * float(weight[c]));
T factor = T(T(1) + scale[index + c]);
T product = T(affine * factor);
out[row * width + c] = T(product + shift[index + c]);
}
""",
)
def norm_modulate(x, weight, shift, scale, indices, eps=1e-5, threads=256):
if x.ndim != 2 or x.shape[-1] != weight.shape[0]:
raise ValueError("Expected [rows,width] and a matching norm weight")
return _kernel(threads)(
inputs=[
x,
weight,
shift,
scale,
indices.astype(mx.uint32),
mx.array([eps], mx.float32),
],
template=[("T", x.dtype), ("TG", threads)],
grid=(threads, x.shape[0], 1),
threadgroup=(threads, 1, 1),
output_shapes=[x.shape],
output_dtypes=[x.dtype],
)[0]
@lru_cache(maxsize=1)
def _qk_kernel():
return mx.fast.metal_kernel(
name="h3_qk_norm_rope",
input_names=["qkv", "qw", "kw", "cosine", "sine", "epsilon"],
output_names=["qout", "kout"],
source=r"""
uint lane = thread_position_in_threadgroup.x;
uint row = threadgroup_position_in_grid.y;
uint heads = qkv_shape[2], dim = qkv_shape[3];
uint token = row / heads, head = row % heads;
uint qb = token * 3 * heads * dim + head * dim;
uint kb = qb + heads * dim;
uint rot = cosine_shape[1];
float qsum = 0, ksum = 0;
for (uint c = lane; c < dim; c += 32) {
float q = float(qkv[qb+c]), k = float(qkv[kb+c]);
qsum += q*q; ksum += k*k;
}
float qi = metal::rsqrt(simd_sum(qsum)/float(dim)+epsilon[0]);
float ki = metal::rsqrt(simd_sum(ksum)/float(dim)+epsilon[0]);
for (uint c = lane; c < dim; c += 32) {
T qn = T(float(qkv[qb+c])*qi*float(qw[c]));
T kn = T(float(qkv[kb+c])*ki*float(kw[c]));
if (c < rot) {
uint partner = (c+rot/2)%rot;
T qother = T(float(qkv[qb+partner])*qi*float(qw[partner]));
T kother = T(float(qkv[kb+partner])*ki*float(kw[partner]));
if (c < rot/2) { qother = -qother; kother = -kother; }
T qcos = T(qn*cosine[token*rot+c]);
T kcos = T(kn*cosine[token*rot+c]);
T qsin = T(qother*sine[token*rot+c]);
T ksin = T(kother*sine[token*rot+c]);
qn = T(qcos+qsin); kn = T(kcos+ksin);
}
qout[row*dim+c] = qn;
kout[row*dim+c] = kn;
}
""",
)
def qk_norm_rope(qkv, qw, kw, cos, sin, eps=1e-5):
"""Fused head RMSNorm/partial RoPE; qkv [S,3,H,D], outputs [1,S,H,D]."""
if qkv.ndim != 4 or qkv.shape[1] != 3:
raise ValueError("QKV must have shape [sequence,3,heads,dim]")
n, _, heads, dim = qkv.shape
shape = (1, n, heads, dim)
rot = cos.shape[-1]
return _qk_kernel()(
inputs=[
qkv,
qw,
kw,
cos.reshape(n, rot),
sin.reshape(n, rot),
mx.array([eps], mx.float32),
],
template=[("T", qkv.dtype)],
grid=(32, n * heads, 1),
threadgroup=(32, 1, 1),
output_shapes=[shape, shape],
output_dtypes=[qkv.dtype, qkv.dtype],
)
|