| """Research kernels. No default dispatch until full-output validation.""" |
|
|
| from functools import lru_cache |
|
|
| import mlx.core as mx |
|
|
| from experiments.v2_mpp import gemm |
|
|
|
|
| class TunedFC2: |
| def __init__(self, layer): |
| self.original = layer |
| self.bias = layer.bias |
| self.weight = layer.weight |
| self.zero = mx.zeros(self.weight.shape[0], self.weight.dtype) |
|
|
| @property |
| def nbytes(self): |
| return self.original.nbytes |
|
|
| def __call__(self, x): |
| if x.ndim != 2 or x.dtype != mx.bfloat16 or self.weight.shape != (5376, 14336): |
| return self.original(x) |
| return gemm( |
| x, |
| self.weight, |
| self.zero if self.bias is None else self.bias, |
| 64, |
| 128, |
| 512, |
| 4, |
| 8, |
| ) |
|
|
|
|
| @lru_cache(None) |
| def _qk(): |
| return mx.fast.metal_kernel( |
| name="h3_r3_qk_native_reduce", |
| 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; |
| #pragma unroll |
| for (uint j = 0; j < 4; ++j) { |
| uint c = lane * 4 + j; |
| float q = float(qkv[qb+c]), k = float(qkv[kb+c]); |
| qsum += q*q; ksum += k*k; |
| } |
| float qi = metal::precise::rsqrt(simd_sum(qsum)/float(dim)+epsilon[0]); |
| float ki = metal::precise::rsqrt(simd_sum(ksum)/float(dim)+epsilon[0]); |
| #pragma unroll |
| for (uint j = 0; j < 4; ++j) { |
| uint c = lane * 4 + j; |
| T qn = T(float(qw[c])*(float(qkv[qb+c])*qi)); |
| T kn = T(float(kw[c])*(float(qkv[kb+c])*ki)); |
| if (c < rot) { |
| uint partner = (c+rot/2)%rot; |
| T qother = T(float(qw[partner])*(float(qkv[qb+partner])*qi)); |
| T kother = T(float(kw[partner])*(float(qkv[kb+partner])*ki)); |
| 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); |
| } |
| uint offset = HEAD_MAJOR ? (head*qkv_shape[0]+token)*dim+c : row*dim+c; |
| qout[offset] = qn; kout[offset] = kn; |
| } |
| """, |
| ) |
|
|
|
|
| def qk_native_reduce(qkv, qw, kw, cos, sin, eps=1e-5, *, head_major=False): |
| n, _, heads, dim = qkv.shape |
| if dim != 128: |
| raise ValueError("Research kernel only supports D128") |
| shape = (1, heads, n, dim) if head_major else (1, n, heads, dim) |
| return _qk()( |
| inputs=[ |
| qkv, |
| qw, |
| kw, |
| cos.reshape(n, -1), |
| sin.reshape(n, -1), |
| mx.array([eps], mx.float32), |
| ], |
| template=[("T", qkv.dtype), ("HEAD_MAJOR", head_major)], |
| grid=(32, n * heads, 1), |
| threadgroup=(32, 1, 1), |
| output_shapes=[shape, shape], |
| output_dtypes=[qkv.dtype, qkv.dtype], |
| ) |
|
|
|
|
| @lru_cache(None) |
| def _norm(): |
| return mx.fast.metal_kernel( |
| name="h3_r3_norm_mod_native_reduce", |
| 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 r = 0; r < width; r += TG*4) { |
| #pragma unroll |
| for (uint j = 0; j < 4; ++j) { |
| uint c = r + tid*4+j; |
| if (c < width) { 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 (warp == 0) partial[lane] = 0; |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| if (lane == 0) partial[warp] = sum; |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| threadgroup float inverse; |
| if (warp == 0) { |
| float s = simd_sum(partial[lane]); |
| if (lane == 0) inverse = metal::precise::rsqrt(s/float(width)+epsilon[0]); |
| } |
| threadgroup_barrier(mem_flags::mem_threadgroup); |
| uint index = indices[row]*width; |
| for (uint r = 0; r < width; r += TG*4) { |
| #pragma unroll |
| for (uint j = 0; j < 4; ++j) { |
| uint c = r+tid*4+j; |
| if (c < width) { |
| T affine = T(float(weight[c])*(float(x[row*width+c])*inverse)); |
| T factor = T(T(1)+scale[index+c]); |
| T product = T(affine*factor); |
| out[row*width+c] = T(product+shift[index+c]); |
| } |
| } |
| } |
| """, |
| ) |
|
|
|
|
| def norm_native_reduce(x, weight, shift, scale, indices, eps=1e-5, threads=1024): |
| return _norm()( |
| 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] |
|
|
