250 lines
7.5 KiB
C
250 lines
7.5 KiB
C
/*
|
|
* Unit test for Q4_K block layout, scale extraction, and dot product.
|
|
* Build: cc -O2 -Wall -DDS4_NO_GPU -DDS4_Q4K_DOT_TEST_MAIN -I. -o tests/test_q4k_dot tests/test_q4k_dot.c -lm -pthread
|
|
* Run: ./tests/test_q4k_dot
|
|
*/
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
|
|
#define QK_K 256
|
|
|
|
typedef struct {
|
|
uint16_t d;
|
|
uint16_t dmin;
|
|
uint8_t scales[12];
|
|
uint8_t qs[QK_K / 2];
|
|
} block_q4_K;
|
|
|
|
typedef struct {
|
|
float d;
|
|
int8_t qs[QK_K];
|
|
int16_t bsums[QK_K / 16];
|
|
} block_q8_K;
|
|
|
|
static inline float f16_to_f32(uint16_t h) {
|
|
uint32_t sign = (h & 0x8000) << 16;
|
|
int32_t exp = (h >> 10) & 0x1F;
|
|
uint32_t frac = h & 0x3FF;
|
|
if (exp == 0) return *(float *)&sign;
|
|
if (exp == 31) return *(float *)(uint32_t[]){ sign | 0x7F800000 | (frac << 13) };
|
|
uint32_t f = sign | ((exp + 127 - 15) << 23) | (frac << 13);
|
|
float out;
|
|
memcpy(&out, &f, sizeof(out));
|
|
return out;
|
|
}
|
|
|
|
static inline void q4_k_get_scale_min(int j, const uint8_t *q, uint8_t *sc, uint8_t *m) {
|
|
if (j < 4) {
|
|
*sc = q[j] & 63;
|
|
*m = q[j + 4] & 63;
|
|
} else {
|
|
*sc = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
|
|
*m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4);
|
|
}
|
|
}
|
|
|
|
/* The corrected dot product, matching ds4.c after the fix. */
|
|
static void vec_dot_q4_K_q8_K(int n, float *s, const block_q4_K *x, const block_q8_K *y) {
|
|
const int nb = n / QK_K;
|
|
float sumf = 0.0f;
|
|
|
|
for (int i = 0; i < nb; i++) {
|
|
const float d = y[i].d * f16_to_f32(x[i].d);
|
|
const float dm = -y[i].d * f16_to_f32(x[i].dmin);
|
|
|
|
const uint8_t *qs = x[i].qs;
|
|
const uint8_t *sc = x[i].scales;
|
|
const int8_t *q8 = y[i].qs;
|
|
|
|
int summs = 0;
|
|
for (int j = 0; j < QK_K / 32; j++) {
|
|
uint8_t sc_val, m_val;
|
|
q4_k_get_scale_min(j, sc, &sc_val, &m_val);
|
|
int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1];
|
|
summs += m_val * gsum;
|
|
}
|
|
|
|
int isum = 0;
|
|
for (int j = 0; j < QK_K / 32; j++) {
|
|
uint8_t sc_val, m_val;
|
|
q4_k_get_scale_min(j, sc, &sc_val, &m_val);
|
|
|
|
const int byte_off = (j >> 1) * 32;
|
|
const int shift = (j & 1) * 4;
|
|
|
|
for (int l = 0; l < 32; l++) {
|
|
isum += ((qs[byte_off + l] >> shift) & 0xF) * (int)q8[j * 32 + l] * sc_val;
|
|
}
|
|
}
|
|
|
|
sumf += d * (float)isum + dm * (float)summs;
|
|
}
|
|
|
|
*s = sumf;
|
|
}
|
|
|
|
/* Reference: fully dequantize Q4_K to float, then dot with Q8_K's dequantized values. */
|
|
static float ref_dot(const block_q4_K *bx, const block_q8_K *by) {
|
|
float x[QK_K];
|
|
const float d = f16_to_f32(bx->d);
|
|
const float dm = f16_to_f32(bx->dmin);
|
|
|
|
for (int j = 0; j < QK_K / 32; j++) {
|
|
uint8_t sc_val, m_val;
|
|
q4_k_get_scale_min(j, bx->scales, &sc_val, &m_val);
|
|
const int byte_off = (j >> 1) * 32;
|
|
const int shift = (j & 1) * 4;
|
|
for (int l = 0; l < 32; l++) {
|
|
int q = (bx->qs[byte_off + l] >> shift) & 0xF;
|
|
x[j * 32 + l] = d * sc_val * (float)q - dm * m_val;
|
|
}
|
|
}
|
|
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < QK_K; i++) sum += x[i] * by->d * (float)by->qs[i];
|
|
return sum;
|
|
}
|
|
|
|
static void fill_q4_K(block_q4_K *bx, uint32_t seed) {
|
|
uint8_t *p = (uint8_t *)bx;
|
|
uint32_t s = seed;
|
|
for (size_t i = 0; i < sizeof(block_q4_K); i++) {
|
|
s = s * 1103515245u + 12345u;
|
|
p[i] = (uint8_t)(s >> 16);
|
|
}
|
|
/* Ensure d/dmin are valid finite f16 */
|
|
bx->d &= 0x7BFF;
|
|
bx->dmin &= 0x7BFF;
|
|
}
|
|
|
|
static void fill_q8_K(block_q8_K *by, uint32_t seed) {
|
|
uint32_t s = seed;
|
|
by->d = ((s & 0xFFFF) / 65536.0f) * 2.0f + 0.01f;
|
|
for (int i = 0; i < QK_K; i++) {
|
|
s = s * 1103515245u + 12345u;
|
|
by->qs[i] = (int8_t)((s >> 16) & 0xFF);
|
|
}
|
|
for (int j = 0; j < QK_K / 16; j++) {
|
|
int32_t sum = 0;
|
|
for (int l = 0; l < 16; l++) sum += (int32_t)by->qs[j * 16 + l];
|
|
by->bsums[j] = (int16_t)sum;
|
|
}
|
|
}
|
|
|
|
static int test_block_sizes(void) {
|
|
int ok = (sizeof(block_q4_K) == 144 && sizeof(block_q8_K) == 292);
|
|
printf(" block sizes: Q4_K=%zu (expect 144), Q8_K=%zu (expect 292): %s\n",
|
|
sizeof(block_q4_K), sizeof(block_q8_K), ok ? "PASS" : "FAIL");
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
static int test_scale_extraction(void) {
|
|
uint8_t scales[12] = {0};
|
|
uint8_t expected_sc[8] = {5, 12, 30, 60, 7, 15, 31, 63};
|
|
uint8_t expected_m[8] = {3, 8, 20, 45, 1, 10, 25, 50};
|
|
|
|
for (int j = 0; j < 4; j++) {
|
|
scales[j] = expected_sc[j];
|
|
scales[j + 4] = expected_m[j];
|
|
}
|
|
for (int j = 4; j < 8; j++) {
|
|
scales[j + 4] = (expected_sc[j] & 0xF) | ((expected_m[j] & 0xF) << 4);
|
|
scales[j - 4] |= ((expected_sc[j] >> 4) << 6);
|
|
scales[j - 0] |= ((expected_m[j] >> 4) << 6);
|
|
}
|
|
|
|
int ok = 1;
|
|
for (int j = 0; j < 8; j++) {
|
|
uint8_t sc, m;
|
|
q4_k_get_scale_min(j, scales, &sc, &m);
|
|
if (sc != expected_sc[j] || m != expected_m[j]) {
|
|
printf(" j=%d: expected sc=%d m=%d, got sc=%d m=%d\n",
|
|
j, expected_sc[j], expected_m[j], sc, m);
|
|
ok = 0;
|
|
}
|
|
}
|
|
|
|
printf(" scale extraction: %s\n", ok ? "PASS" : "FAIL");
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
static int test_dot_reference(void) {
|
|
int ok = 1;
|
|
for (uint32_t seed = 1; seed <= 50; seed++) {
|
|
block_q4_K bx;
|
|
block_q8_K by;
|
|
fill_q4_K(&bx, seed);
|
|
fill_q8_K(&by, seed * 7 + 13);
|
|
|
|
float result = 0.0f;
|
|
vec_dot_q4_K_q8_K(QK_K, &result, &bx, &by);
|
|
float expected = ref_dot(&bx, &by);
|
|
|
|
float err = fabsf(result - expected);
|
|
float rel = fabsf(expected) > 1e-3f ? err / fabsf(expected) : err;
|
|
if (rel > 0.01f) {
|
|
printf(" seed=%u: result=%.6f expected=%.6f rel_err=%.6f\n",
|
|
seed, result, expected, rel);
|
|
ok = 0;
|
|
}
|
|
}
|
|
|
|
printf(" dot vs reference (50 random blocks): %s\n", ok ? "PASS" : "FAIL");
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
/* Test with a hand-crafted known block. */
|
|
static int test_dot_known(void) {
|
|
block_q4_K bx;
|
|
block_q8_K by;
|
|
memset(&bx, 0, sizeof(bx));
|
|
memset(&by, 0, sizeof(by));
|
|
|
|
/* d=1.0 f16=0x3C00, dmin=0, all scales=1 (sc=1, m=0), all qs=0x88 (value 8) */
|
|
bx.d = 0x3C00;
|
|
bx.dmin = 0;
|
|
/* Groups 0-3: sc=1 m=0 */
|
|
for (int j = 0; j < 4; j++) {
|
|
bx.scales[j] = 1;
|
|
bx.scales[j + 4] = 0;
|
|
}
|
|
/* Groups 4-7: sc=1 m=0, high bits=0 */
|
|
for (int j = 4; j < 8; j++) {
|
|
bx.scales[j + 4] = (1 & 0xF) | (0 << 4);
|
|
/* high bits already 0 */
|
|
}
|
|
/* All qs = 0x88: low nibble=8, high nibble=8 */
|
|
memset(bx.qs, 0x88, sizeof(bx.qs));
|
|
|
|
/* Q8: all +1, d=1.0 */
|
|
by.d = 1.0f;
|
|
for (int i = 0; i < QK_K; i++) by.qs[i] = 1;
|
|
for (int j = 0; j < QK_K / 16; j++) by.bsums[j] = 16;
|
|
|
|
/* Expected: d=1, dm=0, isum = sum over 8 groups * 32 values * 8 * 1 * sc(=1) = 8*32*8 = 2048 */
|
|
float result = 0.0f;
|
|
vec_dot_q4_K_q8_K(QK_K, &result, &bx, &by);
|
|
float expected = 2048.0f;
|
|
|
|
int ok = (fabsf(result - expected) < 0.5f);
|
|
printf(" dot known: result=%.1f expected=%.1f: %s\n", result, expected, ok ? "PASS" : "FAIL");
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
int main(void) {
|
|
int failures = 0;
|
|
|
|
printf("Q4_K unit tests:\n");
|
|
failures += test_block_sizes();
|
|
failures += test_scale_extraction();
|
|
failures += test_dot_known();
|
|
failures += test_dot_reference();
|
|
|
|
printf("\n%d/%d tests passed\n", 4 - failures, 4);
|
|
return failures ? 1 : 0;
|
|
}
|