77 lines
2.2 KiB
Metal
77 lines
2.2 KiB
Metal
|
#include <metal_stdlib>
|
||
|
|
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,
|
||
|
|
0.0321, 0.0554, 0.0812, 0.1121,
|
||
|
|
0.1523, 0.2154, 0.2800, 0.3500
|
||
|
|
};
|
||
|
|
|
||
|
|
// Fast Walsh-Hadamard Transform (In-place, SIMD-optimized)
|
||
|
|
// Assumes d=128 (standard head dimension)
|
||
|
|
kernel void kernel_fwht_128(
|
||
|
|
device float* data [[buffer(0)]],
|
||
|
|
uint tid [[thread_position_in_grid]]
|
||
|
|
) {
|
||
|
|
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++) {
|
||
|
|
float x = data[base + j];
|
||
|
|
float y = data[base + j + h];
|
||
|
|
data[base + j] = x + y;
|
||
|
|
data[base + j + h] = x - y;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// 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
|
||
|
|
kernel void kernel_turbo4_dequant(
|
||
|
|
device const uchar* src [[buffer(0)]],
|
||
|
|
device const float* norms [[buffer(1)]],
|
||
|
|
device float* dst [[buffer(2)]],
|
||
|
|
uint tid [[thread_position_in_grid]]
|
||
|
|
) {
|
||
|
|
const uint d = 128;
|
||
|
|
uint base_src = tid * (d / 2);
|
||
|
|
uint base_dst = tid * d;
|
||
|
|
float norm = norms[tid];
|
||
|
|
|
||
|
|
for (uint i = 0; i < d; i++) {
|
||
|
|
uchar packed = src[base_src + (i / 2)];
|
||
|
|
uint idx = (i % 2 == 0) ? (packed & 0x0F) : (packed >> 4);
|
||
|
|
dst[base_dst + i] = 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
|
||
|
|
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)]],
|
||
|
|
uint tid [[thread_position_in_grid]]
|
||
|
|
) {
|
||
|
|
// 1. Dequantize K on the fly
|
||
|
|
// 2. Compute dot product with Q
|
||
|
|
// 3. Store score
|
||
|
|
}
|