352 lines
14 KiB
C
352 lines
14 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_matvec_nx.c — bit-exact gate + micro-bench for the BATCHED (multi-
|
|
* column) matvec kernels fc_matvec_*_nx against their single-column
|
|
* counterparts on one K80 die.
|
|
*
|
|
* The batched-prefill correctness contract is per-COLUMN bit-exactness:
|
|
* column t of fc_matvec_*_nx(x[n][cols]) must be bit-identical (memcmp) to
|
|
* fc_matvec_*(x[t]) — the single-column kernels are already proven against
|
|
* the CPU replica by bench_matvec.c, so equality here chains the proof to
|
|
* the batched path. Weights are random quant bytes with fp16 block scales
|
|
* clamped to a sane exponent (as in bench_matvec.c), repacked via
|
|
* fc_repack_weight exactly like a model upload.
|
|
*
|
|
* Also times nx-vs-looped-single for the model's hot shapes to report the
|
|
* batching win (the single-token kernels are instruction-issue bound on
|
|
* sm_37, so decoding each weight block once per 4 columns should approach
|
|
* a 2-3x effective speedup per pass).
|
|
*
|
|
* Build (from foundry/):
|
|
* gcc -O2 -I. -o cuda/tests/test_matvec_nx cuda/tests/test_matvec_nx.c \
|
|
* cuda/fc_cuda.c -lm -ldl
|
|
* Run: CUDA_VISIBLE_DEVICES=<die> ./cuda/tests/test_matvec_nx
|
|
*/
|
|
|
|
#include "cuda/fc_cuda.h"
|
|
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
|
|
/* Batch-tile width the launcher picks for n columns of quant type t (L6
|
|
* dual-width dispatch: n <= 4 -> the 4-wide tile, n >= 5 -> the 8-wide tile;
|
|
* Q4_K instead crosses over at FC_Q4K_NX8_MIN_N because its min-lifted 8-wide
|
|
* kernel is register-heavy). Only used for the effective-GB/s report; keep in
|
|
* sync with fc_cuda.c. */
|
|
#define Q4K_NX8_MIN_N_DEFAULT 32
|
|
|
|
static int nx_tile(int t, int n)
|
|
{
|
|
const char *e = getenv("FC_NX1_DIRECT");
|
|
if (!(e && e[0] == '0')) {
|
|
if (n == 1)
|
|
return 1; /* L10: n==1 dispatches single-column */
|
|
if (n == 2)
|
|
return 2; /* L10: n==2 uses the 2-wide tile */
|
|
}
|
|
if (t == 12) {
|
|
const char *m = getenv("FC_Q4K_NX8_MIN_N");
|
|
int min_n = m ? atoi(m) : Q4K_NX8_MIN_N_DEFAULT;
|
|
return (min_n > 0 && n >= min_n) ? 8 : 4;
|
|
}
|
|
return n >= 5 ? 8 : 4;
|
|
}
|
|
|
|
static double now_s(void)
|
|
{
|
|
struct timespec ts;
|
|
clock_gettime(CLOCK_MONOTONIC, &ts);
|
|
return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
|
|
}
|
|
|
|
static size_t raw_block_bytes(int t)
|
|
{
|
|
switch (t) {
|
|
case 8: return 34;
|
|
case 11: return 110;
|
|
case 12: return 144;
|
|
case 13: return 176;
|
|
case 14: return 210;
|
|
default: return 0;
|
|
}
|
|
}
|
|
|
|
/* Clamp fp16 block scales to exponent 2^-4..2^-3 (random halves can be
|
|
* inf/nan, which breaks memcmp compares). Mirrors bench_matvec.c. */
|
|
static void sanitize_quant(int t, unsigned char *p, size_t bytes)
|
|
{
|
|
size_t bs = raw_block_bytes(t);
|
|
size_t off;
|
|
for (off = 0; bs && off + bs <= bytes; off += bs) {
|
|
unsigned char *b = p + off;
|
|
unsigned short h = (unsigned short)(0x2C00 | (rand() & 0x3FF));
|
|
unsigned short h2 = (unsigned short)(0x2C00 | (rand() & 0x3FF));
|
|
switch (t) {
|
|
case 11: b[108] = h & 0xFF; b[109] = h >> 8; break;
|
|
case 14: b[208] = h & 0xFF; b[209] = h >> 8; break;
|
|
case 12: case 13:
|
|
b[0] = h & 0xFF; b[1] = h >> 8;
|
|
b[2] = h2 & 0xFF; b[3] = h2 >> 8;
|
|
break;
|
|
case 8: b[0] = h & 0xFF; b[1] = h >> 8; break;
|
|
}
|
|
}
|
|
}
|
|
|
|
static unsigned char *rand_bytes(size_t bytes)
|
|
{
|
|
unsigned char *p = malloc(bytes);
|
|
size_t i;
|
|
for (i = 0; i < bytes; i++)
|
|
p[i] = (unsigned char)(rand() & 0xFF);
|
|
return p;
|
|
}
|
|
|
|
static float *rand_floats(size_t n)
|
|
{
|
|
float *p = malloc(n * sizeof(float));
|
|
size_t i;
|
|
for (i = 0; i < n; i++)
|
|
p[i] = (float)rand() / (float)RAND_MAX - 0.5f;
|
|
return p;
|
|
}
|
|
|
|
/* Random raw quant weight -> repacked device buffer (as a model upload). */
|
|
static fc_buffer *make_weight(int t, long rows, int cols, int *dev_stride)
|
|
{
|
|
size_t nb = (size_t)((t == 8) ? cols / 32 : cols / 256);
|
|
size_t raw_stride = raw_block_bytes(t) * nb;
|
|
unsigned char *raw = rand_bytes(raw_stride * (size_t)rows);
|
|
fc_buffer *b = NULL;
|
|
|
|
sanitize_quant(t, raw, raw_stride * (size_t)rows);
|
|
if (t == 0) { /* F32 rows: no repack */
|
|
free(raw);
|
|
return NULL;
|
|
}
|
|
{
|
|
size_t rs = fc_repack_row_bytes(t, cols);
|
|
unsigned char *tmp = malloc(rs * (size_t)rows);
|
|
if (tmp && fc_repack_weight(t, raw, tmp, rows, cols) == 0)
|
|
b = fc_buffer_from_ptr(tmp, rs * (size_t)rows);
|
|
free(tmp);
|
|
*dev_stride = (int)rs;
|
|
}
|
|
free(raw);
|
|
return b;
|
|
}
|
|
|
|
typedef void (*mv_fn)(fc_cuda_ctx *, fc_buffer *, fc_buffer *, fc_buffer *,
|
|
int, int, int);
|
|
typedef void (*mv_nx_fn)(fc_cuda_ctx *, fc_buffer *, fc_buffer *, fc_buffer *,
|
|
int, int, int, int);
|
|
|
|
/* Verify + bench one type: n-column nx output vs n single-column calls.
|
|
* Single-column results are computed one x column at a time into a per-
|
|
* column device buffer, exactly like the single-token forward would.
|
|
* Returns 0 on bit-exact PASS. */
|
|
static int check_nx(const char *name, int t, mv_fn single, mv_nx_fn nx,
|
|
int rows, int cols, int n, int iters)
|
|
{
|
|
fc_cuda_ctx *ctx = fc_default_ctx();
|
|
int stride = 0;
|
|
fc_buffer *w = NULL;
|
|
int fail = 0, tk, i;
|
|
|
|
if (t == 0) { /* F32 weight */
|
|
float *wh = rand_floats((size_t)rows * cols);
|
|
w = fc_buffer_from_ptr(wh, (size_t)rows * cols * sizeof(float));
|
|
free(wh);
|
|
stride = cols * (int)sizeof(float);
|
|
} else {
|
|
w = make_weight(t, rows, cols, &stride);
|
|
}
|
|
if (!w) {
|
|
printf("%-14s SKIP (weight alloc failed)\n", name);
|
|
return 1;
|
|
}
|
|
|
|
float *xh = rand_floats((size_t)n * cols);
|
|
fc_buffer *xall = fc_buffer_from_ptr(xh, (size_t)n * cols * sizeof(float));
|
|
fc_buffer *x1 = fc_buffer_alloc((size_t)cols * sizeof(float));
|
|
fc_buffer *y1 = fc_buffer_alloc((size_t)rows * sizeof(float));
|
|
fc_buffer *yn = fc_buffer_alloc((size_t)n * rows * sizeof(float));
|
|
float *ref = malloc((size_t)n * rows * sizeof(float));
|
|
float *got = malloc((size_t)n * rows * sizeof(float));
|
|
|
|
/* Reference: n single-column launches. */
|
|
for (tk = 0; tk < n; tk++) {
|
|
fc_buffer_from_cpu(x1, xh + (size_t)tk * cols);
|
|
single(ctx, w, x1, y1, rows, cols, stride);
|
|
fc_sync();
|
|
fc_buffer_to_cpu(y1, ref + (size_t)tk * rows);
|
|
}
|
|
|
|
/* Batched: one nx launch. */
|
|
nx(ctx, w, xall, yn, rows, cols, stride, n);
|
|
fc_sync();
|
|
fc_buffer_to_cpu(yn, got);
|
|
|
|
for (tk = 0; tk < n && !fail; tk++) {
|
|
if (memcmp(ref + (size_t)tk * rows, got + (size_t)tk * rows,
|
|
(size_t)rows * sizeof(float)) != 0) {
|
|
int r;
|
|
for (r = 0; r < rows; r++) {
|
|
float a = ref[(size_t)tk * rows + r];
|
|
float b = got[(size_t)tk * rows + r];
|
|
if (memcmp(&a, &b, 4) != 0) {
|
|
printf("%-14s FAIL col %d row %d: single=%.9g nx=%.9g\n",
|
|
name, tk, r, a, b);
|
|
break;
|
|
}
|
|
}
|
|
fail = 1;
|
|
}
|
|
}
|
|
|
|
if (!fail) {
|
|
/* Bench: nx pass vs n sequential single passes (device-resident x). */
|
|
double t0, dt_nx, dt_1;
|
|
nx(ctx, w, xall, yn, rows, cols, stride, n);
|
|
fc_sync();
|
|
t0 = now_s();
|
|
for (i = 0; i < iters; i++)
|
|
nx(ctx, w, xall, yn, rows, cols, stride, n);
|
|
fc_sync();
|
|
dt_nx = (now_s() - t0) / iters;
|
|
|
|
t0 = now_s();
|
|
for (i = 0; i < iters; i++)
|
|
for (tk = 0; tk < n; tk++)
|
|
single(ctx, w, xall, y1, rows, cols, stride);
|
|
fc_sync();
|
|
dt_1 = (now_s() - t0) / iters;
|
|
|
|
printf("%-14s PASS %6dx%-5d n=%-3d nx %8.3f ms %dx1 %8.3f ms "
|
|
"speedup %.2fx (%6.2f GB/s weight eff.)\n",
|
|
name, rows, cols, n, dt_nx * 1e3, n, dt_1 * 1e3, dt_1 / dt_nx,
|
|
(double)rows * stride * ((n + nx_tile(t, n) - 1) / nx_tile(t, n))
|
|
/ dt_nx / 1e9);
|
|
}
|
|
|
|
free(xh); free(ref); free(got);
|
|
fc_buffer_free(w); fc_buffer_free(xall);
|
|
fc_buffer_free(x1); fc_buffer_free(y1); fc_buffer_free(yn);
|
|
return fail;
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
int fails = 0;
|
|
|
|
if (fc_init() != 0) {
|
|
fprintf(stderr, "no CUDA\n");
|
|
return 1;
|
|
}
|
|
srand(42);
|
|
|
|
printf("== fc_matvec_*_nx vs single-column (bit-exact gate) ==\n");
|
|
/* Model hot shapes; the L6 dual-width dispatch means n <= 4 exercises
|
|
* the 4-wide tile and n >= 5 the 8-wide one — cover both, plus partial
|
|
* tiles (n=2/6/12), one 8-wide tile exactly (n=8 — the K=8
|
|
* batched-decode shape L6 targets), and multi-tile batches. */
|
|
/* L10 n=1: the _nx entry points dispatch the dedicated single-column
|
|
* kernel (FC_NX1_DIRECT, default on) because a 4-wide tile fed one live
|
|
* column runs 4x the FMA chains. Run the suite a SECOND time with
|
|
* FC_NX1_DIRECT=0 to gate the invariant the dispatch rests on — the
|
|
* tiled kernel at n=1 must equal the single-column kernel bit for bit.
|
|
* The n=1 speedup column prints what the dispatch buys. */
|
|
fails += check_nx("q3k n=1", 11, fc_matvec_q3k_ctx, fc_matvec_q3k_nx_ctx,
|
|
4096, 2048, 1, 40);
|
|
fails += check_nx("q4k n=1", 12, fc_matvec_q4k_ctx, fc_matvec_q4k_nx_ctx,
|
|
2048, 512, 1, 40);
|
|
fails += check_nx("q5k n=1", 13, fc_matvec_q5k_ctx, fc_matvec_q5k_nx_ctx,
|
|
8192, 2048, 1, 20);
|
|
fails += check_nx("q6k head n=1", 14, fc_matvec_q6k_ctx,
|
|
fc_matvec_q6k_nx_ctx, 32768, 2048, 1, 10);
|
|
fails += check_nx("q80 n=1", 8, fc_matvec_q80_ctx, fc_matvec_q80_nx_ctx,
|
|
4096, 2048, 1, 40);
|
|
fails += check_nx("f32 n=1", 0, fc_matvec_f32_ctx, fc_matvec_f32_nx_ctx,
|
|
256, 2048, 1, 40);
|
|
fails += check_nx("q3k n=2", 11, fc_matvec_q3k_ctx, fc_matvec_q3k_nx_ctx,
|
|
4096, 2048, 2, 20);
|
|
fails += check_nx("q6k head n=2", 14, fc_matvec_q6k_ctx,
|
|
fc_matvec_q6k_nx_ctx, 32768, 2048, 2, 10);
|
|
fails += check_nx("q4k n=2", 12, fc_matvec_q4k_ctx, fc_matvec_q4k_nx_ctx,
|
|
2048, 512, 2, 20);
|
|
fails += check_nx("q3k n=4", 11, fc_matvec_q3k_ctx, fc_matvec_q3k_nx_ctx,
|
|
4096, 2048, 4, 20);
|
|
fails += check_nx("q3k n=8", 11, fc_matvec_q3k_ctx, fc_matvec_q3k_nx_ctx,
|
|
4096, 2048, 8, 20);
|
|
fails += check_nx("q6k n=4", 14, fc_matvec_q6k_ctx, fc_matvec_q6k_nx_ctx,
|
|
4096, 2048, 4, 10);
|
|
fails += check_nx("q6k n=8", 14, fc_matvec_q6k_ctx, fc_matvec_q6k_nx_ctx,
|
|
4096, 2048, 8, 10);
|
|
fails += check_nx("q3k n=12", 11, fc_matvec_q3k_ctx, fc_matvec_q3k_nx_ctx,
|
|
4096, 2048, 12, 20);
|
|
fails += check_nx("q3k n=6", 11, fc_matvec_q3k_ctx, fc_matvec_q3k_nx_ctx,
|
|
512, 2048, 6, 20);
|
|
fails += check_nx("q3k down n=32", 11, fc_matvec_q3k_ctx,
|
|
fc_matvec_q3k_nx_ctx, 2048, 512, 32, 20);
|
|
fails += check_nx("q4k n=12", 12, fc_matvec_q4k_ctx, fc_matvec_q4k_nx_ctx,
|
|
2048, 512, 12, 20);
|
|
fails += check_nx("q5k n=12", 13, fc_matvec_q5k_ctx, fc_matvec_q5k_nx_ctx,
|
|
8192, 2048, 12, 10);
|
|
fails += check_nx("q6k n=12", 14, fc_matvec_q6k_ctx, fc_matvec_q6k_nx_ctx,
|
|
4096, 2048, 12, 10);
|
|
fails += check_nx("q80 n=12", 8, fc_matvec_q80_ctx, fc_matvec_q80_nx_ctx,
|
|
4096, 2048, 12, 20);
|
|
fails += check_nx("f32 n=12", 0, fc_matvec_f32_ctx, fc_matvec_f32_nx_ctx,
|
|
256, 2048, 12, 20);
|
|
fails += check_nx("q5k n=128", 13, fc_matvec_q5k_ctx, fc_matvec_q5k_nx_ctx,
|
|
8192, 2048, 128, 4);
|
|
|
|
/* PREFILL WIDTHS — the shapes and the n that decide FC_Q4K_NX8_MIN_N.
|
|
* QWEN35_BS_PREFILL_CHUNK is 128 tokens, so n=128 is what every dense
|
|
* Q4_K projection sees during a prefill; the n=12 measurement that
|
|
* disabled fc_matvec_q4k_nx8 outright is two orders of magnitude away
|
|
* from it. Shapes are saber's own Q4_K dense tensors (GGUF census:
|
|
* attn_qkv x28, ffn_down_shexp x38, attn_output x10, attn_v x10 per
|
|
* forward; ffn_down_exps goes through the grouped-MoE path instead).
|
|
* Run the whole suite a SECOND time with FC_Q4K_NX8_MIN_N=0 to gate the
|
|
* 4-wide fallback: both tiles must be bit-exact per column. */
|
|
fails += check_nx("q4k qkv n=128", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 8192, 2048, 128, 10);
|
|
fails += check_nx("q4k dshexp n=128", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 2048, 512, 128, 10);
|
|
fails += check_nx("q4k aout n=128", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 2048, 4096, 128, 10);
|
|
fails += check_nx("q4k av n=128", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 512, 2048, 128, 10);
|
|
fails += check_nx("q4k qkv n=256", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 8192, 2048, 256, 6);
|
|
/* The crossover itself: FC_Q4K_NX8_MIN_N has to hold for EVERY width a
|
|
* prefill can present, because the last chunk is total_tokens % 128. It is
|
|
* not a pure function of n either -- 8-wide loses 22% at n=12 on the narrow
|
|
* 2048x512 shexp shape and wins ~7% at the same-ish n on the wide
|
|
* 8192x2048 qkv shape -- so both shapes are swept across the boundary. */
|
|
fails += check_nx("q4k qkv n=16", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 8192, 2048, 16, 20);
|
|
fails += check_nx("q4k qkv n=24", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 8192, 2048, 24, 20);
|
|
fails += check_nx("q4k qkv n=32", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 8192, 2048, 32, 20);
|
|
fails += check_nx("q4k qkv n=12", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 8192, 2048, 12, 20);
|
|
fails += check_nx("q4k dshexp n=16", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 2048, 512, 16, 40);
|
|
fails += check_nx("q4k dshexp n=32", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 2048, 512, 32, 40);
|
|
fails += check_nx("q4k dshexp n=64", 12, fc_matvec_q4k_ctx,
|
|
fc_matvec_q4k_nx_ctx, 2048, 512, 64, 20);
|
|
|
|
printf(fails ? "RESULT: FAIL\n" : "RESULT: PASS\n");
|
|
fc_free();
|
|
return fails ? 1 : 0;
|
|
}
|