fix: add bounds checking to Metal shader kernels (#57)

Security fixes:
- kernel_fwht_128: Added n_elements param, early exit for invalid tid
- kernel_turbo4_dequant: Added src_size/norms_size validation
- kernel_attention_turbo4: Implemented with full bounds checking
- Added kernel_softmax for attention pipeline
- Added threadgroup_barrier in FWHT for correctness

Closes #57

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;
+            }
+        }
+    }
 }