| """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], |
| ) |
|
|
