80 lines
4.1 KiB
C
80 lines
4.1 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* Validate + benchmark the large-M prefill GEMM fm_gemm_q4_k_f16 against a CPU
|
|
* F32 reference (dequant of the same Q4_K bytes, then GEMM). F16 tolerance.
|
|
*
|
|
* Usage: test_gemm_q4_k_f16 [M] [N] [K] (defaults 1369 x 2048 x 1024)
|
|
*/
|
|
#include "metal/fm_metal.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <time.h>
|
|
|
|
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, mant = h & 0x3FF, out;
|
|
if (exp == 0) { if (!mant) 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 t; clock_gettime(CLOCK_MONOTONIC,&t); return t.tv_sec*1000.0+t.tv_nsec/1e6; }
|
|
static uint32_t rng=0x9e3779b9u; static uint32_t ur(void){rng^=rng<<13;rng^=rng>>17;rng^=rng<<5;return rng;}
|
|
static float fr(void){ return (float)(ur()%2000)/1000.0f-1.0f; }
|
|
#define QK 144u
|
|
|
|
static void deq_block(const uint8_t*blk,float*o){
|
|
float d=f16_to_f32(((const uint16_t*)blk)[0]), 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],d2=d*sc[2*p+1],m2=dmin*mn[2*p+1];const uint8_t*q=qs+p*32;int o0=p*64;
|
|
for(int l=0;l<32;l++){o[o0+l]=d1*(q[l]&0xF)-m1;o[o0+32+l]=d2*(q[l]>>4)-m2;}}
|
|
}
|
|
|
|
int main(int argc,char**argv){
|
|
int M=argc>1?atoi(argv[1]):1369, N=argc>2?atoi(argv[2]):2048, K=argc>3?atoi(argv[3]):1024;
|
|
if(K%256){fprintf(stderr,"K must be multiple of 256\n");return 1;}
|
|
if(fm_init()||!fm_available()){fprintf(stderr,"Metal unavailable\n");return 77;}
|
|
int nblk=K/256; size_t rs=(size_t)nblk*QK, wb=(size_t)N*rs;
|
|
uint8_t*W=malloc(wb);
|
|
for(int r=0;r<N;r++)for(int b=0;b<nblk;b++){uint8_t*blk=W+(size_t)r*rs+(size_t)b*QK;
|
|
((uint16_t*)blk)[0]=f32_to_f16(0.02f+0.01f*(float)(ur()%8));((uint16_t*)blk)[1]=f32_to_f16(0.01f*(float)(ur()%8));
|
|
for(int i=4;i<144;i++)blk[i]=(uint8_t)ur();}
|
|
float*X=malloc((size_t)M*K*4); for(int i=0;i<M*K;i++)X[i]=fr()*0.3f; /* post-RMSNorm scale */
|
|
|
|
float*ref=malloc((size_t)M*N*4);
|
|
double c0=now_ms(); float blk[256];
|
|
for(int r=0;r<N;r++)for(int b=0;b<nblk;b++){deq_block(W+(size_t)r*rs+(size_t)b*QK,blk);
|
|
for(int m=0;m<M;m++){float acc=(b==0)?0.f:ref[(size_t)m*N+r];const float*xm=X+(size_t)m*K+b*256;
|
|
for(int l=0;l<256;l++)acc+=blk[l]*xm[l]; ref[(size_t)m*N+r]=acc;}}
|
|
double cpu_ms=now_ms()-c0;
|
|
|
|
fm_buffer*bW=fm_buffer_from_ptr(W,wb), *bX=fm_buffer_from_ptr(X,(size_t)M*K*4), *bO=fm_buffer_alloc((size_t)M*N*4);
|
|
fm_gemm_q4_k_f16(bW,bX,bO,M,N,K,(int)rs); fm_sync(); /* warm */
|
|
int iters=20; double g0=now_ms();
|
|
for(int i=0;i<iters;i++) fm_gemm_q4_k_f16(bW,bX,bO,M,N,K,(int)rs);
|
|
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<M*N;i++){double e=fabs((double)out[i]-ref[i]);if(e>maxabs)maxabs=e;sumsq+=(double)ref[i]*ref[i];}
|
|
double rms=sqrt(sumsq/((double)M*N)), rel=maxabs/(rms+1e-9);
|
|
printf("gemm_q4_k_f16 M=%d N=%d K=%d\n",M,N,K);
|
|
printf(" err/RMS=%.4g (RMS=%.3g) -> %s\n", rel, rms, rel<2e-2?"PASS":"FAIL");
|
|
printf(" GPU %.2f ms/call CPU-ref(F32) %.0f ms\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<2e-2?0:1;
|
|
}
|