diff --git a/ggml-metal-turbo.metal b/ggml-metal-turbo.metal index 97f0edd0..bbf4c773 100644 --- a/ggml-metal-turbo.metal +++ b/ggml-metal-turbo.metal @@ -2,7 +2,6 @@ using namespace metal; // Lloyd-Max Centroids (4-bit, 16 levels) -// Precomputed for N(0, 1/128) constant float turbo4_centroids[16] = { -0.2154, -0.1523, -0.1121, -0.0812, -0.0554, -0.0321, -0.0105, 0.0105, @@ -10,16 +9,19 @@ constant float turbo4_centroids[16] = { 0.1523, 0.2154, 0.2800, 0.3500 }; -// Fast Walsh-Hadamard Transform (In-place, SIMD-optimized) -// Assumes d=128 (standard head dimension) +// Fast Walsh-Hadamard Transform with bounds checking (Issue #57) kernel void kernel_fwht_128( device float* data [[buffer(0)]], + constant uint& n_elements [[buffer(1)]], uint tid [[thread_position_in_grid]] ) { + if (tid >= n_elements) { + return; + } + const uint d = 128; uint base = tid * d; - // Stage 1-7 (128 = 2^7) for (uint h = 1; h < d; h <<= 1) { for (uint i = 0; i < d; i += (h << 1)) { for (uint j = i; j < i + h; j++) { @@ -29,48 +31,156 @@ kernel void kernel_fwht_128( data[base + j + h] = x - y; } } + threadgroup_barrier(mem_flags::mem_device); } - // Normalize float scale = 1.0 / sqrt(128.0); for (uint i = 0; i < d; i++) { data[base + i] *= scale; } } -// PolarQuant Turbo4 Dequantization (Attention Hot Path) -// Unpacks 4-bit indices, looks up centroids, scales by radius +// Turbo4 Dequantization with bounds checking (Issue #57) kernel void kernel_turbo4_dequant( device const uchar* src [[buffer(0)]], device const float* norms [[buffer(1)]], device float* dst [[buffer(2)]], + constant uint& n_elements [[buffer(3)]], + constant uint& src_size [[buffer(4)]], + constant uint& norms_size [[buffer(5)]], uint tid [[thread_position_in_grid]] ) { + if (tid >= n_elements) { + return; + } + const uint d = 128; uint base_src = tid * (d / 2); uint base_dst = tid * d; + + if (base_src + (d / 2) > src_size) { + return; + } + + if (tid >= norms_size) { + return; + } + float norm = norms[tid]; for (uint i = 0; i < d; i++) { - uchar packed = src[base_src + (i / 2)]; + uint src_idx = base_src + (i / 2); + uint dst_idx = base_dst + i; + + if (src_idx >= src_size || dst_idx >= n_elements * d) { + continue; + } + + uchar packed = src[src_idx]; uint idx = (i % 2 == 0) ? (packed & 0x0F) : (packed >> 4); - dst[base_dst + i] = turbo4_centroids[idx] * norm; + idx = min(idx, 15u); + dst[dst_idx] = turbo4_centroids[idx] * norm; } - - // Note: FWHT is applied separately or fused into attention } -// Fused Attention with TurboQuant (Conceptual) -// This is where the real speed win happens +// Fused Attention with bounds checking (Issue #57) kernel void kernel_attention_turbo4( device const float* q [[buffer(0)]], device const uchar* k_packed [[buffer(1)]], device const float* k_norms [[buffer(2)]], device float* scores [[buffer(3)]], constant uint& d [[buffer(4)]], + constant uint& n_queries [[buffer(5)]], + constant uint& n_keys [[buffer(6)]], + constant uint& k_packed_size [[buffer(7)]], + constant uint& k_norms_size [[buffer(8)]], uint tid [[thread_position_in_grid]] ) { - // 1. Dequantize K on the fly - // 2. Compute dot product with Q - // 3. Store score + if (tid >= n_queries * n_keys) { + return; + } + + uint query_idx = tid / n_keys; + uint key_idx = tid % n_keys; + + if (query_idx >= n_queries || key_idx >= n_keys) { + return; + } + + float score = 0.0; + uint base_q = query_idx * d; + uint base_k_src = key_idx * (d / 2); + + if (base_k_src + (d / 2) > k_packed_size) { + scores[tid] = 0.0; + return; + } + + if (key_idx >= k_norms_size) { + scores[tid] = 0.0; + return; + } + + float k_norm = k_norms[key_idx]; + + for (uint i = 0; i < d; i++) { + if (base_q + i >= n_queries * d) { + break; + } + + float q_val = q[base_q + i]; + + uint src_idx = base_k_src + (i / 2); + if (src_idx >= k_packed_size) { + break; + } + + uchar packed = k_packed[src_idx]; + uint idx = (i % 2 == 0) ? (packed & 0x0F) : (packed >> 4); + idx = min(idx, 15u); + float k_val = turbo4_centroids[idx] * k_norm; + + score += q_val * k_val; + } + + score /= sqrt(float(d)); + scores[tid] = score; +} + +// Softmax kernel with bounds checking +kernel void kernel_softmax( + device float* scores [[buffer(0)]], + constant uint& n_rows [[buffer(1)]], + constant uint& n_cols [[buffer(2)]], + uint row [[thread_position_in_grid]] +) { + if (row >= n_rows) { + return; + } + + uint base = row * n_cols; + uint total = n_rows * n_cols; + + float max_val = -INFINITY; + for (uint i = 0; i < n_cols; i++) { + if (base + i < total) { + max_val = max(max_val, scores[base + i]); + } + } + + float sum = 0.0; + for (uint i = 0; i < n_cols; i++) { + if (base + i < total) { + scores[base + i] = exp(scores[base + i] - max_val); + sum += scores[base + i]; + } + } + + if (sum > 0.0) { + for (uint i = 0; i < n_cols; i++) { + if (base + i < total) { + scores[base + i] /= sum; + } + } + } }