150 lines
6.0 KiB
C
150 lines
6.0 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* Validate + benchmark the Q4_K x B GEMM (fm_sgemm_q4_k) against a CPU
|
|
* reference that dequantizes the same super-block bytes. Since both sides read
|
|
* identical Q4_K bytes, the weight need not be "meaningfully" quantized — only
|
|
* structurally valid — to check the GPU kernel matches the dequant math.
|
|
*
|
|
* Usage: test_gemm_q4_k [rows] [cols] [B]
|
|
*/
|
|
|
|
#include "metal/fm_metal.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <time.h>
|
|
|
|
/* Minimal fp32 -> fp16 (round to nearest even, no NaN/Inf handling needed for
|
|
* the small magnitudes we synthesize). */
|
|
static uint16_t f32_to_f16(float f) {
|
|
uint32_t x; memcpy(&x, &f, 4);
|
|
uint32_t sign = (x >> 16) & 0x8000u;
|
|
int32_t exp = (int32_t)((x >> 23) & 0xFF) - 127 + 15;
|
|
uint32_t mant = x & 0x7FFFFFu;
|
|
if (exp <= 0) return (uint16_t)sign;
|
|
if (exp >= 31) return (uint16_t)(sign | 0x7C00u);
|
|
return (uint16_t)(sign | ((uint32_t)exp << 10) | (mant >> 13));
|
|
}
|
|
static float f16_to_f32(uint16_t h) {
|
|
uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
|
|
uint32_t exp = (h >> 10) & 0x1F;
|
|
uint32_t mant = h & 0x3FF;
|
|
uint32_t out;
|
|
if (exp == 0) { if (mant == 0) { out = sign; } else {
|
|
exp = 127 - 15 + 1; while (!(mant & 0x400)) { mant <<= 1; exp--; }
|
|
mant &= 0x3FF; out = sign | (exp << 23) | (mant << 13); } }
|
|
else if (exp == 31) out = sign | 0x7F800000u | (mant << 13);
|
|
else out = sign | ((exp - 15 + 127) << 23) | (mant << 13);
|
|
float f; memcpy(&f, &out, 4); return f;
|
|
}
|
|
|
|
static double now_ms(void) {
|
|
struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts);
|
|
return ts.tv_sec * 1000.0 + ts.tv_nsec / 1e6;
|
|
}
|
|
|
|
static uint32_t rng = 0x1234567u;
|
|
static uint32_t urand(void) { rng ^= rng << 13; rng ^= rng >> 17; rng ^= rng << 5; return rng; }
|
|
static float frand(void) { return (float)(urand() % 2000) / 1000.0f - 1.0f; }
|
|
|
|
/* Q4_K super-block = 144 bytes: half d, half dmin, uchar scales[12], uchar qs[128]. */
|
|
#define QK_BYTES 144u
|
|
|
|
/* Dequantize block b of row into out[256], mirroring the kernel. */
|
|
static void deq_block(const uint8_t *blk, float *out) {
|
|
float d = f16_to_f32(((const uint16_t *)blk)[0]);
|
|
float dmin = f16_to_f32(((const uint16_t *)blk)[1]);
|
|
const uint8_t *sr = blk + 4, *qs = blk + 16;
|
|
uint8_t sc[8], mn[8];
|
|
for (int j = 0; j < 4; j++) { sc[j] = sr[j] & 63; mn[j] = sr[j+4] & 63; }
|
|
for (int j = 4; j < 8; j++) {
|
|
sc[j] = (sr[j+4] & 0xF) | ((sr[j-4] >> 6) << 4);
|
|
mn[j] = (sr[j+4] >> 4) | ((sr[j] >> 6) << 4);
|
|
}
|
|
for (int p = 0; p < 4; p++) {
|
|
float d1 = d * sc[2*p], m1 = dmin * mn[2*p];
|
|
float d2 = d * sc[2*p+1], m2 = dmin * mn[2*p+1];
|
|
const uint8_t *q = qs + p * 32;
|
|
int o = p * 64;
|
|
for (int l = 0; l < 32; l++) {
|
|
out[o + l] = d1 * (q[l] & 0xF) - m1;
|
|
out[o + 32 + l] = d2 * (q[l] >> 4) - m2;
|
|
}
|
|
}
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
int rows = argc > 1 ? atoi(argv[1]) : 2048;
|
|
int cols = argc > 2 ? atoi(argv[2]) : 1024;
|
|
int B = argc > 3 ? atoi(argv[3]) : 16;
|
|
if (cols % 256) { fprintf(stderr, "cols must be a multiple of 256\n"); return 1; }
|
|
|
|
if (fm_init() != 0 || !fm_available()) { fprintf(stderr, "Metal unavailable\n"); return 77; }
|
|
|
|
int nblk = cols / 256;
|
|
size_t row_stride = (size_t)nblk * QK_BYTES;
|
|
size_t wbytes = (size_t)rows * row_stride;
|
|
uint8_t *W = malloc(wbytes);
|
|
for (int r = 0; r < rows; r++) {
|
|
for (int b = 0; b < nblk; b++) {
|
|
uint8_t *blk = W + (size_t)r * row_stride + (size_t)b * QK_BYTES;
|
|
((uint16_t *)blk)[0] = f32_to_f16(0.02f + 0.01f * (float)(urand() % 8));
|
|
((uint16_t *)blk)[1] = f32_to_f16(0.01f * (float)(urand() % 8));
|
|
for (int i = 4; i < 144; i++) blk[i] = (uint8_t)urand();
|
|
}
|
|
}
|
|
float *X = malloc((size_t)B * cols * sizeof(float));
|
|
for (int i = 0; i < B * cols; i++) X[i] = frand();
|
|
|
|
/* CPU reference: out_ref[b, r] = sum_c dequant(W[r]) * X[b]. */
|
|
float *ref = malloc((size_t)B * rows * sizeof(float));
|
|
double c0 = now_ms();
|
|
float blk[256];
|
|
for (int r = 0; r < rows; r++) {
|
|
for (int b = 0; b < nblk; b++) {
|
|
deq_block(W + (size_t)r * row_stride + (size_t)b * QK_BYTES, blk);
|
|
for (int s = 0; s < B; s++) {
|
|
float acc = (b == 0) ? 0.0f : ref[(size_t)s * rows + r];
|
|
const float *xb = X + (size_t)s * cols + b * 256;
|
|
for (int l = 0; l < 256; l++) acc += blk[l] * xb[l];
|
|
ref[(size_t)s * rows + r] = acc;
|
|
}
|
|
}
|
|
}
|
|
double cpu_ms = now_ms() - c0;
|
|
|
|
/* GPU. */
|
|
fm_buffer *bW = fm_buffer_from_ptr(W, wbytes);
|
|
fm_buffer *bX = fm_buffer_from_ptr(X, (size_t)B * cols * sizeof(float));
|
|
fm_buffer *bO = fm_buffer_alloc((size_t)B * rows * sizeof(float));
|
|
/* warm up (compile pipeline + page in) */
|
|
fm_sgemm_q4_k(bW, bX, bO, rows, cols, (int)row_stride, B); fm_sync();
|
|
|
|
int iters = 50;
|
|
double g0 = now_ms();
|
|
for (int i = 0; i < iters; i++) { fm_sgemm_q4_k(bW, bX, bO, rows, cols, (int)row_stride, B); }
|
|
fm_sync();
|
|
double gpu_ms = (now_ms() - g0) / iters;
|
|
|
|
float *out = fm_buffer_data(bO);
|
|
double maxabs = 0, sumsq = 0;
|
|
for (int i = 0; i < B * rows; i++) {
|
|
double e = fabs((double)out[i] - ref[i]);
|
|
if (e > maxabs) maxabs = e;
|
|
sumsq += (double)ref[i] * ref[i];
|
|
}
|
|
double rms = sqrt(sumsq / (B * rows));
|
|
double rel_to_scale = maxabs / (rms + 1e-12); /* err vs typical magnitude */
|
|
|
|
printf("gemm_q4_k rows=%d cols=%d B=%d\n", rows, cols, B);
|
|
printf(" correctness: max abs err = %.4g, output RMS = %.4g, err/RMS = %.4g -> %s\n",
|
|
maxabs, rms, rel_to_scale, (rel_to_scale < 1e-3) ? "PASS" : "FAIL");
|
|
printf(" GPU: %.3f ms/call (CPU ref build: %.1f ms, unoptimized)\n", gpu_ms, cpu_ms);
|
|
|
|
fm_buffer_free(bW); fm_buffer_free(bX); fm_buffer_free(bO);
|
|
free(W); free(X); free(ref);
|
|
return (rel_to_scale < 1e-3) ? 0 : 1;
|
|
}
|