79 lines
2.3 KiB
C++
79 lines
2.3 KiB
C++
#include "llama-turbo.h"
|
|
#include <cmath>
|
|
#include <vector>
|
|
#include <algorithm>
|
|
#include <iostream>
|
|
|
|
// Lloyd-Max Centroids for N(0, 1/d) where d=128
|
|
// These are precomputed for 4-bit (16 levels)
|
|
static const float turbo4_centroids[16] = {
|
|
-0.2154f, -0.1523f, -0.1121f, -0.0812f,
|
|
-0.0554f, -0.0321f, -0.0105f, 0.0105f,
|
|
0.0321f, 0.0554f, 0.0812f, 0.1121f,
|
|
0.1523f, 0.2154f, 0.2800f, 0.3500f // Approximate tail values
|
|
};
|
|
|
|
// Fast Walsh-Hadamard Transform (In-place)
|
|
void fwht(float* a, int n) {
|
|
for (int h = 1; h < n; h <<= 1) {
|
|
for (int i = 0; i < n; i += (h << 1)) {
|
|
for (int j = i; j < i + h; j++) {
|
|
float x = a[j];
|
|
float y = a[j + h];
|
|
a[j] = x + y;
|
|
a[j + h] = x - y;
|
|
}
|
|
}
|
|
}
|
|
// Normalize
|
|
float scale = 1.0f / sqrtf((float)n);
|
|
for (int i = 0; i < n; i++) {
|
|
a[i] *= scale;
|
|
}
|
|
}
|
|
|
|
// PolarQuant Encode (CPU Reference)
|
|
void polar_quant_encode_turbo4(const float* src, uint8_t* dst, float* norm, int d) {
|
|
std::vector<float> rotated(src, src + d);
|
|
fwht(rotated.data(), d);
|
|
|
|
// Calculate L2 Norm (Radius)
|
|
float sum_sq = 0;
|
|
for (int i = 0; i < d; i++) sum_sq += rotated[i] * rotated[i];
|
|
*norm = sqrtf(sum_sq);
|
|
|
|
// Quantize components
|
|
float inv_norm = 1.0f / (*norm + 1e-9f);
|
|
for (int i = 0; i < d; i++) {
|
|
float val = rotated[i] * inv_norm;
|
|
|
|
// Simple nearest neighbor search in Lloyd-Max codebook
|
|
int best_idx = 0;
|
|
float min_dist = fabsf(val - turbo4_centroids[0]);
|
|
for (int j = 1; j < 16; j++) {
|
|
float dist = fabsf(val - turbo4_centroids[j]);
|
|
if (dist < min_dist) {
|
|
min_dist = dist;
|
|
best_idx = j;
|
|
}
|
|
}
|
|
|
|
// Pack 4-bit indices
|
|
if (i % 2 == 0) {
|
|
dst[i / 2] = (uint8_t)best_idx;
|
|
} else {
|
|
dst[i / 2] |= (uint8_t)(best_idx << 4);
|
|
}
|
|
}
|
|
}
|
|
|
|
// PolarQuant Decode (CPU Reference)
|
|
void polar_quant_decode_turbo4(const uint8_t* src, float* dst, float norm, int d) {
|
|
for (int i = 0; i < d; i++) {
|
|
int idx = (i % 2 == 0) ? (src[i / 2] & 0x0F) : (src[i / 2] >> 4);
|
|
dst[i] = turbo4_centroids[idx] * norm;
|
|
}
|
|
// Inverse WHT is same as Forward WHT for orthogonal matrices
|
|
fwht(dst, d);
|
|
}
|