Salvatore Sanfilippo
06f7135d0f
Guard native cross-session attention, grouped expert, and Q8 batch GEMM tests with __APPLE__. CUDA retains its supported scalar, prefill, and recurrent kernel oracles instead of failing to link against Metal-only entry points.
3719 lines
214 KiB
C
3719 lines
214 KiB
C
/* GPU kernel tests for Qwen3.8-Flash-Next: every kernel
|
|
* at the release and mini model shapes against a double-precision reference.
|
|
* Build: make test-qwen4-kernels */
|
|
|
|
#include <math.h>
|
|
#include <stdbool.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
#include <sys/mman.h>
|
|
|
|
#include "ds4.h"
|
|
#include "ds4_gpu.h"
|
|
|
|
bool ds4_log_is_tty(FILE *fp) {
|
|
(void)fp;
|
|
return false;
|
|
}
|
|
|
|
static uint32_t g_rng = 0x9e3779b9u;
|
|
|
|
static float frand(void) {
|
|
g_rng ^= g_rng << 13;
|
|
g_rng ^= g_rng >> 17;
|
|
g_rng ^= g_rng << 5;
|
|
return ((float)(g_rng & 0xffffffu) / 8388608.0f) - 1.0f;
|
|
}
|
|
|
|
static void require_ok(int ok, const char *what) {
|
|
if (!ok) {
|
|
fprintf(stderr, "%s failed\n", what);
|
|
exit(1);
|
|
}
|
|
}
|
|
|
|
static void check_close(const char *what, const float *got, const double *ref, uint64_t n, double tol) {
|
|
double worst = 0.0, scale = 1e-6;
|
|
uint64_t worst_i = 0;
|
|
for (uint64_t i = 0; i < n; i++) {
|
|
if (!isfinite(got[i])) {
|
|
fprintf(stderr, "%s: non-finite value at %llu\n", what, (unsigned long long)i);
|
|
exit(1);
|
|
}
|
|
const double d = fabs((double)got[i] - ref[i]);
|
|
if (d > worst) { worst = d; worst_i = i; }
|
|
if (fabs(ref[i]) > scale) scale = fabs(ref[i]);
|
|
}
|
|
if (worst > tol * scale) {
|
|
fprintf(stderr, "%s: max|d| %.3e (rel %.3e) at %llu: got %.6f ref %.6f\n",
|
|
what, worst, worst / scale, (unsigned long long)worst_i, got[worst_i], ref[worst_i]);
|
|
exit(1);
|
|
}
|
|
printf(" %-44s ok max|d|=%.2e (scale %.2e)\n", what, worst, scale);
|
|
}
|
|
|
|
static uint16_t f32_to_f16(float f) {
|
|
union { float f; uint32_t u; } v = { f };
|
|
const uint32_t sign = (v.u >> 16) & 0x8000u;
|
|
int32_t exp = (int32_t)((v.u >> 23) & 0xffu) - 127 + 15;
|
|
uint32_t mant = v.u & 0x7fffffu;
|
|
if (exp <= 0) {
|
|
if (exp < -10) return (uint16_t)sign;
|
|
mant |= 0x800000u;
|
|
const uint32_t shift = (uint32_t)(14 - exp);
|
|
uint32_t half = mant >> shift;
|
|
if ((mant >> (shift - 1)) & 1u) half++;
|
|
return (uint16_t)(sign | half);
|
|
}
|
|
if (exp >= 31) return (uint16_t)(sign | 0x7c00u);
|
|
uint32_t half = sign | ((uint32_t)exp << 10) | (mant >> 13);
|
|
if (mant & 0x1000u) half++;
|
|
return (uint16_t)half;
|
|
}
|
|
|
|
static float f16_to_f32(uint16_t h) {
|
|
const uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
|
|
uint32_t exp = (h >> 10) & 0x1fu;
|
|
uint32_t mant = h & 0x3ffu;
|
|
union { uint32_t u; float f; } v;
|
|
if (exp == 0) {
|
|
if (mant == 0) { v.u = sign; return v.f; }
|
|
exp = 127 - 15 + 1;
|
|
while (!(mant & 0x400u)) { mant <<= 1; exp--; }
|
|
mant &= 0x3ffu;
|
|
v.u = sign | (exp << 23) | (mant << 13);
|
|
return v.f;
|
|
}
|
|
if (exp == 31) { v.u = sign | 0x7f800000u | (mant << 13); return v.f; }
|
|
v.u = sign | ((exp + 127 - 15) << 23) | (mant << 13);
|
|
return v.f;
|
|
}
|
|
|
|
static double sigmoid_d(double x) { return x >= 0 ? 1.0 / (1.0 + exp(-x)) : exp(x) / (1.0 + exp(x)); }
|
|
static double silu_d(double x) { return x * sigmoid_d(x); }
|
|
static double softplus_d(double x) { return x > 20.0 ? x : log1p(exp(x)); }
|
|
|
|
/* ---- weight arena ---- */
|
|
|
|
/* anonymous mmap standing in for the model map; weights are appended, never freed */
|
|
typedef struct {
|
|
uint8_t *base;
|
|
uint64_t size;
|
|
uint64_t used;
|
|
} arena_t;
|
|
|
|
static uint64_t arena_alloc(arena_t *a, uint64_t bytes) {
|
|
const uint64_t off = (a->used + 63u) & ~63ull;
|
|
if (off + bytes > a->size) {
|
|
fprintf(stderr, "arena exhausted\n");
|
|
exit(1);
|
|
}
|
|
a->used = off + bytes;
|
|
return off;
|
|
}
|
|
|
|
/* f32 weights with a double shadow for the reference */
|
|
static uint64_t arena_f32(arena_t *a, uint64_t n, double **shadow, float lo, float hi) {
|
|
const uint64_t off = arena_alloc(a, n * sizeof(float));
|
|
float *w = (float *)(a->base + off);
|
|
*shadow = malloc(n * sizeof(double));
|
|
for (uint64_t i = 0; i < n; i++) {
|
|
w[i] = lo + (hi - lo) * (0.5f * frand() + 0.5f);
|
|
(*shadow)[i] = w[i];
|
|
}
|
|
return off;
|
|
}
|
|
|
|
static uint64_t arena_f16(arena_t *a, uint64_t n, double **shadow, float scale) {
|
|
const uint64_t off = arena_alloc(a, n * sizeof(uint16_t));
|
|
uint16_t *w = (uint16_t *)(a->base + off);
|
|
*shadow = malloc(n * sizeof(double));
|
|
for (uint64_t i = 0; i < n; i++) {
|
|
w[i] = f32_to_f16(scale * frand());
|
|
(*shadow)[i] = f16_to_f32(w[i]);
|
|
}
|
|
return off;
|
|
}
|
|
|
|
static uint64_t arena_bf16(arena_t *a, uint64_t n, double **shadow, float scale) {
|
|
const uint64_t off = arena_alloc(a, n * sizeof(uint16_t));
|
|
uint16_t *w = (uint16_t *)(a->base + off);
|
|
*shadow = malloc(n * sizeof(double));
|
|
for (uint64_t i = 0; i < n; i++) {
|
|
const float v = scale * frand();
|
|
uint32_t u;
|
|
memcpy(&u, &v, 4);
|
|
w[i] = (uint16_t)(u >> 16);
|
|
const uint32_t back = (uint32_t)w[i] << 16;
|
|
float r;
|
|
memcpy(&r, &back, 4);
|
|
(*shadow)[i] = r;
|
|
}
|
|
return off;
|
|
}
|
|
|
|
/* q4_0 rows: 18-byte blocks of 32 (f16 scale, 16 nibble bytes; low nibbles first) */
|
|
static uint64_t arena_q4_0(arena_t *a, uint64_t rows, uint64_t cols, double **shadow, float scale) {
|
|
const uint64_t blocks = cols / 32;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 18u);
|
|
uint8_t *w = a->base + off;
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
float vals[32];
|
|
float amax = 0.0f, maxv = 0.0f;
|
|
for (int j = 0; j < 32; j++) {
|
|
vals[j] = scale * frand();
|
|
if (fabsf(vals[j]) > amax) { amax = fabsf(vals[j]); maxv = vals[j]; }
|
|
}
|
|
const float d = maxv / -8.0f;
|
|
const uint16_t dh = f32_to_f16(d);
|
|
const float dq = f16_to_f32(dh);
|
|
const float id = dq ? 1.0f / dq : 0.0f;
|
|
uint8_t *blk = w + (r * blocks + b) * 18u;
|
|
memcpy(blk, &dh, 2);
|
|
for (int j = 0; j < 16; j++) {
|
|
int q0 = (int)(vals[j] * id + 8.5f), q1 = (int)(vals[j + 16] * id + 8.5f);
|
|
if (q0 < 0) q0 = 0;
|
|
if (q0 > 15) q0 = 15;
|
|
if (q1 < 0) q1 = 0;
|
|
if (q1 > 15) q1 = 15;
|
|
blk[2 + j] = (uint8_t)(q0 | (q1 << 4));
|
|
(*shadow)[r * cols + b * 32 + j] = dq * (q0 - 8);
|
|
(*shadow)[r * cols + b * 32 + 16 + j] = dq * (q1 - 8);
|
|
}
|
|
}
|
|
}
|
|
return off;
|
|
}
|
|
|
|
/* q4_K rows: 144-byte super-blocks of 256 (d, dmin, 6-bit scales/mins, nibbles) */
|
|
static uint64_t arena_q4_K(arena_t *a, uint64_t rows, uint64_t cols, double **shadow, float scale) {
|
|
const uint64_t blocks = cols / 256;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 144u);
|
|
uint8_t *w = a->base + off;
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
uint8_t *blk = w + (r * blocks + b) * 144u;
|
|
/* per-group scale/min in 6 bits, block d/dmin as f16 */
|
|
const float d = scale / 63.0f / 15.0f, dmin = d;
|
|
const uint16_t dh = f32_to_f16(d), mh = f32_to_f16(dmin);
|
|
const float dq = f16_to_f32(dh), mq = f16_to_f32(mh);
|
|
memcpy(blk, &dh, 2);
|
|
memcpy(blk + 2, &mh, 2);
|
|
uint8_t sc[8], mn[8];
|
|
for (int g = 0; g < 8; g++) {
|
|
sc[g] = (uint8_t)(1 + (int)(62.0f * (0.5f * frand() + 0.5f)));
|
|
mn[g] = (uint8_t)((int)(63.0f * (0.5f * frand() + 0.5f)));
|
|
}
|
|
uint8_t *s = blk + 4;
|
|
for (int g = 0; g < 4; g++) { s[g] = sc[g] & 63; s[g + 4] = mn[g] & 63; }
|
|
for (int g = 4; g < 8; g++) {
|
|
s[g + 4] = (uint8_t)((sc[g] & 0xF) | ((mn[g] & 0xF) << 4));
|
|
s[g - 4] |= (uint8_t)((sc[g] >> 4) << 6);
|
|
s[g] |= (uint8_t)((mn[g] >> 4) << 6);
|
|
}
|
|
memset(blk + 16, 0, 128);
|
|
for (int g = 0; g < 8; g++) {
|
|
for (int j = 0; j < 32; j++) {
|
|
const int q = (int)(15.0f * (0.5f * frand() + 0.5f));
|
|
blk[16 + (g >> 1) * 32 + j] |= (uint8_t)(q << ((g & 1) * 4));
|
|
(*shadow)[r * cols + b * 256 + g * 32 + j] = dq * sc[g] * q - mq * mn[g];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return off;
|
|
}
|
|
|
|
/* q4_K rows with dyadic dequantized values (d a power of two, mins 0), so
|
|
* d*sc*q is exactly representable in half and the weight-rounding share of
|
|
* the error vanishes; the residual is activation rounding plus accumulation */
|
|
static uint64_t arena_q4_K_dyadic(arena_t *a, uint64_t rows, uint64_t cols, double **shadow) {
|
|
const uint64_t blocks = cols / 256;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 144u);
|
|
uint8_t *w = a->base + off;
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
uint8_t *blk = w + (r * blocks + b) * 144u;
|
|
const uint16_t dh = f32_to_f16(ldexpf(1.0f, -9));
|
|
const float dq = f16_to_f32(dh);
|
|
memcpy(blk, &dh, 2);
|
|
memcpy(blk + 2, &dh, 2);
|
|
uint8_t sc[8] = {0};
|
|
for (int g = 0; g < 8; g++) sc[g] = (uint8_t)(1 + (int)(31.0f * (0.5f * frand() + 0.5f)));
|
|
uint8_t *s = blk + 4;
|
|
memset(blk + 4, 0, 12);
|
|
for (int g = 0; g < 4; g++) s[g] = sc[g] & 63;
|
|
for (int g = 4; g < 8; g++) {
|
|
s[g + 4] = (uint8_t)(sc[g] & 0xF);
|
|
s[g - 4] |= (uint8_t)((sc[g] >> 4) << 6);
|
|
}
|
|
memset(blk + 16, 0, 128);
|
|
for (int g = 0; g < 8; g++) {
|
|
for (int j = 0; j < 32; j++) {
|
|
const int q = (int)(15.0f * (0.5f * frand() + 0.5f));
|
|
blk[16 + (g >> 1) * 32 + j] |= (uint8_t)(q << ((g & 1) * 4));
|
|
(*shadow)[r * cols + b * 256 + g * 32 + j] = (double)dq * sc[g] * q;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return off;
|
|
}
|
|
static uint64_t arena_mxfp4(arena_t *a, uint64_t rows, uint64_t cols, double **shadow) {
|
|
|
|
static const double values[16] = {0, .5, 1, 1.5, 2, 3, 4, 6, -0.0, -.5, -1, -1.5, -2, -3, -4, -6};
|
|
const uint64_t blocks = cols / 32;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 17u);
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
uint8_t *blk = a->base + off + (r * blocks + b) * 17u;
|
|
blk[0] = (uint8_t)(118u + (r + b) % 4u);
|
|
const double scale = ldexp(1.0, (int)blk[0] - 127);
|
|
for (uint32_t i = 0; i < 16; i++) {
|
|
const uint32_t lo = (uint32_t)(r + b * 3u + i) & 15u;
|
|
const uint32_t hi = (uint32_t)(r * 7u + b + i * 3u) & 15u;
|
|
blk[1 + i] = (uint8_t)(lo | (hi << 4));
|
|
(*shadow)[r * cols + b * 32 + i] = scale * values[lo];
|
|
(*shadow)[r * cols + b * 32 + i + 16] = scale * values[hi];
|
|
}
|
|
}
|
|
}
|
|
return off;
|
|
}
|
|
|
|
static const uint8_t ksigns_iq2xs[128] = {
|
|
0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15,
|
|
144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159,
|
|
160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175,
|
|
48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63,
|
|
192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207,
|
|
80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95,
|
|
96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111,
|
|
240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,
|
|
};
|
|
static const uint64_t iq2xxs_grid[256] = {
|
|
0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08,
|
|
0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808,
|
|
0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819,
|
|
0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819,
|
|
0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b,
|
|
0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808,
|
|
0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08,
|
|
0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b,
|
|
0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819,
|
|
0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08,
|
|
0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808,
|
|
0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08,
|
|
0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808,
|
|
0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808,
|
|
0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919,
|
|
0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819,
|
|
0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08,
|
|
0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908,
|
|
0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819,
|
|
0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808,
|
|
0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808,
|
|
0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908,
|
|
0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808,
|
|
0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08,
|
|
0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819,
|
|
0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819,
|
|
0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819,
|
|
0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908,
|
|
0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19,
|
|
0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819,
|
|
0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b,
|
|
0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808,
|
|
0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908,
|
|
0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08,
|
|
0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08,
|
|
0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908,
|
|
0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819,
|
|
0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808,
|
|
0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808,
|
|
0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19,
|
|
0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819,
|
|
0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919,
|
|
0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b,
|
|
0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08,
|
|
0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808,
|
|
0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908,
|
|
0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b,
|
|
0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819,
|
|
0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08,
|
|
0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08,
|
|
0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808,
|
|
0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b,
|
|
0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b,
|
|
0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908,
|
|
0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819,
|
|
0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808,
|
|
0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908,
|
|
0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b,
|
|
0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808,
|
|
0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b,
|
|
0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b,
|
|
0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808,
|
|
0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19,
|
|
0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,
|
|
};
|
|
|
|
/* QWEN4_DUMP_QUANT=DIR: write the raw blocks and the shadow values for offline checks */
|
|
static void dump_quant(const char *type, const uint8_t *raw, uint64_t raw_bytes, const double *shadow, uint64_t n) {
|
|
const char *dir = getenv("QWEN4_DUMP_QUANT");
|
|
if (!dir) return;
|
|
char path[512];
|
|
snprintf(path, sizeof(path), "%s/%s.bin", dir, type);
|
|
FILE *f = fopen(path, "wb");
|
|
if (f) { fwrite(raw, 1, raw_bytes, f); fclose(f); }
|
|
snprintf(path, sizeof(path), "%s/%s.shadow", dir, type);
|
|
f = fopen(path, "wb");
|
|
if (f) { fwrite(shadow, sizeof(double), n, f); fclose(f); }
|
|
}
|
|
|
|
/* q2_K rows: 84-byte super-blocks of 256 (16 scale/min nibbles, 64 packed bytes, d, dmin) */
|
|
static uint64_t arena_q2_K(arena_t *a, uint64_t rows, uint64_t cols, double **shadow, float scale) {
|
|
const uint64_t blocks = cols / 256;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 84u);
|
|
uint8_t *w = a->base + off;
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
uint8_t *blk = w + (r * blocks + b) * 84u;
|
|
const uint16_t dh = f32_to_f16(scale / 15.0f / 3.0f), mh = dh;
|
|
const float dq = f16_to_f32(dh), mq = f16_to_f32(mh);
|
|
memcpy(blk + 80, &dh, 2);
|
|
memcpy(blk + 82, &mh, 2);
|
|
for (int g = 0; g < 16; g++) {
|
|
const int sc = (int)(15.0f * (0.5f * frand() + 0.5f)), mn = (int)(15.0f * (0.5f * frand() + 0.5f));
|
|
blk[g] = (uint8_t)(sc | (mn << 4));
|
|
}
|
|
memset(blk + 16, 0, 64);
|
|
for (int g = 0; g < 16; g++) {
|
|
const int q_base = 32 * (g / 8) + 16 * (g & 1), shift = ((g / 2) & 3) * 2;
|
|
for (int l = 0; l < 16; l++) {
|
|
const int q = (int)(3.0f * (0.5f * frand() + 0.5f)) & 3;
|
|
blk[16 + q_base + l] |= (uint8_t)(q << shift);
|
|
(*shadow)[r * cols + b * 256 + g * 16 + l] = dq * (blk[g] & 0xF) * q - mq * (blk[g] >> 4);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
dump_quant("q2_K", w, rows * blocks * 84u, *shadow, rows * cols);
|
|
return off;
|
|
}
|
|
|
|
/* iq2_xxs rows: 66-byte super-blocks of 256 (d, then per 32: 4 grid indices, 4x7 sign bits, 4-bit scale) */
|
|
static uint64_t arena_iq2_xxs(arena_t *a, uint64_t rows, uint64_t cols, double **shadow, float scale) {
|
|
const uint64_t blocks = cols / 256;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 66u);
|
|
uint8_t *w = a->base + off;
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
uint8_t *blk = w + (r * blocks + b) * 66u;
|
|
const uint16_t dh = f32_to_f16(scale / 16.0f);
|
|
const float dq = f16_to_f32(dh);
|
|
memcpy(blk, &dh, 2);
|
|
for (int ib32 = 0; ib32 < 8; ib32++) {
|
|
uint32_t aux_g = 0, aux_s = 0;
|
|
for (int j = 0; j < 4; j++) {
|
|
aux_g |= (uint32_t)(rand() & 0xFF) << (8 * j);
|
|
aux_s |= (uint32_t)(rand() & 0x7F) << (7 * j);
|
|
}
|
|
aux_s |= (uint32_t)(rand() & 0xF) << 28;
|
|
uint16_t q2[4] = { (uint16_t)aux_g, (uint16_t)(aux_g >> 16), (uint16_t)aux_s, (uint16_t)(aux_s >> 16) };
|
|
memcpy(blk + 2 + ib32 * 8, q2, 8);
|
|
const float dl = dq * (0.5f + (float)(aux_s >> 28)) * 0.25f;
|
|
for (int j = 0; j < 4; j++) {
|
|
const uint8_t *grid = (const uint8_t *)(iq2xxs_grid + ((aux_g >> (8 * j)) & 0xFF));
|
|
const uint32_t signs = ksigns_iq2xs[(aux_s >> (7 * j)) & 127];
|
|
for (int i = 0; i < 8; i++) {
|
|
(*shadow)[r * cols + b * 256 + ib32 * 32 + j * 8 + i] = dl * grid[i] * (((signs >> i) & 1) ? -1.0 : 1.0);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
dump_quant("iq2_xxs", w, rows * blocks * 66u, *shadow, rows * cols);
|
|
return off;
|
|
}
|
|
|
|
/* q8_0 rows of `cols` elements (cols % 32 == 0), `rows` rows */
|
|
static uint64_t arena_q8_0(arena_t *a, uint64_t rows, uint64_t cols, double **shadow, float scale) {
|
|
const uint64_t blocks = cols / 32;
|
|
const uint64_t off = arena_alloc(a, rows * blocks * 34u);
|
|
uint8_t *w = a->base + off;
|
|
*shadow = malloc(rows * cols * sizeof(double));
|
|
for (uint64_t r = 0; r < rows; r++) {
|
|
for (uint64_t b = 0; b < blocks; b++) {
|
|
float vals[32];
|
|
float amax = 0.0f;
|
|
for (int j = 0; j < 32; j++) {
|
|
vals[j] = scale * frand();
|
|
if (fabsf(vals[j]) > amax) amax = fabsf(vals[j]);
|
|
}
|
|
const float d = amax / 127.0f;
|
|
const uint16_t dh = f32_to_f16(d);
|
|
const float dq = f16_to_f32(dh);
|
|
uint8_t *blk = w + (r * blocks + b) * 34u;
|
|
memcpy(blk, &dh, 2);
|
|
for (int j = 0; j < 32; j++) {
|
|
int q = (int)lrintf(d > 0 ? vals[j] / d : 0.0f);
|
|
if (q > 127) q = 127;
|
|
if (q < -127) q = -127;
|
|
((int8_t *)blk)[2 + j] = (int8_t)q;
|
|
(*shadow)[r * cols + b * 32 + j] = (double)dq * q;
|
|
}
|
|
}
|
|
}
|
|
return off;
|
|
}
|
|
|
|
static ds4_gpu_tensor *upload(const float *data, uint64_t n) {
|
|
ds4_gpu_tensor *t = ds4_gpu_tensor_alloc(n * sizeof(float));
|
|
require_ok(t != NULL, "tensor alloc");
|
|
if (data) require_ok(ds4_gpu_tensor_write(t, 0, data, n * sizeof(float)), "tensor write");
|
|
else require_ok(ds4_gpu_tensor_fill_f32(t, 0.0f, n), "tensor fill");
|
|
return t;
|
|
}
|
|
|
|
static float *download(const ds4_gpu_tensor *t, uint64_t n) {
|
|
float *out = malloc(n * sizeof(float));
|
|
require_ok(ds4_gpu_tensor_read(t, 0, out, n * sizeof(float)), "tensor read");
|
|
return out;
|
|
}
|
|
|
|
static void check_tensor(const char *what, const ds4_gpu_tensor *t, const double *ref, uint64_t n, double tol) {
|
|
float *got = download(t, n);
|
|
check_close(what, got, ref, n, tol);
|
|
free(got);
|
|
}
|
|
|
|
static float *rand_vec(uint64_t n, float scale) {
|
|
float *v = malloc(n * sizeof(float));
|
|
for (uint64_t i = 0; i < n; i++) v[i] = scale * frand();
|
|
return v;
|
|
}
|
|
|
|
/* ---- hyper-connections ---- */
|
|
|
|
static void test_hc(arena_t *a, uint32_t E, uint32_t rank, uint32_t T, uint32_t wtype) {
|
|
const bool f16 = wtype == 1u, q8 = wtype == 8u;
|
|
const uint32_t hc = 4, dim = E * hc, CH = DS4_QWEN4_HC_CHUNKS;
|
|
const float eps = 1e-6f;
|
|
double *g_gamma, *g_down, *g_up, *g_inj;
|
|
const uint64_t gamma_off = arena_f32(a, dim, &g_gamma, 0.5f, 1.5f);
|
|
const uint64_t down_off = q8 ? arena_q8_0(a, rank, dim, &g_down, 0.05f)
|
|
: f16 ? arena_f16(a, (uint64_t)rank * dim, &g_down, 0.05f)
|
|
: arena_f32(a, (uint64_t)rank * dim, &g_down, -0.05f, 0.05f);
|
|
/* q8 up rows need rank % 32; smaller ranks keep f16 like the converter does */
|
|
const uint32_t up_type = q8 && (rank % 32) == 0 ? 8u : q8 ? 1u : wtype;
|
|
const uint64_t up_off = up_type == 8u ? arena_q8_0(a, dim, rank, &g_up, 0.2f)
|
|
: up_type == 1u ? arena_f16(a, (uint64_t)dim * rank, &g_up, 0.2f)
|
|
: arena_f32(a, (uint64_t)dim * rank, &g_up, -0.2f, 0.2f);
|
|
const uint64_t inj_off = q8 ? arena_q8_0(a, hc, dim, &g_inj, 0.05f)
|
|
: f16 ? arena_f16(a, (uint64_t)hc * dim, &g_inj, 0.05f)
|
|
: arena_f32(a, (uint64_t)hc * dim, &g_inj, -0.05f, 0.05f);
|
|
float *R = rand_vec((uint64_t)T * dim, 1.0f);
|
|
float *blk = rand_vec((uint64_t)T * E, 1.0f);
|
|
|
|
double *xn = malloc((uint64_t)T * dim * sizeof(double));
|
|
double *lo = malloc((uint64_t)T * rank * sizeof(double));
|
|
double *part = malloc((uint64_t)T * hc * CH * hc * sizeof(double));
|
|
double *mixed = malloc((uint64_t)T * E * sizeof(double));
|
|
double *R2 = malloc((uint64_t)T * dim * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < hc; s++) {
|
|
double ss = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) ss += (double)R[t * dim + s * E + i] * R[t * dim + s * E + i];
|
|
const double inv = 1.0 / sqrt(ss / E + eps);
|
|
for (uint32_t i = 0; i < E; i++) xn[t * dim + s * E + i] = R[t * dim + s * E + i] * inv * g_gamma[s * E + i];
|
|
}
|
|
for (uint32_t r = 0; r < rank; r++) {
|
|
double acc = 0.0;
|
|
for (uint32_t i = 0; i < dim; i++) acc += g_down[(uint64_t)r * dim + i] * xn[t * dim + i];
|
|
lo[t * rank + r] = acc; /* raw: gate_mix applies silu(./hc) */
|
|
}
|
|
double inj[4];
|
|
const uint32_t per = (E + CH - 1) / CH;
|
|
for (uint32_t j = 0; j < hc; j++) {
|
|
double tot = 0.0;
|
|
for (uint32_t s = 0; s < hc; s++) {
|
|
for (uint32_t c = 0; c < CH; c++) {
|
|
double acc = 0.0;
|
|
for (uint32_t i = c * per; i < E && i < (c + 1) * per; i++)
|
|
acc += g_inj[(uint64_t)j * dim + s * E + i] * xn[t * dim + s * E + i];
|
|
part[((uint64_t)t * hc * CH + s * CH + c) * hc + j] = acc;
|
|
tot += acc;
|
|
}
|
|
}
|
|
inj[j] = 2.0 * sigmoid_d(tot / hc);
|
|
}
|
|
for (uint32_t d = 0; d < E; d++) {
|
|
double m = 0.0;
|
|
for (uint32_t s = 0; s < hc; s++) {
|
|
double acc = 0.0;
|
|
for (uint32_t r = 0; r < rank; r++) acc += g_up[(uint64_t)(s * E + d) * rank + r] * silu_d(lo[t * rank + r] / hc);
|
|
m += sigmoid_d(acc) * xn[t * dim + s * E + d];
|
|
}
|
|
mixed[t * E + d] = m / hc;
|
|
}
|
|
for (uint32_t s = 0; s < hc; s++)
|
|
for (uint32_t d = 0; d < E; d++)
|
|
R2[t * dim + s * E + d] = R[t * dim + s * E + d] + inj[s] * blk[t * E + d];
|
|
}
|
|
|
|
ds4_gpu_tensor *gR = upload(R, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *gxn = upload(NULL, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *glo = upload(NULL, (uint64_t)T * rank);
|
|
ds4_gpu_tensor *ginj = upload(NULL, (uint64_t)T * hc * CH * hc);
|
|
ds4_gpu_tensor *gmixed = upload(NULL, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gblk = upload(blk, (uint64_t)T * E);
|
|
require_ok(ds4_gpu_qwen4_hc_norm_tensor(gxn, ginj, gR, a->base, a->size, gamma_off, inj_off, wtype, T, E, hc, hc, eps),
|
|
"hc norm");
|
|
#ifdef __APPLE__
|
|
require_ok(q8 ? ds4_gpu_matmul_q8_0_tensor(glo, a->base, a->size, down_off, dim, rank, gxn, T)
|
|
: f16 ? ds4_gpu_matmul_f16_tensor(glo, a->base, a->size, down_off, dim, rank, gxn, T)
|
|
: ds4_gpu_matmul_f32_tensor(glo, a->base, a->size, down_off, dim, rank, gxn, T), "hc down gemv");
|
|
#else
|
|
require_ok(ds4_gpu_qwen4_dense_mm_tensor(glo, gxn, a->base, a->size,
|
|
down_off, wtype, T, dim, rank), "hc down projection");
|
|
#endif
|
|
require_ok(ds4_gpu_qwen4_hc_gate_mix_tensor(gmixed, gxn, glo, a->base, a->size, up_off, up_type, T, E, hc, rank),
|
|
"hc gate mix");
|
|
require_ok(ds4_gpu_qwen4_hc_combine_tensor(gR, gblk, ginj, T, E, hc), "hc combine");
|
|
char name[96];
|
|
const char *tname = q8 ? "q8_0" : f16 ? "f16" : "f32";
|
|
snprintf(name, sizeof(name), "hc E=%u rank=%u T=%u %s: xn", E, rank, T, tname);
|
|
check_tensor(name, gxn, xn, (uint64_t)T * dim, 1e-5);
|
|
snprintf(name, sizeof(name), "hc E=%u rank=%u T=%u %s: lowrank", E, rank, T, tname);
|
|
check_tensor(name, glo, lo, (uint64_t)T * rank, 2e-4); /* f16 low-rank weights */
|
|
snprintf(name, sizeof(name), "hc E=%u rank=%u T=%u %s: inject partials", E, rank, T, tname);
|
|
check_tensor(name, ginj, part, (uint64_t)T * hc * CH * hc, 2e-5);
|
|
snprintf(name, sizeof(name), "hc E=%u rank=%u T=%u %s: mixed", E, rank, T, tname);
|
|
check_tensor(name, gmixed, mixed, (uint64_t)T * E, 2e-5);
|
|
snprintf(name, sizeof(name), "hc E=%u rank=%u T=%u %s: combine", E, rank, T, tname);
|
|
check_tensor(name, gR, R2, (uint64_t)T * dim, 2e-5);
|
|
ds4_gpu_tensor_free(gblk); ds4_gpu_tensor_free(gmixed); ds4_gpu_tensor_free(ginj);
|
|
ds4_gpu_tensor_free(glo); ds4_gpu_tensor_free(gxn); ds4_gpu_tensor_free(gR);
|
|
free(R2); free(mixed); free(part); free(lo); free(xn); free(blk); free(R);
|
|
free(g_gamma); free(g_down); free(g_up); free(g_inj);
|
|
}
|
|
|
|
/* ---- gated delta net ---- */
|
|
|
|
static void gdn_reference(uint32_t Hk, uint32_t Hv, uint32_t D, uint32_t K, uint32_t T,
|
|
const float *qkv_in, const float *z, const float *ab_in,
|
|
const double *conv_w, const double *ssm_a, const double *dt, const double *norm_w,
|
|
double *state, double *hist, double *out) {
|
|
const uint32_t kd = Hk * D, vd = Hv * D, C = 2 * kd + vd;
|
|
double *conv = malloc(C * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t c = 0; c < C; c++) {
|
|
double acc = conv_w[c * K + (K - 1)] * qkv_in[t * C + c];
|
|
for (uint32_t k = 0; k + 1 < K; k++) acc += conv_w[c * K + k] * hist[k * C + c];
|
|
conv[c] = silu_d(acc);
|
|
}
|
|
for (uint32_t k = 0; k + 2 < K; k++) memcpy(hist + k * C, hist + (k + 1) * C, C * sizeof(double));
|
|
for (uint32_t c = 0; c < C; c++) hist[(K - 2) * C + c] = qkv_in[t * C + c];
|
|
for (uint32_t h = 0; h < Hk; h++) {
|
|
double sq = 0.0, sk = 0.0;
|
|
for (uint32_t i = 0; i < D; i++) { sq += conv[h * D + i] * conv[h * D + i]; sk += conv[kd + h * D + i] * conv[kd + h * D + i]; }
|
|
const double qs = 1.0 / sqrt(sq + 1e-6) / sqrt((double)D), ks = 1.0 / sqrt(sk + 1e-6);
|
|
for (uint32_t i = 0; i < D; i++) { conv[h * D + i] *= qs; conv[kd + h * D + i] *= ks; }
|
|
}
|
|
for (uint32_t j = 0; j < Hv; j++) {
|
|
const uint32_t kh = j % Hk;
|
|
const double g = exp(ssm_a[j] * softplus_d((double)ab_in[t * 2 * Hv + j] + dt[j]));
|
|
const double beta = sigmoid_d(ab_in[t * 2 * Hv + Hv + j]);
|
|
double *S = state + (uint64_t)j * D * D; /* [dv][dk] */
|
|
const double *q = conv + kh * D, *k = conv + kd + kh * D, *v = conv + 2 * kd + j * D;
|
|
double o[128];
|
|
for (uint32_t dv = 0; dv < D; dv++) {
|
|
double u = 0.0;
|
|
for (uint32_t dk = 0; dk < D; dk++) { S[dv * D + dk] *= g; u += S[dv * D + dk] * k[dk]; }
|
|
const double delta = (v[dv] - u) * beta;
|
|
double acc = 0.0;
|
|
for (uint32_t dk = 0; dk < D; dk++) { S[dv * D + dk] += k[dk] * delta; acc += S[dv * D + dk] * q[dk]; }
|
|
o[dv] = acc;
|
|
}
|
|
double ss = 0.0;
|
|
for (uint32_t dv = 0; dv < D; dv++) ss += o[dv] * o[dv];
|
|
const double r = 1.0 / sqrt(ss / D + 1e-6);
|
|
for (uint32_t dv = 0; dv < D; dv++)
|
|
out[((uint64_t)t * Hv + j) * D + dv] = o[dv] * r * norm_w[dv] * sigmoid_d(z[((uint64_t)t * Hv + j) * D + dv]);
|
|
}
|
|
}
|
|
free(conv);
|
|
}
|
|
|
|
static void test_gdn(arena_t *a, uint32_t Hk, uint32_t Hv, uint32_t D, uint32_t T) {
|
|
const uint32_t K = 4, kd = Hk * D, vd = Hv * D, C = 2 * kd + vd;
|
|
double *conv_w, *ssm_a, *dt, *norm_w;
|
|
const uint64_t conv_off = arena_f32(a, (uint64_t)C * K, &conv_w, -0.5f, 0.5f);
|
|
const uint64_t a_off = arena_f32(a, Hv, &ssm_a, -8.0f, -0.1f);
|
|
const uint64_t dt_off = arena_f32(a, Hv, &dt, 0.2f, 1.5f);
|
|
const uint64_t norm_off = arena_f32(a, D, &norm_w, 0.8f, 1.2f);
|
|
float *qkv = rand_vec((uint64_t)T * C, 1.0f);
|
|
float *z = rand_vec((uint64_t)T * vd, 1.0f);
|
|
/* raw alpha/beta come from Q8 projections of a random input so the
|
|
* fused front kernel can be checked against the same reference */
|
|
const uint32_t E = 2560;
|
|
double *alpha_w, *beta_w;
|
|
const uint64_t alpha_off = arena_q8_0(a, Hv, E, &alpha_w, 0.05f);
|
|
const uint64_t beta_off = arena_q8_0(a, Hv, E, &beta_w, 0.05f);
|
|
float *mixed = rand_vec((uint64_t)T * E, 1.0f);
|
|
float *ab = malloc((uint64_t)T * 2 * Hv * sizeof(float));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t j = 0; j < Hv; j++) {
|
|
double sa = 0.0, sb = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) {
|
|
sa += alpha_w[(uint64_t)j * E + i] * mixed[(uint64_t)t * E + i];
|
|
sb += beta_w[(uint64_t)j * E + i] * mixed[(uint64_t)t * E + i];
|
|
}
|
|
ab[(uint64_t)t * 2 * Hv + j] = (float)sa;
|
|
ab[(uint64_t)t * 2 * Hv + Hv + j] = (float)sb;
|
|
}
|
|
}
|
|
float *av = malloc((uint64_t)T * Hv * sizeof(float)), *bv = malloc((uint64_t)T * Hv * sizeof(float));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
memcpy(av + (uint64_t)t * Hv, ab + (uint64_t)t * 2 * Hv, Hv * sizeof(float));
|
|
memcpy(bv + (uint64_t)t * Hv, ab + (uint64_t)t * 2 * Hv + Hv, Hv * sizeof(float));
|
|
}
|
|
float *state0 = rand_vec((uint64_t)Hv * D * D, 0.1f);
|
|
float *hist0 = rand_vec((uint64_t)(K - 1) * C, 1.0f);
|
|
|
|
double *ref_state = malloc((uint64_t)Hv * D * D * sizeof(double));
|
|
double *ref_hist = malloc((uint64_t)(K - 1) * C * sizeof(double));
|
|
double *ref_out = malloc((uint64_t)T * vd * sizeof(double));
|
|
for (uint64_t i = 0; i < (uint64_t)Hv * D * D; i++) ref_state[i] = state0[i];
|
|
for (uint64_t i = 0; i < (uint64_t)(K - 1) * C; i++) ref_hist[i] = hist0[i];
|
|
double *snap_state = malloc((uint64_t)Hv * D * D * sizeof(double));
|
|
double *snap_hist = malloc((uint64_t)(K - 1) * C * sizeof(double));
|
|
memcpy(snap_state, ref_state, (uint64_t)Hv * D * D * sizeof(double));
|
|
memcpy(snap_hist, ref_hist, (uint64_t)(K - 1) * C * sizeof(double));
|
|
gdn_reference(Hk, Hv, D, K, 1, qkv, z, ab, conv_w, ssm_a, dt, norm_w, snap_state, snap_hist, ref_out);
|
|
gdn_reference(Hk, Hv, D, K, T, qkv, z, ab, conv_w, ssm_a, dt, norm_w, ref_state, ref_hist, ref_out);
|
|
|
|
char name[96];
|
|
/* (a) one chunk of T tokens, (b) T single-token calls, (c) fused front
|
|
* kernel (conv + alpha/beta + prep) in pairs of tokens */
|
|
for (int mode = 0; mode < 3; mode++) {
|
|
ds4_gpu_tensor *gqkv = upload(qkv, (uint64_t)T * C);
|
|
ds4_gpu_tensor *gz = upload(z, (uint64_t)T * vd);
|
|
ds4_gpu_tensor *ga = upload(av, (uint64_t)T * Hv);
|
|
ds4_gpu_tensor *gb = upload(bv, (uint64_t)T * Hv);
|
|
ds4_gpu_tensor *gstate = upload(state0, (uint64_t)Hv * D * D);
|
|
ds4_gpu_tensor *ghist = upload(hist0, (uint64_t)(K - 1) * C);
|
|
ds4_gpu_tensor *gout = upload(NULL, (uint64_t)T * vd);
|
|
ds4_gpu_tensor *gmixed = upload(mixed, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsnap_state = mode == 2 ? upload(NULL, (uint64_t)Hv * D * D) : NULL;
|
|
ds4_gpu_tensor *gsnap_hist = mode == 2 ? upload(NULL, (uint64_t)(K - 1) * C) : NULL;
|
|
const uint32_t base_step = mode == 0 ? T : mode == 1 ? 1 : 2;
|
|
for (uint32_t t0 = 0; t0 < T; t0 += base_step) {
|
|
const uint32_t step = t0 + base_step <= T ? base_step : T - t0;
|
|
ds4_gpu_tensor *vqkv = ds4_gpu_tensor_view(gqkv, (uint64_t)t0 * C * 4, (uint64_t)step * C * 4);
|
|
ds4_gpu_tensor *va = ds4_gpu_tensor_view(ga, (uint64_t)t0 * Hv * 4, (uint64_t)step * Hv * 4);
|
|
ds4_gpu_tensor *vb = ds4_gpu_tensor_view(gb, (uint64_t)t0 * Hv * 4, (uint64_t)step * Hv * 4);
|
|
ds4_gpu_tensor *vz = ds4_gpu_tensor_view(gz, (uint64_t)t0 * vd * 4, (uint64_t)step * vd * 4);
|
|
ds4_gpu_tensor *vout = ds4_gpu_tensor_view(gout, (uint64_t)t0 * vd * 4, (uint64_t)step * vd * 4);
|
|
if (mode == 2) {
|
|
ds4_gpu_tensor *vm = ds4_gpu_tensor_view(gmixed, (uint64_t)t0 * E * 4, (uint64_t)step * E * 4);
|
|
require_ok(ds4_gpu_qwen4_gdn_front_tensor(vqkv, ghist, vm, va, vb, a->base, a->size, conv_off, alpha_off, beta_off,
|
|
a_off, dt_off, 8u, step, Hk, Hv, D, K, E,
|
|
t0 == 0 ? gsnap_hist : NULL, 0u, NULL, 0u), "gdn front");
|
|
ds4_gpu_tensor_free(vm);
|
|
} else {
|
|
require_ok(ds4_gpu_qwen4_conv_stream_tensor(vqkv, ghist, a->base, a->size, conv_off, step, C, K, true), "gdn conv");
|
|
require_ok(ds4_gpu_qwen4_gdn_prep_tensor(vqkv, va, vb, a->base, a->size, a_off, dt_off, step, Hk, Hv, D), "gdn prep");
|
|
}
|
|
require_ok(ds4_gpu_qwen4_gdn_scan_tensor(vout, gstate, vqkv, va, vb, step, Hk, Hv, D,
|
|
t0 == 0 ? gsnap_state : NULL, 0u, NULL, 0u), "gdn scan");
|
|
require_ok(ds4_gpu_qwen4_gdn_out_tensor(vout, vz, a->base, a->size, norm_off, step, Hv, D, 1e-6f), "gdn out");
|
|
ds4_gpu_tensor_free(vout); ds4_gpu_tensor_free(vz); ds4_gpu_tensor_free(vb); ds4_gpu_tensor_free(va); ds4_gpu_tensor_free(vqkv);
|
|
}
|
|
const char *mname = mode == 0 ? "chunk" : mode == 1 ? "decode" : "fused";
|
|
snprintf(name, sizeof(name), "gdn %u/%u/%u T=%u %s: output", Hk, Hv, D, T, mname);
|
|
check_tensor(name, gout, ref_out, (uint64_t)T * vd, 5e-5);
|
|
snprintf(name, sizeof(name), "gdn %u/%u/%u T=%u %s: state", Hk, Hv, D, T, mname);
|
|
check_tensor(name, gstate, ref_state, (uint64_t)Hv * D * D, 5e-5);
|
|
snprintf(name, sizeof(name), "gdn %u/%u/%u T=%u %s: conv history", Hk, Hv, D, T, mname);
|
|
check_tensor(name, ghist, ref_hist, (uint64_t)(K - 1) * C, 1e-6);
|
|
if (mode == 2) {
|
|
/* A rejected second token restores precisely the first row,
|
|
* even after later calls have advanced the live state. */
|
|
snprintf(name, sizeof(name), "gdn %u/%u/%u T=%u: first-row state snapshot", Hk, Hv, D, T);
|
|
check_tensor(name, gsnap_state, snap_state, (uint64_t)Hv * D * D, 5e-5);
|
|
snprintf(name, sizeof(name), "gdn %u/%u/%u T=%u: first-row history snapshot", Hk, Hv, D, T);
|
|
check_tensor(name, gsnap_hist, snap_hist, (uint64_t)(K - 1) * C, 1e-6);
|
|
}
|
|
ds4_gpu_tensor_free(gsnap_hist); ds4_gpu_tensor_free(gsnap_state);
|
|
ds4_gpu_tensor_free(gmixed);
|
|
ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(ghist); ds4_gpu_tensor_free(gstate);
|
|
ds4_gpu_tensor_free(gb); ds4_gpu_tensor_free(ga); ds4_gpu_tensor_free(gz); ds4_gpu_tensor_free(gqkv);
|
|
}
|
|
free(bv); free(av);
|
|
free(snap_hist); free(snap_state);
|
|
free(ref_out); free(ref_hist); free(ref_state); free(hist0); free(state0); free(ab); free(z); free(qkv);
|
|
free(conv_w); free(ssm_a); free(dt); free(norm_w);
|
|
free(mixed); free(alpha_w); free(beta_w);
|
|
}
|
|
|
|
/* ---- indexer matrix scorer ---- */
|
|
|
|
static void test_idx_score_mm(uint32_t T, uint32_t n, uint32_t pos0) {
|
|
const uint32_t Hi = 4, Di = 128, ratio = 4;
|
|
float *q = rand_vec((uint64_t)T * Hi * Di, 1.0f);
|
|
float *keyf = rand_vec((uint64_t)n * Di, 1.0f);
|
|
uint16_t *keyh = malloc((uint64_t)n * Di * 2);
|
|
for (uint64_t i = 0; i < (uint64_t)n * Di; i++) { keyh[i] = f32_to_f16(keyf[i]); keyf[i] = f16_to_f32(keyh[i]); }
|
|
double *ref = malloc((uint64_t)T * n * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
const uint32_t visible = (pos0 + t + 1) / ratio;
|
|
for (uint32_t b = 0; b < n; b++) {
|
|
double sum = -3.0e38;
|
|
if (b < visible) {
|
|
sum = 0.0;
|
|
for (uint32_t h = 0; h < Hi; h++) {
|
|
double d = 0.0;
|
|
for (uint32_t x = 0; x < Di; x++) d += (double)q[((uint64_t)t * Hi + h) * Di + x] * keyf[(uint64_t)b * Di + x];
|
|
sum += d > 0.0 ? d : 0.0;
|
|
}
|
|
}
|
|
ref[(uint64_t)t * n + b] = sum;
|
|
}
|
|
}
|
|
ds4_gpu_tensor *gq = upload(q, (uint64_t)T * Hi * Di);
|
|
ds4_gpu_tensor *gk = ds4_gpu_tensor_alloc((uint64_t)n * Di * 2);
|
|
ds4_gpu_tensor *gs = upload(NULL, (uint64_t)T * n);
|
|
require_ok(gq && gk && gs && ds4_gpu_tensor_write(gk, 0, keyh, (uint64_t)n * Di * 2) &&
|
|
ds4_gpu_qwen4_idx_score_tensor(gs, NULL, gq, gk, T, n, Hi, Di, pos0, ratio), "idx score mm");
|
|
float *got = download(gs, (uint64_t)T * n);
|
|
/* invisible entries are -3e38 on both sides: compare them exactly by mapping to 0 */
|
|
for (uint64_t i = 0; i < (uint64_t)T * n; i++) {
|
|
if (ref[i] < -1e37) {
|
|
ref[i] = 0.0;
|
|
got[i] = got[i] < -1e37f ? 0.0f : 1.0f;
|
|
}
|
|
}
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "idx score mm T=%u n=%u pos0=%u", T, n, pos0);
|
|
check_close(name, got, ref, (uint64_t)T * n, 4e-3); /* queries rounded to half */
|
|
free(got); ds4_gpu_tensor_free(gs); ds4_gpu_tensor_free(gk); ds4_gpu_tensor_free(gq); free(ref); free(keyh); free(keyf); free(q);
|
|
}
|
|
|
|
/* ---- indexer radix select ---- */
|
|
|
|
#ifdef __APPLE__
|
|
static int cmp_desc_idx(void *ctx, const void *a, const void *b) {
|
|
#else
|
|
static int cmp_desc_idx(const void *a, const void *b, void *ctx) {
|
|
#endif
|
|
const float *sc = ctx;
|
|
const int32_t ia = *(const int32_t *)a, ib = *(const int32_t *)b;
|
|
if (sc[ia] != sc[ib]) return sc[ia] > sc[ib] ? -1 : 1;
|
|
return ia < ib ? -1 : 1;
|
|
}
|
|
|
|
static void test_idx_select(uint32_t T, uint32_t n, uint32_t k, uint32_t visible) {
|
|
float *sc = rand_vec((uint64_t)T * n, 4.0f);
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t b = 0; b < n; b++) {
|
|
float *v = &sc[(uint64_t)t * n + b];
|
|
if (b >= visible) *v = -3.0e38f; /* incomplete blocks */
|
|
else if (*v < 0.0f) *v = 0.0f; /* relu sums */
|
|
else if ((b % 7) == 3) *v = sc[(uint64_t)t * n + (b / 7) * 7]; /* duplicates */
|
|
}
|
|
}
|
|
ds4_gpu_tensor *gs = upload(sc, (uint64_t)T * n);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * k * 4);
|
|
require_ok(gs && gsel && ds4_gpu_qwen4_idx_select_tensor(gsel, gs, NULL, n, T, k), "idx select");
|
|
int32_t *got = malloc((uint64_t)T * k * 4);
|
|
require_ok(ds4_gpu_tensor_read(gsel, 0, got, (uint64_t)T * k * 4), "idx select read");
|
|
int32_t *order = malloc((uint64_t)n * 4);
|
|
uint64_t bad = 0;
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t b = 0; b < n; b++) order[b] = (int32_t)b;
|
|
#ifdef __APPLE__
|
|
qsort_r(order, n, 4, sc + (uint64_t)t * n, cmp_desc_idx);
|
|
#else
|
|
qsort_r(order, n, 4, cmp_desc_idx, sc + (uint64_t)t * n);
|
|
#endif
|
|
/* the k selected must equal the reference top-k as a set */
|
|
int32_t *g = got + (uint64_t)t * k;
|
|
for (uint32_t i = 0; i < k; i++) {
|
|
bool found = false;
|
|
for (uint32_t j = 0; j < k && !found; j++) found = g[j] == order[i];
|
|
bad += !found;
|
|
}
|
|
}
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "idx select T=%u n=%u k=%u visible=%u", T, n, k, visible);
|
|
require_ok(bad == 0, name);
|
|
printf(" %-44s ok\n", name);
|
|
free(order); free(got); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gs); free(sc);
|
|
}
|
|
|
|
/* Prefill and split attention against scalar attention on the same inputs. */
|
|
static void test_attn_mm_keys(uint32_t T, uint32_t pos0, bool sparse, uint32_t k_blocks, uint32_t H, bool split) {
|
|
const uint32_t Hkv = 2, D = 256, ratio = 4, sel_stride = k_blocks * ratio + ratio;
|
|
const uint32_t cap = pos0 + T;
|
|
const uint64_t qn = (uint64_t)T * H * D, kvn = (uint64_t)cap * Hkv * D;
|
|
float *q = rand_vec(qn, 1.0f), *gate = rand_vec(qn, 2.0f);
|
|
_Float16 *kc = malloc(kvn * 2), *vc = malloc(kvn * 2);
|
|
for (uint64_t i = 0; i < kvn; i++) { kc[i] = (_Float16)(frand() - 0.5f); vc[i] = (_Float16)(frand() - 0.5f); }
|
|
int32_t *sel = malloc((uint64_t)T * sel_stride * 4);
|
|
uint32_t *cnt = malloc(T * 4);
|
|
uint32_t *blocks = malloc(k_blocks * sizeof(*blocks));
|
|
require_ok(blocks != NULL, "attention block selection allocation");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
const uint32_t pos = pos0 + t, n_blocks = (pos + 1) / ratio, nb = n_blocks < k_blocks ? n_blocks : k_blocks;
|
|
uint32_t n = 0;
|
|
for (uint32_t i = 0; i < nb; i++) {
|
|
uint32_t b;
|
|
bool dup;
|
|
do {
|
|
b = (uint32_t)((frand() + 1.0f) * 0.5f * n_blocks) % n_blocks;
|
|
dup = false;
|
|
for (uint32_t j = 0; j < i && !dup; j++) dup = blocks[j] == b;
|
|
} while (dup);
|
|
blocks[i] = b;
|
|
}
|
|
for (uint32_t i = 0; i < nb; i++) for (uint32_t r = 0; r < ratio; r++) sel[(uint64_t)t * sel_stride + n++] = (int32_t)(blocks[i] * ratio + r);
|
|
for (uint32_t k = n_blocks * ratio; k <= pos; k++) sel[(uint64_t)t * sel_stride + n++] = (int32_t)k;
|
|
cnt[t] = n;
|
|
}
|
|
free(blocks);
|
|
#ifndef __APPLE__
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
const float amplitude = t % 3 == 0 ? 64 : t % 3 == 1 ? 0.00001f : 1;
|
|
for (uint32_t i = 0; i < H * D; i++) q[(uint64_t)t * H * D + i] *= amplitude;
|
|
if (sparse && T > 2 && t % 7 == 0) cnt[t] = 0;
|
|
else if (sparse && T > 2 && t % 11 == 1) {
|
|
cnt[t] = 1;
|
|
sel[(uint64_t)t * sel_stride] = -1;
|
|
} else if (sparse && cnt[t] && t + 1 < T)
|
|
sel[(uint64_t)t * sel_stride] = (int32_t)(pos0 + t + 1);
|
|
}
|
|
#endif
|
|
ds4_gpu_tensor *gq = upload(q, qn), *ggate = upload(gate, qn);
|
|
ds4_gpu_tensor *gk = ds4_gpu_tensor_alloc(kvn * 2), *gv = ds4_gpu_tensor_alloc(kvn * 2);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * sel_stride * 4), *gcnt = ds4_gpu_tensor_alloc(T * 4);
|
|
ds4_gpu_tensor *go_ref = upload(NULL, qn), *go_new = upload(NULL, qn);
|
|
ds4_gpu_tensor *partial = split ? ds4_gpu_tensor_alloc((uint64_t)T*H*64*(D+2)*4) : NULL;
|
|
require_ok(!split || partial, "attention partial allocation");
|
|
require_ok(gk && gv && gsel && gcnt && ds4_gpu_tensor_write(gk, 0, kc, kvn * 2) && ds4_gpu_tensor_write(gv, 0, vc, kvn * 2) &&
|
|
ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * sel_stride * 4) && ds4_gpu_tensor_write(gcnt, 0, cnt, T * 4), "attn mm setup");
|
|
setenv("DS4_QWEN4_NO_ATTN_MM", "1", 1);
|
|
require_ok(ds4_gpu_qwen4_attn_decode_tensor(go_ref, gq, ggate, gk, gv, gsel, gcnt, NULL, T, H, Hkv, D, pos0, sparse, sel_stride, 0.0625f), "attn reference");
|
|
unsetenv("DS4_QWEN4_NO_ATTN_MM");
|
|
require_ok(ds4_gpu_qwen4_attn_decode_tensor(go_new, gq, ggate, gk, gv, gsel, gcnt, partial, T, H, Hkv, D, pos0, sparse, sel_stride, 0.0625f), "attn mm");
|
|
float *ref = download(go_ref, qn), *got = download(go_new, qn);
|
|
double worst = 0.0, scale = 0.0;
|
|
for (uint64_t i = 0; i < qn; i++) {
|
|
require_ok(isfinite(got[i]) && isfinite(ref[i]), "finite attention output");
|
|
const double d = fabs((double)got[i] - ref[i]);
|
|
if (d > worst) worst = d;
|
|
if (fabs(ref[i]) > scale) scale = fabs(ref[i]);
|
|
}
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "attn mm T=%u H=%u pos0=%u %s%s", T, H, pos0, sparse ? "sparse" : "dense", split ? " split" : "");
|
|
#ifdef __APPLE__
|
|
require_ok(worst <= 4e-3 * scale, name);
|
|
#else
|
|
require_ok(worst <= 3e-5 * scale, name);
|
|
/* Check selected rows against the definition, independently of both GPU
|
|
* kernels. Include the first, middle and last token and KV head group. */
|
|
double *scores = malloc((size_t)cap * sizeof(*scores));
|
|
require_ok(scores != NULL, "attention CPU reference allocation");
|
|
double cpu_error = 0, scalar_error = 0, cpu_scale = 0;
|
|
for (uint32_t ti = 0; ti < 3; ti++) for (uint32_t hi = 0; hi < 3; hi++) {
|
|
const uint32_t t = ti == 0 ? 0 : ti == 1 ? T/2 : T-1;
|
|
const uint32_t h = hi == 0 ? 0 : hi == 1 ? H/2 : H-1;
|
|
const uint32_t kh = h/(H/Hkv), n = sparse ? cnt[t] : pos0+t+1;
|
|
double peak = -INFINITY, denom = 0, acc[256] = {0};
|
|
for (uint32_t j = 0; j < n; j++) {
|
|
const uint32_t pos = sparse ? (uint32_t)sel[(uint64_t)t*sel_stride+j] : j;
|
|
double dot = 0;
|
|
if (pos > pos0+t) { scores[j] = -INFINITY; continue; }
|
|
for (uint32_t d = 0; d < D; d++)
|
|
dot += (double)q[((uint64_t)t*H+h)*D+d] * (double)kc[((uint64_t)pos*Hkv+kh)*D+d];
|
|
scores[j] = dot * 0.0625;
|
|
peak = fmax(peak,scores[j]);
|
|
}
|
|
for (uint32_t j = 0; j < n; j++) {
|
|
const uint32_t pos = sparse ? (uint32_t)sel[(uint64_t)t*sel_stride+j] : j;
|
|
if (pos > pos0+t) continue;
|
|
const double weight = exp(scores[j]-peak);
|
|
denom += weight;
|
|
for (uint32_t d = 0; d < D; d++) acc[d] += weight * (double)vc[((uint64_t)pos*Hkv+kh)*D+d];
|
|
}
|
|
for (uint32_t d = 0; d < D; d++) {
|
|
const uint64_t i = ((uint64_t)t*H+h)*D+d;
|
|
const double expected = denom > 0 ? acc[d]/denom/(1+exp(-(double)gate[i])) : 0;
|
|
cpu_error = fmax(cpu_error,fabs((double)got[i]-expected));
|
|
scalar_error = fmax(scalar_error,fabs((double)ref[i]-expected));
|
|
cpu_scale = fmax(cpu_scale,fabs(expected));
|
|
}
|
|
}
|
|
free(scores);
|
|
require_ok(cpu_error <= 3e-5 * fmax(cpu_scale,1e-8), "attention double reference");
|
|
printf(" %-44s CPU max|d|=%.2e scalar=%.2e (scale %.2e)\n", name, cpu_error, scalar_error, cpu_scale);
|
|
#endif
|
|
if (T <= 8u && !split) require_ok(worst == 0.0, "short attention tails keep decode arithmetic");
|
|
printf(" %-44s ok max|d|=%.2e (scale %.2e)\n", name, worst, scale);
|
|
free(ref); free(got); free(q); free(gate); free(kc); free(vc); free(sel); free(cnt);
|
|
ds4_gpu_tensor_free(gq); ds4_gpu_tensor_free(ggate); ds4_gpu_tensor_free(gk); ds4_gpu_tensor_free(gv);
|
|
ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gcnt); ds4_gpu_tensor_free(go_ref); ds4_gpu_tensor_free(go_new);
|
|
ds4_gpu_tensor_free(partial);
|
|
}
|
|
|
|
static void test_attn_mm(uint32_t T, uint32_t pos0, bool sparse) {
|
|
test_attn_mm_keys(T, pos0, sparse, 6, 24, false);
|
|
}
|
|
|
|
#ifndef __APPLE__
|
|
static void test_attn_groups(void) {
|
|
test_attn_mm_keys(32, 4093, false, 6, 24, false);
|
|
test_attn_mm_keys(33, 8193, true, 511, 24, false);
|
|
test_attn_mm_keys(32, 97, false, 6, 2, false);
|
|
test_attn_mm_keys(33, 100, true, 6, 32, false);
|
|
test_attn_mm_keys(33, 100, true, 6, 34, false);
|
|
test_attn_mm_keys(1, 510, false, 6, 24, true);
|
|
test_attn_mm_keys(1, 511, false, 6, 24, true);
|
|
test_attn_mm_keys(1, 2052, false, 6, 24, true);
|
|
test_attn_mm_keys(2, 8193, true, 511, 24, true);
|
|
test_attn_mm_keys(8, 8193, true, 511, 32, true);
|
|
test_attn_mm_keys(2, 32769, false, 6, 24, true);
|
|
}
|
|
#endif
|
|
|
|
/* ---- PLE ---- */
|
|
|
|
static void test_ple(arena_t *a, uint32_t E, uint32_t T) {
|
|
const uint32_t hc = 4, dim = E * hc, K = 4, dil = 3, H = (K - 1) * dil;
|
|
const float eps = 1e-6f;
|
|
double *gk, *gq, *gc, *cw;
|
|
const uint64_t gk_off = arena_f32(a, dim, &gk, 0.5f, 1.5f);
|
|
const uint64_t gq_off = arena_f32(a, dim, &gq, 0.5f, 1.5f);
|
|
const uint64_t gc_off = arena_f32(a, dim, &gc, 0.5f, 1.5f);
|
|
const uint64_t cw_off = arena_f32(a, (uint64_t)dim * K, &cw, -0.4f, 0.4f);
|
|
float *R = rand_vec((uint64_t)T * dim, 1.0f);
|
|
float *key = rand_vec((uint64_t)T * dim, 1.0f);
|
|
float *value = rand_vec((uint64_t)T * E, 1.0f);
|
|
float *hist0 = rand_vec((uint64_t)H * dim, 1.0f);
|
|
|
|
double *ref_R = malloc((uint64_t)T * dim * sizeof(double));
|
|
double *ref_hist = malloc((uint64_t)H * dim * sizeof(double));
|
|
double *gated = malloc(dim * sizeof(double)), *normed = malloc(dim * sizeof(double));
|
|
for (uint64_t i = 0; i < (uint64_t)H * dim; i++) ref_hist[i] = hist0[i];
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < hc; s++) {
|
|
double sk = 0.0, sr = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) {
|
|
sk += (double)key[t * dim + s * E + i] * key[t * dim + s * E + i];
|
|
sr += (double)R[t * dim + s * E + i] * R[t * dim + s * E + i];
|
|
}
|
|
const double ik = 1.0 / sqrt(sk / E + eps), ir = 1.0 / sqrt(sr / E + eps);
|
|
double dot = 0.0;
|
|
for (uint32_t i = 0; i < E; i++)
|
|
dot += (key[t * dim + s * E + i] * ik * gk[s * E + i]) * (R[t * dim + s * E + i] * ir * gq[s * E + i]);
|
|
double g = dot / sqrt((double)E);
|
|
const double mag = sqrt(fmax(fabs(g), 1e-6));
|
|
g = sigmoid_d(g > 0 ? mag : (g < 0 ? -mag : 0.0));
|
|
double sg = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) { gated[s * E + i] = g * value[t * E + i]; sg += gated[s * E + i] * gated[s * E + i]; }
|
|
const double ig = 1.0 / sqrt(sg / E + eps);
|
|
for (uint32_t i = 0; i < E; i++) normed[s * E + i] = gated[s * E + i] * ig * gc[s * E + i];
|
|
}
|
|
for (uint32_t c = 0; c < dim; c++) {
|
|
double acc = 0.0;
|
|
for (uint32_t k = 0; k < K; k++) {
|
|
const uint32_t back = (K - 1 - k) * dil;
|
|
acc += cw[(uint64_t)c * K + k] * (back == 0 ? normed[c] : ref_hist[(uint64_t)(H - back) * dim + c]);
|
|
}
|
|
ref_R[t * dim + c] = R[t * dim + c] + gated[c] + silu_d(acc);
|
|
}
|
|
memmove(ref_hist, ref_hist + dim, (uint64_t)(H - 1) * dim * sizeof(double));
|
|
memcpy(ref_hist + (uint64_t)(H - 1) * dim, normed, dim * sizeof(double));
|
|
}
|
|
|
|
char name[96];
|
|
for (int mode = 0; mode < 2; mode++) {
|
|
ds4_gpu_tensor *gR = upload(R, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *gkey = upload(key, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *gval = upload(value, (uint64_t)T * E);
|
|
ds4_gpu_tensor *ggated = upload(NULL, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *gnormed = upload(NULL, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *ghist = upload(hist0, (uint64_t)H * dim);
|
|
const uint32_t step = mode == 0 ? T : 1;
|
|
for (uint32_t t0 = 0; t0 < T; t0 += step) {
|
|
ds4_gpu_tensor *vR = ds4_gpu_tensor_view(gR, (uint64_t)t0 * dim * 4, (uint64_t)step * dim * 4);
|
|
ds4_gpu_tensor *vkey = ds4_gpu_tensor_view(gkey, (uint64_t)t0 * dim * 4, (uint64_t)step * dim * 4);
|
|
ds4_gpu_tensor *vval = ds4_gpu_tensor_view(gval, (uint64_t)t0 * E * 4, (uint64_t)step * E * 4);
|
|
ds4_gpu_tensor *vg = ds4_gpu_tensor_view(ggated, (uint64_t)t0 * dim * 4, (uint64_t)step * dim * 4);
|
|
ds4_gpu_tensor *vn = ds4_gpu_tensor_view(gnormed, (uint64_t)t0 * dim * 4, (uint64_t)step * dim * 4);
|
|
require_ok(ds4_gpu_qwen4_ple_gate_tensor(vg, vn, vR, vkey, vval, a->base, a->size, gk_off, gq_off, gc_off,
|
|
step, E, hc, eps), "ple gate");
|
|
require_ok(ds4_gpu_qwen4_ple_conv_tensor(vR, vg, vn, ghist, a->base, a->size, cw_off, 0u, step, dim, K, dil, NULL, 0u, NULL, 0u), "ple conv");
|
|
ds4_gpu_tensor_free(vn); ds4_gpu_tensor_free(vg); ds4_gpu_tensor_free(vval);
|
|
ds4_gpu_tensor_free(vkey); ds4_gpu_tensor_free(vR);
|
|
}
|
|
const char *mname = mode ? "decode" : "chunk";
|
|
snprintf(name, sizeof(name), "ple E=%u T=%u %s: residual", E, T, mname);
|
|
check_tensor(name, gR, ref_R, (uint64_t)T * dim, 2e-5);
|
|
snprintf(name, sizeof(name), "ple E=%u T=%u %s: history", E, T, mname);
|
|
check_tensor(name, ghist, ref_hist, (uint64_t)H * dim, 2e-5);
|
|
ds4_gpu_tensor_free(ghist); ds4_gpu_tensor_free(gnormed); ds4_gpu_tensor_free(ggated);
|
|
ds4_gpu_tensor_free(gval); ds4_gpu_tensor_free(gkey); ds4_gpu_tensor_free(gR);
|
|
}
|
|
free(normed); free(gated); free(ref_hist); free(ref_R); free(hist0); free(value); free(key); free(R);
|
|
free(gk); free(gq); free(gc); free(cw);
|
|
}
|
|
|
|
/* ---- router ---- */
|
|
|
|
static void test_router(arena_t *a, uint32_t NE, uint32_t k, uint32_t T) {
|
|
const uint32_t E = 2560;
|
|
double *gate_w;
|
|
const uint64_t gate_off = arena_f32(a, E, &gate_w, -0.05f, 0.05f);
|
|
float *x = rand_vec((uint64_t)T * E, 1.0f);
|
|
float *logits = rand_vec((uint64_t)T * NE, 3.0f);
|
|
ds4_gpu_tensor *gl = upload(logits, (uint64_t)T * NE);
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * k * 4);
|
|
ds4_gpu_tensor *gw = upload(NULL, (uint64_t)T * k);
|
|
ds4_gpu_tensor *gsg = upload(NULL, T);
|
|
require_ok(ds4_gpu_qwen4_router_topk_tensor(gsel, gw, gl, gx, a->base, a->size, gate_off, 0u, E, gsg, T, NE, k), "router");
|
|
char name[96];
|
|
{
|
|
double *ref = malloc(T * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
double acc = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) acc += gate_w[i] * x[(uint64_t)t * E + i];
|
|
ref[t] = acc;
|
|
}
|
|
snprintf(name, sizeof(name), "router NE=%u k=%u T=%u: shared gate logit", NE, k, T);
|
|
check_tensor(name, gsg, ref, T, 1e-5);
|
|
free(ref);
|
|
}
|
|
ds4_gpu_tensor_free(gsg); ds4_gpu_tensor_free(gx); free(x); free(gate_w);
|
|
int32_t *sel = malloc((uint64_t)T * k * 4);
|
|
require_ok(ds4_gpu_tensor_read(gsel, 0, sel, (uint64_t)T * k * 4), "router read");
|
|
float *w = download(gw, (uint64_t)T * k);
|
|
double worst = 0.0;
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
double *p = malloc(NE * sizeof(double));
|
|
double mx = -1e300, sum = 0.0;
|
|
for (uint32_t e = 0; e < NE; e++) if (logits[t * NE + e] > mx) mx = logits[t * NE + e];
|
|
for (uint32_t e = 0; e < NE; e++) { p[e] = exp(logits[t * NE + e] - mx); sum += p[e]; }
|
|
for (uint32_t e = 0; e < NE; e++) p[e] /= sum;
|
|
double wsum = 0.0;
|
|
int *ref = malloc(k * sizeof(int));
|
|
for (uint32_t i = 0; i < k; i++) {
|
|
int best = -1;
|
|
for (uint32_t e = 0; e < NE; e++) {
|
|
bool used = false;
|
|
for (uint32_t j = 0; j < i; j++) used |= ref[j] == (int)e;
|
|
if (!used && (best < 0 || p[e] > p[best])) best = (int)e;
|
|
}
|
|
ref[i] = best;
|
|
wsum += p[best];
|
|
}
|
|
for (uint32_t i = 0; i < k; i++) {
|
|
if (sel[t * k + i] != ref[i]) {
|
|
fprintf(stderr, "router: token %u slot %u expert %d != %d\n", t, i, sel[t * k + i], ref[i]);
|
|
exit(1);
|
|
}
|
|
const double d = fabs(w[t * k + i] - p[ref[i]] / wsum);
|
|
if (d > worst) worst = d;
|
|
}
|
|
free(ref); free(p);
|
|
}
|
|
if (worst > 1e-6) { fprintf(stderr, "router weights max|d| %.3e\n", worst); exit(1); }
|
|
snprintf(name, sizeof(name), "router NE=%u k=%u T=%u: softmax top-k", NE, k, T);
|
|
printf(" %-44s ok max|d|=%.2e\n", name, worst);
|
|
free(w); free(sel); ds4_gpu_tensor_free(gw); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gl); free(logits);
|
|
}
|
|
|
|
/* ---- attention + indexer ---- */
|
|
|
|
static void rope_ref(double *x, uint32_t n_rot, uint32_t pos, double base) {
|
|
const uint32_t nh = n_rot / 2;
|
|
for (uint32_t i = 0; i < nh; i++) {
|
|
const double th = (double)pos * pow(base, -2.0 * i / n_rot);
|
|
const double c = cos(th), s = sin(th), x0 = x[i], x1 = x[i + nh];
|
|
x[i] = x0 * c - x1 * s;
|
|
x[i + nh] = x0 * s + x1 * c;
|
|
}
|
|
}
|
|
|
|
/* Decode n_pos tokens one at a time through prep / block keys / scoring /
|
|
* expand / attention and compare against a reference that recomputes the
|
|
* selection. k_blocks is small so the sparse path is exercised early. */
|
|
static void test_attention(arena_t *a, uint32_t H, uint32_t Hkv, uint32_t D, uint32_t n_rot,
|
|
uint32_t Hi, uint32_t Di, uint32_t k_blocks, uint32_t n_pos) {
|
|
const uint32_t ratio = 4, cap = n_pos + 4, group = H / Hkv;
|
|
const double base = 1.0e7, eps = 1e-6;
|
|
const float scale = 1.0f / sqrtf((float)D);
|
|
double *gq_w, *gk_w, *giq_w, *gik_w;
|
|
const uint64_t gq_off = arena_f32(a, D, &gq_w, 0.5f, 1.5f);
|
|
const uint64_t gk_off = arena_f32(a, D, &gk_w, 0.5f, 1.5f);
|
|
const uint64_t giq_off = arena_f32(a, Di, &giq_w, 0.5f, 1.5f);
|
|
const uint64_t gik_off = arena_f32(a, Di, &gik_w, 0.5f, 1.5f);
|
|
|
|
float *qg = rand_vec((uint64_t)n_pos * H * 2 * D, 1.0f);
|
|
float *kp = rand_vec((uint64_t)n_pos * Hkv * D, 1.0f);
|
|
float *vp = rand_vec((uint64_t)n_pos * Hkv * D, 1.0f);
|
|
float *iq = rand_vec((uint64_t)n_pos * Hi * Di, 1.0f);
|
|
float *ik = rand_vec((uint64_t)n_pos * Di, 1.0f);
|
|
|
|
double *kc = malloc((uint64_t)cap * Hkv * D * sizeof(double));
|
|
double *vc = malloc((uint64_t)cap * Hkv * D * sizeof(double));
|
|
double *ref_out = malloc((uint64_t)n_pos * H * D * sizeof(double));
|
|
double *qn = malloc((uint64_t)H * D * sizeof(double));
|
|
double *iqn = malloc((uint64_t)Hi * Di * sizeof(double));
|
|
double *score = malloc((uint64_t)(cap / ratio + 1) * sizeof(double));
|
|
int *sel = malloc((uint64_t)cap * sizeof(int));
|
|
double *p = malloc((uint64_t)cap * sizeof(double));
|
|
for (uint32_t pos = 0; pos < n_pos; pos++) {
|
|
for (uint32_t h = 0; h < H; h++) {
|
|
const float *src = qg + ((uint64_t)pos * H + h) * 2 * D;
|
|
double ss = 0.0;
|
|
for (uint32_t i = 0; i < D; i++) ss += (double)src[i] * src[i];
|
|
const double r = 1.0 / sqrt(ss / D + eps);
|
|
for (uint32_t i = 0; i < D; i++) qn[h * D + i] = src[i] * r * gq_w[i];
|
|
rope_ref(qn + h * D, n_rot, pos, base);
|
|
}
|
|
for (uint32_t h = 0; h < Hkv; h++) {
|
|
const float *src = kp + ((uint64_t)pos * Hkv + h) * D;
|
|
double ss = 0.0;
|
|
for (uint32_t i = 0; i < D; i++) ss += (double)src[i] * src[i];
|
|
const double r = 1.0 / sqrt(ss / D + eps);
|
|
double *dst = kc + ((uint64_t)pos * Hkv + h) * D, *vdst = vc + ((uint64_t)pos * Hkv + h) * D;
|
|
const float *vsrc = vp + ((uint64_t)pos * Hkv + h) * D;
|
|
double tmp[256];
|
|
/* rope on the f32 value before the f16 store, as the kernel does */
|
|
for (uint32_t i = 0; i < D; i++) tmp[i] = src[i] * r * gk_w[i];
|
|
rope_ref(tmp, n_rot, pos, base);
|
|
for (uint32_t i = 0; i < D; i++) dst[i] = f16_to_f32(f32_to_f16((float)tmp[i]));
|
|
for (uint32_t i = 0; i < D; i++) vdst[i] = f16_to_f32(f32_to_f16(vsrc[i]));
|
|
}
|
|
for (uint32_t h = 0; h < Hi; h++) {
|
|
const float *src = iq + ((uint64_t)pos * Hi + h) * Di;
|
|
double ss = 0.0;
|
|
for (uint32_t i = 0; i < Di; i++) ss += (double)src[i] * src[i];
|
|
const double r = 1.0 / sqrt(ss / Di + eps);
|
|
for (uint32_t i = 0; i < Di; i++) iqn[h * Di + i] = src[i] * r * giq_w[i];
|
|
rope_ref(iqn + h * Di, n_rot, pos, base);
|
|
}
|
|
const uint32_t n_vis = pos + 1, n_blocks = n_vis / ratio;
|
|
uint32_t n_sel = 0;
|
|
if (n_blocks <= k_blocks) {
|
|
for (uint32_t t = 0; t < n_vis; t++) sel[n_sel++] = (int)t;
|
|
} else {
|
|
for (uint32_t b = 0; b < n_blocks; b++) {
|
|
double key[128], ss = 0.0;
|
|
for (uint32_t d = 0; d < Di; d++) {
|
|
double acc = 0.0;
|
|
for (uint32_t t = 0; t < ratio; t++) acc += ik[((uint64_t)(b * ratio + t)) * Di + d];
|
|
key[d] = acc / ratio;
|
|
ss += key[d] * key[d];
|
|
}
|
|
const double r = 1.0 / sqrt(ss / Di + eps);
|
|
for (uint32_t d = 0; d < Di; d++) key[d] *= r * gik_w[d];
|
|
rope_ref(key, n_rot, b * ratio, base);
|
|
for (uint32_t d = 0; d < Di; d++) key[d] = f16_to_f32(f32_to_f16((float)key[d]));
|
|
double sum = 0.0;
|
|
for (uint32_t h = 0; h < Hi; h++) {
|
|
double dot = 0.0;
|
|
for (uint32_t d = 0; d < Di; d++) dot += iqn[h * Di + d] * key[d];
|
|
if (dot > 0) sum += dot;
|
|
}
|
|
score[b] = sum;
|
|
}
|
|
bool *taken = calloc(n_blocks, sizeof(bool));
|
|
for (uint32_t i = 0; i < k_blocks; i++) {
|
|
int best = -1;
|
|
for (uint32_t b = 0; b < n_blocks; b++) if (!taken[b] && (best < 0 || score[b] > score[best])) best = (int)b;
|
|
taken[best] = true;
|
|
}
|
|
for (uint32_t b = 0; b < n_blocks; b++) if (taken[b]) for (uint32_t t = 0; t < ratio; t++) sel[n_sel++] = (int)(b * ratio + t);
|
|
for (uint32_t t = n_blocks * ratio; t < n_vis; t++) sel[n_sel++] = (int)t;
|
|
free(taken);
|
|
}
|
|
for (uint32_t h = 0; h < H; h++) {
|
|
const uint32_t kvh = h / group;
|
|
double mx = -1e300;
|
|
for (uint32_t i = 0; i < n_sel; i++) {
|
|
double dot = 0.0;
|
|
for (uint32_t d = 0; d < D; d++) dot += qn[h * D + d] * kc[((uint64_t)sel[i] * Hkv + kvh) * D + d];
|
|
p[i] = dot * scale;
|
|
if (p[i] > mx) mx = p[i];
|
|
}
|
|
double sum = 0.0;
|
|
for (uint32_t i = 0; i < n_sel; i++) { p[i] = exp(p[i] - mx); sum += p[i]; }
|
|
const float *gate = qg + ((uint64_t)pos * H + h) * 2 * D + D;
|
|
for (uint32_t d = 0; d < D; d++) {
|
|
double acc = 0.0;
|
|
for (uint32_t i = 0; i < n_sel; i++) acc += p[i] * vc[((uint64_t)sel[i] * Hkv + kvh) * D + d];
|
|
ref_out[((uint64_t)pos * H + h) * D + d] = acc / sum * sigmoid_d(gate[d]);
|
|
}
|
|
}
|
|
}
|
|
|
|
/* GPU: decode token by token */
|
|
const uint32_t sel_stride = k_blocks * ratio + ratio;
|
|
ds4_gpu_tensor *gqg = upload(qg, (uint64_t)n_pos * H * 2 * D);
|
|
ds4_gpu_tensor *gkp = upload(kp, (uint64_t)n_pos * Hkv * D);
|
|
ds4_gpu_tensor *gvp = upload(vp, (uint64_t)n_pos * Hkv * D);
|
|
ds4_gpu_tensor *giq = upload(iq, (uint64_t)n_pos * Hi * Di);
|
|
ds4_gpu_tensor *gik = upload(ik, (uint64_t)n_pos * Di);
|
|
ds4_gpu_tensor *gq = upload(NULL, (uint64_t)H * D);
|
|
ds4_gpu_tensor *ggate = upload(NULL, (uint64_t)H * D);
|
|
ds4_gpu_tensor *gkc = ds4_gpu_tensor_alloc((uint64_t)cap * Hkv * D * 2);
|
|
ds4_gpu_tensor *gvc = ds4_gpu_tensor_alloc((uint64_t)cap * Hkv * D * 2);
|
|
ds4_gpu_tensor *giqo = upload(NULL, (uint64_t)Hi * Di);
|
|
ds4_gpu_tensor *gikc = upload(NULL, (uint64_t)cap * Di);
|
|
ds4_gpu_tensor *gbk = ds4_gpu_tensor_alloc((uint64_t)(cap / ratio + 1) * Di * 2);
|
|
ds4_gpu_tensor *gscore = upload(NULL, (uint64_t)(cap / ratio + 1));
|
|
ds4_gpu_tensor *gselb = ds4_gpu_tensor_alloc((uint64_t)k_blocks * 4);
|
|
ds4_gpu_tensor *gselt = ds4_gpu_tensor_alloc((uint64_t)sel_stride * 4);
|
|
ds4_gpu_tensor *gnsel = ds4_gpu_tensor_alloc(4);
|
|
ds4_gpu_tensor *gout = upload(NULL, (uint64_t)n_pos * H * D);
|
|
ds4_gpu_tensor *gpart = upload(NULL, ds4_gpu_qwen4_attn_part_floats(1, H, D));
|
|
ds4_gpu_tensor *gpos3 = ds4_gpu_tensor_alloc((uint64_t)cap * 16);
|
|
require_ok(gkc && gvc && gbk && gselb && gselt && gnsel && gpart && gpos3, "attention allocs");
|
|
{
|
|
uint32_t *p3 = calloc((size_t)cap * 4, sizeof(uint32_t));
|
|
for (uint32_t p = 0; p < cap; p++) p3[p * 4] = p3[p * 4 + 1] = p3[p * 4 + 2] = p;
|
|
require_ok(ds4_gpu_tensor_write(gpos3, 0, p3, (uint64_t)cap * 16), "pos3 write");
|
|
free(p3);
|
|
}
|
|
for (uint32_t pos = 0; pos < n_pos; pos++) {
|
|
ds4_gpu_tensor *vqg = ds4_gpu_tensor_view(gqg, (uint64_t)pos * H * 2 * D * 4, (uint64_t)H * 2 * D * 4);
|
|
ds4_gpu_tensor *vk = ds4_gpu_tensor_view(gkp, (uint64_t)pos * Hkv * D * 4, (uint64_t)Hkv * D * 4);
|
|
ds4_gpu_tensor *vv = ds4_gpu_tensor_view(gvp, (uint64_t)pos * Hkv * D * 4, (uint64_t)Hkv * D * 4);
|
|
ds4_gpu_tensor *viq = ds4_gpu_tensor_view(giq, (uint64_t)pos * Hi * Di * 4, (uint64_t)Hi * Di * 4);
|
|
ds4_gpu_tensor *vik = ds4_gpu_tensor_view(gik, (uint64_t)pos * Di * 4, (uint64_t)Di * 4);
|
|
ds4_gpu_tensor *vout = ds4_gpu_tensor_view(gout, (uint64_t)pos * H * D * 4, (uint64_t)H * D * 4);
|
|
require_ok(ds4_gpu_qwen4_attn_prep_tensor(gq, ggate, gkc, gvc, giqo, gikc, vqg, vk, vv, viq, vik, gpos3,
|
|
a->base, a->size, gq_off, gk_off, giq_off,
|
|
1, H, Hkv, D, n_rot, Hi, Di, pos, cap, (float)base, (float)eps), "attn prep");
|
|
const uint32_t n_blocks = (pos + 1) / ratio;
|
|
if ((pos + 1) % ratio == 0) {
|
|
require_ok(ds4_gpu_qwen4_idx_block_key_tensor(gbk, gikc, gpos3, a->base, a->size, gik_off, n_blocks - 1, 1,
|
|
ratio, Di, n_rot, (float)base, (float)eps), "block key");
|
|
}
|
|
const bool sparse = n_blocks > k_blocks;
|
|
if (sparse) {
|
|
require_ok(ds4_gpu_qwen4_idx_score_tensor(gscore, NULL, giqo, gbk, 1, n_blocks, Hi, Di, pos, ratio), "idx score");
|
|
require_ok(ds4_gpu_qwen4_idx_select_tensor(gselb, gscore, NULL, n_blocks, 1, k_blocks), "idx select");
|
|
{ /* the radix select must pick the argsort's set (order may differ) */
|
|
ds4_gpu_tensor *gref = ds4_gpu_tensor_alloc((uint64_t)k_blocks * 4);
|
|
require_ok(gref && ds4_gpu_indexer_topk_tensor(gref, gscore, n_blocks, 1, k_blocks), "topk ref");
|
|
int32_t got_sel[64], ref_sel[64];
|
|
require_ok(k_blocks <= 64 && ds4_gpu_tensor_read(gselb, 0, got_sel, k_blocks * 4) &&
|
|
ds4_gpu_tensor_read(gref, 0, ref_sel, k_blocks * 4), "sel read");
|
|
for (uint32_t i = 0; i < k_blocks; i++) {
|
|
bool found = false;
|
|
for (uint32_t j = 0; j < k_blocks; j++) found |= got_sel[j] == ref_sel[i];
|
|
require_ok(found, "radix select matches argsort set");
|
|
}
|
|
ds4_gpu_tensor_free(gref);
|
|
}
|
|
require_ok(ds4_gpu_qwen4_idx_expand_tensor(gselt, gnsel, gselb, 1, k_blocks, ratio, pos, sel_stride), "idx expand");
|
|
}
|
|
/* even positions exercise the key-split partials (main sets DS4_QWEN4_ATTN_SPLIT_KEYS=8) */
|
|
require_ok(ds4_gpu_qwen4_attn_decode_tensor(vout, gq, ggate, gkc, gvc, gselt, gnsel, (pos % 2) ? NULL : gpart,
|
|
1, H, Hkv, D, pos, sparse, sel_stride, scale), "attn decode");
|
|
ds4_gpu_tensor_free(vout); ds4_gpu_tensor_free(vik); ds4_gpu_tensor_free(viq);
|
|
ds4_gpu_tensor_free(vv); ds4_gpu_tensor_free(vk); ds4_gpu_tensor_free(vqg);
|
|
}
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "attention H=%u/%u D=%u k=%u n=%u", H, Hkv, D, k_blocks, n_pos);
|
|
check_tensor(name, gout, ref_out, (uint64_t)n_pos * H * D, 5e-4); /* half KV; block order differs */
|
|
|
|
ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gnsel); ds4_gpu_tensor_free(gselt); ds4_gpu_tensor_free(gselb);
|
|
ds4_gpu_tensor_free(gscore); ds4_gpu_tensor_free(gbk); ds4_gpu_tensor_free(gikc); ds4_gpu_tensor_free(giqo);
|
|
ds4_gpu_tensor_free(gvc); ds4_gpu_tensor_free(gkc); ds4_gpu_tensor_free(ggate); ds4_gpu_tensor_free(gq);
|
|
ds4_gpu_tensor_free(gik); ds4_gpu_tensor_free(giq); ds4_gpu_tensor_free(gvp); ds4_gpu_tensor_free(gkp); ds4_gpu_tensor_free(gqg);
|
|
free(p); free(sel); free(score); free(iqn); free(qn); free(ref_out); free(vc); free(kc);
|
|
free(ik); free(iq); free(vp); free(kp); free(qg);
|
|
free(gq_w); free(gk_w); free(giq_w); free(gik_w);
|
|
}
|
|
|
|
#ifdef __APPLE__
|
|
/* Native cross-session row kernels are Metal-only; CUDA uses ordered rows. */
|
|
static void same_bytes(const char *what, uint32_t row, const ds4_gpu_tensor *ta, uint64_t offa,
|
|
const ds4_gpu_tensor *tb, uint64_t offb, uint64_t bytes) {
|
|
uint8_t *a = malloc(bytes), *b = malloc(bytes);
|
|
require_ok(ds4_gpu_tensor_read(ta, offa, a, bytes) && ds4_gpu_tensor_read(tb, offb, b, bytes), "rows read");
|
|
for (uint64_t i = 0; i < bytes; i++) {
|
|
if (a[i] != b[i]) {
|
|
fprintf(stderr, "attention rows: %s differs for row %u at byte %llu (%02x vs %02x)\n",
|
|
what, row, (unsigned long long)i, a[i], b[i]);
|
|
exit(1);
|
|
}
|
|
}
|
|
free(a); free(b);
|
|
}
|
|
|
|
/* The rows kernels of the decode batch must reproduce the per-row dispatches
|
|
* bit for bit: the caches written, the block key, scores, selection, token
|
|
* list and output. Rows: dense, dense completing a block, sparse completing
|
|
* a block on the plain selector's width, sparse on the prefiltered width. */
|
|
static void test_attention_rows(arena_t *a) {
|
|
const uint32_t H = 8, Hkv = 2, D = 256, n_rot = 64, Hi = 4, Di = 128, k_blocks = 4, ratio = 4, R = 4;
|
|
const uint32_t sparse_pos = (k_blocks + 1) * ratio - 1;
|
|
const uint32_t pos_of[4] = { 5, 11, sparse_pos, 140 };
|
|
const double base = 1.0e7, eps = 1e-6;
|
|
const float scale = 1.0f / sqrtf((float)D);
|
|
const uint32_t sel_stride = k_blocks * ratio + ratio;
|
|
const uint64_t q_dim = (uint64_t)H * D, kv_dim = (uint64_t)Hkv * D, iq_dim = (uint64_t)Hi * Di;
|
|
double *gq_w, *gk_w, *giq_w, *gik_w;
|
|
const uint64_t gq_off = arena_f32(a, D, &gq_w, 0.5f, 1.5f);
|
|
const uint64_t gk_off = arena_f32(a, D, &gk_w, 0.5f, 1.5f);
|
|
const uint64_t giq_off = arena_f32(a, Di, &giq_w, 0.5f, 1.5f);
|
|
const uint64_t gik_off = arena_f32(a, Di, &gik_w, 0.5f, 1.5f);
|
|
uint32_t n_block_stride = 0, cap_max = 0;
|
|
for (uint32_t r = 0; r < R; r++) {
|
|
if ((pos_of[r] + 1) / ratio > n_block_stride) n_block_stride = (pos_of[r] + 1) / ratio;
|
|
if (pos_of[r] + 4 > cap_max) cap_max = pos_of[r] + 4;
|
|
}
|
|
const uint32_t n_tiles = (n_block_stride + 7) / 8;
|
|
|
|
float *qg = rand_vec(R * 2 * q_dim, 1.0f), *kp = rand_vec(R * kv_dim, 1.0f), *vp = rand_vec(R * kv_dim, 1.0f);
|
|
float *iq = rand_vec(R * iq_dim, 1.0f), *ik = rand_vec(R * Di, 1.0f);
|
|
ds4_gpu_tensor *gqg = upload(qg, R * 2 * q_dim), *gkp = upload(kp, R * kv_dim), *gvp = upload(vp, R * kv_dim);
|
|
ds4_gpu_tensor *giq = upload(iq, R * iq_dim), *gik = upload(ik, R * Di);
|
|
ds4_gpu_tensor *q[2], *gate[2], *iqn[2], *out[2], *score[2], *tile_max[2], *selb[2], *selt[2], *nsel[2];
|
|
ds4_gpu_tensor *kc[2][4], *vc[2][4], *ikc[2][4], *bk[2][4], *pos3[4];
|
|
for (int p = 0; p < 2; p++) {
|
|
q[p] = upload(NULL, R * q_dim); gate[p] = upload(NULL, R * q_dim); iqn[p] = upload(NULL, R * iq_dim);
|
|
out[p] = upload(NULL, R * q_dim); score[p] = upload(NULL, (uint64_t)R * n_block_stride);
|
|
tile_max[p] = ds4_gpu_tensor_alloc((uint64_t)R * n_tiles * 4);
|
|
selb[p] = ds4_gpu_tensor_alloc((uint64_t)R * k_blocks * 4);
|
|
selt[p] = ds4_gpu_tensor_alloc((uint64_t)R * sel_stride * 4);
|
|
nsel[p] = ds4_gpu_tensor_alloc((uint64_t)R * 4);
|
|
require_ok(tile_max[p] && selb[p] && selt[p] && nsel[p], "attention rows allocs");
|
|
}
|
|
/* identical random histories for both paths */
|
|
for (uint32_t r = 0; r < R; r++) {
|
|
const uint32_t cap = pos_of[r] + 4, n_bk = cap / ratio + 1;
|
|
uint16_t *kh = malloc(cap * kv_dim * 2), *vh = malloc(cap * kv_dim * 2), *bh = malloc((uint64_t)n_bk * Di * 2);
|
|
float *ikh = rand_vec((uint64_t)cap * Di, 1.0f);
|
|
uint32_t *p3 = calloc((size_t)cap * 4, sizeof(uint32_t));
|
|
for (uint64_t i = 0; i < cap * kv_dim; i++) { kh[i] = f32_to_f16(frand()); vh[i] = f32_to_f16(frand()); }
|
|
for (uint64_t i = 0; i < (uint64_t)n_bk * Di; i++) bh[i] = f32_to_f16(frand());
|
|
for (uint32_t p = 0; p < cap; p++) p3[p * 4] = p3[p * 4 + 1] = p3[p * 4 + 2] = p;
|
|
pos3[r] = ds4_gpu_tensor_alloc((uint64_t)cap * 16);
|
|
require_ok(pos3[r] && ds4_gpu_tensor_write(pos3[r], 0, p3, (uint64_t)cap * 16), "rows pos3");
|
|
for (int p = 0; p < 2; p++) {
|
|
kc[p][r] = ds4_gpu_tensor_alloc(cap * kv_dim * 2); vc[p][r] = ds4_gpu_tensor_alloc(cap * kv_dim * 2);
|
|
ikc[p][r] = ds4_gpu_tensor_alloc((uint64_t)cap * Di * 4); bk[p][r] = ds4_gpu_tensor_alloc((uint64_t)n_bk * Di * 2);
|
|
require_ok(kc[p][r] && vc[p][r] && ikc[p][r] && bk[p][r] &&
|
|
ds4_gpu_tensor_write(kc[p][r], 0, kh, cap * kv_dim * 2) &&
|
|
ds4_gpu_tensor_write(vc[p][r], 0, vh, cap * kv_dim * 2) &&
|
|
ds4_gpu_tensor_write(ikc[p][r], 0, ikh, (uint64_t)cap * Di * 4) &&
|
|
ds4_gpu_tensor_write(bk[p][r], 0, bh, (uint64_t)n_bk * Di * 2), "rows caches");
|
|
}
|
|
free(kh); free(vh); free(bh); free(ikh); free(p3);
|
|
}
|
|
|
|
/* path A: the per-row dispatches on row views */
|
|
ds4_gpu_tensor *part1 = upload(NULL, ds4_gpu_qwen4_attn_part_floats(1, H, D));
|
|
for (uint32_t r = 0; r < R; r++) {
|
|
const uint32_t pos = pos_of[r], n_blocks = (pos + 1) / ratio;
|
|
const bool sparse = pos >= sparse_pos;
|
|
ds4_gpu_tensor *v[16] = {
|
|
ds4_gpu_tensor_view(gqg, r * 2 * q_dim * 4, 2 * q_dim * 4), ds4_gpu_tensor_view(gkp, r * kv_dim * 4, kv_dim * 4),
|
|
ds4_gpu_tensor_view(gvp, r * kv_dim * 4, kv_dim * 4), ds4_gpu_tensor_view(giq, r * iq_dim * 4, iq_dim * 4),
|
|
ds4_gpu_tensor_view(gik, (uint64_t)r * Di * 4, Di * 4), ds4_gpu_tensor_view(q[0], r * q_dim * 4, q_dim * 4),
|
|
ds4_gpu_tensor_view(gate[0], r * q_dim * 4, q_dim * 4), ds4_gpu_tensor_view(iqn[0], r * iq_dim * 4, iq_dim * 4),
|
|
ds4_gpu_tensor_view(out[0], r * q_dim * 4, q_dim * 4),
|
|
ds4_gpu_tensor_view(score[0], (uint64_t)r * n_block_stride * 4, (uint64_t)n_block_stride * 4),
|
|
ds4_gpu_tensor_view(tile_max[0], (uint64_t)r * n_tiles * 4, (uint64_t)n_tiles * 4),
|
|
ds4_gpu_tensor_view(selb[0], (uint64_t)r * k_blocks * 4, (uint64_t)k_blocks * 4),
|
|
ds4_gpu_tensor_view(selt[0], (uint64_t)r * sel_stride * 4, (uint64_t)sel_stride * 4),
|
|
ds4_gpu_tensor_view(nsel[0], (uint64_t)r * 4, 4), NULL, NULL };
|
|
for (int i = 0; i < 14; i++) require_ok(v[i] != NULL, "rows views");
|
|
require_ok(ds4_gpu_qwen4_attn_prep_tensor(v[5], v[6], kc[0][r], vc[0][r], v[7], ikc[0][r], v[0], v[1], v[2], v[3], v[4],
|
|
pos3[r], a->base, a->size, gq_off, gk_off, giq_off, 1, H, Hkv, D, n_rot, Hi, Di,
|
|
pos, pos + 4, (float)base, (float)eps), "rows: prep");
|
|
if ((pos + 1) % ratio == 0) {
|
|
require_ok(ds4_gpu_qwen4_idx_block_key_tensor(bk[0][r], ikc[0][r], pos3[r], a->base, a->size, gik_off, n_blocks - 1, 1,
|
|
ratio, Di, n_rot, (float)base, (float)eps), "rows: block key");
|
|
}
|
|
if (sparse) {
|
|
require_ok(ds4_gpu_qwen4_idx_score_tensor(v[9], v[10], v[7], bk[0][r], 1, n_blocks, Hi, Di, pos, ratio), "rows: score");
|
|
require_ok(ds4_gpu_qwen4_idx_select_tensor(v[11], v[9], v[10], n_blocks, 1, k_blocks), "rows: select");
|
|
require_ok(ds4_gpu_qwen4_idx_expand_tensor(v[12], v[13], v[11], 1, k_blocks, ratio, pos, sel_stride), "rows: expand");
|
|
}
|
|
require_ok(ds4_gpu_qwen4_attn_decode_tensor(v[8], v[5], v[6], kc[0][r], vc[0][r], v[12], v[13], part1,
|
|
1, H, Hkv, D, pos, sparse, sel_stride, scale), "rows: decode");
|
|
for (int i = 0; i < 14; i++) ds4_gpu_tensor_free(v[i]);
|
|
}
|
|
|
|
/* path B: the rows kernels */
|
|
ds4_gpu_tensor *table = ds4_gpu_tensor_alloc((uint64_t)R * DS4_GPU_QWEN4_ATTN_ROW_BYTES);
|
|
ds4_gpu_tensor *partR = upload(NULL, ds4_gpu_qwen4_attn_part_floats(R, H, D));
|
|
ds4_gpu_qwen4_attn_row rows[4];
|
|
for (uint32_t r = 0; r < R; r++) {
|
|
rows[r].k_cache = kc[1][r]; rows[r].v_cache = vc[1][r]; rows[r].ik_cache = ikc[1][r]; rows[r].block_key = bk[1][r];
|
|
rows[r].pos3 = pos3[r]; rows[r].pos = pos_of[r]; rows[r].use_sel = pos_of[r] >= sparse_pos;
|
|
}
|
|
require_ok(table && ds4_gpu_qwen4_attn_rows_stage(table, 0, rows, R, ratio), "rows: stage");
|
|
require_ok(ds4_gpu_qwen4_attn_prep_rows_tensor(q[1], gate[1], iqn[1], gqg, gkp, gvp, giq, gik, table, 0, rows, R,
|
|
a->base, a->size, gq_off, gk_off, giq_off, H, Hkv, D, n_rot, Hi, Di,
|
|
(float)base, (float)eps), "rows: prep rows");
|
|
require_ok(ds4_gpu_qwen4_idx_block_key_rows_tensor(table, 0, rows, R, a->base, a->size, gik_off, ratio, Di, n_rot,
|
|
(float)base, (float)eps), "rows: block key rows");
|
|
require_ok(ds4_gpu_qwen4_idx_score_rows_tensor(score[1], tile_max[1], iqn[1], table, 0, rows, R, n_block_stride, Hi, Di, ratio),
|
|
"rows: score rows");
|
|
require_ok(ds4_gpu_qwen4_idx_select_rows_tensor(selb[1], score[1], tile_max[1], table, 0, rows, R, n_block_stride, k_blocks),
|
|
"rows: select rows");
|
|
require_ok(ds4_gpu_qwen4_idx_expand_rows_tensor(selt[1], nsel[1], selb[1], table, 0, R, k_blocks, ratio, sel_stride),
|
|
"rows: expand rows");
|
|
require_ok(ds4_gpu_qwen4_attn_decode_rows_tensor(out[1], q[1], gate[1], selt[1], nsel[1], partR, table, 0, rows, R,
|
|
H, Hkv, D, sel_stride, scale), "rows: decode rows");
|
|
|
|
same_bytes("q", R, q[0], 0, q[1], 0, R * q_dim * 4);
|
|
same_bytes("gate", R, gate[0], 0, gate[1], 0, R * q_dim * 4);
|
|
same_bytes("indexer q", R, iqn[0], 0, iqn[1], 0, R * iq_dim * 4);
|
|
same_bytes("output", R, out[0], 0, out[1], 0, R * q_dim * 4);
|
|
for (uint32_t r = 0; r < R; r++) {
|
|
const uint32_t pos = pos_of[r], cap = pos + 4, n_blocks = (pos + 1) / ratio;
|
|
same_bytes("k cache", r, kc[0][r], 0, kc[1][r], 0, cap * kv_dim * 2);
|
|
same_bytes("v cache", r, vc[0][r], 0, vc[1][r], 0, cap * kv_dim * 2);
|
|
same_bytes("indexer k cache", r, ikc[0][r], 0, ikc[1][r], 0, (uint64_t)cap * Di * 4);
|
|
same_bytes("block keys", r, bk[0][r], 0, bk[1][r], 0, (uint64_t)(cap / ratio + 1) * Di * 2);
|
|
if (pos < sparse_pos) continue;
|
|
uint32_t n_sel_a = 0;
|
|
require_ok(ds4_gpu_tensor_read(nsel[0], (uint64_t)r * 4, &n_sel_a, 4), "rows: n_sel read");
|
|
same_bytes("scores", r, score[0], (uint64_t)r * n_block_stride * 4, score[1], (uint64_t)r * n_block_stride * 4,
|
|
(uint64_t)n_blocks * 4);
|
|
same_bytes("selected blocks", r, selb[0], (uint64_t)r * k_blocks * 4, selb[1], (uint64_t)r * k_blocks * 4,
|
|
(uint64_t)k_blocks * 4);
|
|
same_bytes("selected count", r, nsel[0], (uint64_t)r * 4, nsel[1], (uint64_t)r * 4, 4);
|
|
same_bytes("selected tokens", r, selt[0], (uint64_t)r * sel_stride * 4, selt[1], (uint64_t)r * sel_stride * 4,
|
|
(uint64_t)n_sel_a * 4);
|
|
}
|
|
printf(" attention rows: dense, block-completing and sparse rows byte-exact against the per-row kernels\n");
|
|
|
|
ds4_gpu_tensor_free(partR); ds4_gpu_tensor_free(table); ds4_gpu_tensor_free(part1);
|
|
for (uint32_t r = 0; r < R; r++) {
|
|
ds4_gpu_tensor_free(pos3[r]);
|
|
for (int p = 0; p < 2; p++) {
|
|
ds4_gpu_tensor_free(kc[p][r]); ds4_gpu_tensor_free(vc[p][r]);
|
|
ds4_gpu_tensor_free(ikc[p][r]); ds4_gpu_tensor_free(bk[p][r]);
|
|
}
|
|
}
|
|
for (int p = 0; p < 2; p++) {
|
|
ds4_gpu_tensor_free(q[p]); ds4_gpu_tensor_free(gate[p]); ds4_gpu_tensor_free(iqn[p]); ds4_gpu_tensor_free(out[p]);
|
|
ds4_gpu_tensor_free(score[p]); ds4_gpu_tensor_free(tile_max[p]); ds4_gpu_tensor_free(selb[p]);
|
|
ds4_gpu_tensor_free(selt[p]); ds4_gpu_tensor_free(nsel[p]);
|
|
}
|
|
ds4_gpu_tensor_free(gik); ds4_gpu_tensor_free(giq); ds4_gpu_tensor_free(gvp); ds4_gpu_tensor_free(gkp); ds4_gpu_tensor_free(gqg);
|
|
free(ik); free(iq); free(vp); free(kp); free(qg);
|
|
free(gq_w); free(gk_w); free(giq_w); free(gik_w);
|
|
}
|
|
|
|
#endif
|
|
|
|
/* ---- routed experts ---- */
|
|
|
|
static uint64_t arena_tier(arena_t *a, uint32_t wtype, uint64_t rows, uint64_t cols, double **shadow) {
|
|
if (wtype == 39u) return arena_mxfp4(a, rows, cols, shadow);
|
|
return wtype == 12u ? arena_q4_K(a, rows, cols, shadow, 0.05f) :
|
|
wtype == 10u ? arena_q2_K(a, rows, cols, shadow, 0.05f) :
|
|
wtype == 16u ? arena_iq2_xxs(a, rows, cols, shadow, 0.05f) : arena_q8_0(a, rows, cols, shadow, 0.05f);
|
|
}
|
|
|
|
static void check_exact_f32(const char *what, const float *got, const float *ref, uint64_t n);
|
|
|
|
/* The vector scorer must reproduce the scalar scorer bit for bit and emit
|
|
* the selector's tile keys; the prefiltered decode select must reproduce
|
|
* the full-row select bit for bit, including tie order at both ends. */
|
|
static void test_idx_prefilter(void) {
|
|
const uint32_t Hi = 4, Di = 128, k = 512;
|
|
const uint32_t sizes[3] = { 4104u, 16384u, 65536u };
|
|
for (uint32_t si = 0; si < 3; si++) {
|
|
for (uint32_t T = 1; T <= 2; T++) {
|
|
const uint32_t n = sizes[si], n_tiles = (n + 7u) / 8u;
|
|
const uint32_t ratio = 4, pos0 = ratio * (n - 3u) + 1u; /* the last tokens see n-2.. blocks */
|
|
float *q = rand_vec((uint64_t)T * Hi * Di, 1.0f);
|
|
float *keyf = rand_vec((uint64_t)n * Di, 1.0f);
|
|
uint16_t *keyh = malloc((uint64_t)n * Di * 2);
|
|
for (uint64_t i = 0; i < (uint64_t)n * Di; i++) keyh[i] = f32_to_f16(keyf[i]);
|
|
ds4_gpu_tensor *gq = upload(q, (uint64_t)T * Hi * Di);
|
|
ds4_gpu_tensor *gk = ds4_gpu_tensor_alloc((uint64_t)n * Di * 2);
|
|
ds4_gpu_tensor *gs0 = upload(NULL, (uint64_t)T * n), *gs1 = upload(NULL, (uint64_t)T * n);
|
|
ds4_gpu_tensor *gtm = ds4_gpu_tensor_alloc((uint64_t)T * n_tiles * 4);
|
|
require_ok(gq && gk && gs0 && gs1 && gtm && ds4_gpu_tensor_write(gk, 0, keyh, (uint64_t)n * Di * 2), "prefilter setup");
|
|
setenv("DS4_QWEN4_IDX_SCORE_VEC", "0", 1);
|
|
require_ok(ds4_gpu_qwen4_idx_score_tensor(gs0, NULL, gq, gk, T, n, Hi, Di, pos0, ratio), "scalar score");
|
|
setenv("DS4_QWEN4_IDX_SCORE_VEC", "1", 1);
|
|
require_ok(ds4_gpu_qwen4_idx_score_tensor(gs1, gtm, gq, gk, T, n, Hi, Di, pos0, ratio), "vector score");
|
|
float *s0 = download(gs0, (uint64_t)T * n), *s1 = download(gs1, (uint64_t)T * n);
|
|
require_ok(memcmp(s0, s1, (uint64_t)T * n * 4) == 0, "vector scorer matches the scalar scorer");
|
|
uint32_t *tm = malloc((uint64_t)T * n_tiles * 4);
|
|
require_ok(ds4_gpu_tensor_read(gtm, 0, tm, (uint64_t)T * n_tiles * 4), "tile max read");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t tile = 0; tile < n_tiles; tile++) {
|
|
uint32_t m = 0;
|
|
for (uint32_t j = 0; j < 8 && tile * 8 + j < n; j++) {
|
|
const float v = fmaxf(s1[(uint64_t)t * n + tile * 8 + j], 0.0f);
|
|
uint32_t key; memcpy(&key, &v, 4);
|
|
if (key > m) m = key;
|
|
}
|
|
require_ok(tm[(uint64_t)t * n_tiles + tile] == m, "tile maximum matches the score keys");
|
|
}
|
|
}
|
|
/* selection on the real scores, then on rows with ties at both ends */
|
|
for (uint32_t variant = 0; variant < 2; variant++) {
|
|
if (variant) {
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t b = 0; b < n; b++) {
|
|
float *v = &s1[(uint64_t)t * n + b];
|
|
if (*v < -1e37f) continue;
|
|
if ((b % 5) == 2) *v = s1[(uint64_t)t * n + (b / 5) * 5]; /* duplicates */
|
|
if ((b % 1000) == 999) *v = 64.0f; /* top ties across tiles */
|
|
if ((b % 3) == 1) *v = 0.0f; /* zeros: bottom ties */
|
|
}
|
|
}
|
|
require_ok(ds4_gpu_tensor_write(gs1, 0, s1, (uint64_t)T * n * 4), "tie rows upload");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t tile = 0; tile < n_tiles; tile++) {
|
|
uint32_t m = 0;
|
|
for (uint32_t j = 0; j < 8 && tile * 8 + j < n; j++) {
|
|
const float v = fmaxf(s1[(uint64_t)t * n + tile * 8 + j], 0.0f);
|
|
uint32_t key; memcpy(&key, &v, 4);
|
|
if (key > m) m = key;
|
|
}
|
|
tm[(uint64_t)t * n_tiles + tile] = m;
|
|
}
|
|
}
|
|
require_ok(ds4_gpu_tensor_write(gtm, 0, tm, (uint64_t)T * n_tiles * 4), "tie tiles upload");
|
|
}
|
|
ds4_gpu_tensor *gfull = ds4_gpu_tensor_alloc((uint64_t)T * k * 4), *gpre = ds4_gpu_tensor_alloc((uint64_t)T * k * 4);
|
|
setenv("DS4_QWEN4_IDX_PREFILTER", "0", 1);
|
|
require_ok(gfull && gpre && ds4_gpu_qwen4_idx_select_tensor(gfull, gs1, gtm, n, T, k), "full select");
|
|
setenv("DS4_QWEN4_IDX_PREFILTER", "1", 1);
|
|
require_ok(ds4_gpu_qwen4_idx_select_tensor(gpre, gs1, gtm, n, T, k), "prefiltered select");
|
|
unsetenv("DS4_QWEN4_IDX_PREFILTER");
|
|
int32_t *full = malloc((uint64_t)T * k * 4), *pre = malloc((uint64_t)T * k * 4);
|
|
require_ok(ds4_gpu_tensor_read(gfull, 0, full, (uint64_t)T * k * 4) &&
|
|
ds4_gpu_tensor_read(gpre, 0, pre, (uint64_t)T * k * 4), "select read");
|
|
require_ok(memcmp(full, pre, (uint64_t)T * k * 4) == 0, "prefiltered select matches the full select");
|
|
/* A scalar-scorer override leaves tile maxima stale. Poison
|
|
* them and require the full selector fallback to remain exact. */
|
|
setenv("DS4_QWEN4_IDX_SCORE_VEC", "0", 1);
|
|
setenv("DS4_QWEN4_IDX_PREFILTER", "1", 1);
|
|
memset(tm, 0, (uint64_t)T * n_tiles * 4);
|
|
for (uint32_t t = 0; t < T; t++)
|
|
for (uint32_t tile = 0; tile < k; tile++) tm[(uint64_t)t * n_tiles + tile] = 0x7f7fffffu;
|
|
require_ok(ds4_gpu_tensor_write(gtm, 0, tm, (uint64_t)T * n_tiles * 4), "stale tiles upload");
|
|
require_ok(ds4_gpu_qwen4_idx_select_tensor(gpre, gs1, gtm, n, T, k) &&
|
|
ds4_gpu_tensor_read(gpre, 0, pre, (uint64_t)T * k * 4), "scalar scorer fallback");
|
|
require_ok(memcmp(full, pre, (uint64_t)T * k * 4) == 0, "scalar override ignores stale tile maxima");
|
|
setenv("DS4_QWEN4_IDX_SCORE_VEC", "1", 1);
|
|
unsetenv("DS4_QWEN4_IDX_PREFILTER");
|
|
free(full); free(pre);
|
|
ds4_gpu_tensor_free(gfull); ds4_gpu_tensor_free(gpre);
|
|
}
|
|
unsetenv("DS4_QWEN4_IDX_SCORE_VEC");
|
|
printf(" idx prefilter n=%u T=%u: vector scores, tile keys and selections byte-exact\n", n, T);
|
|
free(q); free(keyf); free(keyh); free(s0); free(s1); free(tm);
|
|
ds4_gpu_tensor_free(gq); ds4_gpu_tensor_free(gk); ds4_gpu_tensor_free(gs0); ds4_gpu_tensor_free(gs1); ds4_gpu_tensor_free(gtm);
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Routed gate/up and down types can differ; shared experts stay Q8_0
|
|
* for quantized cases and F32 for the F32 case. */
|
|
static void test_moe_types(arena_t *a, uint32_t NE, uint32_t slots, uint32_t E, uint32_t F,
|
|
uint32_t T, uint32_t wtype, uint32_t dtype) {
|
|
/* Nonzero random padding ensures kernels ignore the physical tail. */
|
|
const uint32_t DF = dtype == 10u ? (F + 255u) / 256u * 256u : F;
|
|
double *gate_w, *up_w, *down_w, *sg_w, *su_w, *sd_w;
|
|
uint64_t gate_off, up_off, down_off, sg_off, su_off, sd_off;
|
|
const bool q8 = wtype != 0u;
|
|
const char *tier_name = wtype == 12u ? "q4_K" : wtype == 10u ? "q2_K" : wtype == 16u ? "iq2_xxs" : wtype == 39u ? "mxfp4" : q8 ? "q8_0" : "f32";
|
|
if (q8) {
|
|
gate_off = arena_tier(a, wtype, (uint64_t)NE * F, E, &gate_w);
|
|
up_off = arena_tier(a, wtype, (uint64_t)NE * F, E, &up_w);
|
|
down_off = arena_tier(a, dtype, (uint64_t)NE * E, DF, &down_w);
|
|
sg_off = arena_q8_0(a, F, E, &sg_w, 0.05f);
|
|
su_off = arena_q8_0(a, F, E, &su_w, 0.05f);
|
|
sd_off = arena_q8_0(a, E, F, &sd_w, 0.05f);
|
|
} else {
|
|
gate_off = arena_f32(a, (uint64_t)NE * F * E, &gate_w, -0.05f, 0.05f);
|
|
up_off = arena_f32(a, (uint64_t)NE * F * E, &up_w, -0.05f, 0.05f);
|
|
down_off = arena_f32(a, (uint64_t)NE * E * F, &down_w, -0.05f, 0.05f);
|
|
sg_off = arena_f32(a, (uint64_t)F * E, &sg_w, -0.05f, 0.05f);
|
|
su_off = arena_f32(a, (uint64_t)F * E, &su_w, -0.05f, 0.05f);
|
|
sd_off = arena_f32(a, (uint64_t)E * F, &sd_w, -0.05f, 0.05f);
|
|
}
|
|
const uint32_t n_out = slots + 1; /* + shared expert */
|
|
float *x = rand_vec((uint64_t)T * E, 1.0f);
|
|
float *sgate = rand_vec(T, 2.0f);
|
|
int32_t *sel = malloc((uint64_t)T * slots * 4);
|
|
float *w = malloc((uint64_t)T * slots * 4);
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < slots; s++) {
|
|
sel[t * slots + s] = (int32_t)((t * 7u + s * 3u) % NE);
|
|
w[t * slots + s] = 0.5f * frand() + 0.6f;
|
|
}
|
|
}
|
|
double *mid = malloc((uint64_t)T * n_out * F * sizeof(double));
|
|
double *part = malloc((uint64_t)T * n_out * E * sizeof(double));
|
|
double *out = malloc((uint64_t)T * E * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t d = 0; d < E; d++) out[t * E + d] = 0.0;
|
|
for (uint32_t s = 0; s < n_out; s++) {
|
|
const bool shared = s == slots;
|
|
const uint32_t e = shared ? 0 : (uint32_t)sel[t * slots + s];
|
|
const double *gw = shared ? sg_w : gate_w + (uint64_t)e * F * E;
|
|
const double *uw = shared ? su_w : up_w + (uint64_t)e * F * E;
|
|
const uint32_t stride = shared ? F : DF;
|
|
const double *dw = shared ? sd_w : down_w + (uint64_t)e * E * DF;
|
|
for (uint32_t f = 0; f < F; f++) {
|
|
double g = 0.0, u = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) {
|
|
g += gw[(uint64_t)f * E + i] * x[t * E + i];
|
|
u += uw[(uint64_t)f * E + i] * x[t * E + i];
|
|
}
|
|
mid[((uint64_t)t * n_out + s) * F + f] = silu_d(g) * u;
|
|
}
|
|
const double wgt = shared ? sigmoid_d(sgate[t]) : w[t * slots + s];
|
|
for (uint32_t d = 0; d < E; d++) {
|
|
double acc = 0.0;
|
|
for (uint32_t f = 0; f < F; f++) acc += dw[(uint64_t)d * stride + f] * mid[((uint64_t)t * n_out + s) * F + f];
|
|
part[((uint64_t)t * n_out + s) * E + d] = acc;
|
|
out[t * E + d] += wgt * acc;
|
|
}
|
|
}
|
|
}
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * slots * 4);
|
|
require_ok(ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * slots * 4), "sel write");
|
|
ds4_gpu_tensor *gw = upload(w, (uint64_t)T * slots);
|
|
ds4_gpu_tensor *gmid = upload(NULL, (uint64_t)T * n_out * F);
|
|
ds4_gpu_tensor *gpart = upload(NULL, (uint64_t)T * n_out * E);
|
|
ds4_gpu_tensor *gsg = upload(sgate, T);
|
|
ds4_gpu_tensor *gout = upload(NULL, (uint64_t)T * E);
|
|
const uint32_t shared_type = q8 ? 8u : 0u;
|
|
require_ok(ds4_gpu_qwen4_moe_mid_tensor(gmid, gx, gsel, a->base, a->size, gate_off, up_off, wtype, NE, T, slots, E, F,
|
|
sg_off, su_off, shared_type), "moe mid");
|
|
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size, down_off, dtype, NE, T, slots, F, E,
|
|
sd_off, shared_type), "moe down");
|
|
if ((dtype == 10u || dtype == 39u) && getenv("DS4_TEST_QWEN4_MV_EXACT")) {
|
|
/* Specialized decode row kernels must match the generic dispatch bit
|
|
* for bit across every NR/NSG geometry; Q4_K gate/up keep their own
|
|
* ordered kernel, so a Q4_K/MXFP4 pair exercises the down rows. */
|
|
const uint64_t nm = (uint64_t)T * n_out * F, np = (uint64_t)T * n_out * E;
|
|
float *bm = malloc(nm * sizeof(float)), *bp = malloc(np * sizeof(float));
|
|
float *am = malloc(nm * sizeof(float)), *ap = malloc(np * sizeof(float));
|
|
require_ok(bm && bp && am && ap, "MoE exact allocation");
|
|
for (uint32_t mode = 0; mode < 9; mode++) {
|
|
const char *nr[] = {"2", "1", "1", "2", "2", "4", "4", "1", "2"};
|
|
const char *nsg[] = {"4", "4", "8", "4", "8", "4", "8", "16", "16"};
|
|
setenv("DS4_QWEN4_MOE_MV_SPECIALIZE", mode ? "1" : "0", 1);
|
|
setenv("DS4_QWEN4_MOE_MV_NR", nr[mode], 1);
|
|
setenv("DS4_QWEN4_MOE_MV_NSG", nsg[mode], 1);
|
|
require_ok(ds4_gpu_qwen4_moe_mid_tensor(gmid, gx, gsel, a->base, a->size,
|
|
gate_off, up_off, wtype, NE, T, slots, E, F, sg_off, su_off, shared_type), "exact mid");
|
|
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size,
|
|
down_off, dtype, NE, T, slots, F, E, sd_off, shared_type), "exact down");
|
|
require_ok(ds4_gpu_tensor_read(gmid, 0, am, nm * sizeof(float)) &&
|
|
ds4_gpu_tensor_read(gpart, 0, ap, np * sizeof(float)), "exact read");
|
|
if (!mode) { memcpy(bm, am, nm * sizeof(float)); memcpy(bp, ap, np * sizeof(float)); }
|
|
else { check_exact_f32(dtype == 10u ? "specialized IQ2 mid" : "Q4K mid beside specialized down", am, bm, nm);
|
|
check_exact_f32(dtype == 10u ? "specialized Q2 down" : "specialized MXFP4 down", ap, bp, np); }
|
|
}
|
|
unsetenv("DS4_QWEN4_MOE_MV_SPECIALIZE");
|
|
unsetenv("DS4_QWEN4_MOE_MV_NR");
|
|
unsetenv("DS4_QWEN4_MOE_MV_NSG");
|
|
free(bm); free(bp); free(am); free(ap);
|
|
}
|
|
if (dtype == 39u && getenv("DS4_TEST_QWEN4_MV_EXACT")) {
|
|
/* The prefetched MXFP4 down rows must match the plain kernel byte for
|
|
* bit, shared Q8 slot included, at the default and generic geometries. */
|
|
const uint64_t np = (uint64_t)T * n_out * E;
|
|
float *bp = malloc(np * sizeof(float)), *ap = malloc(np * sizeof(float));
|
|
require_ok(bp && ap, "down prefetch allocation");
|
|
for (uint32_t spec = 0; spec < 2u; spec++) {
|
|
setenv("DS4_QWEN4_MOE_MV_SPECIALIZE", spec ? "1" : "0", 1);
|
|
for (uint32_t mode = 0; mode < 2u; mode++) {
|
|
setenv("DS4_QWEN4_MOE_DOWN_PREFETCH", mode ? "1" : "0", 1);
|
|
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size,
|
|
down_off, dtype, NE, T, slots, F, E, sd_off, shared_type), "down prefetch dispatch");
|
|
require_ok(ds4_gpu_tensor_read(gpart, 0, mode ? ap : bp, np * sizeof(float)), "down prefetch read");
|
|
}
|
|
check_exact_f32(spec ? "prefetched MXFP4 down, specialized" : "prefetched MXFP4 down, generic", ap, bp, np);
|
|
}
|
|
unsetenv("DS4_QWEN4_MOE_MV_SPECIALIZE");
|
|
unsetenv("DS4_QWEN4_MOE_DOWN_PREFETCH");
|
|
free(bp); free(ap);
|
|
}
|
|
if (dtype == 10u) {
|
|
require_ok(!ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size, down_off,
|
|
dtype, NE, T, slots, F + 1u, E, 0, UINT32_MAX),
|
|
"Q2_K down rejects a partial activation group");
|
|
}
|
|
char name[96];
|
|
const uint32_t CH = DS4_QWEN4_HC_CHUNKS;
|
|
float *R0 = rand_vec((uint64_t)T * 4 * E, 1.0f);
|
|
float *injv = rand_vec((uint64_t)T * 4 * CH * 4, 1.0f);
|
|
double *R_ref = malloc((uint64_t)T * 4 * E * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++)
|
|
for (uint32_t s = 0; s < 4; s++) {
|
|
double tot = 0.0;
|
|
for (uint32_t src = 0; src < 4 * CH; src++) tot += injv[((uint64_t)t * 4 * CH + src) * 4 + s];
|
|
const double wgt = 2.0 * sigmoid_d(tot / 4.0);
|
|
for (uint32_t d = 0; d < E; d++)
|
|
R_ref[((uint64_t)t * 4 + s) * E + d] = (double)R0[((uint64_t)t * 4 + s) * E + d] + wgt * out[t * E + d];
|
|
}
|
|
ds4_gpu_tensor *gR = upload(R0, (uint64_t)T * 4 * E);
|
|
ds4_gpu_tensor *ginj = upload(injv, (uint64_t)T * 4 * CH * 4);
|
|
require_ok(ds4_gpu_qwen4_moe_reduce_tensor(gout, gpart, gw, gsg, NULL, gR, ginj, T, slots, n_out, E, 4), "moe reduce");
|
|
const char *dname = dtype == 10u ? "q2_K" : dtype == 39u ? "mxfp4" : q8 ? "q8_0" : "f32";
|
|
snprintf(name, sizeof(name), "moe %s E=%u F=%u slots=%u T=%u: reduce+combine", dname, E, F, slots, T);
|
|
check_tensor(name, gR, R_ref, (uint64_t)T * 4 * E, 2e-5);
|
|
ds4_gpu_tensor_free(ginj); ds4_gpu_tensor_free(gR); free(R_ref); free(injv); free(R0);
|
|
snprintf(name, sizeof(name), "moe %s E=%u F=%u slots=%u T=%u: mid (+shared)", tier_name, E, F, slots, T);
|
|
check_tensor(name, gmid, mid, (uint64_t)T * n_out * F, 2e-5);
|
|
snprintf(name, sizeof(name), "moe %s E=%u F=%u slots=%u T=%u: down (+shared)", dname, E, F, slots, T);
|
|
check_tensor(name, gpart, part, (uint64_t)T * n_out * E, 2e-5);
|
|
snprintf(name, sizeof(name), "moe %s E=%u F=%u slots=%u T=%u: out", dname, E, F, slots, T);
|
|
check_tensor(name, gout, out, (uint64_t)T * E, 2e-5);
|
|
if (q8 && (E % 64) == 0 && (F % 64) == 0) {
|
|
/* prefill path: lists + tiled GEMMs over routed slots only (f16 tiles) */
|
|
ds4_gpu_tensor *glists = ds4_gpu_tensor_alloc((uint64_t)NE * T * 4);
|
|
ds4_gpu_tensor *gcounts = ds4_gpu_tensor_alloc((uint64_t)NE * 4);
|
|
ds4_gpu_tensor *gmid2 = upload(NULL, (uint64_t)T * slots * F);
|
|
ds4_gpu_tensor *gpart2 = upload(NULL, (uint64_t)T * slots * E);
|
|
require_ok(ds4_gpu_qwen4_moe_build_lists_tensor(glists, gcounts, gsel, T, slots, NE, T), "moe lists");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_mid_tensor(gmid2, gx, glists, gcounts, a->base, a->size, gate_off, up_off, wtype,
|
|
NE, T, slots, slots, E, F, T), "moe mm mid");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_down_tensor(gpart2, gmid2, glists, gcounts, a->base, a->size, down_off, dtype,
|
|
NE, T, slots, slots, F, E, T), "moe mm down");
|
|
double *mid_r = malloc((uint64_t)T * slots * F * sizeof(double));
|
|
double *part_r = malloc((uint64_t)T * slots * E * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++)
|
|
for (uint32_t s = 0; s < slots; s++) {
|
|
memcpy(mid_r + ((uint64_t)t * slots + s) * F, mid + ((uint64_t)t * n_out + s) * F, F * sizeof(double));
|
|
memcpy(part_r + ((uint64_t)t * slots + s) * E, part + ((uint64_t)t * n_out + s) * E, E * sizeof(double));
|
|
}
|
|
int32_t *cnt = malloc((uint64_t)NE * 4);
|
|
require_ok(ds4_gpu_tensor_read(gcounts, 0, cnt, (uint64_t)NE * 4), "counts read");
|
|
int total = 0;
|
|
for (uint32_t e = 0; e < NE; e++) total += cnt[e];
|
|
if (total != (int)(T * slots)) { fprintf(stderr, "moe lists: %d pairs != %u\n", total, T * slots); exit(1); }
|
|
free(cnt);
|
|
snprintf(name, sizeof(name), "moe mm %s E=%u F=%u slots=%u T=%u: mid", tier_name, E, F, slots, T);
|
|
check_tensor(name, gmid2, mid_r, (uint64_t)T * slots * F, 3e-3);
|
|
snprintf(name, sizeof(name), "moe mm %s E=%u F=%u slots=%u T=%u: down", tier_name, E, F, slots, T);
|
|
check_tensor(name, gpart2, part_r, (uint64_t)T * slots * E, 3e-3);
|
|
free(part_r); free(mid_r);
|
|
ds4_gpu_tensor_free(gpart2); ds4_gpu_tensor_free(gmid2); ds4_gpu_tensor_free(gcounts); ds4_gpu_tensor_free(glists);
|
|
}
|
|
ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gsg); ds4_gpu_tensor_free(gpart);
|
|
ds4_gpu_tensor_free(gmid); ds4_gpu_tensor_free(gw); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gx);
|
|
free(out); free(part); free(mid); free(w); free(sel); free(sgate); free(x);
|
|
free(gate_w); free(up_w); free(down_w); free(sg_w); free(su_w); free(sd_w);
|
|
}
|
|
|
|
static void test_moe(arena_t *a, uint32_t NE, uint32_t slots, uint32_t E, uint32_t F, uint32_t T, uint32_t wtype) {
|
|
test_moe_types(a, NE, slots, E, F, T, wtype, wtype ? 8u : 0u);
|
|
}
|
|
|
|
static void check_exact_f32(const char *what, const float *got, const float *ref, uint64_t n) {
|
|
for (uint64_t i = 0; i < n; i++) {
|
|
uint32_t gb, rb;
|
|
memcpy(&gb, got + i, sizeof(gb));
|
|
memcpy(&rb, ref + i, sizeof(rb));
|
|
/* Integer bits keep the finite check valid under -ffast-math. */
|
|
if ((gb & 0x7f800000u) == 0x7f800000u ||
|
|
(rb & 0x7f800000u) == 0x7f800000u || gb != rb) {
|
|
fprintf(stderr, "%s: non-exact value at %llu: got %.9g (0x%08x), ref %.9g (0x%08x)\n",
|
|
what, (unsigned long long)i, got[i], gb, ref[i], rb);
|
|
exit(1);
|
|
}
|
|
}
|
|
}
|
|
|
|
static void test_decode_fusions(arena_t *a) {
|
|
const uint32_t hc = 4u, guard = 17u;
|
|
/* Check the whole residual -> normalization -> injection boundary,
|
|
* including writes at the end of each buffer and rejected aliasing. */
|
|
const uint32_t widths[] = {64u, 2560u};
|
|
for (uint32_t wi = 0; wi < 2u; ++wi) {
|
|
const uint32_t e = widths[wi], dim = hc * e, ni = hc * DS4_QWEN4_HC_CHUNKS * hc;
|
|
double *shadow = NULL;
|
|
const uint64_t gamma = arena_f32(a, dim, &shadow, .8f, 1.2f); free(shadow);
|
|
const uint64_t inject = arena_f16(a, (uint64_t)hc * dim, &shadow, .1f); free(shadow);
|
|
float *r = rand_vec(dim, 1.f), *blk = rand_vec(e, 1.f), *inj = rand_vec(ni, .1f);
|
|
ds4_gpu_tensor *old = upload(r, dim), *block = upload(blk, e), *oldinj = upload(inj, ni);
|
|
ds4_gpu_tensor *seq = upload(NULL, dim + guard), *next = upload(NULL, dim + guard);
|
|
ds4_gpu_tensor *xs = upload(NULL, dim + guard), *xf = upload(NULL, dim + guard);
|
|
ds4_gpu_tensor *is = upload(NULL, ni + guard), *ifused = upload(NULL, ni + guard);
|
|
require_ok(ds4_gpu_tensor_fill_f32(seq, 17.25f, dim + guard) && ds4_gpu_tensor_write(seq, 0, r, dim * 4u), "combine input");
|
|
ds4_gpu_tensor *buffers[] = {next, xs, xf, is, ifused};
|
|
for (uint32_t i = 0; i < 5; ++i) require_ok(ds4_gpu_tensor_fill_f32(buffers[i], 17.25f, (i < 3 ? dim : ni) + guard), "fusion guard fill");
|
|
require_ok(ds4_gpu_begin_commands() && ds4_gpu_qwen4_hc_combine_tensor(seq, block, oldinj, 1, e, hc) &&
|
|
ds4_gpu_qwen4_hc_norm_tensor(xs, is, seq, a->base, a->size, gamma, inject, 1, 1, e, hc, hc, 1.e-6f) &&
|
|
ds4_gpu_qwen4_hc_combine_norm_tensor(next, block, oldinj, xf, ifused, old, a->base, a->size,
|
|
gamma, inject, 1, 1, e, hc, hc, 1.e-6f) && ds4_gpu_end_commands(), "combined normalization");
|
|
float *expected = malloc((dim + guard) * sizeof(float)), *actual = malloc((dim + guard) * sizeof(float));
|
|
ds4_gpu_tensor *left[] = {seq, xs, is}, *right[] = {next, xf, ifused};
|
|
for (uint32_t i = 0; i < 3; ++i) {
|
|
uint32_t n = (i == 2 ? ni : dim) + guard;
|
|
require_ok(ds4_gpu_tensor_read(left[i], 0, expected, n * 4u) && ds4_gpu_tensor_read(right[i], 0, actual, n * 4u), "fusion read");
|
|
check_exact_f32("combine norm output and guard", actual, expected, n);
|
|
}
|
|
require_ok(!ds4_gpu_qwen4_hc_combine_norm_tensor(old, block, oldinj, xf, ifused, old,
|
|
a->base, a->size, gamma, inject, 1, 1, e, hc, hc, 1.e-6f), "reject residual alias");
|
|
require_ok(!ds4_gpu_qwen4_hc_combine_norm_tensor(next, block, oldinj, xf, oldinj, old,
|
|
a->base, a->size, gamma, inject, 1, 1, e, hc, hc, 1.e-6f), "reject injection alias");
|
|
ds4_gpu_tensor *all[] = {old, block, oldinj, seq, next, xs, xf, is, ifused};
|
|
for (uint32_t i = 0; i < 9; ++i) ds4_gpu_tensor_free(all[i]);
|
|
free(expected); free(actual); free(r); free(blk); free(inj);
|
|
}
|
|
|
|
/* Concatenated projection: unequal output sizes and short/real input widths. */
|
|
const uint32_t widths_q8[] = {32u, 2560u}, rows_a[] = {2u, 10240u}, rows_b[] = {48u, 6144u};
|
|
for (uint32_t i = 0; i < 2u; ++i) {
|
|
uint32_t k = widths_q8[i], na = rows_a[i], nb = rows_b[i];
|
|
double *shadow = NULL;
|
|
uint64_t wa = arena_q8_0(a, na, k, &shadow, .1f); free(shadow);
|
|
uint64_t wb = arena_q8_0(a, nb, k, &shadow, .1f); free(shadow);
|
|
float *input = rand_vec(k, 1.f); ds4_gpu_tensor *x = upload(input, k);
|
|
ds4_gpu_tensor *ra = upload(NULL, na + guard), *rb = upload(NULL, nb + guard);
|
|
ds4_gpu_tensor *ca = upload(NULL, na + guard), *cb = upload(NULL, nb + guard);
|
|
require_ok(ds4_gpu_tensor_fill_f32(ra, 17.25f, na + guard) && ds4_gpu_tensor_fill_f32(ca, 17.25f, na + guard) &&
|
|
ds4_gpu_tensor_fill_f32(rb, 17.25f, nb + guard) && ds4_gpu_tensor_fill_f32(cb, 17.25f, nb + guard), "Q8 guards");
|
|
require_ok(ds4_gpu_begin_commands() && ds4_gpu_qwen4_matmul_q8_0_tensor(ra, a->base, a->size, wa, k, na, x, 1) &&
|
|
ds4_gpu_qwen4_matmul_q8_0_tensor(rb, a->base, a->size, wb, k, nb, x, 1) &&
|
|
ds4_gpu_qwen4_q8_pair_tensor(ca, cb, a->base, a->size, wa, wb, k, na, nb, x, 1) && ds4_gpu_end_commands(), "Q8 concatenated dispatch");
|
|
uint32_t cap = (na > nb ? na : nb) + guard;
|
|
float *expected = malloc(cap * 4u), *actual = malloc(cap * 4u);
|
|
require_ok(ds4_gpu_tensor_read(ra, 0, expected, (na + guard) * 4u) && ds4_gpu_tensor_read(ca, 0, actual, (na + guard) * 4u), "Q8 first read");
|
|
check_exact_f32("Q8 concatenated first output", actual, expected, na + guard);
|
|
require_ok(ds4_gpu_tensor_read(rb, 0, expected, (nb + guard) * 4u) && ds4_gpu_tensor_read(cb, 0, actual, (nb + guard) * 4u), "Q8 second read");
|
|
check_exact_f32("Q8 concatenated second output", actual, expected, nb + guard);
|
|
require_ok(!ds4_gpu_qwen4_q8_pair_tensor(ca, cb, a->base, a->size, wa, wb, k, na - 1u, nb, x, 1), "Q8 odd shape fallback");
|
|
ds4_gpu_tensor_free(x); ds4_gpu_tensor_free(ra); ds4_gpu_tensor_free(rb); ds4_gpu_tensor_free(ca); ds4_gpu_tensor_free(cb);
|
|
free(input); free(expected); free(actual);
|
|
}
|
|
printf("Qwen decode fusions: exact residual/injection and Q8 checks passed\n");
|
|
}
|
|
|
|
/* Match host greedy selection at SIMD/chunk boundaries and on exceptional values. */
|
|
static void test_qwen4_argmax(void) {
|
|
const uint32_t sizes[] = {1u, 31u, 32u, 33u, 255u, 256u, 257u, 4095u, 4096u, 4097u, 248320u};
|
|
for (uint32_t shape = 0; shape < sizeof(sizes) / sizeof(sizes[0]); shape++) {
|
|
const uint32_t n = sizes[shape], chunks = (n + 4095u) / 4096u, guard = 7u;
|
|
float *v = rand_vec(n, 10.0f);
|
|
ds4_gpu_tensor *x = upload(NULL, n), *tmp = upload(NULL, chunks * 2u + guard);
|
|
ds4_gpu_tensor *out = upload(NULL, 1u + guard);
|
|
for (uint32_t mode = 0; mode < 8u; mode++) {
|
|
if (mode == 1u) { for (uint32_t j = 0; j < n; j++) v[j] = -3.0f; v[n / 2u] = v[n - 1u] = 11.0f; }
|
|
if (mode == 2u) {
|
|
const uint32_t neg_inf = 0xff800000u, nan = 0x7fc00000u;
|
|
for (uint32_t j = 0; j < n; j++) memcpy(v + j, &neg_inf, sizeof(neg_inf));
|
|
memcpy(v + n / 2u, &nan, sizeof(nan));
|
|
}
|
|
if (mode == 3u) { for (uint32_t j = 0; j < n; j++) v[j] = -2.0e30f; }
|
|
if (mode == 4u) { for (uint32_t j = 0; j < n; j++) v[j] = j & 1u ? -0.0f : 0.0f; }
|
|
if (mode == 5u) {
|
|
const uint32_t inf = 0x7f800000u;
|
|
memcpy(v + n / 2u, &inf, sizeof(inf)); memcpy(v + n - 1u, &inf, sizeof(inf));
|
|
}
|
|
if (mode == 6u) { for (uint32_t j = 0; j < n; j++) v[j] = -1.0e30f; }
|
|
if (mode == 7u) {
|
|
const uint32_t nan = 0x7fc00000u;
|
|
for (uint32_t j = 0; j < n; j++) memcpy(v + j, &nan, sizeof(nan));
|
|
}
|
|
uint32_t expected = 0;
|
|
float best = -1.0e30f;
|
|
for (uint32_t j = 0; j < n; j++) {
|
|
uint32_t bits; memcpy(&bits, v + j, sizeof(bits));
|
|
if ((bits & 0x7fffffffu) > 0x7f800000u) continue;
|
|
if (v[j] > best) { best = v[j]; expected = j; }
|
|
}
|
|
require_ok(ds4_gpu_tensor_write(x, 0, v, n * sizeof(float)) &&
|
|
ds4_gpu_tensor_fill_f32(tmp, 17.25f, chunks * 2u + guard) &&
|
|
ds4_gpu_tensor_fill_f32(out, 17.25f, 1u + guard) &&
|
|
ds4_gpu_begin_commands() && ds4_gpu_qwen4_argmax_tensor(out, tmp, x, n) &&
|
|
ds4_gpu_end_commands(), "Qwen argmax dispatch");
|
|
uint32_t got;
|
|
require_ok(ds4_gpu_tensor_read(out, 0, &got, sizeof(got)) && got == expected, "Qwen argmax index");
|
|
float tail[7];
|
|
require_ok(ds4_gpu_tensor_read(out, sizeof(uint32_t), tail, sizeof(tail)), "argmax output guard read");
|
|
for (uint32_t j = 0; j < guard; j++) require_ok(tail[j] == 17.25f, "argmax output guard");
|
|
require_ok(ds4_gpu_tensor_read(tmp, chunks * 8u, tail, sizeof(tail)), "argmax scratch guard read");
|
|
for (uint32_t j = 0; j < guard; j++) require_ok(tail[j] == 17.25f, "argmax scratch guard");
|
|
}
|
|
require_ok(!ds4_gpu_qwen4_argmax_tensor(out, tmp, x, 0u) &&
|
|
!ds4_gpu_qwen4_argmax_tensor(out, tmp, x, n + 1u), "argmax rejects invalid sizes");
|
|
ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(tmp); ds4_gpu_tensor_free(x); free(v);
|
|
}
|
|
printf("Qwen predictor argmax: CPU indices, ties, special values and guards passed\n");
|
|
}
|
|
|
|
/* The paired mixer must preserve both token rows and its partial-group guard. */
|
|
/* The prefetched F16 gate/mix must reproduce the plain kernel byte for byte
|
|
* on every lane path: ranks that skip the eight-term rounds, that end in a
|
|
* remainder, and the production 320, for one and several tokens. */
|
|
static void test_hc_mix_prefetch(arena_t *a) {
|
|
const uint32_t ranks[] = {8u, 72u, 136u, 200u, 320u}, tokens[] = {1u, 3u, 2u};
|
|
for (uint32_t ir = 0; ir < 5u; ir++) {
|
|
for (uint32_t it = 0; it < 3u; it++) {
|
|
const uint32_t rank = ranks[ir], T = tokens[it], E = rank == 320u ? 2560u : 96u;
|
|
const uint64_t n = (uint64_t)T * E, guard = 13u;
|
|
double *shadow = NULL;
|
|
const uint64_t off = arena_f16(a, (uint64_t)4u * E * rank, &shadow, 0.3f);
|
|
free(shadow);
|
|
float *xn = rand_vec((uint64_t)T * 4u * E, 1.0f), *lo = rand_vec((uint64_t)T * rank, 6.0f);
|
|
float *ref = malloc((n + guard) * sizeof(float)), *got = malloc((n + guard) * sizeof(float));
|
|
require_ok(ref && got, "HC prefetch readback allocation");
|
|
ds4_gpu_tensor *gx = upload(xn, (uint64_t)T * 4u * E), *gl = upload(lo, (uint64_t)T * rank);
|
|
ds4_gpu_tensor *go = upload(NULL, n + guard);
|
|
for (uint32_t mode = 0; mode < 2u; mode++) {
|
|
require_ok(setenv("DS4_QWEN4_HC_MIX_PREFETCH", mode ? "1" : "0", 1) == 0, "HC prefetch override");
|
|
require_ok(ds4_gpu_tensor_fill_f32(go, 127.25f, n + guard) &&
|
|
ds4_gpu_qwen4_hc_gate_mix_tensor(go, gx, gl, a->base, a->size, off, 1u, T, E, 4u, rank) &&
|
|
ds4_gpu_tensor_read(go, 0, mode ? got : ref, (n + guard) * sizeof(float)), "HC prefetch dispatch/read");
|
|
}
|
|
check_exact_f32("HC prefetch reference finite", ref, ref, n + guard);
|
|
check_exact_f32("HC prefetch rows and guard", got, ref, n + guard);
|
|
for (uint64_t j = n; j < n + guard; j++) require_ok(got[j] == 127.25f, "HC prefetch output guard");
|
|
ds4_gpu_tensor_free(go); ds4_gpu_tensor_free(gl); ds4_gpu_tensor_free(gx);
|
|
free(got); free(ref); free(lo); free(xn);
|
|
}
|
|
}
|
|
unsetenv("DS4_QWEN4_HC_MIX_PREFETCH");
|
|
printf("HC prefetched mixers (single and paired): exact against the plain kernels on all lane paths\n");
|
|
}
|
|
|
|
|
|
/* The few-row matvec simdgroup count is rows-per-threadgroup only: Q8 and
|
|
* F16 outputs must match the default byte for byte at 1/2/4/8 groups for the
|
|
* verify-row shapes (T 2 and 3, odd row counts). */
|
|
static void test_mv_ext_groups(arena_t *a) {
|
|
const uint32_t T_list[] = {2u, 3u}, rows_list[] = {640u, 641u, 2560u};
|
|
const char *groups[] = {"2", "1", "4", "8"};
|
|
for (uint32_t wt = 0; wt < 2u; wt++) {
|
|
for (uint32_t it = 0; it < 2u; it++) {
|
|
for (uint32_t ir = 0; ir < 3u; ir++) {
|
|
const uint32_t T = T_list[it], rows = rows_list[ir], in_dim = 2560u;
|
|
const uint64_t n = (uint64_t)T * rows, guard = 9u;
|
|
double *sh = NULL;
|
|
const uint64_t off = wt ? arena_f16(a, (uint64_t)rows * in_dim, &sh, 0.05f)
|
|
: arena_q8_0(a, rows, in_dim, &sh, 0.05f);
|
|
free(sh);
|
|
float *x = rand_vec(n ? (uint64_t)T * in_dim : 1u, 1.0f);
|
|
float *ref = malloc((n + guard) * sizeof(float)), *got = malloc((n + guard) * sizeof(float));
|
|
require_ok(ref && got, "mv_ext groups allocation");
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * in_dim), *go = upload(NULL, n + guard);
|
|
for (uint32_t mode = 0; mode < 4u; mode++) {
|
|
require_ok(setenv("DS4_METAL_MV_EXT_NSG", groups[mode], 1) == 0, "mv_ext groups override");
|
|
require_ok(ds4_gpu_tensor_fill_f32(go, 127.25f, n + guard) &&
|
|
(wt ? ds4_gpu_matmul_f16_tensor(go, a->base, a->size, off, in_dim, rows, gx, T)
|
|
: ds4_gpu_qwen4_matmul_q8_0_tensor(go, a->base, a->size, off, in_dim, rows, gx, T)) &&
|
|
ds4_gpu_tensor_read(go, 0, mode ? got : ref, (n + guard) * sizeof(float)), "mv_ext groups dispatch/read");
|
|
if (mode) check_exact_f32(wt ? "F16 mv_ext groups and guard" : "Q8 mv_ext groups and guard", got, ref, n + guard);
|
|
else check_exact_f32("mv_ext groups reference finite", ref, ref, n + guard);
|
|
}
|
|
ds4_gpu_tensor_free(go); ds4_gpu_tensor_free(gx);
|
|
free(got); free(ref); free(x);
|
|
}
|
|
}
|
|
}
|
|
unsetenv("DS4_METAL_MV_EXT_NSG");
|
|
printf("few-row matvec simdgroup counts: Q8 and F16 verify-row shapes exact at 1/2/4/8 groups\n");
|
|
}
|
|
|
|
#ifdef __APPLE__
|
|
/* The grouped decode-batch kernels must reproduce the per-token kernels bit
|
|
* for bit under heavy expert reuse (sixteen rows over eight experts). */
|
|
static void test_moe_grouped(arena_t *a) {
|
|
const uint32_t NE = 8, slots = 6, E = 2560, F = 640, T = 16, cap = 64;
|
|
double *gate_w, *up_w, *down_w;
|
|
const uint64_t gate_off = arena_q4_K(a, (uint64_t)NE * F, E, &gate_w, 0.05f);
|
|
const uint64_t up_off = arena_q4_K(a, (uint64_t)NE * F, E, &up_w, 0.05f);
|
|
const uint64_t down_off = arena_mxfp4(a, (uint64_t)NE * E, F, &down_w);
|
|
free(gate_w); free(up_w); free(down_w);
|
|
float *x = rand_vec((uint64_t)T * E, 1.0f);
|
|
int32_t *sel = malloc((uint64_t)T * slots * 4);
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < slots; s++) sel[t * slots + s] = (int32_t)((t * 7u + s * 3u) % NE);
|
|
}
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * slots * 4);
|
|
ds4_gpu_tensor *glists = ds4_gpu_tensor_alloc((uint64_t)NE * cap * 4), *gcounts = ds4_gpu_tensor_alloc((uint64_t)NE * 4);
|
|
ds4_gpu_tensor *gmid[2] = { upload(NULL, (uint64_t)T * slots * F), upload(NULL, (uint64_t)T * slots * F) };
|
|
ds4_gpu_tensor *gpart[2] = { upload(NULL, (uint64_t)T * slots * E), upload(NULL, (uint64_t)T * slots * E) };
|
|
require_ok(gsel && glists && gcounts && ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * slots * 4), "grouped sel");
|
|
require_ok(ds4_gpu_qwen4_moe_mid_tensor(gmid[0], gx, gsel, a->base, a->size, gate_off, up_off, 12u, NE, T, slots, E, F,
|
|
0, 0, UINT32_MAX), "grouped: per-token mid");
|
|
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart[0], gmid[0], gsel, a->base, a->size, down_off, 39u, NE, T, slots, F, E,
|
|
0, UINT32_MAX), "grouped: per-token down");
|
|
require_ok(ds4_gpu_qwen4_moe_build_lists_tensor(glists, gcounts, gsel, T, slots, NE, cap), "grouped: lists");
|
|
require_ok(ds4_gpu_qwen4_moe_mid_grouped_tensor(gmid[1], gx, gsel, glists, gcounts, cap, a->base, a->size, gate_off,
|
|
up_off, 12u, NE, T, slots, E, F), "grouped: mid");
|
|
require_ok(ds4_gpu_qwen4_moe_down_grouped_tensor(gpart[1], gmid[0], gsel, glists, gcounts, cap, a->base, a->size,
|
|
down_off, 39u, NE, T, slots, F, E), "grouped: down");
|
|
const uint64_t nm = (uint64_t)T * slots * F, np = (uint64_t)T * slots * E;
|
|
float *am = download(gmid[0], nm), *bm = download(gmid[1], nm);
|
|
float *ap = download(gpart[0], np), *bp = download(gpart[1], np);
|
|
check_exact_f32("grouped Q4_K mid", bm, am, nm);
|
|
check_exact_f32("grouped MXFP4 down", bp, ap, np);
|
|
printf(" MoE grouped kernels: mid and down byte-exact against the per-token kernels under expert reuse\n");
|
|
free(bp); free(ap); free(bm); free(am);
|
|
for (int i = 0; i < 2; i++) { ds4_gpu_tensor_free(gpart[i]); ds4_gpu_tensor_free(gmid[i]); }
|
|
ds4_gpu_tensor_free(gcounts); ds4_gpu_tensor_free(glists); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gx);
|
|
free(sel); free(x);
|
|
}
|
|
|
|
#endif
|
|
|
|
static void test_hc_pair_groups(arena_t *a) {
|
|
const uint32_t types[] = {1u, 0u, 8u}, widths[] = {9u, 64u, 2560u};
|
|
const char *groups[] = {"4", "1", "2", "8", "16"};
|
|
for (uint32_t it = 0; it < 3u; it++) {
|
|
for (uint32_t iw = 0; iw < 3u; iw++) {
|
|
const uint32_t type = types[it], E = widths[iw], rank = iw == 1u ? 32u : 320u;
|
|
const uint64_t n = 2u * E, guard = 17u;
|
|
double *shadow = NULL;
|
|
const uint64_t off = type == 8u ? arena_q8_0(a, 4u * E, rank, &shadow, 0.2f) :
|
|
type == 1u ? arena_f16(a, (uint64_t)4u * E * rank, &shadow, 0.2f) :
|
|
arena_f32(a, (uint64_t)4u * E * rank, &shadow, -0.2f, 0.2f);
|
|
free(shadow);
|
|
float *xn = rand_vec(8u * E, 1.0f), *lo = rand_vec(2u * rank, 1.0f);
|
|
float *ref = malloc((n + guard) * sizeof(float)), *got = malloc((n + guard) * sizeof(float));
|
|
require_ok(ref && got, "HC pair readback allocation");
|
|
ds4_gpu_tensor *gx = upload(xn, 8u * E), *gl = upload(lo, 2u * rank);
|
|
ds4_gpu_tensor *go = upload(NULL, n + guard);
|
|
for (uint32_t mode = 0; mode < 5u; mode++) {
|
|
require_ok(setenv("DS4_QWEN4_HC_PAIR_NSG", groups[mode], 1) == 0, "HC pair override");
|
|
require_ok(ds4_gpu_tensor_fill_f32(go, 127.25f, n + guard) &&
|
|
ds4_gpu_qwen4_hc_gate_mix_tensor(go, gx, gl, a->base, a->size, off, type, 2u, E, 4u, rank) &&
|
|
ds4_gpu_tensor_read(go, 0, got, (n + guard) * sizeof(float)), "HC pair dispatch/read");
|
|
if (!mode) {
|
|
memcpy(ref, got, (n + guard) * sizeof(float));
|
|
check_exact_f32("HC pair reference finite", ref, ref, n + guard);
|
|
} else check_exact_f32("HC pair geometry and guard", got, ref, n + guard);
|
|
for (uint64_t j = n; j < n + guard; j++) require_ok(got[j] == 127.25f, "HC pair output guard");
|
|
}
|
|
ds4_gpu_tensor_free(go); ds4_gpu_tensor_free(gl); ds4_gpu_tensor_free(gx);
|
|
free(got); free(ref); free(lo); free(xn);
|
|
}
|
|
}
|
|
unsetenv("DS4_QWEN4_HC_PAIR_NSG");
|
|
printf("HC paired mixer geometry: all formats and guarded tails exact\n");
|
|
}
|
|
|
|
/* Freeze the stream input and weights across forced old/reuse and automatic
|
|
* dispatch. All three share a command batch and all readbacks follow it. */
|
|
static void test_hc_norm_reuse_case(arena_t *a, uint32_t type, uint32_t E,
|
|
uint32_t hc, uint32_t T, uint32_t n_inject) {
|
|
const uint32_t guard = 5, CH = DS4_QWEN4_HC_CHUNKS;
|
|
const uint64_t dim = (uint64_t)E * hc, n = (uint64_t)T * dim;
|
|
const uint64_t inj_n = (uint64_t)T * hc * CH * n_inject;
|
|
const uint64_t inj_storage_n = inj_n ? inj_n : (uint64_t)T * hc * CH;
|
|
const float sentinel = -1234.5f, eps = 1e-6f;
|
|
uint32_t sentinel_bits;
|
|
memcpy(&sentinel_bits, &sentinel, sizeof(sentinel_bits));
|
|
require_ok(dim % 32u == 0u && n_inject <= 4u, "HC norm reuse fixture shape");
|
|
/* The focused selector can start with an empty arena; keep model offsets
|
|
* nonzero as well as the tensor-view offsets below. Gamma stays FP32. */
|
|
(void)arena_alloc(a, 16u);
|
|
double *shadow;
|
|
const uint64_t gamma_off = arena_f32(a, dim, &shadow, 0.5f, 1.5f);
|
|
free(shadow);
|
|
uint64_t inject_off = 0;
|
|
if (n_inject) {
|
|
inject_off = type == 8u ? arena_q8_0(a, n_inject, dim, &shadow, 0.05f)
|
|
: type == 1u ? arena_f16(a, (uint64_t)n_inject * dim, &shadow, 0.05f)
|
|
: arena_f32(a, (uint64_t)n_inject * dim, &shadow, -0.05f, 0.05f);
|
|
free(shadow);
|
|
}
|
|
float *R = malloc((n + 2u * guard) * sizeof(float));
|
|
require_ok(R != NULL, "HC norm reuse host input allocation");
|
|
for (uint64_t i = 0; i < n + 2u * guard; i++) R[i] = sentinel;
|
|
for (uint32_t t = 0; t < T; t++) for (uint32_t s = 0; s < hc; s++) {
|
|
const float scale = s % 3u == 0u ? 0.125f : s % 3u == 1u ? 4.0f : 32.0f;
|
|
for (uint32_t i = 0; i < E; i++) {
|
|
const uint64_t index = guard + ((uint64_t)t * hc + s) * E + i;
|
|
R[index] = frand() * scale;
|
|
if (t == 0 && s == 0) {
|
|
const uint32_t zero_bits = i & 1u ? 0x80000000u : 0u;
|
|
memcpy(R + index, &zero_bits, sizeof(zero_bits));
|
|
} else if (t == 0 && s == 1) {
|
|
R[index] *= 1e-12f;
|
|
}
|
|
}
|
|
}
|
|
ds4_gpu_tensor *gR_base = upload(R, n + 2u * guard);
|
|
ds4_gpu_tensor *gR = ds4_gpu_tensor_view(gR_base, guard * sizeof(float), n * sizeof(float));
|
|
ds4_gpu_tensor *gxn_base[3], *gxn[3], *ginj_base[3], *ginj[3];
|
|
require_ok(gR != NULL, "HC norm reuse input view");
|
|
for (uint32_t mode = 0; mode < 3; mode++) {
|
|
gxn_base[mode] = upload(NULL, n + 2u * guard);
|
|
ginj_base[mode] = upload(NULL, inj_storage_n + 2u * guard);
|
|
gxn[mode] = ds4_gpu_tensor_view(gxn_base[mode], guard * sizeof(float), n * sizeof(float));
|
|
ginj[mode] = ds4_gpu_tensor_view(ginj_base[mode], guard * sizeof(float), inj_storage_n * sizeof(float));
|
|
require_ok(gxn[mode] && ginj[mode], "HC norm reuse output views");
|
|
require_ok(ds4_gpu_tensor_fill_f32(gxn_base[mode], sentinel, n + 2u * guard) &&
|
|
ds4_gpu_tensor_fill_f32(ginj_base[mode], sentinel, inj_storage_n + 2u * guard),
|
|
"HC norm reuse output sentinels");
|
|
}
|
|
require_ok(ds4_gpu_begin_commands(), "HC norm reuse batch begin");
|
|
for (uint32_t mode = 0; mode < 3; mode++) {
|
|
require_ok(mode == 2u ? unsetenv("DS4_QWEN4_HC_NORM_REUSE") == 0
|
|
: setenv("DS4_QWEN4_HC_NORM_REUSE", mode ? "1" : "0", 1) == 0,
|
|
"select forced or automatic HC norm");
|
|
#ifndef __APPLE__
|
|
/* CUDA uses its original one-token kernel as an independent dispatch
|
|
* control, not the Metal-only environment switch above. */
|
|
if (mode == 0u && T > 8u) {
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
ds4_gpu_tensor *r = ds4_gpu_tensor_view(gR,t*dim*4,dim*4);
|
|
ds4_gpu_tensor *x = ds4_gpu_tensor_view(gxn[mode],t*dim*4,dim*4);
|
|
const uint64_t partials = (uint64_t)hc*CH*n_inject;
|
|
ds4_gpu_tensor *in = n_inject ? ds4_gpu_tensor_view(ginj[mode],t*partials*4,partials*4) : NULL;
|
|
require_ok(r && x && (!n_inject || in),"HC norm decode control views");
|
|
require_ok(ds4_gpu_qwen4_hc_norm_tensor(x,in,r,a->base,a->size,
|
|
gamma_off,inject_off,type,1,E,hc,n_inject,eps),"HC norm decode control");
|
|
ds4_gpu_tensor_free(in); ds4_gpu_tensor_free(x); ds4_gpu_tensor_free(r);
|
|
}
|
|
continue;
|
|
}
|
|
#endif
|
|
require_ok(ds4_gpu_qwen4_hc_norm_tensor(gxn[mode], ginj[mode], gR, a->base, a->size,
|
|
gamma_off, inject_off, type, T, E, hc, n_inject, eps), "HC norm dispatch");
|
|
}
|
|
require_ok(ds4_gpu_end_commands(), "HC norm reuse batch end");
|
|
|
|
char name[160];
|
|
for (uint32_t output = 0; output < 2; output++) {
|
|
const uint64_t values = output ? inj_storage_n : n;
|
|
float *ref = download(output ? ginj_base[0] : gxn_base[0], values + 2u * guard);
|
|
for (uint32_t mode = 1; mode < 3; mode++) {
|
|
float *got = download(output ? ginj_base[mode] : gxn_base[mode], values + 2u * guard);
|
|
snprintf(name, sizeof(name), "HC norm %s type=%u E=%u hc=%u T=%u inject=%u %s",
|
|
mode == 1u ? "reuse" : "automatic", type, E, hc, T, n_inject,
|
|
output ? "inject partials" : "xn");
|
|
/* Integer bits preserve finite checks under -ffast-math. */
|
|
check_exact_f32(name, got, ref, values + 2u * guard);
|
|
for (uint64_t i = 0; i < values + 2u * guard; i++) {
|
|
const bool outside = i < guard || i >= guard + values;
|
|
const bool unused = output && n_inject == 0;
|
|
if (outside || unused) {
|
|
uint32_t bits;
|
|
memcpy(&bits, got + i, sizeof(bits));
|
|
require_ok(bits == sentinel_bits, "HC norm reuse guards and unused injection output");
|
|
}
|
|
}
|
|
free(got);
|
|
}
|
|
free(ref);
|
|
}
|
|
float *got_R = download(gR_base, n + 2u * guard);
|
|
check_exact_f32("HC norm reuse input remains frozen", got_R, R, n + 2u * guard);
|
|
free(got_R);
|
|
printf(" HC norm reuse/default type=%u E=%u hc=%u T=%u inject=%u: byte-exact xn/inj, guards%s\n",
|
|
type, E, hc, T, n_inject, T <= 2u ? " (decode fallback)" : "");
|
|
for (uint32_t mode = 0; mode < 3; mode++) {
|
|
ds4_gpu_tensor_free(ginj[mode]); ds4_gpu_tensor_free(ginj_base[mode]);
|
|
ds4_gpu_tensor_free(gxn[mode]); ds4_gpu_tensor_free(gxn_base[mode]);
|
|
}
|
|
ds4_gpu_tensor_free(gR); ds4_gpu_tensor_free(gR_base);
|
|
free(R);
|
|
}
|
|
|
|
static void test_hc_norm_reuse(arena_t *a) {
|
|
const char *value = getenv("DS4_QWEN4_HC_NORM_REUSE");
|
|
char *saved = value ? strdup(value) : NULL;
|
|
require_ok(!value || saved != NULL, "save HC norm reuse environment");
|
|
const uint32_t types[] = {0u, 1u, 8u};
|
|
const uint32_t shapes[][4] = {
|
|
{2560u, 4u, 3u, 4u},
|
|
{2560u, 4u, 9u, 0u},
|
|
{2560u, 4u, 17u, 4u},
|
|
{2560u, 4u, 8u, 4u},
|
|
{2560u, 4u, 9u, 4u},
|
|
{260u, 8u, 9u, 3u},
|
|
{4u, 8u, 9u, 4u},
|
|
{260u, 8u, 3u, 3u}, /* final chunk is shorter, stream boundary splits Q8 blocks */
|
|
{40u, 4u, 5u, 1u},
|
|
{260u, 8u, 3u, 2u},
|
|
{4u, 8u, 3u, 4u}, /* chunks four through seven are empty */
|
|
{2560u, 4u, 1u, 4u}, /* decode and MTP keep the original dispatch */
|
|
{2560u, 4u, 2u, 4u},
|
|
};
|
|
for (uint32_t type = 0; type < sizeof(types) / sizeof(types[0]); type++) {
|
|
for (uint32_t shape = 0; shape < sizeof(shapes) / sizeof(shapes[0]); shape++) {
|
|
test_hc_norm_reuse_case(a, types[type], shapes[shape][0], shapes[shape][1],
|
|
shapes[shape][2], shapes[shape][3]);
|
|
}
|
|
}
|
|
/* Full outputs around the automatic threshold: the last case uses
|
|
* reuse by default on M3 Ultra, while other devices retain the old path.
|
|
* Each also compares explicit off/on, so both kernels run on every device. */
|
|
for (uint32_t T = 8191u; T <= 8192u; T++) {
|
|
test_hc_norm_reuse_case(a, 1u, 2560u, 4u, T, 4u);
|
|
}
|
|
require_ok(saved ? setenv("DS4_QWEN4_HC_NORM_REUSE", saved, 1) == 0
|
|
: unsetenv("DS4_QWEN4_HC_NORM_REUSE") == 0, "restore HC norm reuse environment");
|
|
free(saved);
|
|
}
|
|
|
|
/* Compare against the original per-row GPU kernel, not a second copy of the
|
|
* ordered implementation. Ten Q4_K blocks per row exercise accumulation
|
|
* across blocks and the fixture's varied packed six-bit scales/minima.
|
|
* Each rows-per-SIMD setting is checked against the same original output. */
|
|
static void test_q4k_ordered_exact(arena_t *a, uint32_t T, uint32_t F, bool shared) {
|
|
const uint32_t E = 2560, NE = 5, slots = 3, D = 64, guard = 16;
|
|
const uint32_t n_out = slots + (shared ? 1u : 0u);
|
|
const uint32_t shared_type = shared ? 8u : UINT32_MAX;
|
|
const bool downstream = F % 32u == 0u;
|
|
const uint64_t mid_n = (uint64_t)T * n_out * F, part_n = (uint64_t)T * n_out * D;
|
|
const char *env_names[] = {"DS4_QWEN4_NO_Q4K_MID", "DS4_QWEN4_Q4K_MID_NSG", "DS4_QWEN4_Q4K_MID_NR"};
|
|
char *saved_env[3];
|
|
for (uint32_t i = 0; i < 3; i++) {
|
|
const char *v = getenv(env_names[i]);
|
|
saved_env[i] = v ? strdup(v) : NULL;
|
|
require_ok(!v || saved_env[i] != NULL, "save ordered mid environment");
|
|
}
|
|
double *gate_w, *shadow;
|
|
const uint64_t gate_off = arena_q4_K(a, (uint64_t)NE * F, E, &gate_w, 0.05f);
|
|
const uint64_t up_off = arena_q4_K(a, (uint64_t)NE * F, E, &shadow, 0.05f); free(shadow);
|
|
uint64_t sg_off = 0, su_off = 0, down_off = 0, sd_off = 0;
|
|
if (shared) {
|
|
sg_off = arena_q8_0(a, F, E, &shadow, 0.05f); free(shadow);
|
|
su_off = arena_q8_0(a, F, E, &shadow, 0.05f); free(shadow);
|
|
}
|
|
if (downstream) {
|
|
down_off = arena_q8_0(a, (uint64_t)NE * D, F, &shadow, 0.05f); free(shadow);
|
|
if (shared) { sd_off = arena_q8_0(a, D, F, &shadow, 0.05f); free(shadow); }
|
|
}
|
|
float *x = rand_vec((uint64_t)T * E, 1.0f);
|
|
for (uint64_t i = 0; i < (uint64_t)T * E; i++) x[i] *= i % 7u == 0u ? 8.0f : i % 7u == 1u ? 0.125f : 1.0f;
|
|
/* Ensure that the first selected gate takes the negative sigmoid branch. */
|
|
double first_gate = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) first_gate += gate_w[i] * x[i];
|
|
require_ok(first_gate != 0.0, "ordered mid negative gate fixture");
|
|
if (first_gate > 0.0) for (uint32_t i = 0; i < E; i++) x[i] = -x[i];
|
|
free(gate_w);
|
|
int32_t sel[6];
|
|
float weights[6], shared_gate[2] = {-1.75f, 2.25f};
|
|
for (uint32_t t = 0; t < T; t++) for (uint32_t s = 0; s < slots; s++) {
|
|
sel[t * slots + s] = (int32_t)((t * 3u + s * 2u) % NE);
|
|
weights[t * slots + s] = 0.125f * (float)(1u + t + s);
|
|
}
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * slots * sizeof(int32_t));
|
|
require_ok(gsel && ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * slots * sizeof(int32_t)), "ordered mid selected");
|
|
ds4_gpu_tensor *gmid = upload(NULL, mid_n + guard);
|
|
ds4_gpu_tensor *gpart = downstream ? upload(NULL, part_n) : NULL;
|
|
ds4_gpu_tensor *gout = downstream ? upload(NULL, (uint64_t)T * D) : NULL;
|
|
ds4_gpu_tensor *gw = downstream ? upload(weights, (uint64_t)T * slots) : NULL;
|
|
ds4_gpu_tensor *gsg = downstream && shared ? upload(shared_gate, T) : NULL;
|
|
float *ref_mid = NULL, *ref_part = NULL, *ref_out = NULL;
|
|
const float sentinel = -1234.5f;
|
|
const uint32_t rows_per_simd[] = {1u, 2u};
|
|
/* One original reference, one unset-default dispatch, then all NR/NSG
|
|
* combinations. Preserve the existing NSG 1..8 coverage for NR2. */
|
|
for (uint32_t mode = 0; mode < 2u + (sizeof(rows_per_simd) / sizeof(rows_per_simd[0])) * 8u; mode++) {
|
|
const bool reference = mode == 0u, defaults = mode == 1u;
|
|
const uint32_t nr = reference || defaults ? 0u : rows_per_simd[(mode - 2u) / 8u];
|
|
const uint32_t nsg = reference || defaults ? 0u : 1u + (mode - 2u) % 8u;
|
|
if (reference) {
|
|
require_ok(setenv(env_names[0], "1", 1) == 0 && unsetenv(env_names[1]) == 0 &&
|
|
unsetenv(env_names[2]) == 0, "select original Q4K mid reference");
|
|
} else {
|
|
require_ok(unsetenv(env_names[0]) == 0, "enable ordered Q4K mid");
|
|
if (defaults) {
|
|
/* T1 and T2 may choose NR1/NSG8 on M3 Ultra. Both defaults
|
|
* must match the original arithmetic. */
|
|
require_ok(unsetenv(env_names[1]) == 0 && unsetenv(env_names[2]) == 0,
|
|
"select default Q4K mid geometry");
|
|
} else {
|
|
char nsg_value[16], nr_value[16];
|
|
snprintf(nsg_value, sizeof(nsg_value), "%u", nsg);
|
|
snprintf(nr_value, sizeof(nr_value), "%u", nr);
|
|
require_ok(setenv(env_names[1], nsg_value, 1) == 0 && setenv(env_names[2], nr_value, 1) == 0,
|
|
"select Q4K mid NR/NSG geometry");
|
|
}
|
|
}
|
|
require_ok(ds4_gpu_tensor_fill_f32(gmid, sentinel, mid_n + guard), "ordered mid sentinel");
|
|
require_ok(ds4_gpu_qwen4_moe_mid_tensor(gmid, gx, gsel, a->base, a->size, gate_off, up_off, 12u,
|
|
NE, T, slots, E, F, sg_off, su_off, shared_type), "ordered mid dispatch");
|
|
float *got_mid = download(gmid, mid_n + guard);
|
|
for (uint64_t i = mid_n; i < mid_n + guard; i++) require_ok(got_mid[i] == sentinel, "ordered mid output guard");
|
|
char name[144], geometry[32];
|
|
if (reference || defaults) snprintf(geometry, sizeof(geometry), "%s", reference ? "original" : "default");
|
|
else snprintf(geometry, sizeof(geometry), "NR=%u NSG=%u", nr, nsg);
|
|
snprintf(name, sizeof(name), "q4_K ordered T=%u F=%u shared=%u %s mid", T, F, shared, geometry);
|
|
if (reference) {
|
|
ref_mid = got_mid;
|
|
check_exact_f32(name, ref_mid, ref_mid, mid_n);
|
|
} else {
|
|
check_exact_f32(name, got_mid, ref_mid, mid_n + guard);
|
|
free(got_mid);
|
|
}
|
|
if (downstream) {
|
|
require_ok(ds4_gpu_qwen4_moe_down_tensor(gpart, gmid, gsel, a->base, a->size, down_off, 8u,
|
|
NE, T, slots, F, D, sd_off, shared_type), "ordered mid downstream");
|
|
require_ok(ds4_gpu_qwen4_moe_reduce_tensor(gout, gpart, gw, gsg, NULL, NULL, NULL,
|
|
T, slots, n_out, D, 0), "ordered mid reduce");
|
|
float *got_part = download(gpart, part_n), *got_out = download(gout, (uint64_t)T * D);
|
|
if (reference) { ref_part = got_part; ref_out = got_out; }
|
|
else {
|
|
snprintf(name, sizeof(name), "q4_K ordered T=%u shared=%u %s down", T, shared, geometry);
|
|
check_exact_f32(name, got_part, ref_part, part_n);
|
|
snprintf(name, sizeof(name), "q4_K ordered T=%u shared=%u %s reduce", T, shared, geometry);
|
|
check_exact_f32(name, got_out, ref_out, (uint64_t)T * D);
|
|
free(got_part); free(got_out);
|
|
}
|
|
}
|
|
}
|
|
printf(" q4_K ordered T=%u F=%u shared=%u: default and NR 1/2 x NSG 1..8 byte-exact mid%s\n",
|
|
T, F, shared, downstream ? ", down, reduce" : " (odd tail)");
|
|
for (uint32_t i = 0; i < 3; i++) {
|
|
if (saved_env[i]) setenv(env_names[i], saved_env[i], 1); else unsetenv(env_names[i]);
|
|
free(saved_env[i]);
|
|
}
|
|
free(ref_out); free(ref_part); free(ref_mid); free(x);
|
|
ds4_gpu_tensor_free(gsg); ds4_gpu_tensor_free(gw); ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gpart);
|
|
ds4_gpu_tensor_free(gmid); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gx);
|
|
}
|
|
|
|
/* Freeze the GPU-built routing lists across launch-cap comparisons.
|
|
* The hot expert spans 21 tiles, exceeding caps eight and sixteen, and
|
|
* the final 32-token tile contains only one token. */
|
|
/* Q2-pack tiers (iq2_xxs gate/up + q2_K down) on the routed tiles: the
|
|
* simdgroup tiles are the reference; the tensor-op levels are bounded
|
|
* against them and against a double-precision exact reference. */
|
|
static void test_moe_mm_tiles_iq2(arena_t *a) {
|
|
const uint32_t T = 641, E = 256, F = 256, NE = 4, slots = 2, n_out = 3, list_cap = T + 7, guard = 16;
|
|
const uint64_t mid_n = (uint64_t)T * n_out * F, part_n = (uint64_t)T * n_out * E;
|
|
const char *env_names[] = {"DS4_QWEN4_MOE_MID_TILES", "DS4_QWEN4_MOE_DOWN_TILES"};
|
|
char *saved_env[3];
|
|
for (uint32_t i = 0; i < 3; i++) {
|
|
const char *names[] = {env_names[0], env_names[1], "DS4_QWEN4_MOE_MM_NAX"};
|
|
const char *v = getenv(names[i]);
|
|
saved_env[i] = v ? strdup(v) : NULL;
|
|
require_ok(!v || saved_env[i] != NULL, "save Q2 tile environment");
|
|
}
|
|
setenv("DS4_QWEN4_MOE_MM_NAX", "0", 1); /* simdgroup reference */
|
|
double *gate_shadow, *up_shadow, *down_shadow;
|
|
const uint64_t gate_off = arena_tier(a, 16u, (uint64_t)NE * F, E, &gate_shadow);
|
|
const uint64_t up_off = arena_tier(a, 16u, (uint64_t)NE * F, E, &up_shadow);
|
|
const uint64_t down_off = arena_tier(a, 10u, (uint64_t)NE * E, F, &down_shadow);
|
|
float *x = rand_vec((uint64_t)T * E, 2.0f);
|
|
int32_t *sel = malloc((uint64_t)T * slots * sizeof(int32_t));
|
|
require_ok(sel != NULL, "Q2 tile selection allocation");
|
|
for (uint32_t t = 0; t < T; t++) { sel[t * slots] = 0; sel[t * slots + 1] = (int32_t)(1u + t % 2u); }
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * slots * sizeof(int32_t));
|
|
ds4_gpu_tensor *glists = ds4_gpu_tensor_alloc((uint64_t)NE * list_cap * sizeof(int32_t));
|
|
ds4_gpu_tensor *gcounts = ds4_gpu_tensor_alloc(NE * sizeof(int32_t));
|
|
require_ok(gsel && glists && gcounts && ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * slots * sizeof(int32_t)),
|
|
"Q2 tile routing setup");
|
|
require_ok(ds4_gpu_qwen4_moe_build_lists_tensor(glists, gcounts, gsel, T, slots, NE, list_cap), "Q2 tile frozen lists");
|
|
double *mid_exact = malloc(mid_n * sizeof(double)), *part_exact = malloc(part_n * sizeof(double));
|
|
require_ok(mid_exact && part_exact, "Q2 tile exact reference allocation");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < slots; s++) {
|
|
const uint32_t e = (uint32_t)sel[t * slots + s];
|
|
const float *xr = x + (uint64_t)t * E;
|
|
for (uint32_t f = 0; f < F; f++) {
|
|
const double *gr = gate_shadow + ((uint64_t)e * F + f) * E;
|
|
const double *ur = up_shadow + ((uint64_t)e * F + f) * E;
|
|
double g = 0.0, u = 0.0;
|
|
for (uint32_t k = 0; k < E; k++) { g += gr[k] * (double)xr[k]; u += ur[k] * (double)xr[k]; }
|
|
mid_exact[((uint64_t)t * n_out + s) * F + f] = (g / (1.0 + exp(-g))) * u;
|
|
}
|
|
for (uint32_t d = 0; d < E; d++) {
|
|
const double *dr = down_shadow + ((uint64_t)e * E + d) * F;
|
|
double p = 0.0;
|
|
for (uint32_t k = 0; k < F; k++) p += dr[k] * mid_exact[((uint64_t)t * n_out + s) * F + k];
|
|
part_exact[((uint64_t)t * n_out + s) * E + d] = p;
|
|
}
|
|
}
|
|
}
|
|
ds4_gpu_tensor *gmid = upload(NULL, mid_n + guard);
|
|
ds4_gpu_tensor *gpart = upload(NULL, part_n + guard);
|
|
const float sentinel = -1234.5f;
|
|
float *ref_mid = NULL, *ref_part = NULL;
|
|
const uint32_t nax_levels[] = {0, 1, 2, 5};
|
|
for (uint32_t li = 0; li < sizeof(nax_levels) / sizeof(nax_levels[0]); li++) {
|
|
const uint32_t nax = nax_levels[li];
|
|
char nax_str[4]; snprintf(nax_str, sizeof(nax_str), "%u", nax);
|
|
setenv("DS4_QWEN4_MOE_MM_NAX", nax_str, 1);
|
|
for (uint32_t i = 0; i < 2; i++) setenv(env_names[i], "8", 1);
|
|
require_ok(ds4_gpu_tensor_fill_f32(gmid, sentinel, mid_n + guard) &&
|
|
ds4_gpu_tensor_fill_f32(gpart, sentinel, part_n + guard), "Q2 tile sentinels");
|
|
char name[96];
|
|
const bool have_nax = ds4_gpu_qwen4_moe_mm_mid_tensor(gmid, gx, glists, gcounts, a->base, a->size, gate_off, up_off,
|
|
16u, NE, T, slots, n_out, E, F, list_cap);
|
|
if (!have_nax && nax != 0) { printf(" Q2 tiles nax=%u: unavailable, skipped\n", nax); continue; }
|
|
require_ok(have_nax, "Q2 tile mid dispatch");
|
|
float *got_mid = download(gmid, mid_n + guard);
|
|
for (uint64_t i = mid_n; i < mid_n + guard; i++) require_ok(got_mid[i] == sentinel, "Q2 tile mid tail guard");
|
|
for (uint32_t t = 0; t < T; t++) for (uint32_t f = 0; f < F; f++)
|
|
require_ok(got_mid[((uint64_t)t * n_out + slots) * F + f] == sentinel, "Q2 tile reserved mid slot");
|
|
double worst = 0.0, scale = 0.0;
|
|
for (uint64_t i = 0; i < mid_n; i++) {
|
|
if (got_mid[i] == sentinel) continue;
|
|
if (ref_mid) { const double d = fabs((double)got_mid[i] - ref_mid[i]); if (d > worst) worst = d; }
|
|
if (fabs(got_mid[i]) > scale) scale = fabs(got_mid[i]);
|
|
}
|
|
double eworst = 0.0, esum = 0.0;
|
|
for (uint64_t i = 0; i < mid_n; i++) {
|
|
if (got_mid[i] == sentinel) continue;
|
|
const double d = fabs((double)got_mid[i] - mid_exact[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" Q2 tiles nax=%u mid vs exact: max=%.3e mean=%.3e\n", nax, eworst, esum / (double)mid_n);
|
|
if (ref_mid) {
|
|
snprintf(name, sizeof(name), "Q2 tile mid nax=%u within 2e-3 of simdgroup", nax);
|
|
require_ok(worst <= 2e-3 * scale, name);
|
|
printf(" Q2 tiles nax=%u mid vs simdgroup: max|d|=%.3e (scale %.3e)\n", nax, worst, scale);
|
|
} else { ref_mid = got_mid; got_mid = NULL; }
|
|
free(got_mid);
|
|
require_ok(ds4_gpu_qwen4_moe_mm_down_tensor(gpart, gmid, glists, gcounts, a->base, a->size, down_off,
|
|
10u, NE, T, slots, n_out, F, E, list_cap), "Q2 tile down dispatch");
|
|
float *got_part = download(gpart, part_n + guard);
|
|
for (uint64_t i = part_n; i < part_n + guard; i++) require_ok(got_part[i] == sentinel, "Q2 tile down tail guard");
|
|
worst = 0.0; scale = 0.0;
|
|
for (uint64_t i = 0; i < part_n; i++) {
|
|
if (got_part[i] == sentinel) continue;
|
|
if (ref_part) { const double d = fabs((double)got_part[i] - ref_part[i]); if (d > worst) worst = d; }
|
|
if (fabs(got_part[i]) > scale) scale = fabs(got_part[i]);
|
|
}
|
|
eworst = 0.0; esum = 0.0;
|
|
for (uint64_t i = 0; i < part_n; i++) {
|
|
if (got_part[i] == sentinel) continue;
|
|
const double d = fabs((double)got_part[i] - part_exact[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" Q2 tiles nax=%u down vs exact: max=%.3e mean=%.3e\n", nax, eworst, esum / (double)part_n);
|
|
if (ref_part) {
|
|
snprintf(name, sizeof(name), "Q2 tile down nax=%u within 2e-3 of simdgroup", nax);
|
|
require_ok(worst <= 2e-3 * scale, name);
|
|
printf(" Q2 tiles nax=%u down vs simdgroup: max|d|=%.3e (scale %.3e)\n", nax, worst, scale);
|
|
} else { ref_part = got_part; got_part = NULL; }
|
|
free(got_part);
|
|
}
|
|
for (uint32_t i = 0; i < 3; i++) {
|
|
const char *names[] = {env_names[0], env_names[1], "DS4_QWEN4_MOE_MM_NAX"};
|
|
if (saved_env[i]) { setenv(names[i], saved_env[i], 1); free(saved_env[i]); } else unsetenv(names[i]);
|
|
}
|
|
free(x); free(sel); free(mid_exact); free(part_exact); free(ref_mid); free(ref_part); free(gate_shadow); free(up_shadow); free(down_shadow);
|
|
ds4_gpu_tensor_free(gx); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(glists); ds4_gpu_tensor_free(gcounts);
|
|
ds4_gpu_tensor_free(gmid); ds4_gpu_tensor_free(gpart);
|
|
}
|
|
|
|
static void test_moe_mm_tiles_exact(arena_t *a, uint32_t down_type) {
|
|
const uint32_t T = 641, E = 256, F = 256, NE = 4, slots = 2, n_out = 3, list_cap = T + 7, guard = 16;
|
|
const uint64_t mid_n = (uint64_t)T * n_out * F, part_n = (uint64_t)T * n_out * E;
|
|
const char *env_names[] = {"DS4_QWEN4_MOE_MID_TILES", "DS4_QWEN4_MOE_DOWN_TILES", "DS4_QWEN4_MOE_MM_NAX"};
|
|
char *saved_env[3];
|
|
for (uint32_t i = 0; i < 3; i++) {
|
|
const char *v = getenv(env_names[i]);
|
|
saved_env[i] = v ? strdup(v) : NULL;
|
|
require_ok(!v || saved_env[i] != NULL, "save MoE tile caps environment");
|
|
}
|
|
setenv("DS4_QWEN4_MOE_MM_NAX", "0", 1); /* caps/nt8 sections pin the simdgroup tiles */
|
|
double *gate_shadow, *up_shadow, *down_shadow;
|
|
const uint64_t gate_off = arena_q4_K(a, (uint64_t)NE * F, E, &gate_shadow, 0.05f);
|
|
const uint64_t up_off = arena_q4_K(a, (uint64_t)NE * F, E, &up_shadow, 0.05f);
|
|
const uint64_t down_off = arena_tier(a, down_type, (uint64_t)NE * E, F, &down_shadow);
|
|
float *x = rand_vec((uint64_t)T * E, 2.0f);
|
|
int32_t *sel = malloc((uint64_t)T * slots * sizeof(int32_t));
|
|
require_ok(sel != NULL, "MoE tile caps selection allocation");
|
|
for (uint32_t t = 0; t < T; t++) { sel[t * slots] = 0; sel[t * slots + 1] = (int32_t)(1u + t % 2u); }
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *gsel = ds4_gpu_tensor_alloc((uint64_t)T * slots * sizeof(int32_t));
|
|
ds4_gpu_tensor *glists = ds4_gpu_tensor_alloc((uint64_t)NE * list_cap * sizeof(int32_t));
|
|
ds4_gpu_tensor *gcounts = ds4_gpu_tensor_alloc(NE * sizeof(int32_t));
|
|
require_ok(gsel && glists && gcounts && ds4_gpu_tensor_write(gsel, 0, sel, (uint64_t)T * slots * sizeof(int32_t)),
|
|
"MoE tile caps routing setup");
|
|
require_ok(ds4_gpu_qwen4_moe_build_lists_tensor(glists, gcounts, gsel, T, slots, NE, list_cap), "MoE tile caps frozen lists");
|
|
int32_t counts[4];
|
|
require_ok(ds4_gpu_tensor_read(gcounts, 0, counts, sizeof(counts)), "MoE tile caps counts read");
|
|
require_ok(counts[0] == (int32_t)T && counts[1] == (int32_t)((T + 1u) / 2u) && counts[2] == (int32_t)(T / 2u) && counts[3] == 0,
|
|
"MoE tile caps hot, partial, and empty experts");
|
|
double *mid_exact = NULL, *part_exact = NULL;
|
|
if (down_type == 39u) {
|
|
/* full-precision reference: double products of the dequantized weights
|
|
* and the float activations, accumulated in double; part takes the
|
|
* exact mid, so mid and down tile paths are judged end to end */
|
|
mid_exact = malloc((uint64_t)T * n_out * F * sizeof(double));
|
|
part_exact = malloc(part_n * sizeof(double));
|
|
require_ok(mid_exact && part_exact, "MoE nax exact reference allocation");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < slots; s++) {
|
|
const uint32_t e = (uint32_t)sel[t * slots + s];
|
|
const float *xr = x + (uint64_t)t * E;
|
|
for (uint32_t f = 0; f < F; f++) {
|
|
const double *gr = gate_shadow + ((uint64_t)e * F + f) * E;
|
|
const double *ur = up_shadow + ((uint64_t)e * F + f) * E;
|
|
double g = 0.0, u = 0.0;
|
|
for (uint32_t k = 0; k < E; k++) { g += gr[k] * (double)xr[k]; u += ur[k] * (double)xr[k]; }
|
|
mid_exact[((uint64_t)t * n_out + s) * F + f] = (g / (1.0 + exp(-g))) * u;
|
|
}
|
|
for (uint32_t d = 0; d < E; d++) {
|
|
const double *dr = down_shadow + ((uint64_t)e * E + d) * F;
|
|
double p = 0.0;
|
|
for (uint32_t k = 0; k < F; k++) p += dr[k] * mid_exact[((uint64_t)t * n_out + s) * F + k];
|
|
part_exact[((uint64_t)t * n_out + s) * E + d] = p;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
ds4_gpu_tensor *gmid = upload(NULL, mid_n + guard);
|
|
ds4_gpu_tensor *gpart = upload(NULL, part_n + guard);
|
|
float *ref_mid = NULL, *ref_part = NULL;
|
|
const float sentinel = -1234.5f;
|
|
const uint32_t caps[] = {8, 1, 16, 32};
|
|
for (uint32_t mode = 0; mode < sizeof(caps) / sizeof(caps[0]); mode++) {
|
|
char value[16];
|
|
snprintf(value, sizeof(value), "%u", caps[mode]);
|
|
for (uint32_t i = 0; i < 2; i++) setenv(env_names[i], value, 1);
|
|
require_ok(ds4_gpu_tensor_fill_f32(gmid, sentinel, mid_n + guard) &&
|
|
ds4_gpu_tensor_fill_f32(gpart, sentinel, part_n + guard), "MoE tile caps sentinels");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_mid_tensor(gmid, gx, glists, gcounts, a->base, a->size, gate_off, up_off,
|
|
12u, NE, T, slots, n_out, E, F, list_cap), "MoE tile caps mid dispatch");
|
|
float *got_mid = download(gmid, mid_n + guard);
|
|
for (uint64_t i = mid_n; i < mid_n + guard; i++) require_ok(got_mid[i] == sentinel, "MoE tile caps mid tail guard");
|
|
for (uint32_t t = 0; t < T; t++) for (uint32_t f = 0; f < F; f++)
|
|
require_ok(got_mid[((uint64_t)t * n_out + slots) * F + f] == sentinel, "MoE tile caps reserved mid slot");
|
|
char name[128];
|
|
snprintf(name, sizeof(name), "MoE tile caps mid T=%u down=%u cap=%u", T, down_type, caps[mode]);
|
|
if (mode == 0) {
|
|
ref_mid = got_mid; check_exact_f32(name, ref_mid, ref_mid, mid_n);
|
|
uint64_t h = 1469598103934665603ull;
|
|
for (uint64_t i = 0; i < mid_n; i++) { uint32_t u; memcpy(&u, &got_mid[i], 4); h = (h ^ u) * 1099511628211ull; }
|
|
if (mid_exact) {
|
|
double eworst = 0.0, esum = 0.0;
|
|
for (uint64_t i = 0; i < mid_n; i++) {
|
|
if (ref_mid[i] == sentinel) continue;
|
|
const double d = fabs((double)ref_mid[i] - mid_exact[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" MoE simdgroup mid vs exact: max=%.3e mean=%.3e\n", eworst, esum / (double)mid_n);
|
|
}
|
|
}
|
|
else { check_exact_f32(name, got_mid, ref_mid, mid_n + guard); free(got_mid); }
|
|
require_ok(ds4_gpu_qwen4_moe_mm_down_tensor(gpart, gmid, glists, gcounts, a->base, a->size, down_off,
|
|
down_type, NE, T, slots, n_out, F, E, list_cap), "MoE tile caps down dispatch");
|
|
float *got_part = download(gpart, part_n + guard);
|
|
for (uint64_t i = part_n; i < part_n + guard; i++) require_ok(got_part[i] == sentinel, "MoE tile caps down tail guard");
|
|
for (uint32_t t = 0; t < T; t++) for (uint32_t e = 0; e < E; e++)
|
|
require_ok(got_part[((uint64_t)t * n_out + slots) * E + e] == sentinel, "MoE tile caps reserved down slot");
|
|
snprintf(name, sizeof(name), "MoE tile caps down T=%u type=%u cap=%u", T, down_type, caps[mode]);
|
|
if (mode == 0) {
|
|
ref_part = got_part; check_exact_f32(name, ref_part, ref_part, part_n);
|
|
uint64_t h = 1469598103934665603ull;
|
|
for (uint64_t i = 0; i < part_n; i++) { uint32_t u; memcpy(&u, &got_part[i], 4); h = (h ^ u) * 1099511628211ull; }
|
|
if (part_exact) {
|
|
double eworst = 0.0, esum = 0.0;
|
|
for (uint64_t i = 0; i < part_n; i++) {
|
|
if (ref_part[i] == sentinel) continue;
|
|
const double d = fabs((double)ref_part[i] - part_exact[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" MoE simdgroup down vs exact: max=%.3e mean=%.3e\n", eworst, esum / (double)part_n);
|
|
}
|
|
}
|
|
else { check_exact_f32(name, got_part, ref_part, part_n + guard); free(got_part); }
|
|
}
|
|
{ /* 64-token gate/up tiles (with their 8/16/32-token tails) must match cap 8 */
|
|
for (uint32_t i = 0; i < 2; i++) setenv(env_names[i], "8", 1);
|
|
setenv("DS4_QWEN4_MOE_MID_NT", "8", 1);
|
|
setenv("DS4_QWEN4_MOE_TAILS", "1", 1);
|
|
require_ok(ds4_gpu_tensor_fill_f32(gmid, sentinel, mid_n + guard), "MoE mid nt8 sentinels");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_mid_tensor(gmid, gx, glists, gcounts, a->base, a->size, gate_off, up_off,
|
|
12u, NE, T, slots, n_out, E, F, list_cap), "MoE mid nt8 dispatch");
|
|
float *got_mid = download(gmid, mid_n + guard);
|
|
check_exact_f32("MoE mid nt8", got_mid, ref_mid, mid_n + guard);
|
|
free(got_mid);
|
|
unsetenv("DS4_QWEN4_MOE_MID_NT");
|
|
unsetenv("DS4_QWEN4_MOE_TAILS");
|
|
}
|
|
if (down_type == 39u) for (uint32_t nax = 1; nax <= 5; nax++) {
|
|
/* tensor-op tiles of 32/64 tokens, half (1/2), float (3/4) or
|
|
* compensated (5) activation operands: cooperative accumulation;
|
|
* bound the drift against the simdgroup tiles and print the
|
|
* exact-reference error */
|
|
char nax_str[4]; snprintf(nax_str, sizeof(nax_str), "%u", nax);
|
|
for (uint32_t i = 0; i < 2; i++) setenv(env_names[i], "8", 1);
|
|
setenv("DS4_QWEN4_MOE_MM_NAX", nax_str, 1);
|
|
require_ok(ds4_gpu_tensor_fill_f32(gmid, sentinel, mid_n + guard) &&
|
|
ds4_gpu_tensor_fill_f32(gpart, sentinel, part_n + guard), "MoE nax sentinels");
|
|
if (ds4_gpu_qwen4_moe_mm_mid_tensor(gmid, gx, glists, gcounts, a->base, a->size, gate_off, up_off,
|
|
12u, NE, T, slots, n_out, E, F, list_cap)) {
|
|
float *got_mid = download(gmid, mid_n + guard);
|
|
double worst = 0.0, scale = 0.0;
|
|
for (uint64_t i = 0; i < mid_n; i++) {
|
|
if (ref_mid[i] == sentinel) continue;
|
|
const double d = fabs((double)got_mid[i] - ref_mid[i]);
|
|
if (d > worst) worst = d;
|
|
if (fabs(ref_mid[i]) > scale) scale = fabs(ref_mid[i]);
|
|
}
|
|
for (uint64_t i = mid_n; i < mid_n + guard; i++) require_ok(got_mid[i] == sentinel, "MoE nax mid tail guard");
|
|
require_ok(worst <= 2e-3 * scale, "MoE nax mid within 2e-3 of the simdgroup tiles");
|
|
{
|
|
uint64_t h = 1469598103934665603ull;
|
|
for (uint64_t i = 0; i < mid_n; i++) { uint32_t u; memcpy(&u, &got_mid[i], 4); h = (h ^ u) * 1099511628211ull; }
|
|
printf(" MoE nax=%u mid: max|d|=%.3e (scale %.3e) hash=%016llx\n", nax, worst, scale, (unsigned long long)h);
|
|
}
|
|
{
|
|
double eworst = 0.0, esum = 0.0;
|
|
for (uint64_t i = 0; i < mid_n; i++) {
|
|
if (ref_mid[i] == sentinel) continue;
|
|
const double d = fabs((double)got_mid[i] - mid_exact[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" MoE nax=%u mid vs exact: max=%.3e mean=%.3e\n", nax, eworst, esum / (double)mid_n);
|
|
}
|
|
require_ok(ds4_gpu_qwen4_moe_mm_down_tensor(gpart, gmid, glists, gcounts, a->base, a->size, down_off,
|
|
down_type, NE, T, slots, n_out, F, E, list_cap), "MoE nax down dispatch");
|
|
float *got_part = download(gpart, part_n + guard);
|
|
worst = 0.0; scale = 0.0;
|
|
for (uint64_t i = 0; i < part_n; i++) {
|
|
if (ref_part[i] == sentinel) continue;
|
|
const double d = fabs((double)got_part[i] - ref_part[i]);
|
|
if (d > worst) worst = d;
|
|
if (fabs(ref_part[i]) > scale) scale = fabs(ref_part[i]);
|
|
}
|
|
for (uint64_t i = part_n; i < part_n + guard; i++) require_ok(got_part[i] == sentinel, "MoE nax down tail guard");
|
|
require_ok(worst <= 2e-3 * scale, "MoE nax down within 2e-3 of the simdgroup tiles");
|
|
{
|
|
uint64_t h = 1469598103934665603ull;
|
|
for (uint64_t i = 0; i < part_n; i++) { uint32_t u; memcpy(&u, &got_part[i], 4); h = (h ^ u) * 1099511628211ull; }
|
|
printf(" MoE nax=%u down: max|d|=%.3e (scale %.3e) hash=%016llx\n", nax, worst, scale, (unsigned long long)h);
|
|
}
|
|
{
|
|
double eworst = 0.0, esum = 0.0;
|
|
for (uint64_t i = 0; i < part_n; i++) {
|
|
if (ref_part[i] == sentinel) continue;
|
|
const double d = fabs((double)got_part[i] - part_exact[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" MoE nax=%u down vs exact: max=%.3e mean=%.3e\n", nax, eworst, esum / (double)part_n);
|
|
}
|
|
free(got_mid); free(got_part);
|
|
} else {
|
|
printf(" MoE nax: tensor API unavailable, skipped\n");
|
|
}
|
|
unsetenv("DS4_QWEN4_MOE_MM_NAX");
|
|
}
|
|
if (down_type == 39u) {
|
|
/* error decomposition on a dyadic Q4_K fixture: the dequantized
|
|
* weights are exactly representable in half, so the simdgroup path's
|
|
* residual is the activation rounding plus fp32 accumulation, and the
|
|
* compensated tensor path's residual is its accumulation alone */
|
|
double *dgs, *dus;
|
|
const uint64_t dgate_off = arena_q4_K_dyadic(a, (uint64_t)NE * F, E, &dgs);
|
|
const uint64_t dup_off = arena_q4_K_dyadic(a, (uint64_t)NE * F, E, &dus);
|
|
double *dmid = malloc((uint64_t)T * n_out * F * sizeof(double));
|
|
require_ok(dmid != NULL, "MoE dyadic reference allocation");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t s = 0; s < slots; s++) {
|
|
const uint32_t e2 = (uint32_t)sel[t * slots + s];
|
|
const float *xr = x + (uint64_t)t * E;
|
|
for (uint32_t f = 0; f < F; f++) {
|
|
const double *gr = dgs + ((uint64_t)e2 * F + f) * E;
|
|
const double *ur = dus + ((uint64_t)e2 * F + f) * E;
|
|
double g = 0.0, u = 0.0;
|
|
for (uint32_t k = 0; k < E; k++) { g += gr[k] * (double)xr[k]; u += ur[k] * (double)xr[k]; }
|
|
dmid[((uint64_t)t * n_out + s) * F + f] = (g / (1.0 + exp(-g))) * u;
|
|
}
|
|
}
|
|
}
|
|
for (uint32_t level = 0; level < 3; level++) {
|
|
double eworst = 0.0, esum = 0.0;
|
|
setenv("DS4_QWEN4_MOE_MM_NAX", level == 0 ? "0" : level == 1 ? "2" : "5", 1);
|
|
require_ok(ds4_gpu_tensor_fill_f32(gmid, sentinel, mid_n + guard), "MoE dyadic sentinels");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_mid_tensor(gmid, gx, glists, gcounts, a->base, a->size, dgate_off, dup_off,
|
|
12u, NE, T, slots, n_out, E, F, list_cap), "MoE dyadic mid dispatch");
|
|
float *got = download(gmid, mid_n + guard);
|
|
for (uint64_t i = 0; i < mid_n; i++) {
|
|
if (ref_mid[i] == sentinel) continue;
|
|
const double d = fabs((double)got[i] - dmid[i]);
|
|
if (d > eworst) eworst = d;
|
|
esum += d;
|
|
}
|
|
printf(" MoE dyadic mid %s: max=%.3e mean=%.3e\n",
|
|
level == 0 ? "simdgroup" : level == 1 ? "nax=2" : "nax=5", eworst, esum / (double)mid_n);
|
|
free(got);
|
|
}
|
|
unsetenv("DS4_QWEN4_MOE_MM_NAX");
|
|
free(dmid); free(dgs); free(dus);
|
|
}
|
|
printf(" MoE tile caps T=%u Q4_K/%s: caps 1,16,32 and 64-token gate/up tiles byte-exact mid/down vs cap8\n",
|
|
T, down_type == 39u ? "mxfp4" : "q8_0");
|
|
for (uint32_t i = 0; i < 3; i++) {
|
|
if (saved_env[i]) setenv(env_names[i], saved_env[i], 1); else unsetenv(env_names[i]);
|
|
free(saved_env[i]);
|
|
}
|
|
free(ref_part); free(ref_mid); free(sel); free(x);
|
|
free(gate_shadow); free(up_shadow); free(down_shadow); free(mid_exact); free(part_exact);
|
|
ds4_gpu_tensor_free(gpart); ds4_gpu_tensor_free(gmid); ds4_gpu_tensor_free(gcounts);
|
|
ds4_gpu_tensor_free(glists); ds4_gpu_tensor_free(gsel); ds4_gpu_tensor_free(gx);
|
|
}
|
|
|
|
/* MTP input staging: cat rows [rms(e)*g_e | 0] and [0 | rms(R_s)*g_h_s]
|
|
* (full-row or per-stream RMS), then R_out = proj[0] + proj[1+s]. */
|
|
static void test_mtp(arena_t *a, uint32_t E, uint32_t hc) {
|
|
double *g_e, *g_h;
|
|
const uint64_t g_e_off = arena_f32(a, E, &g_e, 0.5f, 1.5f);
|
|
const uint64_t g_h_off = arena_f32(a, (uint64_t)hc * E, &g_h, 0.5f, 1.5f);
|
|
float *e = rand_vec(E, 1.0f);
|
|
float *R = rand_vec((uint64_t)hc * E, 1.0f);
|
|
float *proj = rand_vec((uint64_t)(hc + 1u) * E, 1.0f);
|
|
double *cat = calloc((uint64_t)(hc + 1u) * 2u * E, sizeof(double));
|
|
double *R_ref = malloc((uint64_t)hc * E * sizeof(double));
|
|
const double eps = 1e-6;
|
|
double ss = 0.0;
|
|
for (uint32_t i = 0; i < E; i++) ss += (double)e[i] * e[i];
|
|
double inv = 1.0 / sqrt(ss / E + eps);
|
|
for (uint32_t i = 0; i < E; i++) cat[i] = e[i] * inv * g_e[i];
|
|
double full = 0.0;
|
|
for (uint32_t i = 0; i < hc * E; i++) full += (double)R[i] * R[i];
|
|
inv = 1.0 / sqrt(full / ((double)hc * E) + eps);
|
|
for (uint32_t s = 0; s < hc; s++) {
|
|
for (uint32_t i = 0; i < E; i++) {
|
|
cat[(uint64_t)(s + 1u) * 2u * E + E + i] = R[s * E + i] * inv * g_h[s * E + i];
|
|
R_ref[s * E + i] = (double)proj[i] + proj[(s + 1u) * E + i];
|
|
}
|
|
}
|
|
ds4_gpu_tensor *ge = upload(e, E);
|
|
ds4_gpu_tensor *gR = upload(R, (uint64_t)hc * E);
|
|
ds4_gpu_tensor *gcat = upload(NULL, (uint64_t)(hc + 1u) * 2u * E);
|
|
ds4_gpu_tensor *gproj = upload(proj, (uint64_t)(hc + 1u) * E);
|
|
ds4_gpu_tensor *gout = upload(NULL, (uint64_t)hc * E);
|
|
require_ok(ds4_gpu_qwen4_mtp_stage_tensor(gcat, ge, gR, a->base, a->size, g_e_off, g_h_off, E, hc, (float)eps), "mtp stage");
|
|
require_ok(ds4_gpu_qwen4_mtp_combine_tensor(gout, gproj, E, hc), "mtp combine");
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "mtp stage E=%u hc=%u", E, hc);
|
|
check_tensor(name, gcat, cat, (uint64_t)(hc + 1u) * 2u * E, 1e-5);
|
|
snprintf(name, sizeof(name), "mtp combine E=%u hc=%u", E, hc);
|
|
check_tensor(name, gout, R_ref, (uint64_t)hc * E, 1e-6);
|
|
ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gproj); ds4_gpu_tensor_free(gcat); ds4_gpu_tensor_free(gR); ds4_gpu_tensor_free(ge);
|
|
free(R_ref); free(cat); free(proj); free(R); free(e); free(g_e); free(g_h);
|
|
}
|
|
|
|
/* A split immediately below the large-prefill boundary uses the original
|
|
* scan geometry. Compare every output and state value across that boundary,
|
|
* including the first-token snapshot used by speculative verification. */
|
|
static void test_gdn_prefill_dispatch(void) {
|
|
const uint32_t Hk = 16u, Hv = 48u, D = 128u;
|
|
const uint64_t C = (2u * Hk + Hv) * D, V = Hv * D, S = V * D;
|
|
for (uint32_t T = 8191u; T <= 8193u; T++) {
|
|
float *qkv = rand_vec(T * C, 0.05f);
|
|
float *initial = rand_vec(S, 0.01f);
|
|
ds4_gpu_tensor *gqkv = upload(qkv, T * C);
|
|
ds4_gpu_tensor *ga = upload(NULL, (uint64_t)T * Hv);
|
|
ds4_gpu_tensor *gb = upload(NULL, (uint64_t)T * Hv);
|
|
ds4_gpu_tensor *gs = upload(initial, S);
|
|
ds4_gpu_tensor *go = upload(NULL, T * V);
|
|
ds4_gpu_tensor *gsnap = upload(NULL, S);
|
|
require_ok(ds4_gpu_tensor_fill_f32(ga, 0.95f, (uint64_t)T * Hv) &&
|
|
ds4_gpu_tensor_fill_f32(gb, 0.25f, (uint64_t)T * Hv), "GDN prefill gates");
|
|
free(qkv);
|
|
float *ref_out = NULL, *ref_state = NULL, *ref_snap = NULL;
|
|
for (uint32_t mode = 0; mode < 2u; mode++) {
|
|
require_ok(ds4_gpu_tensor_write(gs, 0, initial, S * sizeof(float)), "GDN prefill reset");
|
|
const uint32_t chunk = mode == 0u ? T - 2u : T;
|
|
for (uint32_t pos = 0; pos < T; pos += chunk) {
|
|
const uint32_t n = T - pos < chunk ? T - pos : chunk;
|
|
ds4_gpu_tensor *q = ds4_gpu_tensor_view(gqkv, pos * C * sizeof(float), n * C * sizeof(float));
|
|
ds4_gpu_tensor *a = ds4_gpu_tensor_view(ga, (uint64_t)pos * Hv * sizeof(float), (uint64_t)n * Hv * sizeof(float));
|
|
ds4_gpu_tensor *b = ds4_gpu_tensor_view(gb, (uint64_t)pos * Hv * sizeof(float), (uint64_t)n * Hv * sizeof(float));
|
|
ds4_gpu_tensor *o = ds4_gpu_tensor_view(go, pos * V * sizeof(float), n * V * sizeof(float));
|
|
require_ok(q && a && b && o, "GDN prefill views");
|
|
require_ok(ds4_gpu_qwen4_gdn_scan_tensor(o, gs, q, a, b, n, Hk, Hv, D,
|
|
pos == 0u ? gsnap : NULL, 0u, NULL, 0u), "GDN prefill scan");
|
|
ds4_gpu_tensor_free(q); ds4_gpu_tensor_free(a);
|
|
ds4_gpu_tensor_free(b); ds4_gpu_tensor_free(o);
|
|
}
|
|
float *out = download(go, T * V), *state = download(gs, S), *snap = download(gsnap, S);
|
|
if (mode == 0u) {
|
|
ref_out = out; ref_state = state; ref_snap = snap;
|
|
} else {
|
|
require_ok(memcmp(out, ref_out, T * V * sizeof(float)) == 0 &&
|
|
memcmp(state, ref_state, S * sizeof(float)) == 0 &&
|
|
memcmp(snap, ref_snap, S * sizeof(float)) == 0,
|
|
"GDN prefill byte-exact output/state/snapshot");
|
|
free(out); free(state); free(snap);
|
|
}
|
|
}
|
|
printf(" GDN prefill T=%u: byte-exact split/batch output, state and snapshot\n", T);
|
|
free(ref_out); free(ref_state); free(ref_snap); free(initial);
|
|
ds4_gpu_tensor_free(gqkv); ds4_gpu_tensor_free(ga); ds4_gpu_tensor_free(gb);
|
|
ds4_gpu_tensor_free(gs); ds4_gpu_tensor_free(go); ds4_gpu_tensor_free(gsnap);
|
|
}
|
|
}
|
|
|
|
static double bench_now(void) {
|
|
struct timespec ts;
|
|
clock_gettime(CLOCK_MONOTONIC, &ts);
|
|
return (double)ts.tv_sec + 1e-9 * (double)ts.tv_nsec;
|
|
}
|
|
|
|
static double bench_hc_norm_reuse_batch(arena_t *a, uint32_t type, uint32_t T,
|
|
uint64_t gamma_off, uint64_t inject_off,
|
|
ds4_gpu_tensor *gR, ds4_gpu_tensor *gxn,
|
|
ds4_gpu_tensor *ginj, uint32_t mode, uint32_t calls) {
|
|
require_ok(setenv("DS4_QWEN4_HC_NORM_REUSE", mode ? "1" : "0", 1) == 0, "select timed HC norm kernel");
|
|
const double t0 = bench_now();
|
|
require_ok(ds4_gpu_begin_commands(), "HC norm timing batch begin");
|
|
for (uint32_t i = 0; i < calls; i++) {
|
|
require_ok(ds4_gpu_qwen4_hc_norm_tensor(gxn, ginj, gR, a->base, a->size,
|
|
gamma_off, inject_off, type, T, 2560u, 4u, 4u, 1e-6f), "timed HC norm dispatch");
|
|
}
|
|
/* end_commands waits for GPU completion; the returned wall time includes
|
|
* CPU encoding and the batch's final wait, but no weight/input allocation. */
|
|
require_ok(ds4_gpu_end_commands(), "HC norm timing batch end and wait");
|
|
return 1e6 * (bench_now() - t0) / calls;
|
|
}
|
|
|
|
static double hc_norm_median7(const double *samples) {
|
|
double sorted[7];
|
|
memcpy(sorted, samples, sizeof(sorted));
|
|
for (uint32_t i = 1; i < 7; i++) {
|
|
const double value = sorted[i];
|
|
uint32_t j = i;
|
|
while (j && sorted[j - 1u] > value) { sorted[j] = sorted[j - 1u]; j--; }
|
|
sorted[j] = value;
|
|
}
|
|
return sorted[3];
|
|
}
|
|
|
|
/* Optional scheduling screen; excluded from the default correctness path.
|
|
* Each type reuses its weight matrices at every batch size. Both variants get
|
|
* warmups, then seven paired observations with alternating execution order. */
|
|
static void bench_hc_norm_reuse(arena_t *a) {
|
|
const char *value = getenv("DS4_QWEN4_HC_NORM_REUSE");
|
|
char *saved = value ? strdup(value) : NULL;
|
|
require_ok(!value || saved != NULL, "save HC norm timing environment");
|
|
const uint32_t types[] = {0u, 1u, 8u}, tokens[] = {33u, 128u, 512u, 1024u, 2048u, 4096u, 8192u};
|
|
const uint64_t dim = 2560u * 4u;
|
|
printf("HC norm timing: E=2560 hc=4 inject=4; microseconds/call including batch encoding and final GPU wait\n");
|
|
for (uint32_t type_index = 0; type_index < 3; type_index++) {
|
|
const uint32_t type = types[type_index];
|
|
double *shadow;
|
|
const uint64_t gamma_off = arena_f32(a, dim, &shadow, 0.5f, 1.5f);
|
|
free(shadow);
|
|
const uint64_t inject_off = type == 8u ? arena_q8_0(a, 4u, dim, &shadow, 0.05f)
|
|
: type == 1u ? arena_f16(a, 4u * dim, &shadow, 0.05f)
|
|
: arena_f32(a, 4u * dim, &shadow, -0.05f, 0.05f);
|
|
free(shadow);
|
|
for (uint32_t size = 0; size < sizeof(tokens) / sizeof(tokens[0]); size++) {
|
|
const uint32_t T = tokens[size], calls = T < 512u ? 32u : T < 8192u ? 8u : 2u;
|
|
float *R = rand_vec((uint64_t)T * dim, 1.0f);
|
|
ds4_gpu_tensor *gR = upload(R, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *gxn = upload(NULL, (uint64_t)T * dim);
|
|
ds4_gpu_tensor *ginj = upload(NULL, (uint64_t)T * 4u * DS4_QWEN4_HC_CHUNKS * 4u);
|
|
free(R);
|
|
for (uint32_t mode = 0; mode < 2; mode++) {
|
|
(void)bench_hc_norm_reuse_batch(a, type, T, gamma_off, inject_off, gR, gxn, ginj, mode, calls);
|
|
}
|
|
double samples[2][7];
|
|
for (uint32_t rep = 0; rep < 7; rep++) for (uint32_t order = 0; order < 2; order++) {
|
|
const uint32_t mode = (rep + order) & 1u;
|
|
samples[mode][rep] = bench_hc_norm_reuse_batch(a, type, T, gamma_off, inject_off,
|
|
gR, gxn, ginj, mode, calls);
|
|
}
|
|
const double baseline = hc_norm_median7(samples[0]), reuse = hc_norm_median7(samples[1]);
|
|
printf(" HC norm timing type=%u T=%u calls=%u median baseline=%.3f us reuse=%.3f us speedup=%+.2f%%\n",
|
|
type, T, calls, baseline, reuse, 100.0 * (baseline / reuse - 1.0));
|
|
for (uint32_t mode = 0; mode < 2; mode++) {
|
|
printf(" %s samples_us:", mode ? "reuse" : "baseline");
|
|
for (uint32_t rep = 0; rep < 7; rep++) printf(" %.3f", samples[mode][rep]);
|
|
printf("\n");
|
|
}
|
|
ds4_gpu_tensor_free(ginj); ds4_gpu_tensor_free(gxn); ds4_gpu_tensor_free(gR);
|
|
}
|
|
}
|
|
require_ok(saved ? setenv("DS4_QWEN4_HC_NORM_REUSE", saved, 1) == 0
|
|
: unsetenv("DS4_QWEN4_HC_NORM_REUSE") == 0, "restore HC norm timing environment");
|
|
free(saved);
|
|
}
|
|
|
|
typedef int (*bench_fn)(void *ud);
|
|
|
|
static double bench_run(const char *name, bench_fn fn, void *ud, uint32_t reps) {
|
|
const char *only = getenv("QWEN4_BENCH_ONLY");
|
|
if (only && only[0] && !strstr(name, only)) return 0.0;
|
|
if (only && only[0]) { for (uint32_t i = 0; i < 20; i++) fn(ud); } /* single-bench runs: warm the clocks */
|
|
require_ok(ds4_gpu_begin_commands(), "begin");
|
|
require_ok(fn(ud), name);
|
|
require_ok(ds4_gpu_end_commands(), "end");
|
|
require_ok(ds4_gpu_synchronize(), "sync");
|
|
const double t0 = bench_now();
|
|
require_ok(ds4_gpu_begin_commands(), "begin");
|
|
for (uint32_t r = 0; r < reps; r++) require_ok(fn(ud), name);
|
|
require_ok(ds4_gpu_end_commands(), "end");
|
|
require_ok(ds4_gpu_synchronize(), "sync");
|
|
const double us = 1e6 * (bench_now() - t0) / reps;
|
|
printf(" %-44s %8.1f us\n", name, us);
|
|
return us;
|
|
}
|
|
|
|
typedef struct {
|
|
arena_t *a;
|
|
uint64_t off[14];
|
|
ds4_gpu_tensor *t[43];
|
|
uint32_t n[3];
|
|
} bench_ctx;
|
|
|
|
static int bench_q8_small(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_q8_0_tensor(c->t[1], c->a->base, c->a->size, c->off[0], 2560, 48, c->t[0], 1); }
|
|
static int bench_q8_big(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_q8_0_tensor(c->t[1], c->a->base, c->a->size, c->off[1], 2560, 6144, c->t[0], 1); }
|
|
static int bench_q8_big2(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_q8_0_tensor(c->t[1], c->a->base, c->a->size, c->off[1], 2560, 6144, c->t[0], 2); }
|
|
static int bench_combine(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_hc_combine_tensor(c->t[2], c->t[0], c->t[3], 1, 2560, 4); }
|
|
static int bench_hc_norm(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_hc_norm_tensor(c->t[4], c->t[3], c->t[2], c->a->base, c->a->size, c->off[2], c->off[4], 1u, 1, 2560, 4, 4, 1e-6f); }
|
|
static int bench_hc_down(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f16_tensor(c->t[5], c->a->base, c->a->size, c->off[3], 10240, 320, c->t[4], 1); }
|
|
static int bench_hc_mix(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_hc_gate_mix_tensor(c->t[0], c->t[4], c->t[5], c->a->base, c->a->size, c->off[5], 1u, 1, 2560, 4, 320); }
|
|
static int bench_hc_mix2(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_hc_gate_mix_tensor(c->t[0], c->t[4], c->t[5], c->a->base, c->a->size, c->off[5], 1u, 2, 2560, 4, 320); }
|
|
static int bench_attn(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_attn_decode_tensor(c->t[6], c->t[7], c->t[7], c->t[8], c->t[9], NULL, NULL, c->n[1] ? c->t[10] : NULL,
|
|
1, 24, 2, 256, c->n[0], false, 0, 0.0625f);
|
|
}
|
|
static int bench_gdn_front(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_gdn_front_tensor(c->t[11], c->t[12], c->t[0], c->t[13], c->t[14], c->a->base, c->a->size,
|
|
c->off[6], c->off[0], c->off[0], c->off[7], c->off[7], 8u, 1, 16, 48, 128, 4, 2560, NULL, 0u, NULL, 0u);
|
|
}
|
|
static int bench_gdn_unfused(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_matmul_q8_0_tensor(c->t[13], c->a->base, c->a->size, c->off[0], 2560, 48, c->t[0], 1) &&
|
|
ds4_gpu_matmul_q8_0_tensor(c->t[14], c->a->base, c->a->size, c->off[0], 2560, 48, c->t[0], 1) &&
|
|
ds4_gpu_qwen4_conv_stream_tensor(c->t[11], c->t[12], c->a->base, c->a->size, c->off[6], 1, 10240, 4, true) &&
|
|
ds4_gpu_qwen4_gdn_prep_tensor(c->t[11], c->t[13], c->t[14], c->a->base, c->a->size, c->off[7], c->off[7], 1, 16, 48, 128);
|
|
}
|
|
static int bench_moe_mid(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_moe_mid_tensor(c->t[1], c->t[0], c->t[16], c->a->base, c->a->size, c->off[8], c->off[9], 8u, 16, 1, 10, 2560, 640, c->off[8], c->off[9], 8u);
|
|
}
|
|
static int bench_moe_mid_q4k(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_moe_mid_tensor(c->t[1], c->t[0], c->t[16], c->a->base, c->a->size,
|
|
c->off[12], c->off[13], 12u, 16, c->n[0], 10, 2560, 640,
|
|
c->off[8], c->off[9], 8u);
|
|
}
|
|
static int bench_moe_down(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_moe_down_tensor(c->t[17], c->t[1], c->t[16], c->a->base, c->a->size, c->off[10], 8u, 16, 1, 10, 640, 2560, c->off[10], 8u);
|
|
}
|
|
static int bench_router_gemv(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f32_tensor(c->t[17], c->a->base, c->a->size, c->off[11], 2560, 512, c->t[0], 1); }
|
|
static int bench_router_topk(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_router_topk_tensor(c->t[16], c->t[18], c->t[17], c->t[0], c->a->base, c->a->size, c->off[2], 0u, 2560, c->t[18], 1, 512, 10);
|
|
}
|
|
static int bench_router_topk_nogate(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_router_topk_tensor(c->t[16], c->t[18], c->t[17], NULL, NULL, 0, 0, 0u, 0, NULL, 1, 512, 10);
|
|
}
|
|
static int bench_router_topk_k1(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_router_topk_tensor(c->t[16], c->t[18], c->t[17], NULL, NULL, 0, 0, 0u, 0, NULL, 1, 512, 1);
|
|
}
|
|
/* prefill-sized (T=256) variants */
|
|
static int bench_p_router_f32(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f32_tensor(c->t[19], c->a->base, c->a->size, c->off[11], 2560, 512, c->t[18], 256); }
|
|
static int bench_p_router_mm(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_dense_mm_tensor(c->t[19], c->t[18], c->a->base, c->a->size, c->off[11], 0u, 256, 2560, 512); }
|
|
static int bench_p_topk(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_router_topk_tensor(c->t[16], c->t[1], c->t[19], c->t[18], c->a->base, c->a->size, c->off[2], 0u, 2560, c->t[17], 256, 512, 10); }
|
|
static int bench_p_hc_norm(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_hc_norm_tensor(c->t[1], c->t[3], c->t[19], c->a->base, c->a->size, c->off[2], c->off[4], 1u, 256, 2560, 4, 4, 1e-6f); }
|
|
static int bench_p_hc_down_f16(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f16_tensor(c->t[17], c->a->base, c->a->size, c->off[3], 10240, 320, c->t[1], 256); }
|
|
static int bench_p_hc_down_mm(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_dense_mm_tensor(c->t[17], c->t[1], c->a->base, c->a->size, c->off[3], 1u, 256, 10240, 320); }
|
|
static int bench_p_hc_up_f16(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f16_tensor(c->t[1], c->a->base, c->a->size, c->off[5], 320, 10240, c->t[17], 256); }
|
|
static int bench_p_hc_up_mm(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_dense_mm_tensor(c->t[1], c->t[17], c->a->base, c->a->size, c->off[5], 1u, 256, 320, 10240); }
|
|
static int bench_p_hc_mix_rows(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_hc_mix_rows_tensor(c->t[19], c->t[1], c->t[1], 256, 2560, 4); }
|
|
static int bench_p_q8_gemm(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_q8_0_tensor(c->t[1], c->a->base, c->a->size, c->off[1], 2560, 6144, c->t[18], 256); }
|
|
static int bench_p_q8_mm(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_dense_mm_tensor(c->t[1], c->t[18], c->a->base, c->a->size, c->off[1], 8u, 256, 2560, 6144); }
|
|
static int bench_p_idx_score(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_idx_score_tensor(c->t[26], NULL, c->t[24], c->t[25], 32, 65536, 4, 128, 262144, 4); }
|
|
static int bench_p_idx_argsort(void *ud) { bench_ctx *c = ud; return ds4_gpu_indexer_topk_tensor(c->t[27], c->t[26], 65536, 32, 512); }
|
|
static int bench_p_idx_select(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_idx_select_tensor(c->t[27], c->t[26], NULL, 65536, 32, 512); }
|
|
static int bench_p_q8_gemm_2k(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_q8_0_tensor(c->t[40], c->a->base, c->a->size, c->off[1], 2560, 6144, c->t[39], 2048); }
|
|
static int bench_p_f16_gemm_2k(void *ud) { bench_ctx *c = ud; return ds4_gpu_matmul_f16_tensor(c->t[41], c->a->base, c->a->size, c->off[3], 10240, 320, c->t[42], 2048); }
|
|
static int bench_p_idx_score_1k(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_idx_score_tensor(c->t[29], NULL, c->t[28], c->t[25], 1024, 65536, 4, 128, 262144 - 1024, 4); }
|
|
static int bench_p_idx_select_1k(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_idx_select_tensor(c->t[30], c->t[29], NULL, 65536, 1024, 512); }
|
|
/* sparse prefill attention: 1024 queries at the end of a 256k context, 512 selected blocks each */
|
|
static int bench_p_attn_sparse(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_attn_decode_tensor(c->t[37], c->t[33], c->t[34], c->t[31], c->t[32], c->n[2] ? c->t[36] : c->t[35], c->t[38], NULL,
|
|
1024, 24, 2, 256, 262144 - 1024, true, 2052, 0.0625f);
|
|
}
|
|
static int bench_p_gdn_r4(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_gdn_scan_tensor(c->t[15], c->t[1], c->t[11], c->t[13], c->t[14], 1024, 16, 48, 128, NULL, 0u, NULL, 0u); }
|
|
static int bench_p_moe_mm(void *ud) {
|
|
bench_ctx *c = ud;
|
|
return ds4_gpu_qwen4_moe_build_lists_tensor(c->t[20], c->t[21], c->t[16], 256, 10, 16, 256) &&
|
|
ds4_gpu_qwen4_moe_mm_mid_tensor(c->t[22], c->t[18], c->t[20], c->t[21], c->a->base, c->a->size, c->off[8], c->off[9], 8u, 16, 256, 10, 10, 2560, 640, 256) &&
|
|
ds4_gpu_qwen4_moe_mm_down_tensor(c->t[23], c->t[22], c->t[20], c->t[21], c->a->base, c->a->size, c->off[10], 8u, 16, 256, 10, 10, 640, 2560, 256);
|
|
}
|
|
/* Production-shaped routed tiles: Q4_K gate/up + MXFP4 down, 10 slots per
|
|
* token. One dense case (32 experts, ~640 pairs each) and one at the
|
|
* density of an 8192-token chunk (256 experts, ~80 pairs each). */
|
|
static int bench_p_moe_mm_q4k_case(bench_ctx *c, uint32_t NE, uint32_t seed0, ds4_gpu_tensor **st, uint64_t *offs) {
|
|
const uint32_t T = 2048, slots = 10, E = 2560, F = 640, cap = NE >= 256 ? 512 : T;
|
|
if (!st[0]) {
|
|
double *sh;
|
|
offs[0] = arena_q4_K(c->a, (uint64_t)NE * F, E, &sh, 0.05f); free(sh);
|
|
offs[1] = arena_q4_K(c->a, (uint64_t)NE * F, E, &sh, 0.05f); free(sh);
|
|
offs[2] = arena_tier(c->a, 39u, (uint64_t)NE * E, F, &sh); free(sh);
|
|
int32_t *s = malloc((uint64_t)T * slots * sizeof(int32_t));
|
|
uint32_t seed = seed0;
|
|
for (uint64_t i = 0; i < (uint64_t)T * slots; i++) { seed = seed * 1664525u + 1013904223u; s[i] = (int32_t)((seed >> 8) % NE); }
|
|
st[0] = ds4_gpu_tensor_alloc((uint64_t)T * slots * sizeof(int32_t));
|
|
require_ok(st[0] && ds4_gpu_tensor_write(st[0], 0, s, (uint64_t)T * slots * sizeof(int32_t)), "moe q4k bench selection");
|
|
free(s);
|
|
st[1] = ds4_gpu_tensor_alloc((uint64_t)NE * cap * sizeof(int32_t));
|
|
st[2] = ds4_gpu_tensor_alloc(NE * sizeof(int32_t));
|
|
float *xv = rand_vec((uint64_t)T * E, 1.0f);
|
|
st[3] = upload(xv, (uint64_t)T * E); free(xv);
|
|
st[4] = upload(NULL, (uint64_t)T * slots * F);
|
|
st[5] = upload(NULL, (uint64_t)T * slots * E);
|
|
require_ok(st[1] && st[2] && st[3] && st[4] && st[5], "moe q4k bench buffers");
|
|
}
|
|
return ds4_gpu_qwen4_moe_build_lists_tensor(st[1], st[2], st[0], T, slots, NE, cap) &&
|
|
ds4_gpu_qwen4_moe_mm_mid_tensor(st[4], st[3], st[1], st[2], c->a->base, c->a->size, offs[0], offs[1], 12u, NE, T, slots, slots, E, F, cap) &&
|
|
ds4_gpu_qwen4_moe_mm_down_tensor(st[5], st[4], st[1], st[2], c->a->base, c->a->size, offs[2], 39u, NE, T, slots, slots, F, E, cap);
|
|
}
|
|
/* Q2-pack-shaped routed tiles: IQ2XXS gate/up + Q2_K down (768-wide ff), 10
|
|
* slots per token. Dense case (32 experts, ~640 pairs each). */
|
|
static int bench_p_moe_mm_iq2_case(bench_ctx *c, uint32_t NE, uint32_t seed0, ds4_gpu_tensor **st, uint64_t *offs) {
|
|
const uint32_t T = 2048, slots = 10, E = 2560, F = 768, cap = NE >= 256 ? 512 : T;
|
|
if (!st[0]) {
|
|
double *sh;
|
|
offs[0] = arena_tier(c->a, 16u, (uint64_t)NE * F, E, &sh); free(sh);
|
|
offs[1] = arena_tier(c->a, 16u, (uint64_t)NE * F, E, &sh); free(sh);
|
|
offs[2] = arena_tier(c->a, 10u, (uint64_t)NE * E, F, &sh); free(sh);
|
|
int32_t *s = malloc((uint64_t)T * slots * sizeof(int32_t));
|
|
uint32_t seed = seed0;
|
|
for (uint64_t i = 0; i < (uint64_t)T * slots; i++) { seed = seed * 1664525u + 1013904223u; s[i] = (int32_t)((seed >> 8) % NE); }
|
|
st[0] = ds4_gpu_tensor_alloc((uint64_t)T * slots * sizeof(int32_t));
|
|
require_ok(st[0] && ds4_gpu_tensor_write(st[0], 0, s, (uint64_t)T * slots * sizeof(int32_t)), "moe iq2 bench selection");
|
|
free(s);
|
|
st[1] = ds4_gpu_tensor_alloc((uint64_t)NE * cap * sizeof(int32_t));
|
|
st[2] = ds4_gpu_tensor_alloc(NE * sizeof(int32_t));
|
|
float *xv = rand_vec((uint64_t)T * E, 1.0f);
|
|
st[3] = upload(xv, (uint64_t)T * E); free(xv);
|
|
st[4] = upload(NULL, (uint64_t)T * slots * F);
|
|
st[5] = upload(NULL, (uint64_t)T * slots * E);
|
|
require_ok(st[1] && st[2] && st[3] && st[4] && st[5], "moe iq2 bench buffers");
|
|
}
|
|
return ds4_gpu_qwen4_moe_build_lists_tensor(st[1], st[2], st[0], T, slots, NE, cap) &&
|
|
ds4_gpu_qwen4_moe_mm_mid_tensor(st[4], st[3], st[1], st[2], c->a->base, c->a->size, offs[0], offs[1], 16u, NE, T, slots, slots, E, F, cap) &&
|
|
ds4_gpu_qwen4_moe_mm_down_tensor(st[5], st[4], st[1], st[2], c->a->base, c->a->size, offs[2], 10u, NE, T, slots, slots, F, E, cap);
|
|
}
|
|
static int bench_p_moe_mm_iq2(void *ud) { static ds4_gpu_tensor *st[6]; static uint64_t offs[3]; return bench_p_moe_mm_iq2_case(ud, 32, 4242u, st, offs); }
|
|
|
|
static int bench_p_moe_mm_q4k(void *ud) { static ds4_gpu_tensor *st[6]; static uint64_t offs[3]; return bench_p_moe_mm_q4k_case(ud, 32, 12345u, st, offs); }
|
|
static int bench_p_moe_mm_q4k_lo(void *ud) { static ds4_gpu_tensor *st[6]; static uint64_t offs[3]; return bench_p_moe_mm_q4k_case(ud, 256, 777u, st, offs); }
|
|
static int bench_gdn_scan(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_gdn_scan_tensor(c->t[15], c->t[1], c->t[11], c->t[13], c->t[14], 1, 16, 48, 128, NULL, 0u, NULL, 0u); }
|
|
static int bench_gdn_scan2(void *ud) { bench_ctx *c = ud; return ds4_gpu_qwen4_gdn_scan_tensor(c->t[15], c->t[1], c->t[11], c->t[13], c->t[14], 2, 16, 48, 128, NULL, 0u, NULL, 0u); }
|
|
|
|
/* QWEN4_BENCH=1: per-dispatch cost of the decode kernels at full-model shapes */
|
|
static void bench_dispatch(arena_t *a) {
|
|
bench_ctx c = { .a = a };
|
|
double *sh;
|
|
c.off[0] = arena_q8_0(a, 48, 2560, &sh, 0.05f); free(sh);
|
|
c.off[1] = arena_q8_0(a, 6144, 2560, &sh, 0.05f); free(sh);
|
|
c.off[2] = arena_f32(a, 10240, &sh, 0.5f, 1.5f); free(sh);
|
|
c.off[3] = arena_f16(a, 320ull * 10240, &sh, 0.05f); free(sh);
|
|
c.off[4] = arena_f16(a, 4ull * 10240, &sh, 0.05f); free(sh);
|
|
c.off[5] = arena_f16(a, 10240ull * 320, &sh, 0.05f); free(sh);
|
|
c.off[6] = arena_f32(a, 10240ull * 4, &sh, -0.5f, 0.5f); free(sh);
|
|
c.off[7] = arena_f32(a, 48, &sh, -2.0f, -0.1f); free(sh);
|
|
c.t[0] = upload(NULL, 2 * 2560);
|
|
c.t[2] = upload(NULL, 10240); /* R */
|
|
c.t[4] = upload(NULL, 2 * 10240); /* xn */
|
|
c.t[5] = upload(NULL, 2 * 320); /* lo */
|
|
c.t[6] = upload(NULL, 24 * 256); /* attn out */
|
|
c.t[7] = upload(NULL, 24 * 256); /* q / gate */
|
|
c.t[8] = ds4_gpu_tensor_alloc(4096ull * 512 * 2); /* k cache f16 */
|
|
c.t[9] = ds4_gpu_tensor_alloc(4096ull * 512 * 2); /* v cache f16 */
|
|
c.t[10] = upload(NULL, ds4_gpu_qwen4_attn_part_floats(2, 24, 256));
|
|
c.t[12] = upload(NULL, 3 * 10240); /* conv state */
|
|
c.off[8] = arena_q8_0(a, 16ull * 640, 2560, &sh, 0.05f); free(sh); /* 16 experts gate */
|
|
c.off[9] = arena_q8_0(a, 16ull * 640, 2560, &sh, 0.05f); free(sh); /* up */
|
|
c.off[10] = arena_q8_0(a, 16ull * 2560, 640, &sh, 0.05f); free(sh); /* down */
|
|
c.off[11] = arena_f32(a, 512ull * 2560, &sh, -0.05f, 0.05f); free(sh);
|
|
c.off[12] = arena_q4_K(a, 16ull * 640, 2560, &sh, 0.05f); free(sh);
|
|
c.off[13] = arena_q4_K(a, 16ull * 640, 2560, &sh, 0.05f); free(sh);
|
|
c.t[18] = upload(NULL, 256ull * 2560); /* x for 256 tokens */
|
|
c.t[19] = upload(NULL, 256ull * 10240); /* wide scratch */
|
|
c.t[1] = upload(NULL, 256ull * 10240); /* scratch (also the GDN state) */
|
|
c.t[3] = upload(NULL, 256ull * 4 * DS4_QWEN4_HC_CHUNKS * 4); /* inj partials */
|
|
c.t[20] = ds4_gpu_tensor_alloc(16ull * 256 * 4);
|
|
c.t[21] = ds4_gpu_tensor_alloc(16 * 4);
|
|
c.t[22] = upload(NULL, 256ull * 10 * 640);
|
|
c.t[23] = upload(NULL, 256ull * 10 * 2560);
|
|
c.t[17] = upload(NULL, 256ull * 2560);
|
|
c.t[11] = upload(NULL, 1024ull * 10240); /* qkv for 1024 tokens */
|
|
c.t[13] = upload(NULL, 1024ull * 48); /* decay */
|
|
c.t[14] = upload(NULL, 1024ull * 48); /* beta */
|
|
c.t[15] = upload(NULL, 1024ull * 48 * 128); /* scan output */
|
|
c.t[24] = upload(NULL, 32ull * 512); /* indexer q, 32 tokens */
|
|
c.t[25] = ds4_gpu_tensor_alloc(65536ull * 128 * 2); /* block keys f16 */
|
|
c.t[26] = upload(NULL, 32ull * 65536); /* scores */
|
|
c.t[27] = ds4_gpu_tensor_alloc(32ull * 512 * 4); /* selected */
|
|
c.t[28] = upload(NULL, 1024ull * 512); /* indexer q, 1024 tokens */
|
|
c.t[29] = upload(NULL, 1024ull * 65536); /* scores, 1024 tokens */
|
|
c.t[30] = ds4_gpu_tensor_alloc(1024ull * 512 * 4); /* selected, 1024 tokens */
|
|
{ /* 256k K/V caches, 1024 queries, selections: random blocks vs blocks from the last 8k tokens */
|
|
const uint64_t kv_n = 262144ull * 512;
|
|
_Float16 *kv = malloc(kv_n * 2);
|
|
for (uint64_t i = 0; i < kv_n; i++) kv[i] = (_Float16)(frand() - 0.5f);
|
|
c.t[31] = ds4_gpu_tensor_alloc(kv_n * 2);
|
|
require_ok(c.t[31] && ds4_gpu_tensor_write(c.t[31], 0, kv, kv_n * 2), "k cache write");
|
|
for (uint64_t i = 0; i < kv_n; i += 7) kv[i] = (_Float16)(frand() - 0.5f);
|
|
c.t[32] = ds4_gpu_tensor_alloc(kv_n * 2);
|
|
require_ok(c.t[32] && ds4_gpu_tensor_write(c.t[32], 0, kv, kv_n * 2), "v cache write");
|
|
free(kv);
|
|
float *qf = malloc(1024ull * 6144 * 4);
|
|
for (int i = 0; i < 1024 * 6144; i++) qf[i] = frand() - 0.5f;
|
|
c.t[33] = upload(qf, 1024ull * 6144);
|
|
c.t[34] = upload(qf, 1024ull * 6144);
|
|
free(qf);
|
|
int32_t *sel = malloc(1024ull * 2052 * 4);
|
|
uint32_t *cnt = malloc(1024 * 4);
|
|
for (int variant = 0; variant < 2; variant++) {
|
|
for (int t = 0; t < 1024; t++) {
|
|
const uint32_t pos = 262144 - 1024 + t, n_blocks = pos / 4;
|
|
const uint32_t lo = variant ? n_blocks - 2048 : 0, span = n_blocks - lo;
|
|
uint32_t blocks[512];
|
|
for (int i = 0; i < 512; i++) {
|
|
uint32_t b;
|
|
bool dup;
|
|
do {
|
|
b = lo + (uint32_t)((frand() + 1.0f) * 0.5f * span) % span;
|
|
dup = false;
|
|
for (int j = 0; j < i && !dup; j++) dup = blocks[j] == b;
|
|
} while (dup);
|
|
blocks[i] = b;
|
|
}
|
|
for (int i = 1; i < 512; i++) {
|
|
const uint32_t b = blocks[i];
|
|
int j = i - 1;
|
|
while (j >= 0 && blocks[j] > b) { blocks[j + 1] = blocks[j]; j--; }
|
|
blocks[j + 1] = b;
|
|
}
|
|
for (int i = 0; i < 512; i++)
|
|
for (int r = 0; r < 4; r++) sel[(uint64_t)t * 2052 + i * 4 + r] = (int32_t)(blocks[i] * 4 + r);
|
|
cnt[t] = 2048;
|
|
}
|
|
c.t[35 + variant] = ds4_gpu_tensor_alloc(1024ull * 2052 * 4);
|
|
require_ok(c.t[35 + variant] && ds4_gpu_tensor_write(c.t[35 + variant], 0, sel, 1024ull * 2052 * 4), "sel write");
|
|
}
|
|
c.t[38] = ds4_gpu_tensor_alloc(1024 * 4);
|
|
require_ok(c.t[38] && ds4_gpu_tensor_write(c.t[38], 0, cnt, 1024 * 4), "cnt write");
|
|
c.t[37] = upload(NULL, 1024ull * 6144);
|
|
free(sel); free(cnt);
|
|
}
|
|
c.t[39] = upload(NULL, 2048ull * 2560); /* x for 2048 tokens */
|
|
c.t[40] = upload(NULL, 2048ull * 6144);
|
|
c.t[41] = upload(NULL, 2048ull * 320);
|
|
c.t[42] = upload(NULL, 2048ull * 10240); /* hc-wide x for 2048 tokens */
|
|
{ /* random queries and keys: the select bench needs a real score spread */
|
|
float *q = malloc(1024ull * 512 * 4);
|
|
for (int i = 0; i < 1024 * 512; i++) q[i] = frand() - 0.5f;
|
|
require_ok(ds4_gpu_tensor_write(c.t[28], 0, q, 1024ull * 512 * 4), "indexer q write");
|
|
free(q);
|
|
_Float16 *k = malloc(65536ull * 128 * 2);
|
|
for (int i = 0; i < 65536 * 128; i++) k[i] = (_Float16)(frand() - 0.5f);
|
|
require_ok(c.t[25] && ds4_gpu_tensor_write(c.t[25], 0, k, 65536ull * 128 * 2), "block key write");
|
|
free(k);
|
|
}
|
|
{ /* expert lists for the T=256 moe benches; the decode benches use the first 10 */
|
|
int32_t *sel = malloc(256 * 10 * 4);
|
|
for (int i = 0; i < 256 * 10; i++) sel[i] = (i * 7) % 16;
|
|
c.t[16] = ds4_gpu_tensor_alloc(256 * 10 * 4);
|
|
require_ok(ds4_gpu_tensor_write(c.t[16], 0, sel, 256 * 10 * 4), "sel write");
|
|
free(sel);
|
|
}
|
|
require_ok(ds4_gpu_begin_commands(), "begin");
|
|
require_ok(ds4_gpu_end_commands(), "end");
|
|
bench_run("hc_combine (tiny)", bench_combine, &c, 200);
|
|
bench_run("q8 gemv 48x2560", bench_q8_small, &c, 200);
|
|
bench_run("q8 gemv 6144x2560 T=1", bench_q8_big, &c, 200);
|
|
bench_run("q8 gemv 6144x2560 T=2", bench_q8_big2, &c, 200);
|
|
bench_run("hc_norm (+inject partials)", bench_hc_norm, &c, 200);
|
|
bench_run("hc down gemv f16 320x10240", bench_hc_down, &c, 200);
|
|
bench_run("hc_gate_mix f16", bench_hc_mix, &c, 200);
|
|
bench_run("hc_gate_mix f16 T=2", bench_hc_mix2, &c, 200);
|
|
c.n[0] = 110; c.n[1] = 0; bench_run("attn_decode pos=110 no split", bench_attn, &c, 100);
|
|
c.n[0] = 110; c.n[1] = 1; bench_run("attn_decode pos=110 split", bench_attn, &c, 100);
|
|
c.n[0] = 2000; c.n[1] = 0; bench_run("attn_decode pos=2000 no split", bench_attn, &c, 50);
|
|
c.n[0] = 2000; c.n[1] = 1; bench_run("attn_decode pos=2000 split", bench_attn, &c, 50);
|
|
bench_run("gdn_front (fused)", bench_gdn_front, &c, 100);
|
|
bench_run("gdn gemv a/b + conv + prep (prefill path)", bench_gdn_unfused, &c, 100);
|
|
bench_run("gdn_scan", bench_gdn_scan, &c, 100);
|
|
bench_run("gdn_scan T=2", bench_gdn_scan2, &c, 100);
|
|
printf(" --- prefill T=256 ---\n");
|
|
bench_run("router f32 512x2560 T=256 (DS4)", bench_p_router_f32, &c, 20);
|
|
bench_run("router f32 512x2560 T=256 (dense mm)", bench_p_router_mm, &c, 20);
|
|
bench_run("router_topk T=256", bench_p_topk, &c, 20);
|
|
bench_run("hc_norm T=256", bench_p_hc_norm, &c, 20);
|
|
bench_run("hc down f16 320x10240 T=256 (DS4)", bench_p_hc_down_f16, &c, 20);
|
|
bench_run("hc down f16 320x10240 T=256 (dense mm)", bench_p_hc_down_mm, &c, 20);
|
|
bench_run("hc up f16 10240x320 T=256 (DS4)", bench_p_hc_up_f16, &c, 20);
|
|
bench_run("hc up f16 10240x320 T=256 (dense mm)", bench_p_hc_up_mm, &c, 20);
|
|
bench_run("hc_mix_rows T=256", bench_p_hc_mix_rows, &c, 20);
|
|
bench_run("q8 gemm 6144x2560 T=256 (DS4)", bench_p_q8_gemm, &c, 20);
|
|
bench_run("q8 gemm 6144x2560 T=256 (dense mm)", bench_p_q8_mm, &c, 20);
|
|
bench_run("moe mm lists+mid+down 16 experts T=256 (all 2560 pairs)", bench_p_moe_mm, &c, 10);
|
|
bench_run("moe mm q4k/mxfp4 32 experts T=2048 x10 slots", bench_p_moe_mm_q4k, &c, 10);
|
|
bench_run("moe mm iq2xxs/q2k 32 experts T=2048 x10 slots", bench_p_moe_mm_iq2, &c, 10);
|
|
bench_run("moe mm q4k/mxfp4 lo 256 experts T=2048 x10 slots", bench_p_moe_mm_q4k_lo, &c, 10);
|
|
{ /* decays in (0,1], betas in (0,1) for the scan benches */
|
|
float *g = malloc(1024 * 48 * 4), *b = malloc(1024 * 48 * 4);
|
|
for (int i = 0; i < 1024 * 48; i++) { g[i] = 0.9f + 0.1f * frand(); b[i] = 0.5f * frand() + 0.25f; }
|
|
require_ok(ds4_gpu_tensor_write(c.t[13], 0, g, 1024 * 48 * 4) && ds4_gpu_tensor_write(c.t[14], 0, b, 1024 * 48 * 4),
|
|
"scan input write");
|
|
free(g); free(b);
|
|
}
|
|
bench_run("gdn scan r4 T=1024", bench_p_gdn_r4, &c, 5);
|
|
bench_run("idx score n=65536 T=32", bench_p_idx_score, &c, 10);
|
|
bench_run("idx argsort top-512 n=65536 T=32", bench_p_idx_argsort, &c, 5);
|
|
bench_run("idx select top-512 n=65536 T=32", bench_p_idx_select, &c, 10);
|
|
bench_run("q8 gemm 6144x2560 T=2048 (DS4)", bench_p_q8_gemm_2k, &c, 5);
|
|
bench_run("hc down f16 320x10240 T=2048 (DS4)", bench_p_f16_gemm_2k, &c, 5);
|
|
bench_run("idx score n=65536 T=1024", bench_p_idx_score_1k, &c, 5);
|
|
bench_run("idx select top-512 n=65536 T=1024", bench_p_idx_select_1k, &c, 5);
|
|
setenv("DS4_QWEN4_NO_ATTN_MM", "1", 1);
|
|
c.n[2] = 0; bench_run("attn sparse T=1024 ctx=256k random (per-token kernel)", bench_p_attn_sparse, &c, 5);
|
|
unsetenv("DS4_QWEN4_NO_ATTN_MM");
|
|
c.n[2] = 0; bench_run("attn sparse T=1024 ctx=256k random blocks", bench_p_attn_sparse, &c, 5);
|
|
c.n[2] = 1; bench_run("attn sparse T=1024 ctx=256k blocks in last 8k", bench_p_attn_sparse, &c, 5);
|
|
bench_run("router gemv f32 512x2560", bench_router_gemv, &c, 100);
|
|
bench_run("router_topk (+gate logit)", bench_router_topk, &c, 100);
|
|
bench_run("router_topk no gate", bench_router_topk_nogate, &c, 100);
|
|
bench_run("router_topk no gate k=1", bench_router_topk_k1, &c, 100);
|
|
bench_run("moe_mid q8 10+1 slots (37 MB)", bench_moe_mid, &c, 100);
|
|
bench_run("moe_down q8 10+1 slots (19 MB)", bench_moe_down, &c, 100);
|
|
c.n[0] = 1;
|
|
bench_run("moe_mid q4_K 10+1 slots T=1", bench_moe_mid_q4k, &c, 100);
|
|
c.n[0] = 2;
|
|
bench_run("moe_mid q4_K 10+1 slots T=2", bench_moe_mid_q4k, &c, 100);
|
|
}
|
|
|
|
/* four projections of one input, mixed weight types (q8_0, bf16, q4_0, f16) */
|
|
static void test_multi_gemv(arena_t *a, uint32_t E, uint32_t T) {
|
|
const uint32_t rows[4] = { 64, 48, 32, 16 };
|
|
const uint32_t types[4] = { E % 256 == 0 ? 12u : 8u, 30u, 2u, 1u };
|
|
double *sh[4];
|
|
uint64_t offs[4];
|
|
offs[0] = E % 256 == 0 ? arena_q4_K(a, rows[0], E, &sh[0], 0.05f) : arena_q8_0(a, rows[0], E, &sh[0], 0.05f);
|
|
offs[1] = arena_bf16(a, (uint64_t)rows[1] * E, &sh[1], 0.05f);
|
|
offs[2] = arena_q4_0(a, rows[2], E, &sh[2], 0.05f);
|
|
offs[3] = arena_f16(a, (uint64_t)rows[3] * E, &sh[3], 0.05f);
|
|
float *x = rand_vec((uint64_t)T * E, 1.0f);
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * E);
|
|
ds4_gpu_tensor *outs[4];
|
|
for (int i = 0; i < 4; i++) outs[i] = upload(NULL, (uint64_t)T * rows[i]);
|
|
require_ok(ds4_gpu_qwen4_multi_gemv_tensor(gx, T, E, 4, outs, a->base, a->size, offs, types, rows), "multi gemv");
|
|
for (int i = 0; i < 4; i++) {
|
|
double *ref = malloc((uint64_t)T * rows[i] * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++)
|
|
for (uint32_t r = 0; r < rows[i]; r++) {
|
|
double acc = 0.0;
|
|
for (uint32_t k = 0; k < E; k++) acc += sh[i][(uint64_t)r * E + k] * x[(uint64_t)t * E + k];
|
|
ref[(uint64_t)t * rows[i] + r] = acc;
|
|
}
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "multi gemv out%d type %u rows %u T=%u", i, types[i], rows[i], T);
|
|
check_tensor(name, outs[i], ref, (uint64_t)T * rows[i], 3e-5);
|
|
free(ref); free(sh[i]); ds4_gpu_tensor_free(outs[i]);
|
|
}
|
|
ds4_gpu_tensor_free(gx); free(x);
|
|
}
|
|
|
|
#ifndef __APPLE__
|
|
/* Check half-operand expert tiles independently at every CUDA tile width.
|
|
* The existing tests above still bound their error versus unrounded weights. */
|
|
static void test_half_expert_tiles(arena_t *a, uint32_t T, uint32_t type, uint32_t dtype, uint32_t F) {
|
|
const uint32_t E = 256, NE = 4, NS = 2, NO = 3, cap = T+7, guard = 16;
|
|
const uint32_t DF = dtype == 10 ? (F+255u)/256u*256u : F;
|
|
double *gw, *uw, *dw;
|
|
uint64_t go = arena_tier(a,type,(uint64_t)NE*F,E,&gw);
|
|
uint64_t uo = arena_tier(a,type,(uint64_t)NE*F,E,&uw);
|
|
uint64_t d = arena_tier(a,dtype,(uint64_t)NE*E,DF,&dw);
|
|
float *x = rand_vec((uint64_t)T*E,1.0f);
|
|
int32_t *sel = malloc((uint64_t)T*NS*4);
|
|
require_ok(sel != NULL,"half tile selections allocation");
|
|
for (uint32_t t = 0; t < T; t++) { sel[t*NS] = 0; sel[t*NS+1] = 1+t%2; }
|
|
ds4_gpu_tensor *gx = upload(x,(uint64_t)T*E);
|
|
ds4_gpu_tensor *gs = ds4_gpu_tensor_alloc((uint64_t)T*NS*4);
|
|
ds4_gpu_tensor *gl = ds4_gpu_tensor_alloc((uint64_t)NE*cap*4);
|
|
ds4_gpu_tensor *gc = ds4_gpu_tensor_alloc(NE*4);
|
|
uint64_t nm = (uint64_t)T*NO*F, np = (uint64_t)T*NO*E;
|
|
ds4_gpu_tensor *gm = upload(NULL,nm+guard), *gp = upload(NULL,np+guard);
|
|
require_ok(gs && gl && gc && ds4_gpu_tensor_write(gs,0,sel,(uint64_t)T*NS*4),"half tile selection upload");
|
|
const float sentinel = -1234.5f;
|
|
require_ok(ds4_gpu_tensor_fill_f32(gm,sentinel,nm+guard) &&
|
|
ds4_gpu_tensor_fill_f32(gp,sentinel,np+guard),"half tile sentinels");
|
|
require_ok(ds4_gpu_qwen4_moe_build_lists_tensor(gl,gc,gs,T,NS,NE,cap),"half tile list build");
|
|
int32_t counts[NE];
|
|
require_ok(ds4_gpu_tensor_read(gc,0,counts,sizeof(counts)),"half tile counts read");
|
|
require_ok(counts[0] == (int32_t)T && counts[1]+counts[2] == (int32_t)T && counts[3] == 0,
|
|
"half tile hot and empty experts");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_mid_tensor(gm,gx,gl,gc,a->base,a->size,go,uo,type,NE,T,NS,NO,E,F,cap),"half tile mid");
|
|
require_ok(ds4_gpu_qwen4_moe_mm_down_tensor(gp,gm,gl,gc,a->base,a->size,d,dtype,NE,T,NS,NO,F,E,cap),"half tile down");
|
|
float *mid = download(gm,nm+guard), *part = download(gp,np+guard);
|
|
for (uint64_t i = nm; i < nm+guard; i++) require_ok(mid[i] == sentinel,"half tile mid tail");
|
|
for (uint64_t i = np; i < np+guard; i++) require_ok(part[i] == sentinel,"half tile down tail");
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
for (uint32_t f = 0; f < F; f++) require_ok(mid[((uint64_t)t*NO+NS)*F+f] == sentinel,"half tile reserved mid slot");
|
|
for (uint32_t e = 0; e < E; e++) require_ok(part[((uint64_t)t*NO+NS)*E+e] == sentinel,"half tile reserved down slot");
|
|
}
|
|
const uint32_t samples[] = {0,1,31,32,63,64,127,128,1023,T-1};
|
|
float got[sizeof(samples)/sizeof(samples[0])*NS*(E+F)];
|
|
double ref[sizeof(got)/sizeof(got[0])];
|
|
uint32_t n = 0;
|
|
for (uint32_t j = 0; j < sizeof(samples)/sizeof(samples[0]); j++) {
|
|
uint32_t t = samples[j];
|
|
if (t >= T) continue;
|
|
for (uint32_t s = 0; s < NS; s++) {
|
|
uint32_t e = sel[t*NS+s];
|
|
for (uint32_t f = 0; f < F; f++) {
|
|
double g = 0, u = 0;
|
|
for (uint32_t k = 0; k < E; k++) {
|
|
double v = (_Float16)x[(uint64_t)t*E+k];
|
|
g += (double)(_Float16)(float)gw[((uint64_t)e*F+f)*E+k]*v;
|
|
u += (double)(_Float16)(float)uw[((uint64_t)e*F+f)*E+k]*v;
|
|
}
|
|
got[n] = mid[((uint64_t)t*NO+s)*F+f]; ref[n++] = silu_d(g)*u;
|
|
}
|
|
for (uint32_t r = 0; r < E; r++) {
|
|
double v = 0;
|
|
for (uint32_t k = 0; k < F; k++)
|
|
v += (double)(_Float16)(float)dw[((uint64_t)e*E+r)*DF+k]*
|
|
(double)(_Float16)mid[((uint64_t)t*NO+s)*F+k];
|
|
got[n] = part[((uint64_t)t*NO+s)*E+r]; ref[n++] = v;
|
|
}
|
|
}
|
|
}
|
|
char name[96];
|
|
snprintf(name,sizeof(name),"half expert reference type=%u/%u T=%u F=%u",type,dtype,T,F);
|
|
check_close(name,got,ref,n,3e-5);
|
|
free(mid); free(part); free(sel); free(x); free(gw); free(uw); free(dw);
|
|
ds4_gpu_tensor_free(gx); ds4_gpu_tensor_free(gs); ds4_gpu_tensor_free(gl);
|
|
ds4_gpu_tensor_free(gc); ds4_gpu_tensor_free(gm); ds4_gpu_tensor_free(gp);
|
|
}
|
|
#endif
|
|
|
|
#ifdef __APPLE__
|
|
/* The decode-batch Q8 GEMM: reference in double, and the same sums the
|
|
* per-token matvec finds, to rounding. */
|
|
static void test_batch_mm_q8(arena_t *a, uint32_t in_dim, uint32_t rows, uint32_t T) {
|
|
double *sh;
|
|
const uint64_t off = arena_q8_0(a, rows, in_dim, &sh, 0.05f);
|
|
float *x = rand_vec((uint64_t)T * in_dim, 1.0f);
|
|
double *ref = malloc((uint64_t)T * rows * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++)
|
|
for (uint32_t r = 0; r < rows; r++) {
|
|
double acc = 0.0;
|
|
for (uint32_t k = 0; k < in_dim; k++) acc += sh[(uint64_t)r * in_dim + k] * x[(uint64_t)t * in_dim + k];
|
|
ref[(uint64_t)t * rows + r] = acc;
|
|
}
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * in_dim);
|
|
ds4_gpu_tensor *gout = upload(NULL, (uint64_t)T * rows), *gmv = upload(NULL, (uint64_t)T * rows);
|
|
require_ok(ds4_gpu_qwen4_batch_mm_q8_tensor(gout, gx, a->base, a->size, off, T, in_dim, rows), "batch mm q8");
|
|
require_ok(ds4_gpu_qwen4_matmul_q8_0_tensor(gmv, a->base, a->size, off, in_dim, rows, gx, T), "batch mm q8 matvec");
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "batch mm q8 %ux%u T=%u", rows, in_dim, T);
|
|
check_tensor(name, gout, ref, (uint64_t)T * rows, 3e-5);
|
|
float *am = download(gout, (uint64_t)T * rows), *bm = download(gmv, (uint64_t)T * rows);
|
|
double worst = 0.0, scale = 0.0;
|
|
for (uint64_t i = 0; i < (uint64_t)T * rows; i++) {
|
|
const double d = fabs((double)am[i] - (double)bm[i]);
|
|
if (d > worst) worst = d;
|
|
if (fabs((double)bm[i]) > scale) scale = fabs((double)bm[i]);
|
|
}
|
|
require_ok(worst <= 2e-6 * scale + 1e-6, "batch mm q8 within rounding of the matvec");
|
|
free(bm); free(am); free(ref); free(x); free(sh);
|
|
ds4_gpu_tensor_free(gmv); ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gx);
|
|
}
|
|
|
|
#endif
|
|
|
|
/* dense tiled GEMM against a double reference for f32, f16 and q8_0 rows */
|
|
static void test_dense_mm(arena_t *a, uint32_t in_dim, uint32_t rows, uint32_t T, uint32_t wtype) {
|
|
double *sh;
|
|
uint64_t off = wtype == 8u ? arena_q8_0(a, rows, in_dim, &sh, 0.05f)
|
|
: wtype == 1u ? arena_f16(a, (uint64_t)rows * in_dim, &sh, 0.05f)
|
|
: arena_f32(a, (uint64_t)rows * in_dim, &sh, -0.05f, 0.05f);
|
|
float *x = rand_vec((uint64_t)T * in_dim, 1.0f);
|
|
double *ref = malloc((uint64_t)T * rows * sizeof(double));
|
|
for (uint32_t t = 0; t < T; t++)
|
|
for (uint32_t r = 0; r < rows; r++) {
|
|
double acc = 0.0;
|
|
for (uint32_t k = 0; k < in_dim; k++) acc += sh[(uint64_t)r * in_dim + k] * x[(uint64_t)t * in_dim + k];
|
|
ref[(uint64_t)t * rows + r] = acc;
|
|
}
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T * in_dim);
|
|
ds4_gpu_tensor *gout = upload(NULL, (uint64_t)T * rows);
|
|
require_ok(ds4_gpu_qwen4_dense_mm_tensor(gout, gx, a->base, a->size, off, wtype, T, in_dim, rows), "dense mm");
|
|
char name[96];
|
|
snprintf(name, sizeof(name), "dense mm type %u %ux%u T=%u", wtype, rows, in_dim, T);
|
|
check_tensor(name, gout, ref, (uint64_t)T * rows, 3e-5);
|
|
free(ref); free(x); free(sh);
|
|
ds4_gpu_tensor_free(gout); ds4_gpu_tensor_free(gx);
|
|
}
|
|
|
|
#ifndef __APPLE__
|
|
/* Cross the former 64 MiB output-tile limit with odd output strides. Reuse
|
|
* 17 independent input rows so the full CPU oracle stays inexpensive. */
|
|
static void test_dense_mm_large(arena_t *a, uint32_t wtype) {
|
|
const uint32_t K = 64, M = 2051, T = 8201, patterns = 17, guard = 64;
|
|
const uint64_t count = (uint64_t)T*M;
|
|
double *weights;
|
|
const uint64_t off = wtype == 8
|
|
? arena_q8_0(a, M, K, &weights, 0.05f)
|
|
: arena_f16(a, (uint64_t)M*K, &weights, 0.05f);
|
|
float *rows = rand_vec((uint64_t)patterns*K, 1.0f);
|
|
float *x = malloc((uint64_t)T*K*sizeof(*x));
|
|
double *dots = calloc((uint64_t)patterns*M, sizeof(*dots));
|
|
double *ref = malloc((count+guard)*sizeof(*ref));
|
|
require_ok(x && dots && ref && rows, "large dense allocations");
|
|
for (uint32_t p = 0; p < patterns; p++)
|
|
for (uint32_t m = 0; m < M; m++)
|
|
for (uint32_t k = 0; k < K; k++)
|
|
dots[(uint64_t)p*M+m] += weights[(uint64_t)m*K+k]*rows[p*K+k];
|
|
for (uint32_t t = 0; t < T; t++) {
|
|
memcpy(x+(uint64_t)t*K, rows+(t%patterns)*K, K*sizeof(*x));
|
|
memcpy(ref+(uint64_t)t*M, dots+(t%patterns)*M, M*sizeof(*ref));
|
|
}
|
|
for (uint32_t i = 0; i < guard; i++) ref[count+i] = 17.25;
|
|
ds4_gpu_tensor *gx = upload(x, (uint64_t)T*K);
|
|
ds4_gpu_tensor *out = upload(NULL, count+guard);
|
|
for (unsigned repeat = 0; repeat < 2; repeat++) {
|
|
require_ok(ds4_gpu_tensor_fill_f32(out, 17.25f, count+guard), "large dense guard fill");
|
|
require_ok(ds4_gpu_qwen4_dense_mm_tensor(out, gx, a->base, a->size,
|
|
off, wtype, T, K, M), "large dense projection");
|
|
check_tensor(wtype == 8 ? "large Q8 dense output" : "large F16 dense output",
|
|
out, ref, count, 3e-5);
|
|
float got_guard[64];
|
|
require_ok(ds4_gpu_tensor_read(out, count*4u, got_guard, sizeof(got_guard)), "large dense guard read");
|
|
for (unsigned i = 0; i < guard; i++) require_ok(got_guard[i] == 17.25f, "large dense output guard");
|
|
}
|
|
ds4_gpu_tensor_free(out); ds4_gpu_tensor_free(gx);
|
|
free(ref); free(dots); free(x); free(rows); free(weights);
|
|
}
|
|
#endif
|
|
|
|
int main(void) {
|
|
arena_t arena;
|
|
arena.size = (uint64_t)1536 << 20;
|
|
arena.base = mmap(NULL, arena.size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
|
|
arena.used = 0;
|
|
if (arena.base == MAP_FAILED) { perror("mmap"); return 1; }
|
|
setenv("DS4_QWEN4_ATTN_SPLIT_KEYS", "8", 1);
|
|
require_ok(ds4_gpu_init(), "GPU initialization");
|
|
require_ok(ds4_gpu_set_model_map(arena.base, arena.size), "model map registration");
|
|
|
|
#ifndef __APPLE__
|
|
if (getenv("DS4_TEST_QWEN4_ATTN_GROUPS")) { test_attn_groups(); return 0; }
|
|
if (getenv("DS4_TEST_QWEN4_DENSE_ONLY")) {
|
|
test_dense_mm_large(&arena, 1u);
|
|
test_dense_mm_large(&arena, 8u);
|
|
test_dense_mm(&arena, 10240, 1700, 32, 8u);
|
|
test_dense_mm(&arena, 320, 10240, 40, 1u);
|
|
test_dense_mm(&arena, 67, 97, 35, 1u);
|
|
test_dense_mm(&arena, 96, 129, 35, 8u);
|
|
return 0;
|
|
}
|
|
if (getenv("DS4_TEST_QWEN4_EXPERT_TILES")) {
|
|
test_half_expert_tiles(&arena,33,16,10,64);
|
|
test_half_expert_tiles(&arena,33,12,39,64);
|
|
test_half_expert_tiles(&arena,2049,16,10,192);
|
|
test_half_expert_tiles(&arena,2049,12,39,192);
|
|
return 0;
|
|
}
|
|
#endif
|
|
if (getenv("DS4_TEST_QWEN4_DECODE_FUSIONS")) { test_decode_fusions(&arena); return 0; }
|
|
if (getenv("DS4_TEST_QWEN4_MV_EXACT")) {
|
|
test_moe_types(&arena, 8, 6, 2560, 640, 1, 16u, 10u);
|
|
test_moe_types(&arena, 8, 6, 2560, 640, 2, 16u, 10u);
|
|
test_moe_types(&arena, 8, 6, 256, 256, 9, 16u, 10u);
|
|
test_moe_types(&arena, 8, 6, 256, 672, 3, 16u, 10u);
|
|
test_moe_types(&arena, 8, 6, 2560, 640, 1, 12u, 39u);
|
|
test_moe_types(&arena, 8, 6, 2560, 640, 2, 12u, 39u);
|
|
test_moe_types(&arena, 8, 6, 256, 256, 9, 12u, 39u);
|
|
test_moe_types(&arena, 8, 6, 256, 672, 3, 12u, 39u);
|
|
printf("all Qwen MoE decode specialization tests passed\n");
|
|
return 0;
|
|
}
|
|
if (getenv("DS4_TEST_QWEN4_IDX_PREFILTER_ONLY")) { test_idx_prefilter(); printf("all qwen4 indexer prefilter tests passed\n"); return 0; }
|
|
const char *q4k_ordered_only = getenv("DS4_TEST_QWEN4_Q4K_ORDERED_ONLY");
|
|
if (q4k_ordered_only && q4k_ordered_only[0] && strcmp(q4k_ordered_only, "0") != 0) {
|
|
test_q4k_ordered_exact(&arena, 1, 640, true);
|
|
test_q4k_ordered_exact(&arena, 2, 640, false);
|
|
test_q4k_ordered_exact(&arena, 1, 641, false);
|
|
test_q4k_ordered_exact(&arena, 2, 641, true);
|
|
printf("all qwen4 ordered Q4K geometry tests passed\n");
|
|
return 0;
|
|
}
|
|
const char *hc_norm_only = getenv("DS4_TEST_QWEN4_HC_NORM_REUSE_ONLY");
|
|
if (hc_norm_only && hc_norm_only[0] && strcmp(hc_norm_only, "0") != 0) {
|
|
test_hc_norm_reuse(&arena);
|
|
printf("all qwen4 HC norm reuse tests passed\n");
|
|
return 0;
|
|
}
|
|
const char *hc_norm_bench = getenv("DS4_TEST_QWEN4_HC_NORM_REUSE_BENCH");
|
|
if (hc_norm_bench && hc_norm_bench[0] && strcmp(hc_norm_bench, "0") != 0) {
|
|
bench_hc_norm_reuse(&arena);
|
|
return 0;
|
|
}
|
|
if (getenv("QWEN4_BENCH")) {
|
|
bench_dispatch(&arena);
|
|
return 0;
|
|
}
|
|
printf("hyper-connections\n");
|
|
test_decode_fusions(&arena);
|
|
test_qwen4_argmax();
|
|
test_hc_pair_groups(&arena);
|
|
test_mv_ext_groups(&arena);
|
|
#ifdef __APPLE__
|
|
test_moe_grouped(&arena);
|
|
#endif
|
|
test_hc_mix_prefetch(&arena);
|
|
test_hc(&arena, 2560, 320, 3, 1u);
|
|
test_hc(&arena, 2560, 320, 2, 1u);
|
|
test_hc(&arena, 2560, 320, 2, 0u);
|
|
test_hc(&arena, 2560, 320, 1, 0u);
|
|
test_hc(&arena, 2560, 320, 2, 8u);
|
|
test_hc(&arena, 2560, 320, 1, 8u);
|
|
test_hc(&arena, 64, 8, 5, 1u);
|
|
test_hc(&arena, 64, 8, 2, 0u);
|
|
test_hc(&arena, 64, 8, 1, 8u);
|
|
test_hc(&arena, 64, 8, 3, 8u);
|
|
printf("gated delta net\n");
|
|
test_gdn(&arena, 16, 48, 128, 5);
|
|
test_gdn(&arena, 16, 48, 128, 40);
|
|
test_gdn(&arena, 16, 48, 128, 200);
|
|
test_idx_score_mm(37, 3001, 11000);
|
|
test_idx_score_mm(3, 70, 100);
|
|
test_idx_select(4, 70001, 512, 69000);
|
|
test_idx_select(3, 600, 512, 599);
|
|
test_idx_select(2, 3000, 512, 520);
|
|
test_attn_mm(40, 0, false);
|
|
test_attn_mm(37, 3000, true);
|
|
test_attn_mm(3, 100, true);
|
|
test_attn_mm(4, 128, false);
|
|
test_attn_mm(8, 128, false);
|
|
test_attn_mm(9, 128, false);
|
|
#ifndef __APPLE__
|
|
test_attn_groups();
|
|
#endif
|
|
test_gdn(&arena, 2, 6, 32, 7);
|
|
test_gdn(&arena, 2, 6, 64, 9);
|
|
test_gdn(&arena, 2, 6, 96, 17);
|
|
printf("ple\n");
|
|
test_ple(&arena, 2560, 3);
|
|
test_ple(&arena, 64, 12);
|
|
printf("router\n");
|
|
test_router(&arena, 512, 10, 3);
|
|
test_router(&arena, 32, 10, 5);
|
|
printf("attention\n");
|
|
test_attention(&arena, 24, 2, 256, 64, 4, 128, 2, 21);
|
|
test_attention(&arena, 4, 2, 32, 8, 4, 32, 2, 30);
|
|
#ifdef __APPLE__
|
|
test_attention_rows(&arena);
|
|
#endif
|
|
printf("routed experts\n");
|
|
test_moe(&arena, 16, 10, 2560, 640, 2, 8u);
|
|
test_moe(&arena, 16, 10, 2560, 640, 1, 12u);
|
|
test_moe(&arena, 16, 10, 2560, 640, 2, 12u);
|
|
test_moe_types(&arena, 16, 10, 2560, 640, 2, 12u, 39u);
|
|
test_moe_types(&arena, 16, 10, 2560, 640, 1, 16u, 10u);
|
|
test_moe_types(&arena, 16, 10, 2560, 640, 37, 16u, 10u);
|
|
test_moe_types(&arena, 8, 6, 256, 256, 9, 16u, 10u);
|
|
test_moe(&arena, 16, 10, 2560, 640, 37, 12u);
|
|
test_moe(&arena, 16, 10, 2560, 640, 100, 12u);
|
|
test_moe(&arena, 16, 10, 2560, 640, 37, 10u);
|
|
test_moe(&arena, 16, 10, 2560, 640, 37, 16u);
|
|
test_moe(&arena, 8, 10, 2560, 640, 1, 0u);
|
|
test_moe(&arena, 32, 10, 64, 32, 3, 8u);
|
|
test_moe(&arena, 32, 10, 64, 32, 3, 0u);
|
|
test_q4k_ordered_exact(&arena, 1, 640, true);
|
|
test_q4k_ordered_exact(&arena, 2, 640, false);
|
|
test_q4k_ordered_exact(&arena, 1, 641, false);
|
|
test_q4k_ordered_exact(&arena, 2, 641, true);
|
|
test_moe_mm_tiles_exact(&arena, 8u);
|
|
test_moe_mm_tiles_exact(&arena, 39u);
|
|
test_moe_mm_tiles_iq2(&arena);
|
|
printf("dense mm\n");
|
|
test_dense_mm(&arena, 2560, 512, 37, 0u);
|
|
#ifdef __APPLE__
|
|
test_batch_mm_q8(&arena, 2560, 640, 16);
|
|
test_batch_mm_q8(&arena, 6144, 2560, 16);
|
|
test_batch_mm_q8(&arena, 2560, 128, 8);
|
|
#endif
|
|
test_dense_mm(&arena, 10240, 320, 33, 1u);
|
|
test_dense_mm(&arena, 320, 10240, 40, 1u);
|
|
test_dense_mm(&arena, 2560, 100, 9, 8u);
|
|
#ifndef __APPLE__
|
|
test_half_expert_tiles(&arena,33,16,10,64);
|
|
test_half_expert_tiles(&arena,2049,16,10,64);
|
|
test_half_expert_tiles(&arena,8193,16,10,64);
|
|
test_half_expert_tiles(&arena,33,12,39,64);
|
|
test_half_expert_tiles(&arena,2049,12,39,64);
|
|
test_half_expert_tiles(&arena,8193,12,39,64);
|
|
test_half_expert_tiles(&arena,2049,16,10,192);
|
|
test_half_expert_tiles(&arena,2049,12,39,192);
|
|
test_dense_mm(&arena, 2560, 100, 37, 8u);
|
|
test_dense_mm(&arena, 10240, 1700, 32, 8u);
|
|
test_dense_mm_large(&arena, 1u);
|
|
test_dense_mm_large(&arena, 8u);
|
|
test_dense_mm(&arena, 67, 97, 35, 1u);
|
|
test_dense_mm(&arena, 96, 129, 35, 8u);
|
|
test_dense_mm(&arena, 32, 7, 1, 8u);
|
|
test_dense_mm(&arena, 96, 9, 1, 8u);
|
|
test_dense_mm(&arena, 68, 9, 1, 1u);
|
|
test_dense_mm(&arena, 67, 7, 1, 1u);
|
|
for (uint32_t T = 2; T <= 8; T++) {
|
|
test_dense_mm(&arena, 96, 9, T, 8u);
|
|
test_dense_mm(&arena, 68, 9, T, 1u);
|
|
test_dense_mm(&arena, 67, 7, T, 1u);
|
|
}
|
|
#endif
|
|
test_dense_mm(&arena, 64, 32, 70, 0u);
|
|
printf("multi gemv\n");
|
|
test_multi_gemv(&arena, 2560, 2);
|
|
test_multi_gemv(&arena, 64, 3);
|
|
printf("mtp\n");
|
|
test_mtp(&arena, 2560, 4);
|
|
test_mtp(&arena, 64, 4);
|
|
test_hc_norm_reuse(&arena);
|
|
test_gdn_prefill_dispatch();
|
|
printf("all qwen4 kernel tests passed\n");
|
|
return 0;
|
|
}
|