github-featured/antirez--ds4
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
antirez--ds4/tests/ds4_test.c
T
Salvatore Sanfilippo 0ad494e514 metal: make exact fast paths automatic and trim experiments
2026-08-09 17:48:01 +02:00

6543 lines
265 KiB
C
Raw Permalink Blame History

#define DS4_SERVER_TEST
#define DS4_SERVER_TEST_NO_MAIN
#include "../ds4_server.c"
#ifndef DS4_NO_GPU
#include "../ds4_gpu.h"
#include <math.h>
bool ds4_test_dspark_cache_window_crop(void);
static ds4_engine *test_engine_fast;
static ds4_engine *test_engine_quality;
static const char *test_model_path(void) {
const char *model_path = getenv("DS4_TEST_MODEL");
return (model_path && model_path[0]) ? model_path : "ds4flash.gguf";
}
static bool test_env_bool(const char *name) {
const char *v = getenv(name);
return v && v[0] && strcmp(v, "0") != 0;
}
static uint32_t test_env_u32(const char *name) {
const char *v = getenv(name);
if (!v || !v[0]) return 0;
char *end = NULL;
unsigned long n = strtoul(v, &end, 10);
if (end == v) return 0;
return n > UINT32_MAX ? UINT32_MAX : (uint32_t)n;
}
static uint64_t test_env_gib(const char *name) {
const char *v = getenv(name);
if (!v || !v[0]) return 0;
char *end = NULL;
unsigned long long n = strtoull(v, &end, 10);
if (end == v || n == 0) return 0;
const uint64_t one_gib = 1024ull * 1024ull * 1024ull;
if (n > UINT64_MAX / one_gib) return UINT64_MAX;
return (uint64_t)n * one_gib;
}
static char *test_save_env(const char *name) {
const char *value = getenv(name);
if (!value) return NULL;
size_t len = strlen(value);
char *copy = malloc(len + 1);
TEST_ASSERT(copy != NULL);
if (!copy) return NULL;
memcpy(copy, value, len + 1);
return copy;
}
static void test_restore_env(const char *name, char *saved) {
if (saved) {
setenv(name, saved, 1);
free(saved);
} else {
unsetenv(name);
}
}
typedef struct {
char *cold_decode;
char *batch_selected_addr;
} test_streaming_prefill_env;
static test_streaming_prefill_env test_force_canonical_streaming_prefill(void) {
test_streaming_prefill_env saved = {
.cold_decode =
test_save_env("DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL"),
.batch_selected_addr =
test_save_env("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR"),
};
if (test_env_bool("DS4_TEST_SSD_STREAMING")) {
setenv("DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL", "1", 1);
setenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR", "1", 1);
}
return saved;
}
static void test_restore_canonical_streaming_prefill(
test_streaming_prefill_env saved) {
test_restore_env("DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL",
saved.cold_decode);
test_restore_env("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR",
saved.batch_selected_addr);
}
static ds4_backend test_model_backend(void) {
const char *backend = getenv("DS4_TEST_BACKEND");
if (backend && !strcmp(backend, "cpu")) return DS4_BACKEND_CPU;
#ifdef __APPLE__
return DS4_BACKEND_METAL;
#else
return DS4_BACKEND_CUDA;
#endif
}
static ds4_engine *test_open_engine(bool quality) {
ds4_engine *engine = NULL;
/* DS4_TEST_MTP loads the MTP head on the fast engine so the speculative
* verify regression can reuse it; draft=4 hits the multi-row verify path. */
const char *mtp = getenv("DS4_TEST_MTP");
ds4_engine_options opt = {
.model_path = test_model_path(),
.backend = test_model_backend(),
.quality = quality,
.ssd_streaming = test_env_bool("DS4_TEST_SSD_STREAMING"),
.ssd_streaming_cold = test_env_bool("DS4_TEST_SSD_STREAMING_COLD"),
.ssd_streaming_cache_experts =
test_env_u32("DS4_TEST_SSD_STREAMING_CACHE_EXPERTS"),
.ssd_streaming_cache_bytes =
test_env_gib("DS4_TEST_SSD_STREAMING_CACHE_GB"),
.ssd_streaming_preload_experts =
test_env_u32("DS4_TEST_SSD_STREAMING_PRELOAD_EXPERTS"),
.mtp_path = (mtp && mtp[0] && !quality) ? mtp : NULL,
.mtp_draft_tokens = (mtp && mtp[0] && !quality) ? 4 : 0,
};
TEST_ASSERT(ds4_engine_open(&engine, &opt) == 0);
return engine;
}
static ds4_engine *test_get_engine(bool quality) {
ds4_engine **slot = quality ? &test_engine_quality : &test_engine_fast;
if (*slot) return *slot;
*slot = test_open_engine(quality);
return *slot;
}
static void test_close_engines(void) {
ds4_engine_close(test_engine_fast);
ds4_engine_close(test_engine_quality);
test_engine_fast = NULL;
test_engine_quality = NULL;
}
static void test_close_engine(bool quality) {
ds4_engine **slot = quality ? &test_engine_quality : &test_engine_fast;
ds4_engine_close(*slot);
*slot = NULL;
}
static uint64_t test_round_up_u64(uint64_t n, uint64_t align) {
return (n + align - 1) & ~(align - 1);
}
static uint32_t test_float_ordered_u32(float value) {
uint32_t bits;
memcpy(&bits, &value, sizeof(bits));
return (bits & 0x80000000u) != 0 ? ~bits : bits | 0x80000000u;
}
typedef struct {
size_t mismatch_count;
uint32_t max_ulp;
float max_abs;
} test_float_compare_stats;
static test_float_compare_stats test_compare_float_bits(
const float *reference,
const float *actual,
size_t count) {
test_float_compare_stats stats = {0};
for (size_t i = 0; i < count; i++) {
if (memcmp(&reference[i], &actual[i], sizeof(float)) != 0) {
stats.mismatch_count++;
}
const uint32_t ref_ordered = test_float_ordered_u32(reference[i]);
const uint32_t actual_ordered = test_float_ordered_u32(actual[i]);
const uint32_t ulp = ref_ordered > actual_ordered
? ref_ordered - actual_ordered
: actual_ordered - ref_ordered;
if (ulp > stats.max_ulp) stats.max_ulp = ulp;
const float abs_error = fabsf(reference[i] - actual[i]);
if (abs_error > stats.max_abs) stats.max_abs = abs_error;
}
return stats;
}
#if defined(__APPLE__)
static const uint32_t test_copy_f32_patterns[] = {
0x00000000u, 0x80000000u, /* signed zero */
0x7f800000u, 0xff800000u, /* infinities */
0x7fc00000u, 0x7fc12345u, 0xffc54321u, 0x7fa00001u, /* NaNs */
0x00000001u, 0x007fffffu, 0x80000001u, 0x807fffffu, /* F32 subnormals */
0x32ffffffu, 0x33000000u, 0x33000001u,
0x337fffffu, 0x33800000u, 0x33800001u, /* minimum F16 subnormal boundary */
0x387fbfffu, 0x387fc000u, 0x387fffffu,
0x38800000u, 0x38800001u, /* maximum subnormal/minimum normal boundary */
0x3f800fffu, 0x3f801000u, 0x3f801001u,
0x3f802fffu, 0x3f803000u, 0x3f803001u, /* round-to-nearest ties */
0x477fdfffu, 0x477fe000u, 0x477fefffu,
0x477ff000u, 0x477ff001u, 0x47800000u, /* maximum/overflow boundary */
0xc77fe000u, 0xbf801000u, 0x3eaaaaabu, 0xbeaaaaabu,
};
static void test_fill_copy_f32_patterns(void *dst, uint32_t n, uint32_t salt) {
uint8_t *bytes = dst;
const uint32_t n_patterns =
(uint32_t)(sizeof(test_copy_f32_patterns) / sizeof(test_copy_f32_patterns[0]));
for (uint32_t i = 0; i < n; i++) {
const uint32_t bits = test_copy_f32_patterns[(i + salt) % n_patterns];
memcpy(bytes + (uint64_t)i * sizeof(bits), &bits, sizeof(bits));
}
}
#endif
static uint16_t test_float_to_f16(float f) {
union {
float f;
uint32_t u;
} v = { .f = f };
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;
uint32_t shift = (uint32_t)(14 - exp);
uint32_t half_mant = mant >> shift;
if ((mant >> (shift - 1)) & 1u) half_mant++;
return (uint16_t)(sign | half_mant);
}
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 test_f16_to_f32(uint16_t h) {
uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
uint32_t exp = (h >> 10) & 0x1fu;
uint32_t mant = h & 0x03ffu;
uint32_t bits;
if (exp == 0) {
if (mant == 0) {
bits = sign;
} else {
exp = 1;
while ((mant & 0x0400u) == 0) {
mant <<= 1;
exp--;
}
mant &= 0x03ffu;
bits = sign | ((exp + 127u - 15u) << 23) | (mant << 13);
}
} else if (exp == 31) {
bits = sign | 0x7f800000u | (mant << 13);
} else {
bits = sign | ((exp + 127u - 15u) << 23) | (mant << 13);
}
float f;
memcpy(&f, &bits, sizeof(f));
return f;
}
static void test_fill_q8_0_weights(uint8_t *weights,
uint32_t in_dim,
uint32_t out_dim,
uint32_t seed) {
const uint32_t blocks = in_dim / 32u;
const uint64_t row_bytes = (uint64_t)blocks * 34u;
for (uint32_t o = 0; o < out_dim; o++) {
uint8_t *row = weights + (uint64_t)o * row_bytes;
for (uint32_t b = 0; b < blocks; b++) {
float vals[32];
float amax = 0.0f;
for (uint32_t i = 0; i < 32; i++) {
const uint32_t k = b * 32u + i;
const int v = (int)((o * 17u + k * 23u + (o ^ k) * 3u +
seed * 29u + ((o + seed) ^ k) * 5u) % 67u) - 33;
vals[i] = (float)v / 96.0f;
float av = fabsf(vals[i]);
if (av > amax) amax = av;
}
const uint16_t scale_bits = test_float_to_f16(amax / 127.0f);
const float scale = test_f16_to_f32(scale_bits);
memcpy(row + b * 34u, &scale_bits, sizeof(scale_bits));
int8_t *qs = (int8_t *)(row + b * 34u + 2u);
for (uint32_t i = 0; i < 32; i++) {
int q = scale != 0.0f ? (int)lrintf(vals[i] / scale) : 0;
if (q > 127) q = 127;
if (q < -128) q = -128;
qs[i] = (int8_t)q;
}
}
}
}
static void test_metal_f16_matvec_fast_nr0_4(void) {
/*
* This is the short regression for the long-context repetition failure.
* Decode uses one-token F16 matvecs for several DS4 projections; the fast
* nr0=4 variant must be numerically equivalent to the plain kernel.
*/
const uint32_t in_dim = 4096;
const uint32_t out_dim = 512;
const uint64_t weight_bytes = (uint64_t)in_dim * out_dim * sizeof(uint16_t);
const uint64_t weight_alloc = test_round_up_u64(weight_bytes, (uint64_t)getpagesize());
void *weights_raw = NULL;
TEST_ASSERT(posix_memalign(&weights_raw, (size_t)getpagesize(), (size_t)weight_alloc) == 0);
if (!weights_raw) return;
uint16_t *weights = weights_raw;
memset(weights, 0, (size_t)weight_alloc);
for (uint32_t o = 0; o < out_dim; o++) {
for (uint32_t i = 0; i < in_dim; i++) {
float w = (float)((int)((o * 3u + i * 5u) % 23u) - 11) / 64.0f;
weights[(uint64_t)o * in_dim + i] = test_float_to_f16(w);
}
}
ds4_gpu_tensor *x = ds4_gpu_tensor_alloc((uint64_t)in_dim * sizeof(float));
ds4_gpu_tensor *out = ds4_gpu_tensor_alloc((uint64_t)out_dim * sizeof(float));
TEST_ASSERT(x != NULL);
TEST_ASSERT(out != NULL);
if (!x || !out) {
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
return;
}
float *x_host = malloc((size_t)in_dim * sizeof(float));
float *out_host = malloc((size_t)out_dim * sizeof(float));
TEST_ASSERT(x_host != NULL);
TEST_ASSERT(out_host != NULL);
if (!x_host || !out_host) {
free(x_host);
free(out_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
return;
}
for (uint32_t i = 0; i < in_dim; i++) {
x_host[i] = (float)((int)(i % 31u) - 15) / 32.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, (uint64_t)in_dim * sizeof(float)) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(weights_raw, weight_alloc) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(ds4_gpu_matmul_f16_tensor(out, weights_raw, weight_alloc, 0,
in_dim, out_dim, x, 1) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(out, 0, out_host, (uint64_t)out_dim * sizeof(float)) != 0);
float max_abs = 0.0f;
for (uint32_t o = 0; o < out_dim; o++) {
float ref = 0.0f;
for (uint32_t i = 0; i < in_dim; i++) {
float w = (float)((int)((o * 3u + i * 5u) % 23u) - 11) / 64.0f;
ref += w * x_host[i];
}
float err = fabsf(out_host[o] - ref);
if (err > max_abs) max_abs = err;
}
TEST_ASSERT(max_abs < 0.02f);
free(x_host);
free(out_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
}
static void test_metal_f16_prefill_matmul(void) {
const uint32_t in_dim = 128;
const uint32_t out_dim = 64;
const uint32_t n_tok = 128;
const uint64_t weight_bytes = (uint64_t)out_dim * in_dim * sizeof(uint16_t);
const uint64_t weight_alloc = test_round_up_u64(weight_bytes, (uint64_t)getpagesize());
const uint64_t x_bytes = (uint64_t)n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tok * out_dim * sizeof(float);
void *weights_raw = NULL;
TEST_ASSERT(posix_memalign(&weights_raw, (size_t)getpagesize(), (size_t)weight_alloc) == 0);
if (!weights_raw) return;
uint16_t *weights = weights_raw;
memset(weights, 0, (size_t)weight_alloc);
for (uint32_t o = 0; o < out_dim; o++) {
for (uint32_t i = 0; i < in_dim; i++) {
const int v = (int)((o * 11u + i * 13u + (o ^ i) * 5u) % 61u) - 30;
weights[(uint64_t)o * in_dim + i] = test_float_to_f16((float)v / 96.0f);
}
}
ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes);
ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(out_bytes);
TEST_ASSERT(x != NULL);
TEST_ASSERT(out != NULL);
if (!x || !out) {
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
return;
}
float *x_host = malloc((size_t)x_bytes);
float *out_host = malloc((size_t)out_bytes);
TEST_ASSERT(x_host != NULL);
TEST_ASSERT(out_host != NULL);
if (!x_host || !out_host) {
free(x_host);
free(out_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
return;
}
for (uint32_t t = 0; t < n_tok; t++) {
for (uint32_t i = 0; i < in_dim; i++) {
const int v = (int)((t * 7u + i * 17u + (t ^ i) * 3u) % 73u) - 36;
x_host[(uint64_t)t * in_dim + i] = (float)v / 80.0f;
}
}
for (uint32_t i = 0; i < n_tok * out_dim; i++) {
out_host[i] = 12345.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(out, 0, out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(weights_raw, weight_alloc) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(ds4_gpu_matmul_f16_tensor(out, weights_raw, weight_alloc, 0,
in_dim, out_dim, x, n_tok) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(out, 0, out_host, out_bytes) != 0);
float max_abs = 0.0f;
float rms = 0.0f;
for (uint32_t t = 0; t < n_tok; t++) {
for (uint32_t o = 0; o < out_dim; o++) {
float ref = 0.0f;
for (uint32_t i = 0; i < in_dim; i++) {
ref += test_f16_to_f32(weights[(uint64_t)o * in_dim + i]) *
x_host[(uint64_t)t * in_dim + i];
}
const float got = out_host[(uint64_t)t * out_dim + o];
TEST_ASSERT(isfinite(got));
const float err = fabsf(got - ref);
if (err > max_abs) max_abs = err;
rms += err * err;
}
}
rms = sqrtf(rms / (float)(n_tok * out_dim));
TEST_ASSERT(max_abs < 0.08f);
TEST_ASSERT(rms < 0.02f);
free(x_host);
free(out_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
}
static void test_metal_q8_0_prefill_matmul(void) {
const uint32_t in_dim = 128;
const uint32_t out_dim = 64;
const uint32_t n_tok = 128;
const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u;
const uint64_t weight_bytes = (uint64_t)out_dim * row_bytes;
const uint64_t weight_alloc = test_round_up_u64(weight_bytes, (uint64_t)getpagesize());
const uint64_t x_bytes = (uint64_t)n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tok * out_dim * sizeof(float);
void *weights_raw = NULL;
TEST_ASSERT(posix_memalign(&weights_raw, (size_t)getpagesize(), (size_t)weight_alloc) == 0);
if (!weights_raw) return;
uint8_t *weights = weights_raw;
memset(weights, 0, (size_t)weight_alloc);
test_fill_q8_0_weights(weights, in_dim, out_dim, 0);
ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes);
ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(out_bytes);
TEST_ASSERT(x != NULL);
TEST_ASSERT(out != NULL);
if (!x || !out) {
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
return;
}
float *x_host = malloc((size_t)x_bytes);
float *out_host = malloc((size_t)out_bytes);
TEST_ASSERT(x_host != NULL);
TEST_ASSERT(out_host != NULL);
if (!x_host || !out_host) {
free(x_host);
free(out_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
return;
}
for (uint32_t t = 0; t < n_tok; t++) {
for (uint32_t i = 0; i < in_dim; i++) {
const int v = (int)((t * 19u + i * 7u + (t ^ i)) % 71u) - 35;
x_host[(uint64_t)t * in_dim + i] = (float)v / 80.0f;
}
}
for (uint32_t i = 0; i < n_tok * out_dim; i++) {
out_host[i] = 12345.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(out, 0, out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(weights_raw, weight_alloc) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor(out, weights_raw, weight_alloc, 0,
in_dim, out_dim, x, n_tok) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(out, 0, out_host, out_bytes) != 0);
float max_abs = 0.0f;
float rms = 0.0f;
for (uint32_t t = 0; t < n_tok; t++) {
for (uint32_t o = 0; o < out_dim; o++) {
const uint8_t *row = weights + (uint64_t)o * row_bytes;
float ref = 0.0f;
for (uint32_t b = 0; b < in_dim / 32u; b++) {
uint16_t scale_bits;
memcpy(&scale_bits, row + b * 34u, sizeof(scale_bits));
const float scale = test_f16_to_f32(scale_bits);
const int8_t *qs = (const int8_t *)(row + b * 34u + 2u);
for (uint32_t i = 0; i < 32; i++) {
ref += scale * (float)qs[i] *
x_host[(uint64_t)t * in_dim + b * 32u + i];
}
}
const float got = out_host[(uint64_t)t * out_dim + o];
TEST_ASSERT(isfinite(got));
const float err = fabsf(got - ref);
if (err > max_abs) max_abs = err;
rms += err * err;
}
}
rms = sqrtf(rms / (float)(n_tok * out_dim));
TEST_ASSERT(max_abs < 0.08f);
TEST_ASSERT(rms < 0.02f);
free(x_host);
free(out_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(out);
free(weights_raw);
}
static void test_metal_pack_slot_rows_f32(void) {
const uint32_t n_rows = 3;
const uint32_t width = 5;
const uint32_t n_slots = 4;
const uint32_t slot_cap = 6;
const uint64_t slot_count = (uint64_t)n_slots * slot_cap * width;
const uint64_t out_count = (uint64_t)n_rows * n_slots * width;
const uint64_t slot_bytes = slot_count * sizeof(float);
const uint64_t out_bytes = out_count * sizeof(float);
ds4_gpu_tensor *slots = ds4_gpu_tensor_alloc(slot_bytes);
ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(out_bytes);
TEST_ASSERT(slots != NULL);
TEST_ASSERT(out != NULL);
if (!slots || !out) {
ds4_gpu_tensor_free(slots);
ds4_gpu_tensor_free(out);
return;
}
float *slot_host = malloc((size_t)slot_bytes);
float *out_host = malloc((size_t)out_bytes);
TEST_ASSERT(slot_host != NULL);
TEST_ASSERT(out_host != NULL);
if (!slot_host || !out_host) {
free(slot_host);
free(out_host);
ds4_gpu_tensor_free(slots);
ds4_gpu_tensor_free(out);
return;
}
for (uint32_t slot = 0; slot < n_slots; slot++) {
for (uint32_t row = 0; row < slot_cap; row++) {
for (uint32_t col = 0; col < width; col++) {
slot_host[((uint64_t)slot * slot_cap + row) * width + col] =
(float)(slot * 1000u + row * 100u + col);
}
}
}
for (uint64_t i = 0; i < out_count; i++) out_host[i] = -1.0f;
TEST_ASSERT(ds4_gpu_tensor_write(slots, 0, slot_host, slot_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(out, 0, out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_pack_slot_rows_f32_tensor(out,
slots,
n_rows,
width,
n_slots,
slot_cap) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(out, 0, out_host, out_bytes) != 0);
for (uint32_t row = 0; row < n_rows; row++) {
for (uint32_t slot = 0; slot < n_slots; slot++) {
for (uint32_t col = 0; col < width; col++) {
const float ref =
slot_host[((uint64_t)slot * slot_cap + row) * width + col];
const float got =
out_host[((uint64_t)row * n_slots + slot) * width + col];
TEST_ASSERT(got == ref);
}
}
}
free(slot_host);
free(out_host);
ds4_gpu_tensor_free(slots);
ds4_gpu_tensor_free(out);
}
static void test_metal_store_raw_kv_batch_wrap(void) {
const uint32_t raw_cap = 5;
const uint32_t head_dim = 3;
const uint32_t n_tokens = 4;
const uint32_t pos0 = 3;
const uint64_t kv_count = (uint64_t)n_tokens * head_dim;
const uint64_t raw_count = (uint64_t)raw_cap * head_dim;
const uint64_t kv_bytes = kv_count * sizeof(float);
const uint64_t raw_bytes = raw_count * sizeof(float);
ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(kv_bytes);
ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes);
TEST_ASSERT(kv != NULL);
TEST_ASSERT(raw != NULL);
if (!kv || !raw) {
ds4_gpu_tensor_free(kv);
ds4_gpu_tensor_free(raw);
return;
}
float kv_host[12];
float raw_host[15];
for (uint32_t t = 0; t < n_tokens; t++) {
for (uint32_t d = 0; d < head_dim; d++) {
kv_host[(uint64_t)t * head_dim + d] = (float)(100u * t + d);
}
}
for (uint64_t i = 0; i < raw_count; i++) raw_host[i] = -1.0f;
TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, kv_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_store_raw_kv_batch_tensor(raw,
kv,
raw_cap,
pos0,
n_tokens,
head_dim) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(raw, 0, raw_host, raw_bytes) != 0);
for (uint32_t t = 0; t < n_tokens; t++) {
const uint32_t row = (pos0 + t) % raw_cap;
for (uint32_t d = 0; d < head_dim; d++) {
const float ref = kv_host[(uint64_t)t * head_dim + d];
const float got = raw_host[(uint64_t)row * head_dim + d];
TEST_ASSERT(got == ref);
}
}
for (uint32_t d = 0; d < head_dim; d++) {
TEST_ASSERT(raw_host[(uint64_t)2u * head_dim + d] == -1.0f);
}
ds4_gpu_tensor_free(kv);
ds4_gpu_tensor_free(raw);
}
static void test_dspark_cache_window_crop(void) {
TEST_ASSERT(ds4_test_dspark_cache_window_crop());
}
static void test_metal_q8_0_decode_pair_exact_case(
uint32_t out0_dim,
uint32_t out1_dim,
uint32_t seed0,
uint32_t seed1) {
/* Exercise the Q-A/KV contract with unequal, odd output extents and
* independently page-aligned model ranges. The paired kernel must be
* bit-identical to two standalone decode matvec dispatches. */
const uint32_t in_dim = 4096;
const uint64_t page = (uint64_t)getpagesize();
const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u;
const uint64_t weight0_bytes = (uint64_t)out0_dim * row_bytes;
const uint64_t weight1_bytes = (uint64_t)out1_dim * row_bytes;
const uint64_t weight1_offset = test_round_up_u64(weight0_bytes, page);
const uint64_t weight_alloc =
test_round_up_u64(weight1_offset + weight1_bytes, page);
void *weights_raw = NULL;
TEST_ASSERT(posix_memalign(&weights_raw, (size_t)page, (size_t)weight_alloc) == 0);
if (!weights_raw) return;
memset(weights_raw, 0, (size_t)weight_alloc);
test_fill_q8_0_weights((uint8_t *)weights_raw, in_dim, out0_dim, seed0);
test_fill_q8_0_weights((uint8_t *)weights_raw + weight1_offset,
in_dim, out1_dim, seed1);
const uint64_t x_bytes = (uint64_t)in_dim * sizeof(float);
const uint64_t out0_bytes = (uint64_t)out0_dim * sizeof(float);
const uint64_t out1_bytes = (uint64_t)out1_dim * sizeof(float);
ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes);
ds4_gpu_tensor *ref0 = ds4_gpu_tensor_alloc(out0_bytes);
ds4_gpu_tensor *ref1 = ds4_gpu_tensor_alloc(out1_bytes);
ds4_gpu_tensor *pair0 = ds4_gpu_tensor_alloc(out0_bytes);
ds4_gpu_tensor *pair1 = ds4_gpu_tensor_alloc(out1_bytes);
TEST_ASSERT(x != NULL);
TEST_ASSERT(ref0 != NULL);
TEST_ASSERT(ref1 != NULL);
TEST_ASSERT(pair0 != NULL);
TEST_ASSERT(pair1 != NULL);
if (!x || !ref0 || !ref1 || !pair0 || !pair1) {
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(ref0);
ds4_gpu_tensor_free(ref1);
ds4_gpu_tensor_free(pair0);
ds4_gpu_tensor_free(pair1);
free(weights_raw);
return;
}
float *x_host = malloc((size_t)x_bytes);
float *ref0_host = malloc((size_t)out0_bytes);
float *ref1_host = malloc((size_t)out1_bytes);
float *pair0_host = malloc((size_t)out0_bytes);
float *pair1_host = malloc((size_t)out1_bytes);
TEST_ASSERT(x_host != NULL);
TEST_ASSERT(ref0_host != NULL);
TEST_ASSERT(ref1_host != NULL);
TEST_ASSERT(pair0_host != NULL);
TEST_ASSERT(pair1_host != NULL);
if (!x_host || !ref0_host || !ref1_host || !pair0_host || !pair1_host) {
free(x_host);
free(ref0_host);
free(ref1_host);
free(pair0_host);
free(pair1_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(ref0);
ds4_gpu_tensor_free(ref1);
ds4_gpu_tensor_free(pair0);
ds4_gpu_tensor_free(pair1);
free(weights_raw);
return;
}
for (uint32_t i = 0; i < in_dim; i++) {
const int v = (int)((i * 29u + (i ^ (i >> 3u)) * 7u) % 127u) - 63;
x_host[i] = (float)v / 72.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(weights_raw, weight_alloc) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor(ref0, weights_raw, weight_alloc, 0,
in_dim, out0_dim, x, 1) != 0);
TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor(ref1, weights_raw, weight_alloc,
weight1_offset,
in_dim, out1_dim, x, 1) != 0);
TEST_ASSERT(ds4_gpu_matmul_q8_0_pair_tensor(pair0, pair1,
weights_raw, weight_alloc,
0, weight1_offset,
in_dim, out0_dim, out1_dim,
x, 1) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(ref0, 0, ref0_host, out0_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(ref1, 0, ref1_host, out1_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(pair0, 0, pair0_host, out0_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(pair1, 0, pair1_host, out1_bytes) != 0);
uint32_t mismatch0 = 0;
uint32_t mismatch1 = 0;
float max_abs0 = 0.0f;
float max_abs1 = 0.0f;
for (uint32_t i = 0; i < out0_dim; i++) {
if (memcmp(&ref0_host[i], &pair0_host[i], sizeof(float)) != 0) mismatch0++;
const float err = fabsf(ref0_host[i] - pair0_host[i]);
if (err > max_abs0) max_abs0 = err;
}
for (uint32_t i = 0; i < out1_dim; i++) {
if (memcmp(&ref1_host[i], &pair1_host[i], sizeof(float)) != 0) mismatch1++;
const float err = fabsf(ref1_host[i] - pair1_host[i]);
if (err > max_abs1) max_abs1 = err;
}
if (mismatch0 != 0 || mismatch1 != 0) {
fprintf(stderr,
"ds4-test: paired Q8_0 exactness mismatches=%u/%u max_abs=%g, %u/%u max_abs=%g\n",
mismatch0, out0_dim, max_abs0,
mismatch1, out1_dim, max_abs1);
}
TEST_ASSERT(mismatch0 == 0);
TEST_ASSERT(mismatch1 == 0);
free(x_host);
free(ref0_host);
free(ref1_host);
free(pair0_host);
free(pair1_host);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(ref0);
ds4_gpu_tensor_free(ref1);
ds4_gpu_tensor_free(pair0);
ds4_gpu_tensor_free(pair1);
free(weights_raw);
}
static void test_metal_q8_0_decode_pair_exact(void) {
/* Cover both possible one-bank tail directions. Distinct seeds ensure a
* mistaken A-for-B weight binding cannot compare equal by construction. */
test_metal_q8_0_decode_pair_exact_case(77, 19, 11, 97);
test_metal_q8_0_decode_pair_exact_case(19, 77, 23, 131);
}
#if defined(__APPLE__)
static void test_metal_f16_compressor_pair_state_store_exact_case(
uint32_t width,
uint32_t ratio,
uint32_t pos,
uint32_t ape_type,
uint32_t seed,
bool test_decode_pack) {
const uint32_t in_dim = 4096u;
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t head_dim = width / coff;
const uint32_t state_rows = coff * ratio;
const bool emit = ((pos + 1u) % ratio) == 0u;
TEST_ASSERT(!test_decode_pack ||
(ratio == 4u && emit &&
(head_dim == 128u || head_dim == 512u)));
const uint64_t page = (uint64_t)getpagesize();
const uint64_t weight_bytes =
(uint64_t)width * in_dim * sizeof(uint16_t);
const uint64_t score_weight_offset =
test_round_up_u64(weight_bytes, page);
const uint64_t ape_offset = test_round_up_u64(
score_weight_offset + weight_bytes, page);
const uint64_t ape_elem_bytes = ape_type == 1u ? 2u : 4u;
const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem_bytes;
const uint64_t norm_offset =
test_round_up_u64(ape_offset + ape_bytes, page);
const uint64_t model_bytes = test_round_up_u64(
norm_offset + (uint64_t)head_dim * sizeof(float), page);
const uint64_t x_bytes = (uint64_t)in_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)width * sizeof(float);
const uint64_t state_count = (uint64_t)state_rows * width;
const uint64_t state_bytes = state_count * sizeof(float);
const uint64_t comp_bytes = (uint64_t)head_dim * sizeof(float);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page,
(size_t)model_bytes) == 0);
ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes);
ds4_gpu_tensor *ref_kv = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *ref_score = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *fused_kv = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *fused_score = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *fused_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *fused_state_score = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *fused_comp = ds4_gpu_tensor_alloc(comp_bytes);
float *x_host = malloc((size_t)x_bytes);
float *ref_kv_host = malloc((size_t)out_bytes);
float *ref_score_host = malloc((size_t)out_bytes);
float *fused_kv_host = malloc((size_t)out_bytes);
float *fused_score_host = malloc((size_t)out_bytes);
float *ref_state_kv_host = malloc((size_t)state_bytes);
float *ref_state_score_host = malloc((size_t)state_bytes);
float *fused_state_kv_host = malloc((size_t)state_bytes);
float *fused_state_score_host = malloc((size_t)state_bytes);
float *ref_comp_host = malloc((size_t)comp_bytes);
float *fused_comp_host = malloc((size_t)comp_bytes);
TEST_ASSERT(model_raw != NULL);
TEST_ASSERT(x != NULL);
TEST_ASSERT(ref_kv != NULL);
TEST_ASSERT(ref_score != NULL);
TEST_ASSERT(fused_kv != NULL);
TEST_ASSERT(fused_score != NULL);
TEST_ASSERT(ref_state_kv != NULL);
TEST_ASSERT(ref_state_score != NULL);
TEST_ASSERT(fused_state_kv != NULL);
TEST_ASSERT(fused_state_score != NULL);
TEST_ASSERT(ref_comp != NULL);
TEST_ASSERT(fused_comp != NULL);
TEST_ASSERT(x_host != NULL);
TEST_ASSERT(ref_kv_host != NULL);
TEST_ASSERT(ref_score_host != NULL);
TEST_ASSERT(fused_kv_host != NULL);
TEST_ASSERT(fused_score_host != NULL);
TEST_ASSERT(ref_state_kv_host != NULL);
TEST_ASSERT(ref_state_score_host != NULL);
TEST_ASSERT(fused_state_kv_host != NULL);
TEST_ASSERT(fused_state_score_host != NULL);
TEST_ASSERT(ref_comp_host != NULL);
TEST_ASSERT(fused_comp_host != NULL);
const bool allocated = model_raw && x && ref_kv && ref_score && fused_kv &&
fused_score && ref_state_kv && ref_state_score && fused_state_kv &&
fused_state_score && ref_comp && fused_comp && x_host && ref_kv_host &&
ref_score_host && fused_kv_host && fused_score_host &&
ref_state_kv_host && ref_state_score_host && fused_state_kv_host &&
fused_state_score_host && ref_comp_host && fused_comp_host;
const char *pair_disable_env =
"DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ";
const char *store_disable_env =
"DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE";
const char *decode_pack_disable_env =
"DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_RATIO4_DECODE_PACK_FUSION";
const char *exact_reduction_disable_env =
"DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_EXACT_REDUCTION_FUSION";
const char *exact_reduction_poison_env =
"DS4_METAL_TEST_POISON_COMPRESSOR_EXACT_REDUCTION_SCRATCH";
char *saved_pair_disable = test_save_env(pair_disable_env);
char *saved_store_disable = test_save_env(store_disable_env);
char *saved_decode_pack_disable =
test_save_env(decode_pack_disable_env);
char *saved_exact_reduction_disable =
test_save_env(exact_reduction_disable_env);
char *saved_exact_reduction_poison =
test_save_env(exact_reduction_poison_env);
test_float_compare_stats kv_stats = {0};
test_float_compare_stats score_stats = {0};
test_float_compare_stats state_kv_stats = {0};
test_float_compare_stats state_score_stats = {0};
test_float_compare_stats comp_stats = {0};
if (allocated) {
memset(model_raw, 0, (size_t)model_bytes);
uint16_t *kv_weights = model_raw;
uint16_t *score_weights =
(uint16_t *)((uint8_t *)model_raw + score_weight_offset);
for (uint32_t o = 0; o < width; o++) {
for (uint32_t i = 0; i < in_dim; i++) {
const int kv_value =
(int)((o * 17u + i * 23u + (o ^ i) * 3u +
seed * 29u) % 67u) - 33;
const int score_value =
(int)((o * 31u + i * 11u + (o ^ (i >> 2u)) * 5u +
seed * 19u) % 71u) - 35;
const uint64_t wi = (uint64_t)o * in_dim + i;
kv_weights[wi] = test_float_to_f16(
(float)kv_value / 96.0f);
score_weights[wi] = test_float_to_f16(
(float)score_value / 104.0f);
}
}
if (ape_type == 1u) {
uint16_t *ape =
(uint16_t *)((uint8_t *)model_raw + ape_offset);
for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) {
const int value =
(int)((i * 13u + (i ^ (i >> 3u)) * 7u +
seed * 17u) % 61u) - 30;
ape[i] = test_float_to_f16((float)value / 80.0f);
}
} else {
float *ape = (float *)((uint8_t *)model_raw + ape_offset);
for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) {
const int value =
(int)((i * 13u + (i ^ (i >> 3u)) * 7u +
seed * 17u) % 61u) - 30;
ape[i] = (float)value / 80.0f;
}
}
float *norm = (float *)((uint8_t *)model_raw + norm_offset);
for (uint32_t i = 0; i < head_dim; i++) {
norm[i] = 0.75f +
(float)((i * 7u + seed * 3u) % 23u) / 64.0f;
}
for (uint32_t i = 0; i < in_dim; i++) {
const int value =
(int)((i * 29u + (i ^ (i >> 4u)) * 9u +
seed * 11u) % 127u) - 63;
x_host[i] = (float)value / 88.0f;
}
for (uint32_t i = 0; i < width; i++) {
const uint32_t poison = 0x7fc00001u + (i & 0x3ffu);
memcpy(ref_kv_host + i, &poison, sizeof(poison));
memcpy(ref_score_host + i, &poison, sizeof(poison));
memcpy(fused_kv_host + i, &poison, sizeof(poison));
memcpy(fused_score_host + i, &poison, sizeof(poison));
}
for (uint64_t i = 0; i < state_count; i++) {
const int kv_value =
(int)((i * 5u + seed * 13u) % 97u) - 48;
const int score_value =
(int)((i * 7u + seed * 5u) % 101u) - 50;
ref_state_kv_host[i] = (float)kv_value / 64.0f;
fused_state_kv_host[i] = ref_state_kv_host[i];
ref_state_score_host[i] = (float)score_value / 72.0f;
fused_state_score_host[i] = ref_state_score_host[i];
}
if (test_decode_pack) {
static const uint32_t edge_bits[8] = {
0x00000000u, 0x80000000u, 0x00000001u, 0x80000001u,
0x3f800000u, 0xbf800000u, 0x42a00000u, 0xc2a00000u,
};
for (uint32_t col = 0; col < 4u; col++) {
for (uint32_t row = 0; row < 8u; row++) {
const uint64_t state_col =
(row >= 4u ? head_dim : 0u) + col;
const uint64_t state_index =
(uint64_t)row * width + state_col;
const uint32_t score_bits =
edge_bits[(row + col) & 7u];
const uint32_t kv_bits =
edge_bits[(7u - row + col) & 7u];
memcpy(ref_state_score_host + state_index,
&score_bits, sizeof(score_bits));
memcpy(fused_state_score_host + state_index,
&score_bits, sizeof(score_bits));
memcpy(ref_state_kv_host + state_index,
&kv_bits, sizeof(kv_bits));
memcpy(fused_state_kv_host + state_index,
&kv_bits, sizeof(kv_bits));
}
}
}
for (uint32_t i = 0; i < head_dim; i++) {
const uint32_t poison = 0x7fc01001u + (i & 0x3ffu);
memcpy(ref_comp_host + i, &poison, sizeof(poison));
memcpy(fused_comp_host + i, &poison, sizeof(poison));
}
TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_kv, 0, ref_kv_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_score, 0, ref_score_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_kv, 0, fused_kv_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_score, 0, fused_score_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_score, 0, ref_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, fused_state_kv_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, fused_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_comp, 0, ref_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_comp, 0, fused_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(unsetenv(pair_disable_env) == 0);
TEST_ASSERT(unsetenv(store_disable_env) == 0);
TEST_ASSERT(setenv(exact_reduction_disable_env, "1", 1) == 0);
TEST_ASSERT(unsetenv(exact_reduction_poison_env) == 0);
if (test_decode_pack) {
TEST_ASSERT(setenv(decode_pack_disable_env, "1", 1) == 0);
}
TEST_ASSERT(ds4_gpu_matmul_f16_pair_tensor(
ref_kv, ref_score, model_raw, model_bytes,
0, score_weight_offset, in_dim, width, x, 1) != 0);
TEST_ASSERT(ds4_gpu_compressor_update_tensor(
ref_kv, ref_score, ref_state_kv, ref_state_score,
ref_comp, model_raw, model_bytes, ape_offset, ape_type,
norm_offset, 0, head_dim, ratio, pos, 0, 0, 0,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f,
1.0e-6f, false, test_decode_pack, false) != 0);
TEST_ASSERT(ds4_gpu_matmul_f16_pair_compressor_store_tensor(
fused_kv, fused_score,
fused_state_kv, fused_state_score,
model_raw, model_bytes, 0, score_weight_offset,
ape_offset, ape_type, in_dim, width, x,
ratio, pos) == 1);
if (test_decode_pack) {
TEST_ASSERT(unsetenv(decode_pack_disable_env) == 0);
TEST_ASSERT(unsetenv(exact_reduction_disable_env) == 0);
TEST_ASSERT(setenv(exact_reduction_poison_env, "1", 1) == 0);
}
TEST_ASSERT(ds4_gpu_compressor_update_tensor(
fused_kv, fused_score,
fused_state_kv, fused_state_score,
fused_comp, model_raw, model_bytes,
ape_offset, ape_type, norm_offset, 0,
head_dim, ratio, pos, 0, 0, 0,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f,
1.0e-6f, true, test_decode_pack, false) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_kv, 0, ref_kv_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_score, 0, ref_score_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_kv, 0, fused_kv_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_score, 0, fused_score_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_score, 0, ref_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_state_kv, 0, fused_state_kv_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_state_score, 0, fused_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_comp, 0, ref_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_comp, 0, fused_comp_host, comp_bytes) != 0);
kv_stats = test_compare_float_bits(
ref_kv_host, fused_kv_host, width);
score_stats = test_compare_float_bits(
ref_score_host, fused_score_host, width);
state_kv_stats = test_compare_float_bits(
ref_state_kv_host, fused_state_kv_host, (size_t)state_count);
state_score_stats = test_compare_float_bits(
ref_state_score_host, fused_state_score_host,
(size_t)state_count);
comp_stats = test_compare_float_bits(
ref_comp_host, fused_comp_host, head_dim);
}
test_restore_env(pair_disable_env, saved_pair_disable);
test_restore_env(store_disable_env, saved_store_disable);
test_restore_env(decode_pack_disable_env, saved_decode_pack_disable);
test_restore_env(exact_reduction_disable_env,
saved_exact_reduction_disable);
test_restore_env(
exact_reduction_poison_env, saved_exact_reduction_poison);
fprintf(stderr,
"ds4-test: compressor pair state-store exact width=%u ratio=%u "
"pos=%u emit=%u ape=%s decode_pack=%u exact_reduce=%u "
"proj=%zu/%zu state=%zu/%zu "
"comp=%zu max_ulp=%u/%u/%u/%u/%u\n",
width, ratio, pos, emit ? 1u : 0u,
ape_type == 1u ? "f16" : "f32",
test_decode_pack ? 1u : 0u,
test_decode_pack ? 1u : 0u,
kv_stats.mismatch_count, score_stats.mismatch_count,
state_kv_stats.mismatch_count,
state_score_stats.mismatch_count,
comp_stats.mismatch_count,
kv_stats.max_ulp, score_stats.max_ulp,
state_kv_stats.max_ulp, state_score_stats.max_ulp,
comp_stats.max_ulp);
TEST_ASSERT(kv_stats.mismatch_count == 0);
TEST_ASSERT(score_stats.mismatch_count == 0);
TEST_ASSERT(state_kv_stats.mismatch_count == 0);
TEST_ASSERT(state_score_stats.mismatch_count == 0);
TEST_ASSERT(comp_stats.mismatch_count == 0);
free(fused_comp_host);
free(ref_comp_host);
free(fused_state_score_host);
free(fused_state_kv_host);
free(ref_state_score_host);
free(ref_state_kv_host);
free(fused_score_host);
free(fused_kv_host);
free(ref_score_host);
free(ref_kv_host);
free(x_host);
ds4_gpu_tensor_free(fused_comp);
ds4_gpu_tensor_free(ref_comp);
ds4_gpu_tensor_free(fused_state_score);
ds4_gpu_tensor_free(fused_state_kv);
ds4_gpu_tensor_free(ref_state_score);
ds4_gpu_tensor_free(ref_state_kv);
ds4_gpu_tensor_free(fused_score);
ds4_gpu_tensor_free(fused_kv);
ds4_gpu_tensor_free(ref_score);
ds4_gpu_tensor_free(ref_kv);
ds4_gpu_tensor_free(x);
free(model_raw);
}
static void test_metal_f16_compressor_pair_state_store_exact(void) {
test_metal_f16_compressor_pair_state_store_exact_case(
256, 4, 8, 0, 17, false);
test_metal_f16_compressor_pair_state_store_exact_case(
256, 4, 11, 1, 23, true);
test_metal_f16_compressor_pair_state_store_exact_case(
1024, 4, 11, 1, 29, true);
test_metal_f16_compressor_pair_state_store_exact_case(
512, 128, 255, 1, 43, false);
}
static void test_metal_compressor_ape_add_exact_case(
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens,
uint32_t ape_type,
uint32_t seed,
bool test_pack_fusion) {
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint32_t n_comp = n_tokens / ratio;
const uint64_t input_count = (uint64_t)n_tokens * width;
const uint64_t state_count = (uint64_t)state_rows * width;
const uint64_t comp_count = (uint64_t)n_comp * head_dim;
const uint64_t input_bytes = input_count * sizeof(float);
const uint64_t state_bytes = state_count * sizeof(float);
const uint64_t comp_bytes = comp_count * sizeof(float);
const uint64_t ape_elem_bytes = ape_type == 1u ? 2u : 4u;
const uint64_t ape_bytes = (uint64_t)width * ratio * ape_elem_bytes;
const uint64_t page = (uint64_t)getpagesize();
const uint64_t norm_offset = test_round_up_u64(ape_bytes, page);
const uint64_t model_bytes = test_round_up_u64(
norm_offset + (uint64_t)head_dim * sizeof(float), page);
ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(input_bytes);
ds4_gpu_tensor *sc = ds4_gpu_tensor_alloc(input_bytes);
ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *fused_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *fused_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *fused_state_score = ds4_gpu_tensor_alloc(state_bytes);
TEST_ASSERT(kv != NULL);
TEST_ASSERT(sc != NULL);
TEST_ASSERT(ref_comp != NULL);
TEST_ASSERT(fused_comp != NULL);
TEST_ASSERT(ref_state_kv != NULL);
TEST_ASSERT(ref_state_score != NULL);
TEST_ASSERT(fused_state_kv != NULL);
TEST_ASSERT(fused_state_score != NULL);
float *kv_host = malloc((size_t)input_bytes);
float *sc_host = malloc((size_t)input_bytes);
float *ref_comp_host = malloc((size_t)comp_bytes);
float *fused_comp_host = malloc((size_t)comp_bytes);
float *ref_state_kv_host = malloc((size_t)state_bytes);
float *ref_state_score_host = malloc((size_t)state_bytes);
float *fused_state_kv_host = malloc((size_t)state_bytes);
float *fused_state_score_host = malloc((size_t)state_bytes);
const uint64_t poison_count = input_count > state_count ?
input_count : state_count;
float *poison_host = test_pack_fusion ?
malloc((size_t)(poison_count * sizeof(float))) : NULL;
TEST_ASSERT(kv_host != NULL);
TEST_ASSERT(sc_host != NULL);
TEST_ASSERT(ref_comp_host != NULL);
TEST_ASSERT(fused_comp_host != NULL);
TEST_ASSERT(ref_state_kv_host != NULL);
TEST_ASSERT(ref_state_score_host != NULL);
TEST_ASSERT(fused_state_kv_host != NULL);
TEST_ASSERT(fused_state_score_host != NULL);
TEST_ASSERT(!test_pack_fusion || poison_host != NULL);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page,
(size_t)model_bytes) == 0);
const bool allocated = kv && sc && ref_comp && fused_comp && ref_state_kv &&
ref_state_score && fused_state_kv && fused_state_score && kv_host &&
sc_host && ref_comp_host && fused_comp_host && ref_state_kv_host &&
ref_state_score_host && fused_state_kv_host && fused_state_score_host &&
(!test_pack_fusion || poison_host) && model_raw;
const char *disable_env = "DS4_METAL_DISABLE_COMPRESSOR_APE_ADD";
const char *pack_disable_env =
"DS4_METAL_DISABLE_COMPRESSOR_RATIO4_PACK_FUSION";
char *saved_disable = test_save_env(disable_env);
char *saved_pack_disable = test_save_env(pack_disable_env);
test_float_compare_stats comp_stats = {0};
test_float_compare_stats state_kv_stats = {0};
test_float_compare_stats state_score_stats = {0};
if (allocated) {
memset(model_raw, 0, (size_t)model_bytes);
if (ape_type == 1u) {
uint16_t *ape = model_raw;
for (uint64_t i = 0; i < (uint64_t)width * ratio; i++) {
const int value = (int)((i * 17u + (i ^ (i >> 4u)) * 5u +
seed * 13u) % 127u) - 63;
ape[i] = test_float_to_f16((float)value / 80.0f);
}
} else {
float *ape = model_raw;
for (uint64_t i = 0; i < (uint64_t)width * ratio; i++) {
const int value = (int)((i * 17u + (i ^ (i >> 4u)) * 5u +
seed * 13u) % 127u) - 63;
ape[i] = (float)value / 80.0f;
}
}
float *norm = (float *)((uint8_t *)model_raw + norm_offset);
for (uint32_t i = 0; i < head_dim; i++) {
norm[i] = 0.75f + (float)((i * 7u + seed) % 19u) / 64.0f;
}
for (uint64_t i = 0; i < input_count; i++) {
const int kv_value = (int)((i * 29u + (i ^ (i >> 5u)) * 3u +
seed * 23u) % 181u) - 90;
const int sc_value = (int)((i * 31u + (i ^ (i >> 3u)) * 11u +
seed * 17u) % 173u) - 86;
kv_host[i] = (float)kv_value / 112.0f;
sc_host[i] = (float)sc_value / 96.0f;
}
kv_host[0] = -0.0f;
sc_host[0] = -0.0f;
// Exercise exact-add edge values in the first active APE row. Using
// bit patterns avoids host fast-math rewriting signed zeros or
// subnormals before the legacy and fused Metal paths see them.
const uint32_t cutoff = (n_tokens / ratio) * ratio;
const uint32_t edge_token = cutoff < n_tokens ? cutoff : cutoff - ratio;
const uint64_t edge_ape =
(uint64_t)((pos0 + edge_token) % ratio) * width;
float *edge_score = sc_host + (uint64_t)edge_token * width;
static const uint16_t edge_f16[] = {
0x0000u, 0x8000u, 0x0001u, 0x8001u,
0x3c00u, 0xbc00u, 0x7bffu, 0xfbffu,
};
static const uint32_t edge_f32[] = {
0x00000000u, 0x80000000u, 0x00000001u, 0x80000001u,
0x3f800000u, 0xbf800000u, 0x7f7fffffu, 0xff7fffffu,
};
static const uint32_t edge_score_f16[] = {
0x80000000u, 0x00000000u, 0x00000000u, 0x80000000u,
0xbf800000u, 0x3f800000u, 0xc77fe000u, 0x477fe000u,
};
static const uint32_t edge_score_f32[] = {
0x80000000u, 0x00000000u, 0x00000000u, 0x80000000u,
0xbf800000u, 0x3f800000u, 0xff7fffffu, 0x7f7fffffu,
};
if (ape_type == 1u) {
uint16_t *ape = model_raw;
memcpy(ape + edge_ape, edge_f16, sizeof(edge_f16));
for (uint32_t i = 0; i < 8u; i++) {
memcpy(edge_score + i, edge_score_f16 + i, sizeof(uint32_t));
}
} else {
uint32_t *ape = model_raw;
memcpy(ape + edge_ape, edge_f32, sizeof(edge_f32));
for (uint32_t i = 0; i < 8u; i++) {
memcpy(edge_score + i, edge_score_f32 + i, sizeof(uint32_t));
}
}
for (uint64_t i = 0; i < state_count; i++) {
ref_state_kv_host[i] = 1234.0f;
fused_state_kv_host[i] = 1234.0f;
ref_state_score_host[i] = -1234.0f;
fused_state_score_host[i] = -1234.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, fused_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_score, 0, ref_state_score_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, fused_state_score_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0);
ds4_gpu_set_quality(false);
if (test_pack_fusion) {
TEST_ASSERT(unsetenv(disable_env) == 0);
} else {
TEST_ASSERT(setenv(disable_env, "1", 1) == 0);
}
TEST_ASSERT(setenv(pack_disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_compressor_prefill_tensor(
ref_comp, ref_state_kv, ref_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, ratio, pos0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0);
if (test_pack_fusion) {
// Overwrite every persistent pack cell with qNaN payloads through
// the legacy replay path. This makes an omitted candidate write
// observable even for plane-zero padding, which normal legacy
// packing would otherwise leave at the correct 0/-inf values.
for (uint64_t i = 0; i < poison_count; i++) {
const uint32_t bits =
0x7fc00001u + (uint32_t)(i & 0x3ffu);
memcpy(poison_host + i, &bits, sizeof(bits));
}
TEST_ASSERT(ds4_gpu_tensor_write(
kv, 0, poison_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
sc, 0, poison_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, poison_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, poison_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor(
fused_comp, fused_state_kv, fused_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, 0, n_comp * ratio, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
kv, 0, kv_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
sc, 0, sc_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, fused_state_kv_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, fused_state_score_host,
state_bytes) != 0);
TEST_ASSERT(unsetenv(pack_disable_env) == 0);
} else {
TEST_ASSERT(unsetenv(disable_env) == 0);
}
TEST_ASSERT(ds4_gpu_compressor_prefill_tensor(
fused_comp, fused_state_kv, fused_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, ratio, pos0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0);
if (test_pack_fusion) {
TEST_ASSERT(ds4_gpu_tensor_read(
kv, 0, poison_host, input_bytes) != 0);
TEST_ASSERT(memcmp(kv_host, poison_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(ds4_gpu_tensor_read(
sc, 0, poison_host, input_bytes) != 0);
TEST_ASSERT(memcmp(sc_host, poison_host,
(size_t)input_bytes) == 0);
}
TEST_ASSERT(ds4_gpu_tensor_read(
ref_comp, 0, ref_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_comp, 0, fused_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_state_kv, 0, fused_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_score, 0, ref_state_score_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_state_score, 0, fused_state_score_host,
state_bytes) != 0);
comp_stats = test_compare_float_bits(
ref_comp_host, fused_comp_host, (size_t)comp_count);
state_kv_stats = test_compare_float_bits(
ref_state_kv_host, fused_state_kv_host, (size_t)state_count);
state_score_stats = test_compare_float_bits(
ref_state_score_host, fused_state_score_host, (size_t)state_count);
}
test_restore_env(disable_env, saved_disable);
test_restore_env(pack_disable_env, saved_pack_disable);
fprintf(stderr,
"ds4-test: compressor %s exact head=%u ratio=%u pos=%u "
"tokens=%u ape=%s comp=%zu/%llu state_kv=%zu/%llu "
"state_score=%zu/%llu max_ulp=%u/%u/%u\n",
test_pack_fusion ? "ratio4 pack" : "APE add",
head_dim, ratio, pos0, n_tokens, ape_type == 1u ? "f16" : "f32",
comp_stats.mismatch_count, (unsigned long long)comp_count,
state_kv_stats.mismatch_count, (unsigned long long)state_count,
state_score_stats.mismatch_count, (unsigned long long)state_count,
comp_stats.max_ulp, state_kv_stats.max_ulp,
state_score_stats.max_ulp);
TEST_ASSERT(comp_stats.mismatch_count == 0);
TEST_ASSERT(state_kv_stats.mismatch_count == 0);
TEST_ASSERT(state_score_stats.mismatch_count == 0);
free(model_raw);
free(poison_host);
free(fused_state_score_host);
free(fused_state_kv_host);
free(ref_state_score_host);
free(ref_state_kv_host);
free(fused_comp_host);
free(ref_comp_host);
free(sc_host);
free(kv_host);
ds4_gpu_tensor_free(fused_state_score);
ds4_gpu_tensor_free(fused_state_kv);
ds4_gpu_tensor_free(ref_state_score);
ds4_gpu_tensor_free(ref_state_kv);
ds4_gpu_tensor_free(fused_comp);
ds4_gpu_tensor_free(ref_comp);
ds4_gpu_tensor_free(sc);
ds4_gpu_tensor_free(kv);
}
static void test_metal_compressor_ape_add_exact(void) {
test_metal_compressor_ape_add_exact_case(128, 4, 3, 16, 1, 7, false);
test_metal_compressor_ape_add_exact_case(512, 4, 0, 16, 1, 13, false);
test_metal_compressor_ape_add_exact_case(512, 4, 1, 14, 0, 19, false);
test_metal_compressor_ape_add_exact_case(512, 128, 127, 257, 1, 31, false);
}
static void test_metal_compressor_ratio4_pack_exact(void) {
test_metal_compressor_ape_add_exact_case(128, 4, 0, 4, 1, 37, true);
test_metal_compressor_ape_add_exact_case(512, 4, 0, 8, 1, 41, true);
test_metal_compressor_ape_add_exact_case(512, 4, 1, 14, 0, 43, true);
}
static void test_metal_compressor_ratio4_replay_pack_exact_case(
uint32_t head_dim,
uint32_t n_tokens,
uint32_t seed) {
const uint32_t width = 2u * head_dim;
const uint32_t n_comp = n_tokens / 4u;
const uint64_t input_count = (uint64_t)n_tokens * width;
const uint64_t state_count = (uint64_t)8u * width;
const uint64_t comp_count = (uint64_t)n_comp * head_dim;
const uint64_t input_bytes = input_count * sizeof(float);
const uint64_t state_bytes = state_count * sizeof(float);
const uint64_t comp_bytes = comp_count * sizeof(float);
const uint64_t ape_bytes = (uint64_t)4u * width * sizeof(uint16_t);
const uint64_t page = (uint64_t)getpagesize();
const uint64_t norm_offset = test_round_up_u64(ape_bytes, page);
const uint64_t model_bytes = test_round_up_u64(
norm_offset + (uint64_t)head_dim * sizeof(float), page);
ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(input_bytes);
ds4_gpu_tensor *sc = ds4_gpu_tensor_alloc(input_bytes);
ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *fused_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *fused_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *fused_state_score = ds4_gpu_tensor_alloc(state_bytes);
float *kv_host = malloc((size_t)input_bytes);
float *sc_host = malloc((size_t)input_bytes);
float *state_kv_host = malloc((size_t)state_bytes);
float *state_score_host = malloc((size_t)state_bytes);
float *source_after_host = malloc((size_t)input_bytes);
float *ref_comp_host = malloc((size_t)comp_bytes);
float *fused_comp_host = malloc((size_t)comp_bytes);
float *ref_state_kv_host = malloc((size_t)state_bytes);
float *fused_state_kv_host = malloc((size_t)state_bytes);
float *ref_state_score_host = malloc((size_t)state_bytes);
float *fused_state_score_host = malloc((size_t)state_bytes);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(
&model_raw, (size_t)page, (size_t)model_bytes) == 0);
TEST_ASSERT(kv != NULL);
TEST_ASSERT(sc != NULL);
TEST_ASSERT(ref_comp != NULL);
TEST_ASSERT(fused_comp != NULL);
TEST_ASSERT(ref_state_kv != NULL);
TEST_ASSERT(ref_state_score != NULL);
TEST_ASSERT(fused_state_kv != NULL);
TEST_ASSERT(fused_state_score != NULL);
TEST_ASSERT(kv_host != NULL);
TEST_ASSERT(sc_host != NULL);
TEST_ASSERT(state_kv_host != NULL);
TEST_ASSERT(state_score_host != NULL);
TEST_ASSERT(source_after_host != NULL);
TEST_ASSERT(ref_comp_host != NULL);
TEST_ASSERT(fused_comp_host != NULL);
TEST_ASSERT(ref_state_kv_host != NULL);
TEST_ASSERT(fused_state_kv_host != NULL);
TEST_ASSERT(ref_state_score_host != NULL);
TEST_ASSERT(fused_state_score_host != NULL);
TEST_ASSERT(model_raw != NULL);
const char *ape_disable_env = "DS4_METAL_DISABLE_COMPRESSOR_APE_ADD";
const char *pack_disable_env =
"DS4_METAL_DISABLE_COMPRESSOR_RATIO4_PACK_FUSION";
char *saved_ape_disable = test_save_env(ape_disable_env);
char *saved_pack_disable = test_save_env(pack_disable_env);
test_float_compare_stats comp_stats = {0};
test_float_compare_stats state_kv_stats = {0};
test_float_compare_stats state_score_stats = {0};
const bool allocated = kv && sc && ref_comp && fused_comp &&
ref_state_kv && ref_state_score && fused_state_kv &&
fused_state_score && kv_host && sc_host && state_kv_host &&
state_score_host && source_after_host && ref_comp_host && fused_comp_host &&
ref_state_kv_host && fused_state_kv_host && ref_state_score_host &&
fused_state_score_host && model_raw;
if (allocated) {
memset(model_raw, 0, (size_t)model_bytes);
uint16_t *ape = model_raw;
for (uint64_t i = 0; i < (uint64_t)4u * width; i++) {
const int value =
(int)((i * 19u + (i ^ (i >> 3u)) * 7u + seed * 11u) %
113u) - 56;
ape[i] = test_float_to_f16((float)value / 72.0f);
}
float *norm = (float *)((uint8_t *)model_raw + norm_offset);
for (uint32_t i = 0; i < head_dim; i++) {
norm[i] = 0.875f + (float)((i * 5u + seed) % 17u) / 64.0f;
}
for (uint64_t i = 0; i < input_count; i++) {
const int kv_value =
(int)((i * 31u + (i ^ (i >> 4u)) * 5u + seed * 13u) %
193u) - 96;
const int sc_value =
(int)((i * 37u + (i ^ (i >> 5u)) * 9u + seed * 17u) %
181u) - 90;
kv_host[i] = (float)kv_value / 104.0f;
sc_host[i] = (float)sc_value / 88.0f;
}
for (uint64_t i = 0; i < state_count; i++) {
const int kv_value =
(int)((i * 23u + (i ^ (i >> 2u)) * 3u + seed * 29u) %
167u) - 83;
const int sc_value =
(int)((i * 41u + (i ^ (i >> 6u)) * 11u + seed * 7u) %
157u) - 78;
state_kv_host[i] = (float)kv_value / 80.0f;
state_score_host[i] = (float)sc_value / 92.0f;
}
const uint32_t negative_zero = 0x80000000u;
memcpy(kv_host, &negative_zero, sizeof(negative_zero));
memcpy(state_kv_host, &negative_zero, sizeof(negative_zero));
TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_kv, 0, state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_score, 0, state_score_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, state_score_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(unsetenv(ape_disable_env) == 0);
TEST_ASSERT(setenv(pack_disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor(
ref_comp, ref_state_kv, ref_state_score, kv, sc,
model_raw, model_bytes, 0, 1, norm_offset, 0,
head_dim, 0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0);
// Poison every persistent pack cell through the legacy full-fill path
// so a missing candidate write cannot inherit the reference value.
for (uint64_t i = 0; i < input_count; i++) {
const uint32_t bits = 0x7fc00001u + (uint32_t)(i & 0x3ffu);
memcpy(kv_host + i, &bits, sizeof(bits));
memcpy(sc_host + i, &bits, sizeof(bits));
}
for (uint64_t i = 0; i < state_count; i++) {
const uint32_t bits = 0x7fc00401u + (uint32_t)(i & 0x3ffu);
memcpy(ref_state_kv_host + i, &bits, sizeof(bits));
memcpy(ref_state_score_host + i, &bits, sizeof(bits));
}
TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, ref_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, ref_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor(
fused_comp, fused_state_kv, fused_state_score, kv, sc,
model_raw, model_bytes, 0, 1, norm_offset, 0,
head_dim, 0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0);
// Restore finite sources and the candidate's original replay state.
for (uint64_t i = 0; i < input_count; i++) {
const int kv_value =
(int)((i * 31u + (i ^ (i >> 4u)) * 5u + seed * 13u) %
193u) - 96;
const int sc_value =
(int)((i * 37u + (i ^ (i >> 5u)) * 9u + seed * 17u) %
181u) - 90;
kv_host[i] = (float)kv_value / 104.0f;
sc_host[i] = (float)sc_value / 88.0f;
}
memcpy(kv_host, &negative_zero, sizeof(negative_zero));
TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_kv, 0, state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_state_score, 0, state_score_host,
state_bytes) != 0);
TEST_ASSERT(unsetenv(pack_disable_env) == 0);
TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor(
fused_comp, fused_state_kv, fused_state_score, kv, sc,
model_raw, model_bytes, 0, 1, norm_offset, 0,
head_dim, 0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
kv, 0, source_after_host, input_bytes) != 0);
TEST_ASSERT(memcmp(kv_host, source_after_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(ds4_gpu_tensor_read(
sc, 0, source_after_host, input_bytes) != 0);
TEST_ASSERT(memcmp(sc_host, source_after_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_comp, 0, ref_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_comp, 0, fused_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_state_kv, 0, fused_state_kv_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_score, 0, ref_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_state_score, 0, fused_state_score_host,
state_bytes) != 0);
comp_stats = test_compare_float_bits(
ref_comp_host, fused_comp_host, (size_t)comp_count);
state_kv_stats = test_compare_float_bits(
ref_state_kv_host, fused_state_kv_host, (size_t)state_count);
state_score_stats = test_compare_float_bits(
ref_state_score_host, fused_state_score_host,
(size_t)state_count);
}
test_restore_env(ape_disable_env, saved_ape_disable);
test_restore_env(pack_disable_env, saved_pack_disable);
fprintf(stderr,
"ds4-test: compressor ratio4 replay pack exact head=%u "
"tokens=%u comp=%zu/%llu state_kv=%zu/%llu "
"state_score=%zu/%llu max_ulp=%u/%u/%u\n",
head_dim, n_tokens,
comp_stats.mismatch_count, (unsigned long long)comp_count,
state_kv_stats.mismatch_count, (unsigned long long)state_count,
state_score_stats.mismatch_count, (unsigned long long)state_count,
comp_stats.max_ulp, state_kv_stats.max_ulp,
state_score_stats.max_ulp);
TEST_ASSERT(comp_stats.mismatch_count == 0);
TEST_ASSERT(state_kv_stats.mismatch_count == 0);
TEST_ASSERT(state_score_stats.mismatch_count == 0);
free(model_raw);
free(fused_state_score_host);
free(ref_state_score_host);
free(fused_state_kv_host);
free(ref_state_kv_host);
free(fused_comp_host);
free(ref_comp_host);
free(source_after_host);
free(state_score_host);
free(state_kv_host);
free(sc_host);
free(kv_host);
ds4_gpu_tensor_free(fused_state_score);
ds4_gpu_tensor_free(fused_state_kv);
ds4_gpu_tensor_free(ref_state_score);
ds4_gpu_tensor_free(ref_state_kv);
ds4_gpu_tensor_free(fused_comp);
ds4_gpu_tensor_free(ref_comp);
ds4_gpu_tensor_free(sc);
ds4_gpu_tensor_free(kv);
}
static void test_metal_compressor_ratio4_replay_pack_exact(void) {
test_metal_compressor_ratio4_replay_pack_exact_case(128, 4, 47);
test_metal_compressor_ratio4_replay_pack_exact_case(512, 8, 53);
}
static void test_metal_compressor_ratio4_direct_pool_exact_case(
uint32_t head_dim,
uint32_t pos0,
uint32_t n_tokens,
uint32_t ape_type,
bool replay,
uint32_t seed) {
const uint32_t ratio = 4u;
const uint32_t width = 2u * head_dim;
const uint32_t state_rows = 8u;
const uint32_t n_comp = n_tokens / ratio;
const uint64_t input_count = (uint64_t)n_tokens * width;
const uint64_t state_count = (uint64_t)state_rows * width;
const uint64_t comp_count = (uint64_t)n_comp * head_dim;
const uint64_t input_bytes = input_count * sizeof(float);
const uint64_t state_bytes = state_count * sizeof(float);
const uint64_t comp_bytes = comp_count * sizeof(float);
const uint64_t ape_elem_bytes = ape_type == 1u ? 2u : 4u;
const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem_bytes;
const uint64_t page = (uint64_t)getpagesize();
const uint64_t norm_offset = test_round_up_u64(ape_bytes, page);
const uint64_t model_bytes = test_round_up_u64(
norm_offset + (uint64_t)head_dim * sizeof(float), page);
TEST_ASSERT(n_comp != 0);
TEST_ASSERT(!replay || ((pos0 & 3u) == 0u && (n_tokens & 3u) == 0u));
ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(input_bytes);
ds4_gpu_tensor *sc = ds4_gpu_tensor_alloc(input_bytes);
ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *direct_comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *direct_state_kv = ds4_gpu_tensor_alloc(state_bytes);
ds4_gpu_tensor *direct_state_score = ds4_gpu_tensor_alloc(state_bytes);
float *kv_host = malloc((size_t)input_bytes);
float *sc_host = malloc((size_t)input_bytes);
float *source_after_host = malloc((size_t)input_bytes);
float *state_kv_host = malloc((size_t)state_bytes);
float *state_score_host = malloc((size_t)state_bytes);
float *ref_comp_host = malloc((size_t)comp_bytes);
float *direct_comp_host = malloc((size_t)comp_bytes);
float *ref_state_kv_host = malloc((size_t)state_bytes);
float *direct_state_kv_host = malloc((size_t)state_bytes);
float *ref_state_score_host = malloc((size_t)state_bytes);
float *direct_state_score_host = malloc((size_t)state_bytes);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(
&model_raw, (size_t)page, (size_t)model_bytes) == 0);
TEST_ASSERT(kv != NULL);
TEST_ASSERT(sc != NULL);
TEST_ASSERT(ref_comp != NULL);
TEST_ASSERT(direct_comp != NULL);
TEST_ASSERT(ref_state_kv != NULL);
TEST_ASSERT(ref_state_score != NULL);
TEST_ASSERT(direct_state_kv != NULL);
TEST_ASSERT(direct_state_score != NULL);
TEST_ASSERT(kv_host != NULL);
TEST_ASSERT(sc_host != NULL);
TEST_ASSERT(source_after_host != NULL);
TEST_ASSERT(state_kv_host != NULL);
TEST_ASSERT(state_score_host != NULL);
TEST_ASSERT(ref_comp_host != NULL);
TEST_ASSERT(direct_comp_host != NULL);
TEST_ASSERT(ref_state_kv_host != NULL);
TEST_ASSERT(direct_state_kv_host != NULL);
TEST_ASSERT(ref_state_score_host != NULL);
TEST_ASSERT(direct_state_score_host != NULL);
TEST_ASSERT(model_raw != NULL);
const char *ape_disable_env = "DS4_METAL_DISABLE_COMPRESSOR_APE_ADD";
const char *pack_disable_env =
"DS4_METAL_DISABLE_COMPRESSOR_RATIO4_PACK_FUSION";
const char *direct_disable_env =
"DS4_METAL_DISABLE_COMPRESSOR_RATIO4_DIRECT_POOL";
char *saved_ape_disable = test_save_env(ape_disable_env);
char *saved_pack_disable = test_save_env(pack_disable_env);
char *saved_direct_disable = test_save_env(direct_disable_env);
test_float_compare_stats comp_stats = {0};
test_float_compare_stats state_kv_stats = {0};
test_float_compare_stats state_score_stats = {0};
const bool allocated = kv && sc && ref_comp && direct_comp &&
ref_state_kv && ref_state_score && direct_state_kv &&
direct_state_score && kv_host && sc_host && source_after_host &&
state_kv_host && state_score_host && ref_comp_host &&
direct_comp_host && ref_state_kv_host && direct_state_kv_host &&
ref_state_score_host && direct_state_score_host && model_raw;
if (allocated) {
memset(model_raw, 0, (size_t)model_bytes);
if (ape_type == 1u) {
uint16_t *ape = model_raw;
for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) {
const int value =
(int)((i * 17u + (i ^ (i >> 4u)) * 5u + seed * 13u) %
127u) - 63;
ape[i] = test_float_to_f16((float)value / 80.0f);
}
} else {
float *ape = model_raw;
for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) {
const int value =
(int)((i * 17u + (i ^ (i >> 4u)) * 5u + seed * 13u) %
127u) - 63;
ape[i] = (float)value / 80.0f;
}
}
float *norm = (float *)((uint8_t *)model_raw + norm_offset);
for (uint32_t i = 0; i < head_dim; i++) {
norm[i] = 0.75f + (float)((i * 7u + seed) % 19u) / 64.0f;
}
for (uint64_t i = 0; i < input_count; i++) {
const uint32_t token = (uint32_t)(i / width);
const uint32_t col = (uint32_t)(i - (uint64_t)token * width);
const int kv_value =
(int)(((uint64_t)token * 37u + (uint64_t)col * 29u +
(col >= head_dim ? 71u : 3u) + seed * 23u) %
193u) - 96;
const int sc_value =
(int)(((uint64_t)token * 41u + (uint64_t)col * 31u +
(col >= head_dim ? 17u : 83u) + seed * 11u) %
181u) - 90;
kv_host[i] = (float)kv_value / 104.0f;
sc_host[i] = (float)sc_value / 32.0f;
}
const uint32_t negative_zero = 0x80000000u;
memcpy(kv_host, &negative_zero, sizeof(negative_zero));
memcpy(sc_host + width + head_dim, &negative_zero,
sizeof(negative_zero));
for (uint32_t row = 0; row < state_rows; row++) {
for (uint32_t col = 0; col < width; col++) {
const uint64_t i = (uint64_t)row * width + col;
if (replay && row < ratio && col < head_dim) {
const int kv_value =
(int)(((uint64_t)row * 43u + (uint64_t)col * 19u +
seed * 29u) % 167u) - 83;
const int sc_value =
(int)(((uint64_t)row * 47u + (uint64_t)col * 23u +
seed * 7u) % 157u) - 78;
state_kv_host[i] = (float)kv_value / 80.0f;
state_score_host[i] = (float)sc_value / 28.0f;
} else {
const uint32_t kv_bits =
0x7fc00001u + (uint32_t)(i & 0x3ffu);
const uint32_t score_bits =
0x7fc00401u + (uint32_t)(i & 0x3ffu);
memcpy(state_kv_host + i, &kv_bits, sizeof(kv_bits));
memcpy(state_score_host + i, &score_bits,
sizeof(score_bits));
}
}
}
for (uint64_t i = 0; i < comp_count; i++) {
const uint32_t bits = 0x7fc00801u + (uint32_t)(i & 0x3ffu);
memcpy(ref_comp_host + i, &bits, sizeof(bits));
memcpy(direct_comp_host + i, &bits, sizeof(bits));
}
TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_comp, 0, ref_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
direct_comp, 0, direct_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_kv, 0, state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
direct_state_kv, 0, state_kv_host, state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
ref_state_score, 0, state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
direct_state_score, 0, state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(unsetenv(ape_disable_env) == 0);
TEST_ASSERT(setenv(direct_disable_env, "1", 1) == 0);
TEST_ASSERT(unsetenv(pack_disable_env) == 0);
int ref_ok;
if (replay) {
ref_ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor(
ref_comp, ref_state_kv, ref_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, pos0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f);
} else {
ref_ok = ds4_gpu_compressor_prefill_tensor(
ref_comp, ref_state_kv, ref_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, ratio, pos0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f);
}
TEST_ASSERT(ref_ok != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
kv, 0, source_after_host, input_bytes) != 0);
TEST_ASSERT(memcmp(kv_host, source_after_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(ds4_gpu_tensor_read(
sc, 0, source_after_host, input_bytes) != 0);
TEST_ASSERT(memcmp(sc_host, source_after_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(unsetenv(direct_disable_env) == 0);
TEST_ASSERT(setenv(pack_disable_env, "1", 1) == 0);
int direct_ok;
if (replay) {
direct_ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor(
direct_comp, direct_state_kv, direct_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, pos0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f);
} else {
direct_ok = ds4_gpu_compressor_prefill_tensor(
direct_comp, direct_state_kv, direct_state_score, kv, sc,
model_raw, model_bytes, 0, ape_type, norm_offset, 0,
head_dim, ratio, pos0, n_tokens, 0, 0, false,
10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f);
}
TEST_ASSERT(direct_ok != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
kv, 0, source_after_host, input_bytes) != 0);
TEST_ASSERT(memcmp(kv_host, source_after_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(ds4_gpu_tensor_read(
sc, 0, source_after_host, input_bytes) != 0);
TEST_ASSERT(memcmp(sc_host, source_after_host,
(size_t)input_bytes) == 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_comp, 0, ref_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
direct_comp, 0, direct_comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_kv, 0, ref_state_kv_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
direct_state_kv, 0, direct_state_kv_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_state_score, 0, ref_state_score_host,
state_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
direct_state_score, 0, direct_state_score_host,
state_bytes) != 0);
comp_stats = test_compare_float_bits(
ref_comp_host, direct_comp_host, (size_t)comp_count);
state_kv_stats = test_compare_float_bits(
ref_state_kv_host, direct_state_kv_host, (size_t)state_count);
state_score_stats = test_compare_float_bits(
ref_state_score_host, direct_state_score_host,
(size_t)state_count);
}
test_restore_env(ape_disable_env, saved_ape_disable);
test_restore_env(pack_disable_env, saved_pack_disable);
test_restore_env(direct_disable_env, saved_direct_disable);
fprintf(stderr,
"ds4-test: compressor ratio4 direct pool exact mode=%s "
"head=%u pos=%u tokens=%u comp_rows=%u ape=%s "
"comp=%zu/%llu state_kv=%zu/%llu state_score=%zu/%llu "
"max_ulp=%u/%u/%u\n",
replay ? "replay" : "prefill",
head_dim, pos0, n_tokens, n_comp,
ape_type == 1u ? "f16" : "f32",
comp_stats.mismatch_count, (unsigned long long)comp_count,
state_kv_stats.mismatch_count, (unsigned long long)state_count,
state_score_stats.mismatch_count,
(unsigned long long)state_count,
comp_stats.max_ulp, state_kv_stats.max_ulp,
state_score_stats.max_ulp);
TEST_ASSERT(comp_stats.mismatch_count == 0);
TEST_ASSERT(state_kv_stats.mismatch_count == 0);
TEST_ASSERT(state_score_stats.mismatch_count == 0);
free(model_raw);
free(direct_state_score_host);
free(ref_state_score_host);
free(direct_state_kv_host);
free(ref_state_kv_host);
free(direct_comp_host);
free(ref_comp_host);
free(state_score_host);
free(state_kv_host);
free(source_after_host);
free(sc_host);
free(kv_host);
ds4_gpu_tensor_free(direct_state_score);
ds4_gpu_tensor_free(direct_state_kv);
ds4_gpu_tensor_free(ref_state_score);
ds4_gpu_tensor_free(ref_state_kv);
ds4_gpu_tensor_free(direct_comp);
ds4_gpu_tensor_free(ref_comp);
ds4_gpu_tensor_free(sc);
ds4_gpu_tensor_free(kv);
}
static void test_metal_compressor_ratio4_direct_pool_exact(void) {
/* n_comp == 1 deliberately stays on the exact GGML reduction path. */
test_metal_compressor_ratio4_direct_pool_exact_case(
512, 0, 4, 1, false, 59);
test_metal_compressor_ratio4_direct_pool_exact_case(
512, 1, 16, 0, false, 61);
test_metal_compressor_ratio4_direct_pool_exact_case(
128, 3, 14, 1, false, 67);
test_metal_compressor_ratio4_direct_pool_exact_case(
512, 2048, 8, 1, true, 71);
test_metal_compressor_ratio4_direct_pool_exact_case(
128, 12, 12, 0, true, 73);
test_metal_compressor_ratio4_direct_pool_exact_case(
128, 8, 4, 1, true, 79);
}
static void test_metal_inplace_rope_pair_exact(void) {
typedef struct {
uint32_t head_dim;
uint32_t n_rot;
uint32_t n_head;
uint32_t n_tok;
uint32_t pos0;
bool inverse;
float ext_factor;
} rope_case;
static const rope_case cases[] = {
{ 512, 64, 64, 1, UINT32_MAX, false, 1.0f },
{ 512, 64, 1, 1, 2047, false, 1.0f },
{ 512, 64, 64, 1, 65533, true, 1.0f },
{ 128, 64, 64, 1, 37, false, 0.0f },
{ 512, 64, 4, 32, 0, false, 0.0f },
{ 512, 64, 7, 33, 2047, true, 1.0f },
{ 128, 64, 64, 34, 65533, false, 1.0f },
{ 128, 64, 64, 35, 37, true, 0.0f },
{ 128, 64, 4, 35, UINT32_MAX - 16u, true, 1.0f },
};
const char *disable_env = "DS4_METAL_DISABLE_INPLACE_ROPE_PAIR";
const char *shared_disable_env =
"DS4_METAL_DISABLE_SHARED_ROPE_COEFF";
const char *affine_disable_env =
"DS4_METAL_DISABLE_AFFINE_ROPE_PAIR";
char *saved_disable = test_save_env(disable_env);
char *saved_shared_disable = test_save_env(shared_disable_env);
char *saved_affine_disable = test_save_env(affine_disable_env);
size_t total_pair_mismatch = 0;
size_t total_shared_mismatch = 0;
size_t total_affine_mismatch = 0;
size_t total_pair_prefix_mismatch = 0;
size_t total_pair_tail_mismatch = 0;
size_t total_shared_prefix_mismatch = 0;
size_t total_shared_tail_mismatch = 0;
size_t total_affine_prefix_mismatch = 0;
size_t total_affine_tail_mismatch = 0;
size_t total_elements = 0;
for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) {
const rope_case *c = &cases[ci];
const size_t elements =
(size_t)c->n_tok * c->n_head * c->head_dim;
const uint64_t bytes = (uint64_t)elements * sizeof(float);
const uint32_t n_nope = c->head_dim - c->n_rot;
const float freq_base = c->ext_factor != 0.0f ? 160000.0f : 10000.0f;
const float freq_scale = c->ext_factor != 0.0f ? 1.0f / 16.0f : 1.0f;
const uint32_t n_ctx_orig = c->ext_factor != 0.0f ? 65536u : 0u;
float attn_factor = 1.0f;
if (c->ext_factor != 0.0f) {
attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(bytes);
ds4_gpu_tensor *pair_candidate = ds4_gpu_tensor_alloc(bytes);
ds4_gpu_tensor *shared_candidate = ds4_gpu_tensor_alloc(bytes);
ds4_gpu_tensor *affine_candidate = ds4_gpu_tensor_alloc(bytes);
float *input = malloc((size_t)bytes);
float *reference_host = malloc((size_t)bytes);
float *pair_host = malloc((size_t)bytes);
float *shared_host = malloc((size_t)bytes);
float *affine_host = malloc((size_t)bytes);
TEST_ASSERT(reference != NULL);
TEST_ASSERT(pair_candidate != NULL);
TEST_ASSERT(shared_candidate != NULL);
TEST_ASSERT(affine_candidate != NULL);
TEST_ASSERT(input != NULL);
TEST_ASSERT(reference_host != NULL);
TEST_ASSERT(pair_host != NULL);
TEST_ASSERT(shared_host != NULL);
TEST_ASSERT(affine_host != NULL);
const bool allocated = reference && pair_candidate &&
shared_candidate && affine_candidate && input && reference_host &&
pair_host && shared_host && affine_host;
if (allocated) {
for (size_t i = 0; i < elements; i++) {
const uint32_t key =
(uint32_t)(i * 37u + (i ^ (i >> 5u)) * 11u + ci * 101u);
const int value = (int)(key % 4093u) - 2046;
input[i] = (float)value / 1024.0f;
if ((i + ci * 17u) % 257u == 0u) {
const uint32_t negative_zero = 0x80000000u;
memcpy(&input[i], &negative_zero, sizeof(negative_zero));
}
}
TEST_ASSERT(ds4_gpu_tensor_write(reference, 0, input, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
pair_candidate, 0, input, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
shared_candidate, 0, input, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
affine_candidate, 0, input, bytes) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(setenv(affine_disable_env, "1", 1) == 0);
TEST_ASSERT(setenv(disable_env, "1", 1) == 0);
TEST_ASSERT(setenv(shared_disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_rope_tail_tensor(
reference,
c->n_tok,
c->n_head,
c->head_dim,
c->n_rot,
c->pos0,
n_ctx_orig,
c->inverse,
freq_base,
freq_scale,
c->ext_factor,
attn_factor,
32.0f,
1.0f) != 0);
TEST_ASSERT(unsetenv(disable_env) == 0);
TEST_ASSERT(ds4_gpu_rope_tail_tensor(
pair_candidate,
c->n_tok,
c->n_head,
c->head_dim,
c->n_rot,
c->pos0,
n_ctx_orig,
c->inverse,
freq_base,
freq_scale,
c->ext_factor,
attn_factor,
32.0f,
1.0f) != 0);
TEST_ASSERT(unsetenv(shared_disable_env) == 0);
TEST_ASSERT(ds4_gpu_rope_tail_tensor(
shared_candidate,
c->n_tok,
c->n_head,
c->head_dim,
c->n_rot,
c->pos0,
n_ctx_orig,
c->inverse,
freq_base,
freq_scale,
c->ext_factor,
attn_factor,
32.0f,
1.0f) != 0);
TEST_ASSERT(setenv(shared_disable_env, "1", 1) == 0);
TEST_ASSERT(unsetenv(affine_disable_env) == 0);
TEST_ASSERT(ds4_gpu_rope_tail_tensor(
affine_candidate,
c->n_tok,
c->n_head,
c->head_dim,
c->n_rot,
c->pos0,
n_ctx_orig,
c->inverse,
freq_base,
freq_scale,
c->ext_factor,
attn_factor,
32.0f,
1.0f) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
reference, 0, reference_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
pair_candidate, 0, pair_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
shared_candidate, 0, shared_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
affine_candidate, 0, affine_host, bytes) != 0);
const test_float_compare_stats pair_stats =
test_compare_float_bits(
reference_host, pair_host, elements);
const test_float_compare_stats shared_stats =
test_compare_float_bits(
pair_host, shared_host, elements);
const test_float_compare_stats affine_stats =
test_compare_float_bits(
pair_host, affine_host, elements);
size_t pair_prefix_mismatch = 0;
size_t pair_tail_mismatch = 0;
size_t shared_prefix_mismatch = 0;
size_t shared_tail_mismatch = 0;
size_t affine_prefix_mismatch = 0;
size_t affine_tail_mismatch = 0;
for (uint32_t t = 0; t < c->n_tok; t++) {
for (uint32_t h = 0; h < c->n_head; h++) {
const size_t row =
((size_t)t * c->n_head + h) * c->head_dim;
for (uint32_t d = 0; d < n_nope; d++) {
if (memcmp(&input[row + d],
&pair_host[row + d],
sizeof(float)) != 0) {
pair_prefix_mismatch++;
}
if (memcmp(&input[row + d],
&shared_host[row + d],
sizeof(float)) != 0) {
shared_prefix_mismatch++;
}
if (memcmp(&input[row + d],
&affine_host[row + d],
sizeof(float)) != 0) {
affine_prefix_mismatch++;
}
}
for (uint32_t d = n_nope; d < c->head_dim; d++) {
if (memcmp(&reference_host[row + d],
&pair_host[row + d],
sizeof(float)) != 0) {
pair_tail_mismatch++;
}
if (memcmp(&pair_host[row + d],
&shared_host[row + d],
sizeof(float)) != 0) {
shared_tail_mismatch++;
}
if (memcmp(&pair_host[row + d],
&affine_host[row + d],
sizeof(float)) != 0) {
affine_tail_mismatch++;
}
}
}
}
fprintf(stderr,
"ds4-test: in-place RoPE exactness case=%zu "
"shape=%ux%ux%u pos=%u inverse=%d ext=%g "
"pair=%zu/%zu shared=%zu/%zu affine=%zu/%zu "
"pair_prefix=%zu pair_tail=%zu "
"shared_prefix=%zu shared_tail=%zu "
"affine_prefix=%zu affine_tail=%zu "
"pair_max_ulp=%u shared_max_ulp=%u affine_max_ulp=%u "
"pair_max_abs=%g shared_max_abs=%g affine_max_abs=%g\n",
ci,
c->n_tok,
c->n_head,
c->head_dim,
c->pos0,
c->inverse ? 1 : 0,
c->ext_factor,
pair_stats.mismatch_count,
elements,
shared_stats.mismatch_count,
elements,
affine_stats.mismatch_count,
elements,
pair_prefix_mismatch,
pair_tail_mismatch,
shared_prefix_mismatch,
shared_tail_mismatch,
affine_prefix_mismatch,
affine_tail_mismatch,
pair_stats.max_ulp,
shared_stats.max_ulp,
affine_stats.max_ulp,
pair_stats.max_abs,
shared_stats.max_abs,
affine_stats.max_abs);
TEST_ASSERT(pair_stats.mismatch_count == 0);
TEST_ASSERT(shared_stats.mismatch_count == 0);
TEST_ASSERT(affine_stats.mismatch_count == 0);
TEST_ASSERT(pair_prefix_mismatch == 0);
TEST_ASSERT(pair_tail_mismatch == 0);
TEST_ASSERT(shared_prefix_mismatch == 0);
TEST_ASSERT(shared_tail_mismatch == 0);
TEST_ASSERT(affine_prefix_mismatch == 0);
TEST_ASSERT(affine_tail_mismatch == 0);
total_pair_mismatch += pair_stats.mismatch_count;
total_shared_mismatch += shared_stats.mismatch_count;
total_affine_mismatch += affine_stats.mismatch_count;
total_pair_prefix_mismatch += pair_prefix_mismatch;
total_pair_tail_mismatch += pair_tail_mismatch;
total_shared_prefix_mismatch += shared_prefix_mismatch;
total_shared_tail_mismatch += shared_tail_mismatch;
total_affine_prefix_mismatch += affine_prefix_mismatch;
total_affine_tail_mismatch += affine_tail_mismatch;
total_elements += elements;
}
free(affine_host);
free(shared_host);
free(pair_host);
free(reference_host);
free(input);
ds4_gpu_tensor_free(affine_candidate);
ds4_gpu_tensor_free(shared_candidate);
ds4_gpu_tensor_free(pair_candidate);
ds4_gpu_tensor_free(reference);
}
test_restore_env(disable_env, saved_disable);
test_restore_env(shared_disable_env, saved_shared_disable);
test_restore_env(affine_disable_env, saved_affine_disable);
fprintf(stderr,
"ds4-test: in-place RoPE total pair=%zu/%zu shared=%zu/%zu "
"affine=%zu/%zu "
"pair_prefix=%zu pair_tail=%zu "
"shared_prefix=%zu shared_tail=%zu "
"affine_prefix=%zu affine_tail=%zu\n",
total_pair_mismatch,
total_elements,
total_shared_mismatch,
total_elements,
total_affine_mismatch,
total_elements,
total_pair_prefix_mismatch,
total_pair_tail_mismatch,
total_shared_prefix_mismatch,
total_shared_tail_mismatch,
total_affine_prefix_mismatch,
total_affine_tail_mismatch);
TEST_ASSERT(total_pair_mismatch == 0);
TEST_ASSERT(total_shared_mismatch == 0);
TEST_ASSERT(total_affine_mismatch == 0);
TEST_ASSERT(total_pair_prefix_mismatch == 0);
TEST_ASSERT(total_pair_tail_mismatch == 0);
TEST_ASSERT(total_shared_prefix_mismatch == 0);
TEST_ASSERT(total_shared_tail_mismatch == 0);
TEST_ASSERT(total_affine_prefix_mismatch == 0);
TEST_ASSERT(total_affine_tail_mismatch == 0);
}
static void test_metal_contiguous_f32_f16_roundtrip_exact(void) {
typedef struct {
uint32_t n;
uint32_t src_offset;
uint32_t dst_offset;
} copy_case;
static const copy_case cases[] = {
{ 1, 0, 0 },
{ 3, 4, 2 },
{ 4, 16, 8 },
{ 5, 12, 6 },
{ 17, 20, 10 },
{ 65, 4, 2 },
};
const char *env_name = "DS4_METAL_DISABLE_CONTIG_F32_F16_COPY";
char *saved_env = test_save_env(env_name);
size_t half_mismatch = 0;
size_t half_guard_mismatch = 0;
size_t roundtrip_mismatch = 0;
size_t roundtrip_guard_mismatch = 0;
size_t half_total = 0;
size_t roundtrip_total = 0;
for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) {
const uint32_t n = cases[ci].n;
const uint64_t src_bytes = cases[ci].src_offset +
(uint64_t)n * sizeof(float) + 16u;
const uint64_t half_bytes = cases[ci].dst_offset +
(uint64_t)n * sizeof(uint16_t) + 16u;
const uint32_t raw_cap = 3;
const uint32_t raw_row = 1;
const uint64_t raw_bytes =
(uint64_t)raw_cap * n * sizeof(float);
ds4_gpu_tensor *src_base = ds4_gpu_tensor_alloc(src_bytes);
ds4_gpu_tensor *src_view = src_base
? ds4_gpu_tensor_view(src_base,
cases[ci].src_offset,
(uint64_t)n * sizeof(float))
: NULL;
ds4_gpu_tensor *half_ref = ds4_gpu_tensor_alloc(half_bytes);
ds4_gpu_tensor *half_vec = ds4_gpu_tensor_alloc(half_bytes);
ds4_gpu_tensor *raw_ref = ds4_gpu_tensor_alloc(raw_bytes);
ds4_gpu_tensor *raw_vec = ds4_gpu_tensor_alloc(raw_bytes);
TEST_ASSERT(src_base != NULL);
TEST_ASSERT(src_view != NULL);
TEST_ASSERT(half_ref != NULL);
TEST_ASSERT(half_vec != NULL);
TEST_ASSERT(raw_ref != NULL);
TEST_ASSERT(raw_vec != NULL);
uint8_t *src_host = malloc((size_t)src_bytes);
uint16_t *half_init = malloc((size_t)half_bytes);
uint16_t *half_ref_host = malloc((size_t)half_bytes);
uint16_t *half_vec_host = malloc((size_t)half_bytes);
uint32_t *raw_init = malloc((size_t)raw_bytes);
uint32_t *raw_ref_host = malloc((size_t)raw_bytes);
uint32_t *raw_vec_host = malloc((size_t)raw_bytes);
TEST_ASSERT(src_host != NULL);
TEST_ASSERT(half_init != NULL);
TEST_ASSERT(half_ref_host != NULL);
TEST_ASSERT(half_vec_host != NULL);
TEST_ASSERT(raw_init != NULL);
TEST_ASSERT(raw_ref_host != NULL);
TEST_ASSERT(raw_vec_host != NULL);
const bool allocated = src_base && src_view && half_ref && half_vec &&
raw_ref && raw_vec && src_host && half_init && half_ref_host &&
half_vec_host && raw_init && raw_ref_host && raw_vec_host;
if (allocated) {
memset(src_host, 0x6d, (size_t)src_bytes);
test_fill_copy_f32_patterns(src_host + cases[ci].src_offset,
n,
(uint32_t)(ci * 7u));
const size_t half_words = (size_t)(half_bytes / sizeof(uint16_t));
for (size_t i = 0; i < half_words; i++) {
half_init[i] = (uint16_t)(0xa55au ^ (uint16_t)(i * 73u));
}
const size_t raw_words = (size_t)raw_cap * n;
for (size_t i = 0; i < raw_words; i++) {
raw_init[i] = 0x4a000000u + (uint32_t)i;
}
TEST_ASSERT(ds4_gpu_tensor_write(src_base, 0, src_host, src_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(half_ref, 0, half_init, half_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(half_vec, 0, half_init, half_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(raw_ref, 0, raw_init, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(raw_vec, 0, raw_init, raw_bytes) != 0);
TEST_ASSERT(setenv(env_name, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_tensor_copy_f32_to_f16(
half_ref,
cases[ci].dst_offset,
src_base,
cases[ci].src_offset,
n) != 0);
TEST_ASSERT(ds4_gpu_store_raw_kv_tensor(
raw_ref, src_view, raw_cap, raw_row, n) != 0);
TEST_ASSERT(setenv(env_name, "0", 1) == 0);
TEST_ASSERT(ds4_gpu_tensor_copy_f32_to_f16(
half_vec,
cases[ci].dst_offset,
src_base,
cases[ci].src_offset,
n) != 0);
TEST_ASSERT(ds4_gpu_store_raw_kv_tensor(
raw_vec, src_view, raw_cap, raw_row, n) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
half_ref, 0, half_ref_host, half_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
half_vec, 0, half_vec_host, half_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
raw_ref, 0, raw_ref_host, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
raw_vec, 0, raw_vec_host, raw_bytes) != 0);
const size_t half_first = cases[ci].dst_offset / sizeof(uint16_t);
const size_t half_last = half_first + n;
for (size_t i = 0; i < half_words; i++) {
if (half_ref_host[i] == half_vec_host[i]) continue;
if (i >= half_first && i < half_last) {
half_mismatch++;
} else {
half_guard_mismatch++;
}
}
const size_t raw_first = (size_t)raw_row * n;
const size_t raw_last = raw_first + n;
for (size_t i = 0; i < raw_words; i++) {
if (raw_ref_host[i] == raw_vec_host[i]) continue;
if (i >= raw_first && i < raw_last) {
roundtrip_mismatch++;
} else {
roundtrip_guard_mismatch++;
}
}
half_total += n;
roundtrip_total += n;
}
free(raw_vec_host);
free(raw_ref_host);
free(raw_init);
free(half_vec_host);
free(half_ref_host);
free(half_init);
free(src_host);
ds4_gpu_tensor_free(raw_vec);
ds4_gpu_tensor_free(raw_ref);
ds4_gpu_tensor_free(half_vec);
ds4_gpu_tensor_free(half_ref);
ds4_gpu_tensor_free(src_view);
ds4_gpu_tensor_free(src_base);
}
test_restore_env(env_name, saved_env);
fprintf(stderr,
"ds4-test: contiguous conversion exactness "
"f32_f16=%zu/%zu guard=%zu, f16_f32_roundtrip=%zu/%zu guard=%zu\n",
half_mismatch,
half_total,
half_guard_mismatch,
roundtrip_mismatch,
roundtrip_total,
roundtrip_guard_mismatch);
TEST_ASSERT(half_mismatch == 0);
TEST_ASSERT(half_guard_mismatch == 0);
TEST_ASSERT(roundtrip_mismatch == 0);
TEST_ASSERT(roundtrip_guard_mismatch == 0);
}
#endif
#if defined(__APPLE__)
static void test_metal_gathered_kv_stage_exact(void) {
const uint32_t head_dim = 512;
const uint32_t raw_cap = 7;
const uint32_t n_raw = 5;
const uint32_t n_comp = 3;
const uint32_t raw_starts[] = {0, 2, 5, 6};
const uint64_t raw_bytes =
(uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes =
(uint64_t)n_comp * head_dim * sizeof(uint16_t);
const uint64_t payload_bytes =
((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t);
const uint64_t raw_out_bytes =
(uint64_t)n_raw * head_dim * sizeof(uint16_t);
const uint64_t raw_view_offset = 4;
const uint64_t comp_view_offset = 2;
const uint64_t dst_view_offset = 6;
const uint64_t raw_base_bytes = raw_view_offset + raw_bytes + 12;
const uint64_t comp_base_bytes = comp_view_offset + comp_bytes + 14;
const uint64_t dst_base_bytes = dst_view_offset + payload_bytes + 10;
const char *envs[] = {
"DS4_METAL_REQUIRE_GATHERED_KV_STAGE",
"DS4_METAL_DISABLE_CONTIG_F32_F16_COPY",
"DS4_METAL_DISABLE_CONTIG_F16_F16_COPY",
};
char *saved[sizeof(envs)/sizeof(envs[0])];
for (size_t i = 0; i < sizeof(envs)/sizeof(envs[0]); i++) {
saved[i] = test_save_env(envs[i]);
}
ds4_gpu_tensor *raw_base = ds4_gpu_tensor_alloc(raw_base_bytes);
ds4_gpu_tensor *raw = raw_base
? ds4_gpu_tensor_view(raw_base, raw_view_offset, raw_bytes)
: NULL;
ds4_gpu_tensor *comp_base = ds4_gpu_tensor_alloc(comp_base_bytes);
ds4_gpu_tensor *comp = comp_base
? ds4_gpu_tensor_view(comp_base, comp_view_offset, comp_bytes)
: NULL;
ds4_gpu_tensor *ref_base = ds4_gpu_tensor_alloc(dst_base_bytes);
ds4_gpu_tensor *ref = ref_base
? ds4_gpu_tensor_view(ref_base, dst_view_offset, payload_bytes)
: NULL;
ds4_gpu_tensor *fused_base = ds4_gpu_tensor_alloc(dst_base_bytes);
ds4_gpu_tensor *fused = fused_base
? ds4_gpu_tensor_view(fused_base, dst_view_offset, payload_bytes)
: NULL;
TEST_ASSERT(raw_base != NULL);
TEST_ASSERT(raw != NULL);
TEST_ASSERT(comp_base != NULL);
TEST_ASSERT(comp != NULL);
TEST_ASSERT(ref_base != NULL);
TEST_ASSERT(ref != NULL);
TEST_ASSERT(fused_base != NULL);
TEST_ASSERT(fused != NULL);
float *raw_host = malloc((size_t)raw_bytes);
uint16_t *comp_host = malloc((size_t)comp_bytes);
uint8_t *dst_init = malloc((size_t)dst_base_bytes);
uint8_t *ref_host = malloc((size_t)dst_base_bytes);
uint8_t *fused_host = malloc((size_t)dst_base_bytes);
TEST_ASSERT(raw_host != NULL);
TEST_ASSERT(comp_host != NULL);
TEST_ASSERT(dst_init != NULL);
TEST_ASSERT(ref_host != NULL);
TEST_ASSERT(fused_host != NULL);
static const uint16_t half_patterns[] = {
0x0000u, 0x8000u, 0x0001u, 0x03ffu, 0x0400u,
0x3555u, 0x3c00u, 0x3c01u, 0x7bffu, 0xfbffu,
0x7c00u, 0xfc00u, 0x7e00u, 0x7e01u, 0xfe55u,
};
const bool allocated = raw_base && raw && comp_base && comp &&
ref_base && ref && fused_base && fused && raw_host && comp_host &&
dst_init && ref_host && fused_host;
size_t raw_mismatch = 0;
size_t comp_mismatch = 0;
size_t guard_mismatch = 0;
if (allocated) {
for (uint32_t row = 0; row < raw_cap; row++) {
test_fill_copy_f32_patterns(
raw_host + (uint64_t)row * head_dim,
head_dim,
row * 17u + 3u);
}
for (uint64_t i = 0; i < (uint64_t)n_comp * head_dim; i++) {
comp_host[i] = half_patterns[(i * 7u + (i >> 3u)) %
(sizeof(half_patterns)/sizeof(half_patterns[0]))];
}
for (uint64_t i = 0; i < dst_base_bytes; i++) {
dst_init[i] = (uint8_t)(0xa5u ^ (uint8_t)(i * 37u));
}
TEST_ASSERT(ds4_gpu_tensor_write(
raw, 0, raw_host, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
comp, 0, comp_host, comp_bytes) != 0);
for (size_t ci = 0;
ci < sizeof(raw_starts)/sizeof(raw_starts[0]);
ci++) {
TEST_ASSERT(ds4_gpu_tensor_write(
ref_base, 0, dst_init, dst_base_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
fused_base, 0, dst_init, dst_base_bytes) != 0);
TEST_ASSERT(unsetenv(envs[0]) == 0);
TEST_ASSERT(unsetenv(envs[1]) == 0);
TEST_ASSERT(unsetenv(envs[2]) == 0);
ds4_gpu_set_quality(true);
TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor(
ref, raw, raw_cap, raw_starts[ci], n_raw,
comp, 1, n_comp, head_dim) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(setenv(envs[0], "1", 1) == 0);
TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor(
fused, raw, raw_cap, raw_starts[ci], n_raw,
comp, 1, n_comp, head_dim) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_base, 0, ref_host, dst_base_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
fused_base, 0, fused_host,
dst_base_bytes) != 0);
for (uint64_t i = 0; i < dst_base_bytes; i++) {
if (i < dst_view_offset ||
i >= dst_view_offset + payload_bytes) {
if (ref_host[i] != dst_init[i]) guard_mismatch++;
if (fused_host[i] != dst_init[i]) guard_mismatch++;
} else if (ref_host[i] == fused_host[i]) {
continue;
} else if (i < dst_view_offset + raw_out_bytes) {
raw_mismatch++;
} else {
comp_mismatch++;
}
}
}
/* Component-copy diagnostics and quality mode prevent strict
* selection of the gathered kernel. */
TEST_ASSERT(setenv(envs[0], "1", 1) == 0);
TEST_ASSERT(setenv(envs[1], "1", 1) == 0);
TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor(
fused, raw, raw_cap, 5, n_raw,
comp, 1, n_comp, head_dim) == 0);
TEST_ASSERT(unsetenv(envs[1]) == 0);
TEST_ASSERT(setenv(envs[2], "1", 1) == 0);
TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor(
fused, raw, raw_cap, 5, n_raw,
comp, 1, n_comp, head_dim) == 0);
TEST_ASSERT(unsetenv(envs[2]) == 0);
ds4_gpu_set_quality(true);
TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor(
fused, raw, raw_cap, 5, n_raw,
comp, 1, n_comp, head_dim) == 0);
ds4_gpu_set_quality(false);
}
for (size_t i = 0; i < sizeof(envs)/sizeof(envs[0]); i++) {
test_restore_env(envs[i], saved[i]);
}
fprintf(stderr,
"ds4-test: gathered KV staging exact cases=%zu "
"raw_bytes=%zu comp_bytes=%zu guard_bytes=%zu\n",
sizeof(raw_starts)/sizeof(raw_starts[0]),
raw_mismatch, comp_mismatch, guard_mismatch);
TEST_ASSERT(raw_mismatch == 0);
TEST_ASSERT(comp_mismatch == 0);
TEST_ASSERT(guard_mismatch == 0);
free(fused_host);
free(ref_host);
free(dst_init);
free(comp_host);
free(raw_host);
ds4_gpu_tensor_free(fused);
ds4_gpu_tensor_free(fused_base);
ds4_gpu_tensor_free(ref);
ds4_gpu_tensor_free(ref_base);
ds4_gpu_tensor_free(comp);
ds4_gpu_tensor_free(comp_base);
ds4_gpu_tensor_free(raw);
ds4_gpu_tensor_free(raw_base);
}
static void test_metal_contiguous_compressed_f16_attention_exact(void) {
const uint32_t head_dim = 512;
const uint32_t n_head = 2;
const uint32_t raw_cap = 7;
const uint32_t n_raw = 5;
const uint32_t raw_start = 5;
const uint32_t n_comp = 3;
const uint64_t raw_bytes =
(uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes =
(uint64_t)n_comp * head_dim * sizeof(uint16_t);
const uint64_t q_bytes =
(uint64_t)n_head * head_dim * sizeof(float);
const uint64_t page = (uint64_t)getpagesize();
const char *env_name = "DS4_METAL_DISABLE_CONTIG_F16_F16_COPY";
static const uint16_t half_patterns[] = {
0x0000u, 0x8000u, 0x0001u, 0x03ffu, 0x0400u,
0x1001u, 0x3555u, 0x3c00u, 0x3c01u, 0x4000u,
0xbc00u, 0xc000u, 0x7bffu, 0xfbffu,
};
ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes);
ds4_gpu_tensor *comp_base = ds4_gpu_tensor_alloc(comp_bytes + 18u);
ds4_gpu_tensor *comp = comp_base
? ds4_gpu_tensor_view(comp_base, 2u, comp_bytes)
: NULL;
ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *heads_blit = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *heads_compute = ds4_gpu_tensor_alloc(q_bytes);
TEST_ASSERT(raw != NULL);
TEST_ASSERT(comp_base != NULL);
TEST_ASSERT(comp != NULL);
TEST_ASSERT(q != NULL);
TEST_ASSERT(heads_blit != NULL);
TEST_ASSERT(heads_compute != NULL);
float *raw_host = malloc((size_t)raw_bytes);
uint16_t *comp_host = malloc((size_t)comp_bytes);
float *q_host = malloc((size_t)q_bytes);
float *blit_host = malloc((size_t)q_bytes);
float *compute_host = malloc((size_t)q_bytes);
TEST_ASSERT(raw_host != NULL);
TEST_ASSERT(comp_host != NULL);
TEST_ASSERT(q_host != NULL);
TEST_ASSERT(blit_host != NULL);
TEST_ASSERT(compute_host != NULL);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0);
const bool allocated = raw && comp_base && comp && q && heads_blit &&
heads_compute && raw_host && comp_host && q_host && blit_host &&
compute_host && model_raw;
char *saved_env = test_save_env(env_name);
test_float_compare_stats stats = {0};
if (allocated) {
memset(model_raw, 0, (size_t)page);
float *sinks = model_raw;
sinks[0] = -0.375f;
sinks[1] = 0.1875f;
for (uint64_t i = 0; i < (uint64_t)raw_cap * head_dim; i++) {
const int value =
(int)((i * 19u + (i ^ (i >> 5u)) * 7u) % 193u) - 96;
raw_host[i] = (float)value / 128.0f;
}
for (uint64_t i = 0; i < (uint64_t)n_comp * head_dim; i++) {
comp_host[i] = half_patterns[(i * 5u + (i >> 4u)) %
(sizeof(half_patterns) / sizeof(half_patterns[0]))];
}
for (uint32_t i = 0; i < n_head * head_dim; i++) {
const int value = (int)((i * 37u + (i ^ (i >> 3u)) * 11u) % 251u) - 125;
q_host[i] = (float)value / 96.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(comp, 0, comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(setenv(env_name, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
heads_blit,
model_raw,
page,
0,
q,
raw,
n_raw,
raw_cap,
raw_start,
comp,
1,
n_comp,
NULL,
0,
n_head,
head_dim) != 0);
TEST_ASSERT(setenv(env_name, "0", 1) == 0);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
heads_compute,
model_raw,
page,
0,
q,
raw,
n_raw,
raw_cap,
raw_start,
comp,
1,
n_comp,
NULL,
0,
n_head,
head_dim) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
heads_blit, 0, blit_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
heads_compute, 0, compute_host, q_bytes) != 0);
stats = test_compare_float_bits(
blit_host, compute_host, (size_t)n_head * head_dim);
}
test_restore_env(env_name, saved_env);
fprintf(stderr,
"ds4-test: contiguous compressed-F16 staging exactness "
"mismatches=%zu/%u max_ulp=%u max_abs=%g\n",
stats.mismatch_count,
n_head * head_dim,
stats.max_ulp,
stats.max_abs);
TEST_ASSERT(stats.mismatch_count == 0);
free(model_raw);
free(compute_host);
free(blit_host);
free(q_host);
free(comp_host);
free(raw_host);
ds4_gpu_tensor_free(heads_compute);
ds4_gpu_tensor_free(heads_blit);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(comp);
ds4_gpu_tensor_free(comp_base);
ds4_gpu_tensor_free(raw);
}
static void test_metal_persistent_zero_attention_mask_exact_case(
uint32_t raw_cap,
uint32_t n_raw,
uint32_t raw_start,
uint32_t n_comp,
uint32_t seed) {
const uint32_t head_dim = 512;
const uint32_t n_head = 2;
const uint64_t raw_bytes =
(uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes =
(uint64_t)n_comp * head_dim * sizeof(uint16_t);
const uint64_t q_bytes =
(uint64_t)n_head * head_dim * sizeof(float);
const uint64_t mask_bytes = (uint64_t)n_comp * sizeof(float);
const uint64_t page = (uint64_t)getpagesize();
ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes);
ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *comp_mask = ds4_gpu_tensor_alloc(mask_bytes);
ds4_gpu_tensor *legacy = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *persistent = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *masked = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *pad_legacy = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *after_mask = ds4_gpu_tensor_alloc(q_bytes);
TEST_ASSERT(raw != NULL);
TEST_ASSERT(comp != NULL);
TEST_ASSERT(q != NULL);
TEST_ASSERT(comp_mask != NULL);
TEST_ASSERT(legacy != NULL);
TEST_ASSERT(persistent != NULL);
TEST_ASSERT(masked != NULL);
TEST_ASSERT(pad_legacy != NULL);
TEST_ASSERT(after_mask != NULL);
float *raw_host = malloc((size_t)raw_bytes);
uint16_t *comp_host = malloc((size_t)comp_bytes);
float *q_host = malloc((size_t)q_bytes);
float *mask_host = malloc((size_t)mask_bytes);
float *legacy_host = malloc((size_t)q_bytes);
float *persistent_host = malloc((size_t)q_bytes);
float *masked_host = malloc((size_t)q_bytes);
float *pad_legacy_host = malloc((size_t)q_bytes);
float *after_mask_host = malloc((size_t)q_bytes);
TEST_ASSERT(raw_host != NULL);
TEST_ASSERT(comp_host != NULL);
TEST_ASSERT(q_host != NULL);
TEST_ASSERT(mask_host != NULL);
TEST_ASSERT(legacy_host != NULL);
TEST_ASSERT(persistent_host != NULL);
TEST_ASSERT(masked_host != NULL);
TEST_ASSERT(pad_legacy_host != NULL);
TEST_ASSERT(after_mask_host != NULL);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0);
const bool allocated = raw && comp && q && comp_mask && legacy &&
persistent && masked && pad_legacy && after_mask && raw_host &&
comp_host && q_host && mask_host && legacy_host && persistent_host &&
masked_host && pad_legacy_host && after_mask_host &&
model_raw;
const char *disable_env =
"DS4_METAL_DISABLE_PERSISTENT_ZERO_ATTN_MASK";
const char *pad_disable_env =
"DS4_METAL_DISABLE_GATHERED_KV_PAD_FUSION";
const char *shared_pad_disable_env =
"DS4_METAL_DISABLE_SHARED_KV_PAD";
char *saved_disable = test_save_env(disable_env);
char *saved_pad_disable = test_save_env(pad_disable_env);
char *saved_shared_pad_disable = test_save_env(shared_pad_disable_env);
test_float_compare_stats persistent_stats = {0};
test_float_compare_stats masked_stats = {0};
test_float_compare_stats pad_stats = {0};
test_float_compare_stats after_mask_stats = {0};
if (allocated) {
memset(model_raw, 0, (size_t)page);
float *sinks = model_raw;
sinks[0] = -0.3125f;
sinks[1] = 0.21875f;
for (uint64_t i = 0; i < (uint64_t)raw_cap * head_dim; i++) {
const int value = (int)((i * 17u + (i ^ (i >> 5u)) * 11u +
seed * 13u) % 211u) - 105;
raw_host[i] = (float)value / 128.0f;
}
for (uint64_t i = 0; i < (uint64_t)n_comp * head_dim; i++) {
const int value = (int)((i * 23u + (i ^ (i >> 4u)) * 7u +
seed * 19u) % 193u) - 96;
comp_host[i] = test_float_to_f16((float)value / 112.0f);
}
for (uint32_t i = 0; i < n_head * head_dim; i++) {
const int value = (int)((i * 31u + (i ^ (i >> 3u)) * 5u +
seed * 29u) % 227u) - 113;
q_host[i] = (float)value / 104.0f;
}
for (uint32_t i = 0; i < n_comp; i++) {
mask_host[i] = i == 0 ? -8.0f : -(float)(i + 1u) / 8.0f;
}
TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(comp, 0, comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
comp_mask, 0, mask_host, mask_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(setenv(disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
legacy, model_raw, page, 0, q, raw,
n_raw, raw_cap, raw_start, comp, 1, n_comp,
NULL, 0, n_head, head_dim) != 0);
unsetenv(disable_env);
unsetenv(pad_disable_env);
unsetenv(shared_pad_disable_env);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
persistent, model_raw, page, 0, q, raw,
n_raw, raw_cap, raw_start, comp, 1, n_comp,
NULL, 0, n_head, head_dim) != 0);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
masked, model_raw, page, 0, q, raw,
n_raw, raw_cap, raw_start, comp, 1, n_comp,
comp_mask, 1, n_head, head_dim) != 0);
TEST_ASSERT(setenv(pad_disable_env, "1", 1) == 0);
TEST_ASSERT(setenv(shared_pad_disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
pad_legacy, model_raw, page, 0, q, raw,
n_raw, raw_cap, raw_start, comp, 1, n_comp,
comp_mask, 1, n_head, head_dim) != 0);
unsetenv(pad_disable_env);
unsetenv(shared_pad_disable_env);
TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor(
after_mask, model_raw, page, 0, q, raw,
n_raw, raw_cap, raw_start, comp, 1, n_comp,
NULL, 0, n_head, head_dim) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
legacy, 0, legacy_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
persistent, 0, persistent_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
masked, 0, masked_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
pad_legacy, 0, pad_legacy_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
after_mask, 0, after_mask_host, q_bytes) != 0);
persistent_stats = test_compare_float_bits(
legacy_host, persistent_host, (size_t)n_head * head_dim);
masked_stats = test_compare_float_bits(
legacy_host, masked_host, (size_t)n_head * head_dim);
pad_stats = test_compare_float_bits(
pad_legacy_host, masked_host, (size_t)n_head * head_dim);
after_mask_stats = test_compare_float_bits(
legacy_host, after_mask_host, (size_t)n_head * head_dim);
}
test_restore_env(disable_env, saved_disable);
test_restore_env(pad_disable_env, saved_pad_disable);
test_restore_env(shared_pad_disable_env, saved_shared_pad_disable);
fprintf(stderr,
"ds4-test: persistent zero attention mask exact keys=%u "
"candidate=%zu/%u max_ulp=%u masked_diff=%zu/%u "
"pad_fusion=%zu/%u max_ulp=%u "
"after_mask=%zu/%u max_ulp=%u\n",
n_raw + n_comp,
persistent_stats.mismatch_count, n_head * head_dim,
persistent_stats.max_ulp,
masked_stats.mismatch_count, n_head * head_dim,
pad_stats.mismatch_count, n_head * head_dim,
pad_stats.max_ulp,
after_mask_stats.mismatch_count, n_head * head_dim,
after_mask_stats.max_ulp);
TEST_ASSERT(persistent_stats.mismatch_count == 0);
TEST_ASSERT(masked_stats.mismatch_count != 0);
TEST_ASSERT(pad_stats.mismatch_count == 0);
TEST_ASSERT(after_mask_stats.mismatch_count == 0);
free(model_raw);
free(after_mask_host);
free(pad_legacy_host);
free(masked_host);
free(persistent_host);
free(legacy_host);
free(mask_host);
free(q_host);
free(comp_host);
free(raw_host);
ds4_gpu_tensor_free(after_mask);
ds4_gpu_tensor_free(pad_legacy);
ds4_gpu_tensor_free(masked);
ds4_gpu_tensor_free(persistent);
ds4_gpu_tensor_free(legacy);
ds4_gpu_tensor_free(comp_mask);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(comp);
ds4_gpu_tensor_free(raw);
}
static void test_metal_persistent_zero_attention_mask_exact(void) {
test_metal_persistent_zero_attention_mask_exact_case(7, 5, 5, 3, 11);
test_metal_persistent_zero_attention_mask_exact_case(37, 29, 35, 3, 23);
}
typedef enum {
TEST_METAL_PREFILL_MASK_CACHE_RAW = 1,
TEST_METAL_PREFILL_MASK_CACHE_RATIO4 = 2,
TEST_METAL_PREFILL_MASK_CACHE_RATIO128 = 3,
} test_metal_prefill_mask_cache_kind;
typedef struct {
uint32_t n_tokens;
uint32_t n_comp;
uint32_t window;
uint32_t ratio;
} test_metal_prefill_mask_cache_shape;
static int test_metal_zero_prefix_prefill_mask_cache_call(
test_metal_prefill_mask_cache_kind kind,
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw,
const ds4_gpu_tensor *comp,
const ds4_gpu_tensor *comp_mask,
const test_metal_prefill_mask_cache_shape *shape,
bool masked,
uint32_t n_head,
uint32_t head_dim) {
if (kind == TEST_METAL_PREFILL_MASK_CACHE_RAW) {
if (masked) return 0;
return ds4_gpu_attention_prefill_raw_heads_tensor(
heads, model_map, model_size, 0, q, raw,
shape->n_tokens, shape->window, n_head, head_dim);
}
if (masked) {
return ds4_gpu_attention_prefill_masked_mixed_heads_tensor(
heads, model_map, model_size, 0, q, raw, comp, 1, comp_mask,
shape->n_tokens, shape->n_comp, shape->window, shape->ratio,
n_head, head_dim);
}
return ds4_gpu_attention_prefill_static_mixed_heads_tensor(
heads, model_map, model_size, 0, q, raw, comp, 1,
shape->n_tokens, shape->n_comp, shape->window, shape->ratio,
n_head, head_dim);
}
static bool test_metal_zero_prefix_prefill_mask_cache_run(
test_metal_prefill_mask_cache_kind kind,
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw,
const ds4_gpu_tensor *comp,
const ds4_gpu_tensor *comp_mask,
const test_metal_prefill_mask_cache_shape *shape,
bool masked,
uint32_t n_head,
uint32_t head_dim,
float *host) {
const uint64_t bytes =
(uint64_t)shape->n_tokens * n_head * head_dim * sizeof(float);
const int call_ok = test_metal_zero_prefix_prefill_mask_cache_call(
kind, heads, model_map, model_size, q, raw, comp, comp_mask,
shape, masked, n_head, head_dim);
TEST_ASSERT(call_ok != 0);
if (!call_ok) return false;
const int read_ok = ds4_gpu_tensor_read(heads, 0, host, bytes);
TEST_ASSERT(read_ok != 0);
return read_ok != 0;
}
static void test_metal_zero_prefix_prefill_mask_cache_compare(
const float *expected,
const float *actual,
size_t count,
size_t *total_mismatches,
uint32_t *max_ulp) {
const test_float_compare_stats stats =
test_compare_float_bits(expected, actual, count);
*total_mismatches += stats.mismatch_count;
if (stats.max_ulp > *max_ulp) *max_ulp = stats.max_ulp;
}
static void test_metal_zero_prefix_prefill_mask_cache_exact_kind(
test_metal_prefill_mask_cache_kind kind,
uint32_t seed) {
const uint32_t head_dim = 512;
const uint32_t n_head = 1;
const uint32_t max_tokens = 129;
const uint32_t max_comp = 32;
const uint64_t raw_count = (uint64_t)max_tokens * head_dim;
const uint64_t comp_count = (uint64_t)max_comp * head_dim;
const uint64_t q_count = (uint64_t)max_tokens * n_head * head_dim;
const uint64_t mask_count = (uint64_t)max_tokens * max_comp;
const uint64_t raw_bytes = raw_count * sizeof(float);
const uint64_t comp_bytes = comp_count * sizeof(uint16_t);
const uint64_t q_bytes = q_count * sizeof(float);
const uint64_t mask_bytes = mask_count * sizeof(float);
const uint64_t page = (uint64_t)getpagesize();
const test_metal_prefill_mask_cache_shape shape_a = {
.n_tokens = 128,
.n_comp = kind == TEST_METAL_PREFILL_MASK_CACHE_RAW ? 0u :
(kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? 32u : 1u),
.window = 128,
.ratio = kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? 4u :
(kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO128 ? 128u : 0u),
};
const test_metal_prefill_mask_cache_shape shape_b = {
.n_tokens = 129,
.n_comp = kind == TEST_METAL_PREFILL_MASK_CACHE_RAW ? 0u :
(kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? 31u : 2u),
.window = 63,
.ratio = shape_a.ratio,
};
const size_t count_a =
(size_t)shape_a.n_tokens * n_head * head_dim;
const size_t count_b =
(size_t)shape_b.n_tokens * n_head * head_dim;
ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes);
ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_bytes);
ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes);
ds4_gpu_tensor *comp_mask = ds4_gpu_tensor_alloc(mask_bytes);
ds4_gpu_tensor *heads = ds4_gpu_tensor_alloc(q_bytes);
float *raw_host = malloc((size_t)raw_bytes);
uint16_t *comp_host = malloc((size_t)comp_bytes);
float *q_host = calloc((size_t)q_count, sizeof(float));
float *mask_host = malloc((size_t)mask_bytes);
float *ref_a = malloc((size_t)q_bytes);
float *ref_b = malloc((size_t)q_bytes);
float *actual = malloc((size_t)q_bytes);
float *masked_actual = malloc((size_t)q_bytes);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0);
TEST_ASSERT(raw != NULL);
TEST_ASSERT(comp != NULL);
TEST_ASSERT(q != NULL);
TEST_ASSERT(comp_mask != NULL);
TEST_ASSERT(heads != NULL);
TEST_ASSERT(raw_host != NULL);
TEST_ASSERT(comp_host != NULL);
TEST_ASSERT(q_host != NULL);
TEST_ASSERT(mask_host != NULL);
TEST_ASSERT(ref_a != NULL);
TEST_ASSERT(ref_b != NULL);
TEST_ASSERT(actual != NULL);
TEST_ASSERT(masked_actual != NULL);
TEST_ASSERT(model_raw != NULL);
const char *disable_env =
"DS4_METAL_DISABLE_ZERO_PREFIX_PREFILL_MASK_CACHE";
char *saved_disable = test_save_env(disable_env);
size_t total_mismatches = 0;
uint32_t max_ulp = 0;
size_t key_difference = 0;
size_t masked_difference = 0;
const bool allocated = raw && comp && q && comp_mask && heads &&
raw_host && comp_host && q_host && mask_host && ref_a && ref_b &&
actual && masked_actual && model_raw;
if (allocated) {
memset(model_raw, 0, (size_t)page);
((float *)model_raw)[0] = -1.0f;
for (uint32_t row = 0; row < max_tokens; row++) {
for (uint32_t col = 0; col < head_dim; col++) {
const int value = (int)((row * 37u + col * 17u +
(col ^ (col >> 3u)) * 5u +
seed * 13u) % 257u) - 128;
raw_host[(uint64_t)row * head_dim + col] =
(float)value / 256.0f;
}
}
for (uint32_t row = 0; row < max_comp; row++) {
for (uint32_t col = 0; col < head_dim; col++) {
const int value = (int)((row * 29u + col * 11u +
(col ^ (col >> 4u)) * 7u +
seed * 19u) % 193u) - 96;
comp_host[(uint64_t)row * head_dim + col] =
test_float_to_f16(0.375f + (float)value / 384.0f);
}
}
for (uint64_t i = 0; i < mask_count; i++) {
mask_host[i] = -65504.0f;
}
if (shape_a.n_comp != 0) {
for (uint32_t row = 0; row < shape_a.n_tokens; row++) {
const uint32_t visible = (row + 1u) / shape_a.ratio;
for (uint32_t col = 0; col < shape_a.n_comp; col++) {
if (col < visible) {
mask_host[(uint64_t)row * shape_a.n_comp + col] =
(col & 1u) != 0u ? -2.0f : -65504.0f;
}
}
}
}
TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(comp, 0, comp_host, comp_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
comp_mask, 0, mask_host, mask_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(setenv(disable_env, "1", 1) == 0);
const bool have_ref_a =
test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_a, false, n_head, head_dim, ref_a);
const bool have_ref_b =
test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_b, false, n_head, head_dim, ref_b);
if (have_ref_a && have_ref_b) {
const test_float_compare_stats key_stats =
test_compare_float_bits(ref_a, ref_b, count_a);
key_difference = key_stats.mismatch_count;
TEST_ASSERT(key_difference != 0);
}
TEST_ASSERT(unsetenv(disable_env) == 0);
if (test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_a, false, n_head, head_dim, actual) && have_ref_a) {
test_metal_zero_prefix_prefill_mask_cache_compare(
ref_a, actual, count_a, &total_mismatches, &max_ulp);
}
if (test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_a, false, n_head, head_dim, actual) && have_ref_a) {
test_metal_zero_prefix_prefill_mask_cache_compare(
ref_a, actual, count_a, &total_mismatches, &max_ulp);
}
if (kind != TEST_METAL_PREFILL_MASK_CACHE_RAW) {
if (test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_a, true, n_head, head_dim, masked_actual) &&
have_ref_a) {
const test_float_compare_stats masked_stats =
test_compare_float_bits(ref_a, masked_actual, count_a);
masked_difference = masked_stats.mismatch_count;
TEST_ASSERT(masked_difference != 0);
}
if (test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_a, false, n_head, head_dim, actual) && have_ref_a) {
test_metal_zero_prefix_prefill_mask_cache_compare(
ref_a, actual, count_a,
&total_mismatches, &max_ulp);
}
}
if (test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_b, false, n_head, head_dim, actual) && have_ref_b) {
test_metal_zero_prefix_prefill_mask_cache_compare(
ref_b, actual, count_b, &total_mismatches, &max_ulp);
}
if (test_metal_zero_prefix_prefill_mask_cache_run(
kind, heads, model_raw, page, q, raw, comp, comp_mask,
&shape_a, false, n_head, head_dim, actual) && have_ref_a) {
test_metal_zero_prefix_prefill_mask_cache_compare(
ref_a, actual, count_a, &total_mismatches, &max_ulp);
}
}
test_restore_env(disable_env, saved_disable);
const char *kind_name = kind == TEST_METAL_PREFILL_MASK_CACHE_RAW ? "raw" :
(kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? "ratio4" : "ratio128");
fprintf(stderr,
"ds4-test: zero-prefix prefill mask cache %s exact "
"mismatches=%zu max_ulp=%u key_diff=%zu masked_diff=%zu\n",
kind_name, total_mismatches, max_ulp,
key_difference, masked_difference);
TEST_ASSERT(total_mismatches == 0);
TEST_ASSERT(max_ulp == 0);
free(model_raw);
free(masked_actual);
free(actual);
free(ref_b);
free(ref_a);
free(mask_host);
free(q_host);
free(comp_host);
free(raw_host);
ds4_gpu_tensor_free(heads);
ds4_gpu_tensor_free(comp_mask);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(comp);
ds4_gpu_tensor_free(raw);
}
static void test_metal_zero_prefix_prefill_mask_cache_exact(void) {
test_metal_zero_prefix_prefill_mask_cache_exact_kind(
TEST_METAL_PREFILL_MASK_CACHE_RAW, 41);
test_metal_zero_prefix_prefill_mask_cache_exact_kind(
TEST_METAL_PREFILL_MASK_CACHE_RATIO4, 43);
test_metal_zero_prefix_prefill_mask_cache_exact_kind(
TEST_METAL_PREFILL_MASK_CACHE_RATIO128, 47);
}
#endif
#if defined(__APPLE__)
static void test_metal_hc_split_weighted_sum_norm_batch_exact(void) {
/* Compare the batched HC+RMSNorm fusion against the exact two-dispatch
* sequence used by the reference path at DS4's production dimensions. */
const uint32_t n_embd = 7168;
const uint32_t n_hc = 4;
const uint32_t n_rows = 3;
const uint32_t sinkhorn_iters = 20;
const float hc_eps = 1.0e-6f;
const float norm_eps = 1.0e-6f;
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t page = (uint64_t)getpagesize();
const uint64_t scale_offset = 0;
const uint64_t base_offset = test_round_up_u64(3u * sizeof(float), page);
const uint64_t norm_weight_offset =
test_round_up_u64(base_offset + mix_hc * sizeof(float), page);
const uint64_t model_alloc = test_round_up_u64(
norm_weight_offset + (uint64_t)n_embd * sizeof(float), page);
const uint64_t mix_bytes = (uint64_t)n_rows * mix_hc * sizeof(float);
const uint64_t residual_bytes =
(uint64_t)n_rows * n_hc * n_embd * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_rows * n_embd * sizeof(float);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)model_alloc) == 0);
if (!model_raw) return;
memset(model_raw, 0, (size_t)model_alloc);
float *scale = (float *)((uint8_t *)model_raw + scale_offset);
float *base = (float *)((uint8_t *)model_raw + base_offset);
float *norm_weight = (float *)((uint8_t *)model_raw + norm_weight_offset);
scale[0] = 0.625f;
scale[1] = -0.75f;
scale[2] = 0.4375f;
for (uint32_t i = 0; i < mix_hc; i++) {
const int value = (int)((i * 17u + 5u) % 29u) - 14;
base[i] = (float)value / 16.0f;
}
for (uint32_t i = 0; i < n_embd; i++) {
norm_weight[i] = 0.5f + (float)((i * 13u + 7u) % 31u) / 32.0f;
}
ds4_gpu_tensor *mix = ds4_gpu_tensor_alloc(mix_bytes);
ds4_gpu_tensor *residual = ds4_gpu_tensor_alloc(residual_bytes);
ds4_gpu_tensor *ref_split = ds4_gpu_tensor_alloc(mix_bytes);
ds4_gpu_tensor *fused_split = ds4_gpu_tensor_alloc(mix_bytes);
ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *fused_out = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *fused_norm = ds4_gpu_tensor_alloc(out_bytes);
TEST_ASSERT(mix != NULL);
TEST_ASSERT(residual != NULL);
TEST_ASSERT(ref_split != NULL);
TEST_ASSERT(fused_split != NULL);
TEST_ASSERT(ref_out != NULL);
TEST_ASSERT(fused_out != NULL);
TEST_ASSERT(ref_norm != NULL);
TEST_ASSERT(fused_norm != NULL);
if (!mix || !residual || !ref_split || !fused_split ||
!ref_out || !fused_out || !ref_norm || !fused_norm) {
ds4_gpu_tensor_free(mix);
ds4_gpu_tensor_free(residual);
ds4_gpu_tensor_free(ref_split);
ds4_gpu_tensor_free(fused_split);
ds4_gpu_tensor_free(ref_out);
ds4_gpu_tensor_free(fused_out);
ds4_gpu_tensor_free(ref_norm);
ds4_gpu_tensor_free(fused_norm);
free(model_raw);
return;
}
float *mix_host = malloc((size_t)mix_bytes);
float *residual_host = malloc((size_t)residual_bytes);
float *ref_split_host = malloc((size_t)mix_bytes);
float *fused_split_host = malloc((size_t)mix_bytes);
float *ref_out_host = malloc((size_t)out_bytes);
float *fused_out_host = malloc((size_t)out_bytes);
float *ref_norm_host = malloc((size_t)out_bytes);
float *fused_norm_host = malloc((size_t)out_bytes);
TEST_ASSERT(mix_host != NULL);
TEST_ASSERT(residual_host != NULL);
TEST_ASSERT(ref_split_host != NULL);
TEST_ASSERT(fused_split_host != NULL);
TEST_ASSERT(ref_out_host != NULL);
TEST_ASSERT(fused_out_host != NULL);
TEST_ASSERT(ref_norm_host != NULL);
TEST_ASSERT(fused_norm_host != NULL);
if (!mix_host || !residual_host || !ref_split_host || !fused_split_host ||
!ref_out_host || !fused_out_host || !ref_norm_host || !fused_norm_host) {
free(mix_host);
free(residual_host);
free(ref_split_host);
free(fused_split_host);
free(ref_out_host);
free(fused_out_host);
free(ref_norm_host);
free(fused_norm_host);
ds4_gpu_tensor_free(mix);
ds4_gpu_tensor_free(residual);
ds4_gpu_tensor_free(ref_split);
ds4_gpu_tensor_free(fused_split);
ds4_gpu_tensor_free(ref_out);
ds4_gpu_tensor_free(fused_out);
ds4_gpu_tensor_free(ref_norm);
ds4_gpu_tensor_free(fused_norm);
free(model_raw);
return;
}
for (uint32_t row = 0; row < n_rows; row++) {
for (uint32_t i = 0; i < mix_hc; i++) {
const int value =
(int)(((row + 1u) * 19u + i * 11u + (i ^ row) * 3u) % 47u) - 23;
mix_host[(uint64_t)row * mix_hc + i] = (float)value / 9.0f;
}
for (uint32_t hc = 0; hc < n_hc; hc++) {
for (uint32_t d = 0; d < n_embd; d++) {
const uint32_t key =
d * 37u + hc * 173u + row * 997u + ((d >> 3u) ^ (d * 7u));
const int value = (int)(key % 2047u) - 1023;
residual_host[((uint64_t)row * n_hc + hc) * n_embd + d] =
(float)value / 512.0f;
}
}
}
TEST_ASSERT(ds4_gpu_tensor_write(mix, 0, mix_host, mix_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(residual, 0, residual_host, residual_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0);
TEST_ASSERT(ds4_gpu_hc_split_weighted_sum_tensor(
ref_out, ref_split, mix, residual,
model_raw, model_alloc, scale_offset, base_offset,
n_embd, n_hc, sinkhorn_iters, hc_eps) != 0);
TEST_ASSERT(ds4_gpu_rms_norm_weight_rows_tensor(
ref_norm, ref_out, model_raw, model_alloc, norm_weight_offset,
n_embd, n_rows, norm_eps) != 0);
TEST_ASSERT(ds4_gpu_hc_split_weighted_sum_norm_tensor(
fused_out, fused_norm, fused_split, mix, residual,
model_raw, model_alloc, scale_offset, base_offset, norm_weight_offset,
n_embd, n_hc, sinkhorn_iters, hc_eps, norm_eps) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(ref_split, 0, ref_split_host, mix_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(fused_split, 0, fused_split_host, mix_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(ref_out, 0, ref_out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(fused_out, 0, fused_out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(ref_norm, 0, ref_norm_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(fused_norm, 0, fused_norm_host, out_bytes) != 0);
const test_float_compare_stats split_stats = test_compare_float_bits(
ref_split_host, fused_split_host, (size_t)(n_rows * mix_hc));
const test_float_compare_stats out_stats = test_compare_float_bits(
ref_out_host, fused_out_host, (size_t)n_rows * n_embd);
const test_float_compare_stats norm_stats = test_compare_float_bits(
ref_norm_host, fused_norm_host, (size_t)n_rows * n_embd);
fprintf(stderr,
"ds4-test: batch HC+RMSNorm exactness rows=%u "
"split=%zu/%llu max_ulp=%u max_abs=%g, "
"collapse=%zu/%llu max_ulp=%u max_abs=%g, "
"norm=%zu/%llu max_ulp=%u max_abs=%g\n",
n_rows,
split_stats.mismatch_count,
(unsigned long long)(n_rows * mix_hc),
split_stats.max_ulp,
split_stats.max_abs,
out_stats.mismatch_count,
(unsigned long long)((uint64_t)n_rows * n_embd),
out_stats.max_ulp,
out_stats.max_abs,
norm_stats.mismatch_count,
(unsigned long long)((uint64_t)n_rows * n_embd),
norm_stats.max_ulp,
norm_stats.max_abs);
TEST_ASSERT(split_stats.mismatch_count == 0);
TEST_ASSERT(out_stats.mismatch_count == 0);
TEST_ASSERT(norm_stats.mismatch_count == 0);
free(mix_host);
free(residual_host);
free(ref_split_host);
free(fused_split_host);
free(ref_out_host);
free(fused_out_host);
free(ref_norm_host);
free(fused_norm_host);
ds4_gpu_tensor_free(mix);
ds4_gpu_tensor_free(residual);
ds4_gpu_tensor_free(ref_split);
ds4_gpu_tensor_free(fused_split);
ds4_gpu_tensor_free(ref_out);
ds4_gpu_tensor_free(fused_out);
ds4_gpu_tensor_free(ref_norm);
ds4_gpu_tensor_free(fused_norm);
free(model_raw);
}
static void test_metal_output_hc_weights4_exact(void) {
const uint32_t n_hc = 4;
const float eps = 1.0e-6f;
const uint64_t page = (uint64_t)getpagesize();
const uint64_t scale_offset = 0;
const uint64_t base_offset = page;
const uint64_t model_alloc = 2u * page;
const uint64_t bytes = n_hc * sizeof(float);
const char *require_env = "DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4";
char *saved_require = test_save_env(require_env);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(
&model_raw, (size_t)page, (size_t)model_alloc) == 0);
ds4_gpu_tensor *pre = ds4_gpu_tensor_alloc(bytes);
ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(bytes);
ds4_gpu_tensor *candidate = ds4_gpu_tensor_alloc(bytes);
TEST_ASSERT(model_raw != NULL);
TEST_ASSERT(pre != NULL);
TEST_ASSERT(reference != NULL);
TEST_ASSERT(candidate != NULL);
size_t total_mismatch = 0;
uint32_t max_ulp = 0;
const bool allocated = model_raw && pre && reference && candidate;
if (allocated) {
memset(model_raw, 0, (size_t)model_alloc);
float *scale = (float *)((uint8_t *)model_raw + scale_offset);
float *base = (float *)((uint8_t *)model_raw + base_offset);
for (uint32_t ci = 0; ci < 4; ci++) {
float pre_host[4];
float ref_host[4];
float candidate_host[4];
for (uint32_t i = 0; i < 4; i++) {
pre_host[i] =
((float)((int)(ci * 17u + i * 11u) - 23)) / 8.0f;
base[i] =
((float)((int)(ci * 13u + i * 7u) - 19)) / 16.0f;
}
scale[0] = 0.62500012f + (float)ci * 0.125f;
if (ci == 1) {
/* The first lane distinguishes the required two-rounding
* sequence from an illegally contracted multiply-add. */
const uint32_t pre_bits = 0x4620a541u;
const uint32_t scale_bits = 0x462483bau;
const uint32_t base_bits = 0xccce790eu;
memcpy(&pre_host[0], &pre_bits, sizeof(pre_bits));
memcpy(&scale[0], &scale_bits, sizeof(scale_bits));
memcpy(&base[0], &base_bits, sizeof(base_bits));
} else if (ci == 2) {
const uint32_t pre_bits[4] = {
0x00000000u, 0x80000000u,
0x00000001u, 0x80000001u,
};
const uint32_t base_bits[4] = {
0x80000000u, 0x00000000u,
0x00800000u, 0x80800000u,
};
for (uint32_t i = 0; i < 4; i++) {
memcpy(&pre_host[i], &pre_bits[i], sizeof(uint32_t));
memcpy(&base[i], &base_bits[i], sizeof(uint32_t));
}
scale[0] = -1.00000012f;
} else if (ci == 3) {
pre_host[0] = 100.0f;
pre_host[1] = -100.0f;
pre_host[2] = 88.0f;
pre_host[3] = -88.0f;
scale[0] = 1.0f;
for (uint32_t i = 0; i < 4; i++) base[i] = 0.0f;
}
memset(ref_host, 0xa5, sizeof(ref_host));
memset(candidate_host, 0xa5, sizeof(candidate_host));
TEST_ASSERT(ds4_gpu_tensor_write(pre, 0, pre_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
reference, 0, ref_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
candidate, 0, candidate_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0);
TEST_ASSERT(unsetenv(require_env) == 0);
ds4_gpu_test_set_flags(0);
ds4_gpu_set_quality(true);
TEST_ASSERT(ds4_gpu_output_hc_weights_tensor(
reference, pre, model_raw, model_alloc,
scale_offset, base_offset, n_hc, eps) != 0);
ds4_gpu_set_quality(false);
ds4_gpu_test_set_flags(DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4);
TEST_ASSERT(setenv(require_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_output_hc_weights_tensor(
candidate, pre, model_raw, model_alloc,
scale_offset, base_offset, n_hc, eps) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
reference, 0, ref_host, bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
candidate, 0, candidate_host, bytes) != 0);
const test_float_compare_stats stats =
test_compare_float_bits(ref_host, candidate_host, n_hc);
fprintf(stderr,
"ds4-test: output HC weights4 exact case=%u "
"mismatch=%zu/%u max_ulp=%u max_abs=%g\n",
ci, stats.mismatch_count, n_hc,
stats.max_ulp, stats.max_abs);
TEST_ASSERT(stats.mismatch_count == 0);
total_mismatch += stats.mismatch_count;
if (stats.max_ulp > max_ulp) max_ulp = stats.max_ulp;
}
/* Quality mode prevents strict selection of the fast path. */
TEST_ASSERT(setenv(require_env, "1", 1) == 0);
ds4_gpu_set_quality(true);
TEST_ASSERT(ds4_gpu_output_hc_weights_tensor(
candidate, pre, model_raw, model_alloc,
scale_offset, base_offset, n_hc, eps) == 0);
ds4_gpu_set_quality(false);
ds4_gpu_test_set_flags(0);
}
test_restore_env(require_env, saved_require);
fprintf(stderr,
"ds4-test: output HC weights4 total mismatch=%zu/16 max_ulp=%u\n",
total_mismatch, max_ulp);
TEST_ASSERT(total_mismatch == 0);
TEST_ASSERT(max_ulp == 0);
ds4_gpu_tensor_free(candidate);
ds4_gpu_tensor_free(reference);
ds4_gpu_tensor_free(pre);
free(model_raw);
}
static void test_metal_hc_rms_scale_project_f16_exact_shape(
uint32_t in_dim,
uint32_t seed) {
const uint32_t out_dim = 24;
/* One full 32-row matmul tile plus a tail row covers both load paths. */
const uint32_t n_rows = 33;
const float eps = 1.0e-6f;
const uint64_t page = (uint64_t)getpagesize();
const uint64_t weight_offset = page;
const uint64_t weight_bytes =
(uint64_t)in_dim * out_dim * sizeof(uint16_t);
const uint64_t model_alloc = test_round_up_u64(
weight_offset + weight_bytes, page);
const uint64_t x_count = (uint64_t)in_dim * n_rows;
const uint64_t out_count = (uint64_t)out_dim * n_rows;
const uint64_t x_bytes = x_count * sizeof(float);
const uint64_t out_bytes = out_count * sizeof(float);
const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float);
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(
&model_raw, (size_t)page, (size_t)model_alloc) == 0);
ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes);
ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(x_bytes);
ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(out_bytes);
ds4_gpu_tensor *scaled_out = ds4_gpu_tensor_alloc(out_bytes);
/* Deliberately too small for the full-RMS fallback. */
ds4_gpu_tensor *scale_scratch = ds4_gpu_tensor_alloc(scale_bytes);
float *x_host = malloc((size_t)x_bytes);
float *ref_norm_host = malloc((size_t)x_bytes);
float *ref_out_host = malloc((size_t)out_bytes);
float *scaled_out_host = malloc((size_t)out_bytes);
float *scale_host = malloc((size_t)scale_bytes);
float *expected_scale = malloc((size_t)scale_bytes);
TEST_ASSERT(model_raw != NULL);
TEST_ASSERT(x != NULL);
TEST_ASSERT(ref_norm != NULL);
TEST_ASSERT(ref_out != NULL);
TEST_ASSERT(scaled_out != NULL);
TEST_ASSERT(scale_scratch != NULL);
TEST_ASSERT(x_host != NULL);
TEST_ASSERT(ref_norm_host != NULL);
TEST_ASSERT(ref_out_host != NULL);
TEST_ASSERT(scaled_out_host != NULL);
TEST_ASSERT(scale_host != NULL);
TEST_ASSERT(expected_scale != NULL);
const bool allocated = model_raw && x && ref_norm && ref_out &&
scaled_out && scale_scratch && x_host && ref_norm_host &&
ref_out_host && scaled_out_host && scale_host && expected_scale;
test_float_compare_stats scale_stats = {0};
test_float_compare_stats out_stats = {0};
if (allocated) {
memset(model_raw, 0, (size_t)model_alloc);
uint16_t *weights = (uint16_t *)((uint8_t *)model_raw + weight_offset);
for (uint32_t o = 0; o < out_dim; o++) {
for (uint32_t i = 0; i < in_dim; i++) {
const uint32_t key = i * 37u + o * 1009u + seed * 53u +
((i >> 4u) ^ (o * 17u));
uint16_t bits;
if (key % 257u == 0u) {
bits = (key & 1u) ? 0x8000u : 0x0000u;
} else if (key % 263u == 0u) {
bits = (key & 1u) ? 0x8001u : 0x0001u;
} else if (key % 269u == 0u) {
bits = (key & 1u) ? 0x8400u : 0x0400u;
} else {
const int value = (int)(key % 127u) - 63;
bits = test_float_to_f16((float)value / 128.0f);
}
weights[(uint64_t)o * in_dim + i] = bits;
}
}
static const uint32_t rounding_bits[] = {
0x3f800fffu, 0x3f801000u, 0x3f801001u,
0x3f802fffu, 0x3f803000u, 0x3f803001u,
0xbf800fffu, 0xbf801000u, 0xbf801001u,
0x3eaaaaabu, 0xbeaaaaabu,
};
static const int sentinel_exp[] = { -2, -10, 0, -4, -6, -1 };
for (uint32_t row = 0; row < n_rows; row++) {
for (uint32_t i = 0; i < in_dim; i++) {
const uint32_t key = i * 131u + row * 977u + seed * 71u +
((i >> 3u) ^ (row * 29u));
const float sign = (key & 1u) ? -1.0f : 1.0f;
float value;
switch (row % 6u) {
case 0:
value = (float)((int)(key % 4093u) - 2046) / 512.0f;
break;
case 1:
value = sign * ldexpf(
(float)(1u + (key & 7u)) / 8.0f, -19);
break;
case 2:
value = i % 257u == row % 257u
? sign * (16.0f + (float)(key & 7u))
: sign * (float)(1u + (key & 31u)) / 8192.0f;
break;
case 3: {
const uint32_t bits = rounding_bits[
key % (sizeof(rounding_bits) /
sizeof(rounding_bits[0]))];
memcpy(&value, &bits, sizeof(value));
break;
}
case 4:
if ((key & 7u) == 0u) {
const uint32_t bits = (key & 8u) ? 0x80000000u : 0u;
memcpy(&value, &bits, sizeof(value));
} else {
value = sign * ldexpf(
1.0f + (float)(key & 3u) * 0.25f,
(int)((key >> 4u) % 14u) - 10);
}
break;
default:
value = sign * (float)(1u + (key % 251u)) / 256.0f;
break;
}
x_host[(uint64_t)row * in_dim + i] = value;
}
x_host[(uint64_t)row * in_dim] =
ldexpf(1.0f, sentinel_exp[row % 6u]);
}
for (uint64_t i = 0; i < out_count; i++) {
const uint32_t bits = 0x7fc01000u + (uint32_t)(i & 0xfffu);
memcpy(scaled_out_host + i, &bits, sizeof(bits));
}
for (uint32_t row = 0; row < n_rows; row++) {
const uint32_t bits = 0x7fc02000u + row;
memcpy(scale_host + row, &bits, sizeof(bits));
}
TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
scaled_out, 0, scaled_out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_write(
scale_scratch, 0, scale_host, scale_bytes) != 0);
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0);
ds4_gpu_set_quality(false);
int ref_begun = ds4_gpu_begin_commands();
int ref_ok = ref_begun;
if (ref_ok) ref_ok = ds4_gpu_rms_norm_plain_rows_tensor(
ref_norm, x, in_dim, n_rows, eps);
if (ref_ok) ref_ok = ds4_gpu_matmul_f16_tensor(
ref_out, model_raw, model_alloc, weight_offset,
in_dim, out_dim, ref_norm, n_rows);
const int ref_end = ref_begun ? ds4_gpu_end_commands() : 0;
TEST_ASSERT(ref_ok != 0);
TEST_ASSERT(ref_end != 0);
int scaled_begun = ds4_gpu_begin_commands();
int scaled_ok = scaled_begun;
if (scaled_ok) scaled_ok = ds4_gpu_hc_rms_scale_project_f16_tensor(
scaled_out, scale_scratch, model_raw, model_alloc, weight_offset,
in_dim, out_dim, x, n_rows, eps);
const int scaled_end = scaled_begun ? ds4_gpu_end_commands() : 0;
TEST_ASSERT(scaled_ok != 0);
TEST_ASSERT(scaled_end != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_norm, 0, ref_norm_host, x_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_out, 0, ref_out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
scaled_out, 0, scaled_out_host, out_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
scale_scratch, 0, scale_host, scale_bytes) != 0);
for (uint32_t row = 0; row < n_rows; row++) {
expected_scale[row] = ldexpf(
ref_norm_host[(uint64_t)row * in_dim],
-sentinel_exp[row % 6u]);
TEST_ASSERT(isfinite(expected_scale[row]));
TEST_ASSERT(isfinite(scale_host[row]));
}
scale_stats = test_compare_float_bits(
expected_scale, scale_host, n_rows);
out_stats = test_compare_float_bits(
ref_out_host, scaled_out_host, (size_t)out_count);
}
fprintf(stderr,
"ds4-test: HC RMS scale F16 projection exact K=%u rows=%u "
"scale=%zu/%u max_ulp=%u projection=%zu/%llu max_ulp=%u\n",
in_dim, n_rows,
scale_stats.mismatch_count, n_rows, scale_stats.max_ulp,
out_stats.mismatch_count, (unsigned long long)out_count,
out_stats.max_ulp);
TEST_ASSERT(scale_stats.mismatch_count == 0);
TEST_ASSERT(scale_stats.max_ulp == 0);
TEST_ASSERT(out_stats.mismatch_count == 0);
TEST_ASSERT(out_stats.max_ulp == 0);
free(expected_scale);
free(scale_host);
free(scaled_out_host);
free(ref_out_host);
free(ref_norm_host);
free(x_host);
ds4_gpu_tensor_free(scale_scratch);
ds4_gpu_tensor_free(scaled_out);
ds4_gpu_tensor_free(ref_out);
ds4_gpu_tensor_free(ref_norm);
ds4_gpu_tensor_free(x);
free(model_raw);
}
static void test_metal_hc_rms_scale_project_f16_exact(void) {
const char *disable_env = "DS4_METAL_DISABLE_HC_RMS_SCALE_PROJ";
char *saved_disable = test_save_env(disable_env);
TEST_ASSERT(unsetenv(disable_env) == 0);
ds4_gpu_test_set_flags(DS4_GPU_TEST_HC_RMS_SCALE_PROJ);
test_metal_hc_rms_scale_project_f16_exact_shape(16384u, 59u);
test_metal_hc_rms_scale_project_f16_exact_shape(28672u, 61u);
ds4_gpu_test_set_flags(0);
test_restore_env(disable_env, saved_disable);
}
static void test_metal_router_simd_finalize_exact(void) {
typedef struct {
const char *name;
bool has_bias;
uint32_t pattern;
} router_case;
static const router_case cases[] = {
{ "unique", false, 0 },
{ "bias", true, 0 },
{ "top6-ties", false, 1 },
{ "signed-zero-extremes", false, 2 },
{ "clamp-underflow", false, 3 },
{ "sum-rounding", true, 4 },
};
const uint32_t n_expert = 256;
const uint32_t n_used = 6;
const uint32_t modes = 4;
const uint64_t probs_bytes = (uint64_t)n_expert * sizeof(float);
const uint64_t selected_bytes = (uint64_t)n_used * sizeof(int32_t);
const uint64_t weights_bytes = (uint64_t)n_used * sizeof(float);
const uint64_t page = (uint64_t)getpagesize();
const char *disable_env =
"DS4_METAL_DISABLE_PRE_M5_ROUTER_SIMD_FINALIZE";
const char *weights_disable_env =
"DS4_METAL_DISABLE_PRE_M5_ROUTER_SIMD_WEIGHTS_FUSION";
const char *transform_finalize_disable_env =
"DS4_METAL_DISABLE_PRE_M5_ROUTER_TRANSFORM_FINALIZE_FUSION";
const char *select_disable_env =
"DS4_METAL_DISABLE_ROUTER_SELECT_FUSION";
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0);
ds4_gpu_tensor *logits = ds4_gpu_tensor_alloc(probs_bytes);
ds4_gpu_tensor *ref_selected = ds4_gpu_tensor_alloc(selected_bytes);
ds4_gpu_tensor *simd_selected = ds4_gpu_tensor_alloc(selected_bytes);
ds4_gpu_tensor *ref_weights = ds4_gpu_tensor_alloc(weights_bytes);
ds4_gpu_tensor *simd_weights = ds4_gpu_tensor_alloc(weights_bytes);
ds4_gpu_tensor *ref_probs = ds4_gpu_tensor_alloc(probs_bytes);
ds4_gpu_tensor *simd_probs = ds4_gpu_tensor_alloc(probs_bytes);
float *logits_host = malloc((size_t)probs_bytes);
float *ref_probs_host = malloc((size_t)probs_bytes);
float *simd_probs_host = malloc((size_t)probs_bytes);
int32_t ref_selected_host[6];
int32_t simd_selected_host[6];
int32_t unique_selected_host[6];
float ref_weights_host[6];
float simd_weights_host[6];
TEST_ASSERT(model_raw != NULL);
TEST_ASSERT(logits != NULL);
TEST_ASSERT(ref_selected != NULL);
TEST_ASSERT(simd_selected != NULL);
TEST_ASSERT(ref_weights != NULL);
TEST_ASSERT(simd_weights != NULL);
TEST_ASSERT(ref_probs != NULL);
TEST_ASSERT(simd_probs != NULL);
TEST_ASSERT(logits_host != NULL);
TEST_ASSERT(ref_probs_host != NULL);
TEST_ASSERT(simd_probs_host != NULL);
char *saved_disable = test_save_env(disable_env);
char *saved_weights_disable = test_save_env(weights_disable_env);
char *saved_transform_finalize_disable =
test_save_env(transform_finalize_disable_env);
char *saved_select_disable = test_save_env(select_disable_env);
size_t total_selected_mismatch = 0;
size_t total_weights_mismatch = 0;
size_t total_probs_mismatch = 0;
const bool allocated = model_raw && logits && ref_selected &&
simd_selected && ref_weights && simd_weights && ref_probs &&
simd_probs && logits_host && ref_probs_host && simd_probs_host;
if (allocated) {
memset(model_raw, 0, (size_t)page);
float *bias = model_raw;
for (uint32_t i = 0; i < n_used; i++) {
bias[200u + i] = 16.0f - (float)i;
}
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(setenv(transform_finalize_disable_env, "1", 1) == 0);
TEST_ASSERT(unsetenv(select_disable_env) == 0);
for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) {
const router_case *c = &cases[ci];
for (uint32_t i = 0; i < n_expert; i++) {
const int value =
(int)((i * 47u + (i ^ (i >> 3u)) * 13u) % 257u) - 128;
logits_host[i] = (float)value / 32.0f;
}
if (c->pattern == 1) {
static const uint32_t tied[] = {
7u, 19u, 43u, 71u, 103u, 149u, 211u, 239u,
};
for (uint32_t i = 0; i < n_expert; i++) {
logits_host[i] = -7.0f - (float)(i % 17u) / 64.0f;
}
for (size_t i = 0; i < sizeof(tied) / sizeof(tied[0]); i++) {
logits_host[tied[i]] = 1.0f;
}
} else if (c->pattern == 2) {
for (uint32_t i = 0; i < n_expert; i++) {
const uint32_t zero_bits = (i & 1u) ? 0x80000000u : 0u;
memcpy(&logits_host[i], &zero_bits, sizeof(zero_bits));
}
logits_host[3] = 80.0f;
logits_host[17] = 40.0f;
logits_host[61] = 20.0f;
logits_host[127] = -20.0f;
logits_host[193] = -40.0f;
logits_host[251] = -80.0f;
} else if (c->pattern == 3) {
for (uint32_t i = 0; i < n_expert; i++) {
logits_host[i] = -30.0f - (float)(i % 11u);
}
} else if (c->pattern == 4) {
static const float rounding_logits[6] = {
-0.6356699467f,
-0.8182631135f,
-2.7906901836f,
-3.4414808750f,
-3.4991359711f,
-3.1251864433f,
};
for (uint32_t i = 0; i < n_expert; i++) {
logits_host[i] = -20.0f;
}
for (uint32_t i = 0; i < n_used; i++) {
logits_host[200u + i] = rounding_logits[i];
}
}
TEST_ASSERT(ds4_gpu_tensor_write(
logits, 0, logits_host, probs_bytes) != 0);
TEST_ASSERT(setenv(disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_router_select_tensor(
ref_selected, ref_weights, ref_probs,
model_raw, page, 0, 0, 1, 0,
n_expert, n_used, 1.5f, 1, 0,
c->has_bias, false, logits) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_selected, 0, ref_selected_host, selected_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_weights, 0, ref_weights_host, weights_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_probs, 0, ref_probs_host, probs_bytes) != 0);
if (ci == 0) {
memcpy(unique_selected_host,
ref_selected_host,
sizeof(unique_selected_host));
} else if (c->has_bias) {
TEST_ASSERT(memcmp(unique_selected_host,
ref_selected_host,
sizeof(unique_selected_host)) != 0);
}
TEST_ASSERT(unsetenv(disable_env) == 0);
for (uint32_t mode = 0; mode < modes; mode++) {
const bool fused_weights = mode != 0;
const bool fused_transform = mode >= 2;
if (fused_weights) {
TEST_ASSERT(unsetenv(weights_disable_env) == 0);
} else {
TEST_ASSERT(setenv(weights_disable_env, "1", 1) == 0);
}
if (fused_transform) {
TEST_ASSERT(unsetenv(transform_finalize_disable_env) == 0);
} else {
TEST_ASSERT(setenv(
transform_finalize_disable_env, "1", 1) == 0);
}
for (uint32_t i = 0; i < n_expert; i++) {
const uint32_t poison_bits = 0x7fc00001u + i;
memcpy(&simd_probs_host[i],
&poison_bits,
sizeof(poison_bits));
}
TEST_ASSERT(ds4_gpu_tensor_write(
simd_probs, 0, simd_probs_host, probs_bytes) != 0);
TEST_ASSERT(ds4_gpu_router_select_tensor(
simd_selected, simd_weights, simd_probs,
model_raw, page, 0, 0, 1, 0,
n_expert, n_used, 1.5f, 1, 0,
c->has_bias, false, logits) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
simd_selected, 0, simd_selected_host,
selected_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
simd_weights, 0, simd_weights_host,
weights_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
simd_probs, 0, simd_probs_host, probs_bytes) != 0);
size_t selected_mismatch = 0;
for (uint32_t i = 0; i < n_used; i++) {
if (ref_selected_host[i] != simd_selected_host[i]) {
selected_mismatch++;
}
}
const test_float_compare_stats weights_stats =
test_compare_float_bits(
ref_weights_host, simd_weights_host, n_used);
const test_float_compare_stats probs_stats =
test_compare_float_bits(
ref_probs_host, simd_probs_host, n_expert);
fprintf(stderr,
"ds4-test: router SIMD finalize exactness "
"case=%s mode=%s rep=%u selected=%zu/%u weights=%zu/%u "
"probs=%zu/%u max_weight_ulp=%u max_prob_ulp=%u\n",
c->name,
fused_transform ? "fused-transform" :
fused_weights ? "fused-weights" : "split-weights",
fused_transform ? mode - 2u : 0u,
selected_mismatch,
n_used,
weights_stats.mismatch_count,
n_used,
probs_stats.mismatch_count,
n_expert,
weights_stats.max_ulp,
probs_stats.max_ulp);
TEST_ASSERT(selected_mismatch == 0);
TEST_ASSERT(weights_stats.mismatch_count == 0);
TEST_ASSERT(probs_stats.mismatch_count == 0);
total_selected_mismatch += selected_mismatch;
total_weights_mismatch += weights_stats.mismatch_count;
total_probs_mismatch += probs_stats.mismatch_count;
}
}
}
test_restore_env(disable_env, saved_disable);
test_restore_env(weights_disable_env, saved_weights_disable);
test_restore_env(transform_finalize_disable_env,
saved_transform_finalize_disable);
test_restore_env(select_disable_env, saved_select_disable);
fprintf(stderr,
"ds4-test: router SIMD finalize total selected=%zu/%zu "
"weights=%zu/%zu probs=%zu/%zu\n",
total_selected_mismatch,
(sizeof(cases) / sizeof(cases[0])) * modes * (size_t)n_used,
total_weights_mismatch,
(sizeof(cases) / sizeof(cases[0])) * modes * (size_t)n_used,
total_probs_mismatch,
(sizeof(cases) / sizeof(cases[0])) * modes * (size_t)n_expert);
TEST_ASSERT(total_selected_mismatch == 0);
TEST_ASSERT(total_weights_mismatch == 0);
TEST_ASSERT(total_probs_mismatch == 0);
free(simd_probs_host);
free(ref_probs_host);
free(logits_host);
ds4_gpu_tensor_free(simd_probs);
ds4_gpu_tensor_free(ref_probs);
ds4_gpu_tensor_free(simd_weights);
ds4_gpu_tensor_free(ref_weights);
ds4_gpu_tensor_free(simd_selected);
ds4_gpu_tensor_free(ref_selected);
ds4_gpu_tensor_free(logits);
free(model_raw);
}
static void test_metal_router_weights_batch_exact(void) {
typedef struct {
const char *name;
uint32_t n_tokens;
bool has_bias;
bool hash_mode;
uint32_t pattern;
} router_batch_case;
static const router_batch_case cases[] = {
{ "rows2-unique", 2, false, false, 0 },
{ "rows17-bias-ties", 17, true, false, 1 },
{ "rows129-hash-duplicates", 129, false, true, 2 },
{ "rows2048-typical", 2048, false, false, 0 },
{ "rows3-clamp-underflow", 3, false, false, 3 },
};
const uint32_t n_expert = 256;
const uint32_t n_used = 6;
const uint32_t max_tokens = 2048;
const uint32_t hash_rows = 64;
const uint32_t repeats = 2;
const uint64_t logits_bytes =
(uint64_t)max_tokens * n_expert * sizeof(float);
const uint64_t selected_bytes =
(uint64_t)max_tokens * n_used * sizeof(int32_t);
const uint64_t weights_bytes =
(uint64_t)max_tokens * n_used * sizeof(float);
const uint64_t tokens_bytes = (uint64_t)max_tokens * sizeof(int32_t);
const uint64_t page = (uint64_t)getpagesize();
const uint64_t bias_offset = 0;
const uint64_t hash_offset = 2048;
const char *disable_env =
"DS4_METAL_DISABLE_ROUTER_WEIGHTS_BATCH_FUSION";
const char *select_disable_env =
"DS4_METAL_DISABLE_ROUTER_SELECT_FUSION";
void *model_raw = NULL;
TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0);
ds4_gpu_tensor *logits = ds4_gpu_tensor_alloc(logits_bytes);
ds4_gpu_tensor *tokens = ds4_gpu_tensor_alloc(tokens_bytes);
ds4_gpu_tensor *ref_selected = ds4_gpu_tensor_alloc(selected_bytes);
ds4_gpu_tensor *batch_selected = ds4_gpu_tensor_alloc(selected_bytes);
ds4_gpu_tensor *ref_weights = ds4_gpu_tensor_alloc(weights_bytes);
ds4_gpu_tensor *batch_weights = ds4_gpu_tensor_alloc(weights_bytes);
ds4_gpu_tensor *ref_probs = ds4_gpu_tensor_alloc(logits_bytes);
ds4_gpu_tensor *batch_probs = ds4_gpu_tensor_alloc(logits_bytes);
float *logits_host = malloc((size_t)logits_bytes);
int32_t *tokens_host = malloc((size_t)tokens_bytes);
int32_t *ref_selected_host = malloc((size_t)selected_bytes);
int32_t *batch_selected_host = malloc((size_t)selected_bytes);
float *ref_weights_host = malloc((size_t)weights_bytes);
float *batch_weights_host = malloc((size_t)weights_bytes);
float *ref_probs_host = malloc((size_t)logits_bytes);
float *batch_probs_host = malloc((size_t)logits_bytes);
TEST_ASSERT(model_raw != NULL);
TEST_ASSERT(logits != NULL);
TEST_ASSERT(tokens != NULL);
TEST_ASSERT(ref_selected != NULL);
TEST_ASSERT(batch_selected != NULL);
TEST_ASSERT(ref_weights != NULL);
TEST_ASSERT(batch_weights != NULL);
TEST_ASSERT(ref_probs != NULL);
TEST_ASSERT(batch_probs != NULL);
TEST_ASSERT(logits_host != NULL);
TEST_ASSERT(tokens_host != NULL);
TEST_ASSERT(ref_selected_host != NULL);
TEST_ASSERT(batch_selected_host != NULL);
TEST_ASSERT(ref_weights_host != NULL);
TEST_ASSERT(batch_weights_host != NULL);
TEST_ASSERT(ref_probs_host != NULL);
TEST_ASSERT(batch_probs_host != NULL);
char *saved_disable = test_save_env(disable_env);
char *saved_select_disable = test_save_env(select_disable_env);
size_t total_selected_mismatch = 0;
size_t total_weights_mismatch = 0;
size_t total_probs_mismatch = 0;
const bool allocated = model_raw && logits && tokens && ref_selected &&
batch_selected && ref_weights && batch_weights && ref_probs &&
batch_probs && logits_host && tokens_host && ref_selected_host &&
batch_selected_host && ref_weights_host && batch_weights_host &&
ref_probs_host && batch_probs_host;
if (allocated) {
memset(model_raw, 0, (size_t)page);
float *bias = (float *)((uint8_t *)model_raw + bias_offset);
int32_t *hash = (int32_t *)((uint8_t *)model_raw + hash_offset);
for (uint32_t i = 0; i < n_expert; i++) {
bias[i] = (float)((int)((i * 29u) % 67u) - 33) / 16.0f;
}
for (uint32_t row = 0; row < hash_rows; row++) {
const int32_t base = (int32_t)((row * 37u) % n_expert);
hash[(uint64_t)row * n_used + 0u] = base;
hash[(uint64_t)row * n_used + 1u] = base;
hash[(uint64_t)row * n_used + 2u] = (base + 19) % (int32_t)n_expert;
hash[(uint64_t)row * n_used + 3u] = (base + 43) % (int32_t)n_expert;
hash[(uint64_t)row * n_used + 4u] = (base + 43) % (int32_t)n_expert;
hash[(uint64_t)row * n_used + 5u] = (base + 101) % (int32_t)n_expert;
}
for (uint32_t row = 0; row < max_tokens; row++) {
tokens_host[row] = (int32_t)((row * 13u + 7u) % hash_rows);
}
TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0);
ds4_gpu_set_quality(false);
TEST_ASSERT(ds4_gpu_tensor_write(
tokens, 0, tokens_host, tokens_bytes) != 0);
TEST_ASSERT(unsetenv(select_disable_env) == 0);
for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) {
const router_batch_case *c = &cases[ci];
const size_t prob_count = (size_t)c->n_tokens * n_expert;
const size_t route_count = (size_t)c->n_tokens * n_used;
const uint64_t case_probs_bytes = prob_count * sizeof(float);
const uint64_t case_selected_bytes = route_count * sizeof(int32_t);
const uint64_t case_weights_bytes = route_count * sizeof(float);
for (uint32_t row = 0; row < c->n_tokens; row++) {
for (uint32_t expert = 0; expert < n_expert; expert++) {
const size_t index = (size_t)row * n_expert + expert;
const int value =
(int)((row * 131u + expert * 47u +
(expert ^ (expert >> 3u)) * 13u) % 513u) -
256;
logits_host[index] = (float)value / 64.0f;
if (c->pattern == 1) {
logits_host[index] =
-7.0f - (float)((row + expert) % 17u) / 64.0f;
} else if (c->pattern == 3) {
logits_host[index] =
-30.0f - (float)((row * 17u + expert * 29u) % 11u);
}
}
if (c->pattern == 1) {
static const uint32_t tied[] = {
7u, 19u, 43u, 71u, 103u, 149u, 211u, 239u,
};
for (size_t i = 0; i < sizeof(tied) / sizeof(tied[0]); i++) {
logits_host[(size_t)row * n_expert + tied[i]] = 1.0f;
}
}
}
TEST_ASSERT(ds4_gpu_tensor_write(
logits, 0, logits_host, case_probs_bytes) != 0);
TEST_ASSERT(setenv(disable_env, "1", 1) == 0);
TEST_ASSERT(ds4_gpu_router_select_batch_tensor(
ref_selected, ref_weights, ref_probs,
model_raw, page, bias_offset, hash_offset, hash_rows,
1, 0, c->has_bias, c->hash_mode, logits, tokens,
n_expert, n_used, 1.5f, c->n_tokens) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_selected, 0, ref_selected_host, case_selected_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_weights, 0, ref_weights_host, case_weights_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
ref_probs, 0, ref_probs_host, case_probs_bytes) != 0);
if (c->hash_mode) {
for (uint32_t row = 0; row < c->n_tokens; row++) {
const uint32_t hash_row = (uint32_t)tokens_host[row];
for (uint32_t lane = 0; lane < n_used; lane++) {
TEST_ASSERT(
ref_selected_host[(size_t)row * n_used + lane] ==
hash[(uint64_t)hash_row * n_used + lane]);
}
}
}
if (c->pattern == 3) {
for (uint32_t row = 0; row < c->n_tokens; row++) {
float sum = 0.0f;
for (uint32_t lane = 0; lane < n_used; lane++) {
const int32_t expert =
ref_selected_host[(size_t)row * n_used + lane];
sum += ref_probs_host[(size_t)row * n_expert +
(uint32_t)expert];
}
TEST_ASSERT(sum < 6.103515625e-5f);
}
}
TEST_ASSERT(unsetenv(disable_env) == 0);
for (uint32_t rep = 0; rep < repeats; rep++) {
TEST_ASSERT(ds4_gpu_router_select_batch_tensor(
batch_selected, batch_weights, batch_probs,
model_raw, page, bias_offset, hash_offset, hash_rows,
1, 0, c->has_bias, c->hash_mode, logits, tokens,
n_expert, n_used, 1.5f, c->n_tokens) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
batch_selected, 0, batch_selected_host,
case_selected_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
batch_weights, 0, batch_weights_host,
case_weights_bytes) != 0);
TEST_ASSERT(ds4_gpu_tensor_read(
batch_probs, 0, batch_probs_host,
case_probs_bytes) != 0);
size_t selected_mismatch = 0;
for (size_t i = 0; i < route_count; i++) {
if (ref_selected_host[i] != batch_selected_host[i]) {
selected_mismatch++;
}
}
const test_float_compare_stats weights_stats =
test_compare_float_bits(
ref_weights_host, batch_weights_host, route_count);
const test_float_compare_stats probs_stats =
test_compare_float_bits(
ref_probs_host, batch_probs_host, prob_count);
fprintf(stderr,
"ds4-test: router batch weights exactness "
"case=%s rep=%u selected=%zu/%zu weights=%zu/%zu "
"probs=%zu/%zu max_weight_ulp=%u max_prob_ulp=%u\n",
c->name,
rep,
selected_mismatch,
route_count,
weights_stats.mismatch_count,
route_count,
probs_stats.mismatch_count,
prob_count,
weights_stats.max_ulp,
probs_stats.max_ulp);
TEST_ASSERT(selected_mismatch == 0);
TEST_ASSERT(weights_stats.mismatch_count == 0);
TEST_ASSERT(probs_stats.mismatch_count == 0);
total_selected_mismatch += selected_mismatch;
total_weights_mismatch += weights_stats.mismatch_count;
total_probs_mismatch += probs_stats.mismatch_count;
}
}
}
test_restore_env(disable_env, saved_disable);
test_restore_env(select_disable_env, saved_select_disable);
fprintf(stderr,
"ds4-test: router batch weights total selected=%zu "
"weights=%zu probs=%zu\n",
total_selected_mismatch,
total_weights_mismatch,
total_probs_mismatch);
TEST_ASSERT(total_selected_mismatch == 0);
TEST_ASSERT(total_weights_mismatch == 0);
TEST_ASSERT(total_probs_mismatch == 0);
free(batch_probs_host);
free(ref_probs_host);
free(batch_weights_host);
free(ref_weights_host);
free(batch_selected_host);
free(ref_selected_host);
free(tokens_host);
free(logits_host);
ds4_gpu_tensor_free(batch_probs);
ds4_gpu_tensor_free(ref_probs);
ds4_gpu_tensor_free(batch_weights);
ds4_gpu_tensor_free(ref_weights);
ds4_gpu_tensor_free(batch_selected);
ds4_gpu_tensor_free(ref_selected);
ds4_gpu_tensor_free(tokens);
ds4_gpu_tensor_free(logits);
free(model_raw);
}
#endif
static void test_metal_kernel_group(void) {
test_metal_f16_matvec_fast_nr0_4();
test_metal_f16_prefill_matmul();
test_metal_q8_0_prefill_matmul();
test_metal_pack_slot_rows_f32();
test_metal_store_raw_kv_batch_wrap();
test_dspark_cache_window_crop();
test_metal_q8_0_decode_pair_exact();
#if defined(__APPLE__)
test_metal_f16_compressor_pair_state_store_exact();
test_metal_compressor_ape_add_exact();
test_metal_compressor_ratio4_pack_exact();
test_metal_compressor_ratio4_replay_pack_exact();
test_metal_compressor_ratio4_direct_pool_exact();
test_metal_inplace_rope_pair_exact();
test_metal_contiguous_f32_f16_roundtrip_exact();
test_metal_gathered_kv_stage_exact();
test_metal_contiguous_compressed_f16_attention_exact();
test_metal_persistent_zero_attention_mask_exact();
test_metal_zero_prefix_prefill_mask_cache_exact();
test_metal_hc_split_weighted_sum_norm_batch_exact();
test_metal_output_hc_weights4_exact();
test_metal_hc_rms_scale_project_f16_exact();
test_metal_router_simd_finalize_exact();
test_metal_router_weights_batch_exact();
#endif
}
static void test_metal_short_prefill_ratio4(void) {
ds4_engine *engine = test_get_engine(false);
if (!engine) return;
const int tokens[] = {
ds4_token_user(engine),
ds4_token_assistant(engine),
ds4_token_eos(engine),
};
for (size_t i = 0; i < sizeof(tokens) / sizeof(tokens[0]); i++) {
TEST_ASSERT(tokens[i] >= 0);
if (tokens[i] < 0) return;
}
for (size_t n = 1; n <= 3; n++) {
ds4_tokens prompt = {0};
for (size_t i = 0; i < n; i++) {
ds4_tokens_push(&prompt, tokens[i]);
}
TEST_ASSERT(prompt.len == (int)n);
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, 2048) == 0);
if (!session) {
ds4_tokens_free(&prompt);
return;
}
char err[160] = {0};
const int rc = ds4_session_sync(session, &prompt, err, sizeof(err));
if (rc != 0) {
fprintf(stderr, "ds4-test: short prefill failed for %zu token(s): %s\n",
n, err);
}
TEST_ASSERT(rc == 0);
ds4_session_free(session);
ds4_tokens_free(&prompt);
}
}
static char *test_read_file(const char *path) {
FILE *fp = fopen(path, "rb");
if (!fp) return NULL;
if (fseek(fp, 0, SEEK_END) != 0) {
fclose(fp);
return NULL;
}
long len = ftell(fp);
if (len < 0) {
fclose(fp);
return NULL;
}
rewind(fp);
char *s = malloc((size_t)len + 1);
if (!s) {
fclose(fp);
return NULL;
}
size_t nread = fread(s, 1, (size_t)len, fp);
fclose(fp);
if (nread != (size_t)len) {
free(s);
return NULL;
}
s[len] = '\0';
return s;
}
typedef struct {
const char *name;
int number;
} test_long_fact;
static const test_long_fact test_long_facts[] = {
{"Bob", 34},
{"Alice", 52},
{"Clara", 71},
{"Diego", 93},
{"Elena", 16},
{"Felix", 88},
{"Greta", 47},
{"Hugo", 29},
{"Iris", 64},
{"Jonas", 12},
{"Kira", 81},
{"Leo", 39},
{"Marta", 76},
{"Nadia", 23},
{"Owen", 58},
{"Priya", 97},
};
static bool test_is_name_boundary(char c) {
unsigned char uc = (unsigned char)c;
return c == '\0' || !(isalnum(uc) || c == '_');
}
static bool test_parse_assignment_value(const char *p, int *value) {
while (*p == ' ' || *p == '\t') p++;
if (*p != '=') return false;
p++;
while (*p == ' ' || *p == '\t') p++;
if (!isdigit((unsigned char)*p)) return false;
int v = 0;
while (isdigit((unsigned char)*p)) {
v = v * 10 + (*p - '0');
p++;
}
*value = v;
return true;
}
static bool test_output_has_fact(const char *text, const test_long_fact *fact) {
const size_t name_len = strlen(fact->name);
const char *p = text;
bool saw_wrong_assignment = false;
int wrong_value = -1;
while ((p = strstr(p, fact->name)) != NULL) {
const bool before_ok = p == text || test_is_name_boundary(p[-1]);
const bool after_ok = test_is_name_boundary(p[name_len]) ||
p[name_len] == ' ' ||
p[name_len] == '\t' ||
p[name_len] == '=';
if (before_ok && after_ok) {
int value = 0;
if (test_parse_assignment_value(p + name_len, &value)) {
if (value == fact->number) return true;
saw_wrong_assignment = true;
wrong_value = value;
}
}
p += name_len;
}
if (saw_wrong_assignment) {
fprintf(stderr,
"ds4-test: long-context wrong assignment for %s: got %d expected %d\n",
fact->name, wrong_value, fact->number);
} else {
fprintf(stderr,
"ds4-test: long-context missing assignment for %s=%d\n",
fact->name, fact->number);
}
return false;
}
static int test_hex_digit(char c) {
if (c >= '0' && c <= '9') return c - '0';
if (c >= 'a' && c <= 'f') return 10 + c - 'a';
if (c >= 'A' && c <= 'F') return 10 + c - 'A';
return -1;
}
static bool test_hex_to_bytes(const char *hex, unsigned char *out, int cap, int *len) {
int n = 0;
while (*hex && !isspace((unsigned char)*hex)) {
int hi = test_hex_digit(hex[0]);
int lo = test_hex_digit(hex[1]);
if (hi < 0 || lo < 0 || n >= cap) return false;
out[n++] = (unsigned char)((hi << 4) | lo);
hex += 2;
}
*len = n;
return true;
}
static bool test_token_bytes_equal(ds4_engine *engine, int token,
const unsigned char *want, int want_len) {
size_t got_len = 0;
char *got = ds4_token_text(engine, token, &got_len);
bool eq = got && got_len == (size_t)want_len &&
memcmp(got, want, (size_t)want_len) == 0;
free(got);
return eq;
}
static void test_long_prefill_progress(void *ud, const char *event, int current, int total) {
(void)ud;
if (strcmp(event, "prefill_chunk")) return;
if (current == 0 || current == total || current % 8192 == 0) {
fprintf(stderr, "ds4-test: long-context prefill %d/%d\n", current, total);
}
}
static void test_long_story_fact_recall(void) {
const char *prompt_path = getenv("DS4_TEST_LONG_PROMPT");
if (!prompt_path || !prompt_path[0]) {
prompt_path = "tests/long_context_story_prompt.txt";
}
char *prompt_text = test_read_file(prompt_path);
TEST_ASSERT(prompt_text != NULL);
if (!prompt_text) return;
ds4_engine *engine = test_get_engine(false);
if (!engine) {
free(prompt_text);
return;
}
ds4_tokens prompt = {0};
ds4_tokenize_rendered_chat(engine, prompt_text, &prompt);
TEST_ASSERT(prompt.len > 30000);
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, 100000) == 0);
if (!session) {
ds4_tokens_free(&prompt);
free(prompt_text);
return;
}
char err[160];
ds4_session_set_progress(session, test_long_prefill_progress, NULL);
TEST_ASSERT(ds4_session_sync(session, &prompt, err, sizeof(err)) == 0);
ds4_session_set_progress(session, NULL, NULL);
buf out = {0};
uint64_t rng = 12345;
int generated = 0;
bool decode_ok = true;
for (; generated < 350; generated++) {
int token = ds4_session_sample(session, 0.0f, 0, 1.0f, 0.0f, &rng);
if (token == ds4_token_eos(engine)) break;
size_t piece_len = 0;
char *piece = ds4_token_text(engine, token, &piece_len);
buf_append(&out, piece, piece_len);
free(piece);
if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
decode_ok = false;
break;
}
}
const char *text = out.ptr ? out.ptr : "";
TEST_ASSERT(decode_ok);
TEST_ASSERT(generated > 0);
for (size_t i = 0; i < sizeof(test_long_facts) / sizeof(test_long_facts[0]); i++) {
TEST_ASSERT(test_output_has_fact(text, &test_long_facts[i]));
}
buf_free(&out);
ds4_session_free(session);
ds4_tokens_free(&prompt);
free(prompt_text);
}
#define TEST_VEC_MAX_STEPS 16
#define TEST_VEC_MAX_TOP 32
#define TEST_VEC_MAX_TOKEN_BYTES 128
typedef struct {
unsigned char bytes[TEST_VEC_MAX_TOKEN_BYTES];
int len;
float logprob;
} test_vec_top;
typedef struct {
unsigned char selected[TEST_VEC_MAX_TOKEN_BYTES];
int selected_len;
int ntop;
test_vec_top top[TEST_VEC_MAX_TOP];
} test_vec_step;
typedef struct {
char id[96];
char prompt_path[512];
int ctx;
int nsteps;
test_vec_step steps[TEST_VEC_MAX_STEPS];
} test_vec_case;
static char *test_trim_line(char *line) {
while (*line && isspace((unsigned char)*line)) line++;
size_t n = strlen(line);
while (n && isspace((unsigned char)line[n - 1])) line[--n] = '\0';
return line;
}
static bool test_read_vector_case(FILE *fp, test_vec_case *vc) {
char line[2048];
memset(vc, 0, sizeof(*vc));
while (fgets(line, sizeof(line), fp)) {
char *p = test_trim_line(line);
if (!p[0] || p[0] == '#') continue;
if (sscanf(p, "case %95s %d %d %511s",
vc->id, &vc->ctx, &vc->nsteps, vc->prompt_path) == 4) {
TEST_ASSERT(vc->nsteps > 0 && vc->nsteps <= TEST_VEC_MAX_STEPS);
return true;
}
TEST_ASSERT(!"unexpected line before vector case");
}
return false;
}
static bool test_fill_vector_case(FILE *fp, test_vec_case *vc) {
char line[2048];
int step_index = -1;
int top_index = 0;
while (fgets(line, sizeof(line), fp)) {
char *p = test_trim_line(line);
if (!p[0] || p[0] == '#') continue;
if (!strcmp(p, "end")) return true;
if (!strncmp(p, "step ", 5)) {
char hex[TEST_VEC_MAX_TOKEN_BYTES * 2 + 2];
int ntop = 0;
if (sscanf(p, "step %d %257s %d", &step_index, hex, &ntop) != 3) {
TEST_ASSERT(!"bad vector step line");
return false;
}
TEST_ASSERT(step_index >= 0 && step_index < vc->nsteps);
TEST_ASSERT(ntop >= 0 && ntop <= TEST_VEC_MAX_TOP);
vc->steps[step_index].ntop = ntop;
TEST_ASSERT(test_hex_to_bytes(hex,
vc->steps[step_index].selected,
TEST_VEC_MAX_TOKEN_BYTES,
&vc->steps[step_index].selected_len));
top_index = 0;
continue;
}
if (!strncmp(p, "top ", 4)) {
char hex[TEST_VEC_MAX_TOKEN_BYTES * 2 + 2];
float lp = 0.0f;
TEST_ASSERT(step_index >= 0 && step_index < vc->nsteps);
TEST_ASSERT(top_index < vc->steps[step_index].ntop);
if (sscanf(p, "top %257s %f", hex, &lp) != 2) {
TEST_ASSERT(!"bad vector top line");
return false;
}
test_vec_top *top = &vc->steps[step_index].top[top_index++];
top->logprob = lp;
TEST_ASSERT(test_hex_to_bytes(hex, top->bytes,
TEST_VEC_MAX_TOKEN_BYTES, &top->len));
continue;
}
TEST_ASSERT(!"unexpected vector line");
return false;
}
TEST_ASSERT(!"unterminated vector case");
return false;
}
static void test_logprob_vector_case(ds4_engine *engine, const test_vec_case *vc) {
char *prompt_text = test_read_file(vc->prompt_path);
TEST_ASSERT(prompt_text != NULL);
if (!prompt_text) return;
ds4_tokens prompt = {0};
ds4_encode_chat_prompt(engine, "", prompt_text, DS4_THINK_NONE, &prompt);
free(prompt_text);
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, vc->ctx) == 0);
if (!session) {
ds4_tokens_free(&prompt);
return;
}
char err[160];
if (ds4_session_sync(session, &prompt, err, sizeof(err)) != 0) {
fprintf(stderr, "ds4-test: vector %s prefill failed: %s\n", vc->id, err);
TEST_ASSERT(false);
ds4_session_free(session);
ds4_tokens_free(&prompt);
return;
}
ds4_token_score scores[20];
for (int i = 0; i < vc->nsteps; i++) {
const test_vec_step *step = &vc->steps[i];
int nscore = ds4_session_top_logprobs(session, scores, 20);
int token = ds4_session_argmax(session);
if (!test_token_bytes_equal(engine, token, step->selected, step->selected_len)) {
fprintf(stderr, "ds4-test: vector %s step %d selected token mismatch\n",
vc->id, i);
TEST_ASSERT(false);
}
for (int t = 0; t < step->ntop; t++) {
bool found = false;
float local_lp = 0.0f;
for (int j = 0; j < nscore; j++) {
if (scores[j].id < 0) continue;
if (test_token_bytes_equal(engine, scores[j].id,
step->top[t].bytes,
step->top[t].len)) {
found = true;
local_lp = scores[j].logprob;
break;
}
}
if (!found) {
fprintf(stderr, "ds4-test: vector %s step %d official top token missing locally\n",
vc->id, i);
TEST_ASSERT(false);
} else if (fabsf(local_lp - step->top[t].logprob) > 4.0f) {
fprintf(stderr,
"ds4-test: vector %s step %d logprob delta too high: local=%g official=%g\n",
vc->id, i, local_lp, step->top[t].logprob);
TEST_ASSERT(false);
}
}
if (i + 1 < vc->nsteps) {
if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
fprintf(stderr, "ds4-test: vector %s step %d eval failed: %s\n",
vc->id, i, err);
TEST_ASSERT(false);
break;
}
}
}
ds4_session_free(session);
ds4_tokens_free(&prompt);
}
static bool test_logprob_vector_case_disabled(const char *path,
const test_vec_case *vc) {
/*
* This one long-context vector currently matches the public DeepSeek API less
* after adding the official Hadamard+FP4 indexer path. The public official
* implementation and the API appear to disagree here; the official graph has
* slightly lower local perplexity on the A/B check we ran, so DS4 keeps that
* implementation and only excludes this brittle API fixture for now.
*/
return !strcmp(path, "tests/test-vectors/flash-pre-0731/official.vec") &&
!strcmp(vc->id, "long_memory_archive");
}
static void test_official_logprob_vectors_run(const char *case_filter) {
const char *path = getenv("DS4_TEST_VECTOR_FILE");
if (!path || !path[0]) {
path = "tests/test-vectors/flash-0731/official.vec";
}
FILE *fp = fopen(path, "rb");
TEST_ASSERT(fp != NULL);
if (!fp) return;
char *saved_prefill_chunk = test_save_env("DS4_METAL_PREFILL_CHUNK");
char *saved_disable_metal4 = test_save_env("DS4_METAL_DISABLE_METAL4");
test_streaming_prefill_env saved_canonical_streaming_prefill =
test_force_canonical_streaming_prefill();
setenv("DS4_METAL_PREFILL_CHUNK", "2048", 1);
if (getenv("DS4_TEST_LOGPROB_AUTO_METAL") == NULL) {
setenv("DS4_METAL_DISABLE_METAL4", "1", 1);
} else {
unsetenv("DS4_METAL_DISABLE_METAL4");
}
ds4_engine *engine = test_open_engine(false);
if (!engine) {
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
test_restore_env("DS4_METAL_PREFILL_CHUNK", saved_prefill_chunk);
fclose(fp);
return;
}
test_vec_case vc;
int ran = 0;
while (test_read_vector_case(fp, &vc)) {
if (!test_fill_vector_case(fp, &vc)) break;
if (case_filter && case_filter[0] && strcmp(vc.id, case_filter)) {
continue;
}
if (test_logprob_vector_case_disabled(path, &vc)) {
fprintf(stderr, "ds4-test: vector %s skipped (API/official graph mismatch)\n",
vc.id);
continue;
}
fprintf(stderr, "ds4-test: vector %s\n", vc.id);
test_logprob_vector_case(engine, &vc);
ran++;
}
TEST_ASSERT(!case_filter || !case_filter[0] || ran == 1);
ds4_engine_close(engine);
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
test_restore_env("DS4_METAL_PREFILL_CHUNK", saved_prefill_chunk);
fclose(fp);
}
static void test_official_logprob_vectors(void) {
test_official_logprob_vectors_run(NULL);
}
static void test_metal_ssd_streaming_cache_pressure(void) {
#ifndef __APPLE__
fprintf(stderr,
"ds4-test: Metal SSD streaming cache-pressure repro skipped "
"(Metal-only)\n");
#else
/*
* Regression repro for GitHub issue #384.
*
* The bug needs the Metal SSD-streaming decode layer-batch path and a small
* routed-expert cache. Under pressure, a cache entry referenced by an
* already-encoded-but-not-yet-executed layer can be reused for a later
* layer in the same command buffer, producing deterministic wrong logits.
*/
char *saved_streaming = test_save_env("DS4_TEST_SSD_STREAMING");
char *saved_cache_gb = test_save_env("DS4_TEST_SSD_STREAMING_CACHE_GB");
char *saved_cache_experts =
test_save_env("DS4_TEST_SSD_STREAMING_CACHE_EXPERTS");
char *saved_disable_layer_batch =
test_save_env("DS4_METAL_DISABLE_STREAMING_LAYER_BATCH");
char *saved_disable_static_decode =
test_save_env("DS4_METAL_DISABLE_STREAMING_STATIC_DECODE_MAP");
char *saved_one_stage =
test_save_env("DS4_METAL_MOE_ONE_STAGE_PROFILE");
setenv("DS4_TEST_SSD_STREAMING", "1", 1);
setenv("DS4_TEST_SSD_STREAMING_CACHE_GB", "16", 1);
unsetenv("DS4_TEST_SSD_STREAMING_CACHE_EXPERTS");
unsetenv("DS4_METAL_DISABLE_STREAMING_LAYER_BATCH");
unsetenv("DS4_METAL_DISABLE_STREAMING_STATIC_DECODE_MAP");
unsetenv("DS4_METAL_MOE_ONE_STAGE_PROFILE");
fprintf(stderr,
"ds4-test: Metal SSD streaming cache-pressure repro "
"(16GiB cache, layer-batched decode, short_code_completion)\n");
test_official_logprob_vectors_run("short_code_completion");
test_restore_env("DS4_METAL_MOE_ONE_STAGE_PROFILE", saved_one_stage);
test_restore_env("DS4_METAL_DISABLE_STREAMING_STATIC_DECODE_MAP",
saved_disable_static_decode);
test_restore_env("DS4_METAL_DISABLE_STREAMING_LAYER_BATCH",
saved_disable_layer_batch);
test_restore_env("DS4_TEST_SSD_STREAMING_CACHE_EXPERTS",
saved_cache_experts);
test_restore_env("DS4_TEST_SSD_STREAMING_CACHE_GB", saved_cache_gb);
test_restore_env("DS4_TEST_SSD_STREAMING", saved_streaming);
#endif
}
static void test_logits_topk(const float *logits, int n, int *out, int k);
static bool test_topk_contains(const int *top, int k, int id);
#define TEST_LOCAL_GOLDEN_MAX_TOP 128
typedef struct {
int id;
float logit;
} test_local_golden_top;
typedef struct {
char id[96];
char mode[16];
char prompt_path[512];
int ctx;
int frontier;
int ntop;
test_local_golden_top top[TEST_LOCAL_GOLDEN_MAX_TOP];
} test_local_golden_case;
static bool test_read_local_golden_case(FILE *fp, test_local_golden_case *tc) {
char line[2048];
memset(tc, 0, sizeof(*tc));
while (fgets(line, sizeof(line), fp)) {
char *p = test_trim_line(line);
if (!p[0] || p[0] == '#') continue;
if (sscanf(p, "case %95s %15s %d %d %511s %d",
tc->id, tc->mode, &tc->ctx, &tc->frontier,
tc->prompt_path, &tc->ntop) == 6) {
TEST_ASSERT(tc->ctx > tc->frontier);
TEST_ASSERT(tc->frontier > 0);
TEST_ASSERT(tc->ntop > 0 && tc->ntop <= TEST_LOCAL_GOLDEN_MAX_TOP);
return true;
}
TEST_ASSERT(!"unexpected line before local golden case");
return false;
}
return false;
}
static bool test_fill_local_golden_case(FILE *fp, test_local_golden_case *tc) {
char line[2048];
int seen = 0;
while (fgets(line, sizeof(line), fp)) {
char *p = test_trim_line(line);
if (!p[0] || p[0] == '#') continue;
if (!strcmp(p, "end")) {
TEST_ASSERT(seen == tc->ntop);
return seen == tc->ntop;
}
int rank = -1;
int id = -1;
float logit = 0.0f;
if (sscanf(p, "top %d %d %f", &rank, &id, &logit) != 3) {
TEST_ASSERT(!"bad local golden top line");
return false;
}
TEST_ASSERT(rank == seen);
TEST_ASSERT(seen < tc->ntop);
if (seen >= tc->ntop) return false;
tc->top[seen].id = id;
tc->top[seen].logit = logit;
seen++;
}
TEST_ASSERT(!"unterminated local golden case");
return false;
}
static int test_local_golden_overlap(const test_local_golden_case *tc,
const int *cand_top,
int n) {
int overlap = 0;
if (n > tc->ntop) n = tc->ntop;
for (int i = 0; i < n; i++) {
if (test_topk_contains(cand_top, n, tc->top[i].id)) overlap++;
}
return overlap;
}
static float test_local_golden_max_abs(const test_local_golden_case *tc,
const float *cand_logits,
int n) {
float max_abs = 0.0f;
if (n > tc->ntop) n = tc->ntop;
for (int i = 0; i < n; i++) {
const int id = tc->top[i].id;
if (id < 0) continue;
const float abs_delta = fabsf(cand_logits[id] - tc->top[i].logit);
if (abs_delta > max_abs) max_abs = abs_delta;
}
return max_abs;
}
static void test_local_golden_case_run(ds4_engine *engine,
const test_local_golden_case *tc) {
char *prompt_text = test_read_file(tc->prompt_path);
TEST_ASSERT(prompt_text != NULL);
if (!prompt_text) return;
ds4_tokens prompt = {0};
if (!strcmp(tc->mode, "text")) {
ds4_tokenize_text(engine, prompt_text, &prompt);
} else if (!strcmp(tc->mode, "rendered")) {
ds4_tokenize_rendered_chat(engine, prompt_text, &prompt);
} else if (!strcmp(tc->mode, "chat")) {
ds4_encode_chat_prompt(engine, "", prompt_text, DS4_THINK_NONE, &prompt);
} else {
TEST_ASSERT(!"unknown local golden prompt mode");
}
free(prompt_text);
TEST_ASSERT(prompt.len >= tc->frontier);
if (prompt.len < tc->frontier) {
ds4_tokens_free(&prompt);
return;
}
ds4_tokens prefix = {
.v = prompt.v,
.len = tc->frontier,
.cap = tc->frontier,
};
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, tc->ctx) == 0);
if (!session) {
ds4_tokens_free(&prompt);
return;
}
char err[160];
TEST_ASSERT(ds4_session_sync(session, &prefix, err, sizeof(err)) == 0);
const int vocab = ds4_engine_vocab_size(engine);
float *cand_logits = malloc((size_t)vocab * sizeof(cand_logits[0]));
TEST_ASSERT(cand_logits != NULL);
if (cand_logits &&
ds4_session_copy_logits(session, cand_logits, vocab) == vocab) {
int cand_top[TEST_LOCAL_GOLDEN_MAX_TOP];
const int ntop = tc->ntop < TEST_LOCAL_GOLDEN_MAX_TOP ?
tc->ntop : TEST_LOCAL_GOLDEN_MAX_TOP;
test_logits_topk(cand_logits, vocab, cand_top, ntop);
const int top5_overlap = test_local_golden_overlap(tc, cand_top, 5);
const int top20_overlap = test_local_golden_overlap(tc, cand_top, 20);
const int top64_overlap = test_local_golden_overlap(tc, cand_top, 64);
const float top20_max_abs =
test_local_golden_max_abs(tc, cand_logits, 20);
fprintf(stderr,
"ds4-test: local golden %s top1 ref=%d cand=%d "
"top5_overlap=%d/5 top20_overlap=%d/20 top64_overlap=%d/64 "
"top20_max_abs=%g\n",
tc->id, tc->top[0].id, cand_top[0],
top5_overlap, top20_overlap, top64_overlap, top20_max_abs);
/*
* This is intentionally tolerant: it is meant to catch substantial
* backend drift (wrong tiling, skipped work, bad dispatch), not tiny
* floating-point differences from otherwise sane kernel changes.
*/
TEST_ASSERT(cand_top[0] == tc->top[0].id);
TEST_ASSERT(top5_overlap >= 4);
TEST_ASSERT(top20_overlap >= 15);
TEST_ASSERT(top64_overlap >= 40);
TEST_ASSERT(top20_max_abs <= 8.0f);
} else {
TEST_ASSERT(false);
}
free(cand_logits);
ds4_session_free(session);
ds4_tokens_free(&prompt);
}
static void test_local_golden_vectors(void) {
const char *path = getenv("DS4_TEST_LOCAL_GOLDEN_FILE");
if (!path || !path[0]) {
path = "tests/test-vectors/flash-0731/local-golden.vec";
}
FILE *fp = fopen(path, "rb");
TEST_ASSERT(fp != NULL);
if (!fp) return;
char *saved_prefill_chunk = test_save_env("DS4_METAL_PREFILL_CHUNK");
char *saved_disable_metal4 = test_save_env("DS4_METAL_DISABLE_METAL4");
char *saved_moe_tile_max = test_save_env("DS4_METAL_MOE_TILE_MAX");
test_streaming_prefill_env saved_canonical_streaming_prefill =
test_force_canonical_streaming_prefill();
setenv("DS4_METAL_PREFILL_CHUNK", "4096", 1);
setenv("DS4_METAL_DISABLE_METAL4", "1", 1);
unsetenv("DS4_METAL_MOE_TILE_MAX");
ds4_engine *engine = test_open_engine(false);
if (!engine) {
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
test_restore_env("DS4_METAL_MOE_TILE_MAX", saved_moe_tile_max);
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
test_restore_env("DS4_METAL_PREFILL_CHUNK", saved_prefill_chunk);
fclose(fp);
return;
}
test_local_golden_case tc;
while (test_read_local_golden_case(fp, &tc)) {
if (!test_fill_local_golden_case(fp, &tc)) break;
test_local_golden_case_run(engine, &tc);
}
ds4_engine_close(engine);
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
test_restore_env("DS4_METAL_MOE_TILE_MAX", saved_moe_tile_max);
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
test_restore_env("DS4_METAL_PREFILL_CHUNK", saved_prefill_chunk);
fclose(fp);
}
#define TEST_MPP_EQ_MAX_CASES 8
#define TEST_MPP_EQ_TOPK 20
#define TEST_MPP_EQ_TOP5 5
#define TEST_MPP_EQ_DELTAS 5
typedef struct {
char id[96];
int ctx;
int vocab_size;
int gen_steps;
ds4_tokens prompt;
float *ref_logits;
int ref_gen[TEST_VEC_MAX_STEPS];
int ref_gen_len;
} test_mpp_eq_case;
typedef struct {
int ref_top1;
int cand_top1;
int overlap;
int top5_overlap;
int max_rank_delta;
int nonfinite;
float rms;
float max_abs;
float top20_max_abs;
bool same_top1;
bool pass;
} test_mpp_eq_result;
typedef struct {
const char *label;
int cases;
int capture_failures;
int logits_failures;
int greedy_failures;
int top1_mismatches;
int min_overlap;
int min_top5_overlap;
int worst_rank_delta;
float worst_rms;
float worst_max_abs;
float worst_top20_max_abs;
} test_mpp_eq_summary;
static void test_mpp_eq_case_free(test_mpp_eq_case *tc) {
if (!tc) return;
ds4_tokens_free(&tc->prompt);
free(tc->ref_logits);
memset(tc, 0, sizeof(*tc));
}
static void test_logits_topk(const float *logits, int n, int *out, int k) {
for (int i = 0; i < k; i++) out[i] = -1;
for (int id = 0; id < n; id++) {
const float v = logits[id];
if (!isfinite(v)) continue;
for (int j = 0; j < k; j++) {
if (out[j] < 0 || v > logits[out[j]]) {
for (int l = k - 1; l > j; l--) out[l] = out[l - 1];
out[j] = id;
break;
}
}
}
}
static bool test_topk_contains(const int *top, int k, int id) {
for (int i = 0; i < k; i++) {
if (top[i] == id) return true;
}
return false;
}
static int test_topk_rank(const int *top, int k, int id) {
for (int i = 0; i < k; i++) {
if (top[i] == id) return i;
}
return -1;
}
static void test_note_delta(int *ids, float *ref_vals, float *cand_vals,
float *abs_vals, int id, float ref, float cand) {
const float abs_delta = fabsf(cand - ref);
for (int i = 0; i < TEST_MPP_EQ_DELTAS; i++) {
if (ids[i] < 0 || abs_delta > abs_vals[i]) {
for (int j = TEST_MPP_EQ_DELTAS - 1; j > i; j--) {
ids[j] = ids[j - 1];
ref_vals[j] = ref_vals[j - 1];
cand_vals[j] = cand_vals[j - 1];
abs_vals[j] = abs_vals[j - 1];
}
ids[i] = id;
ref_vals[i] = ref;
cand_vals[i] = cand;
abs_vals[i] = abs_delta;
return;
}
}
}
static float test_top_union_max_abs(const float *ref, const float *cand,
const int *ref_top, const int *cand_top, int k) {
float max_abs = 0.0f;
for (int i = 0; i < k; i++) {
if (ref_top[i] >= 0) {
const float d = fabsf(cand[ref_top[i]] - ref[ref_top[i]]);
if (d > max_abs) max_abs = d;
}
if (cand_top[i] >= 0 && !test_topk_contains(ref_top, k, cand_top[i])) {
const float d = fabsf(cand[cand_top[i]] - ref[cand_top[i]]);
if (d > max_abs) max_abs = d;
}
}
return max_abs;
}
/*
* Metal4/TensorOps equivalence is a smoke test, not a demand for bitwise local
* logits. Tensor kernels change precision and reduction order, so the useful
* invariant here is: no NaNs, same first greedy token, and same short greedy
* continuation. Larger logit drift is still printed so it can be compared with
* official API-vector and long-context recall gates.
*/
static test_mpp_eq_result test_compare_mpp_logits(const test_mpp_eq_case *tc,
const float *cand_logits,
bool assert_thresholds) {
int ref_top[TEST_MPP_EQ_TOPK];
int cand_top[TEST_MPP_EQ_TOPK];
test_logits_topk(tc->ref_logits, tc->vocab_size, ref_top, TEST_MPP_EQ_TOPK);
test_logits_topk(cand_logits, tc->vocab_size, cand_top, TEST_MPP_EQ_TOPK);
int overlap = 0;
int top5_overlap = 0;
int max_rank_delta = 0;
for (int i = 0; i < TEST_MPP_EQ_TOPK; i++) {
const int cand_rank = test_topk_rank(cand_top, TEST_MPP_EQ_TOPK, ref_top[i]);
if (ref_top[i] >= 0 && cand_rank >= 0) {
overlap++;
const int rank_delta = abs(cand_rank - i);
if (rank_delta > max_rank_delta) max_rank_delta = rank_delta;
}
if (i < TEST_MPP_EQ_TOP5 &&
ref_top[i] >= 0 &&
test_topk_contains(cand_top, TEST_MPP_EQ_TOP5, ref_top[i])) {
top5_overlap++;
}
}
double sumsq = 0.0;
float max_abs = 0.0f;
int nonfinite = 0;
int delta_ids[TEST_MPP_EQ_DELTAS];
float delta_ref[TEST_MPP_EQ_DELTAS];
float delta_cand[TEST_MPP_EQ_DELTAS];
float delta_abs[TEST_MPP_EQ_DELTAS];
for (int i = 0; i < TEST_MPP_EQ_DELTAS; i++) {
delta_ids[i] = -1;
delta_ref[i] = 0.0f;
delta_cand[i] = 0.0f;
delta_abs[i] = 0.0f;
}
for (int i = 0; i < tc->vocab_size; i++) {
if (!isfinite(tc->ref_logits[i]) || !isfinite(cand_logits[i])) {
nonfinite++;
continue;
}
const float delta = cand_logits[i] - tc->ref_logits[i];
const float abs_delta = fabsf(delta);
if (abs_delta > max_abs) max_abs = abs_delta;
sumsq += (double)delta * (double)delta;
test_note_delta(delta_ids, delta_ref, delta_cand, delta_abs,
(int)i, tc->ref_logits[i], cand_logits[i]);
}
const float rms = (float)sqrt(sumsq / (double)tc->vocab_size);
const float top_abs = test_top_union_max_abs(tc->ref_logits, cand_logits,
ref_top, cand_top, TEST_MPP_EQ_TOPK);
const bool same_top1 = ref_top[0] >= 0 && ref_top[0] == cand_top[0];
test_mpp_eq_result result = {
.ref_top1 = ref_top[0],
.cand_top1 = cand_top[0],
.overlap = overlap,
.top5_overlap = top5_overlap,
.max_rank_delta = max_rank_delta,
.nonfinite = nonfinite,
.rms = rms,
.max_abs = max_abs,
.top20_max_abs = top_abs,
.same_top1 = same_top1,
.pass = nonfinite == 0 && same_top1,
};
fprintf(stderr,
"ds4-test: Tensor equivalence %s top1 ref=%d cand=%d top5_overlap=%d/%d overlap=%d/%d max_rank_delta=%d rms=%g max_abs=%g top20_max_abs=%g\n",
tc->id, ref_top[0], cand_top[0],
top5_overlap, TEST_MPP_EQ_TOP5,
overlap, TEST_MPP_EQ_TOPK,
max_rank_delta, rms, max_abs, top_abs);
fprintf(stderr, "ds4-test: Tensor equivalence %s largest deltas:", tc->id);
for (int i = 0; i < TEST_MPP_EQ_DELTAS && delta_ids[i] >= 0; i++) {
fprintf(stderr, " id=%d ref=%g cand=%g abs=%g",
delta_ids[i], delta_ref[i], delta_cand[i], delta_abs[i]);
}
fputc('\n', stderr);
if (assert_thresholds) {
TEST_ASSERT(nonfinite == 0);
TEST_ASSERT(same_top1);
}
return result;
}
static bool test_mpp_capture(ds4_engine *engine, const test_mpp_eq_case *tc,
float *logits, int *gen, int *gen_len) {
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, tc->ctx) == 0);
if (!session) return false;
char err[160];
bool ok = ds4_session_sync(session, &tc->prompt, err, sizeof(err)) == 0;
TEST_ASSERT(ok);
if (ok) {
ok = ds4_session_copy_logits(session, logits, tc->vocab_size) == tc->vocab_size;
TEST_ASSERT(ok);
}
int n = 0;
while (ok && n < tc->gen_steps) {
const int token = ds4_session_argmax(session);
gen[n++] = token;
if (n < tc->gen_steps && ds4_session_eval(session, token, err, sizeof(err)) != 0) {
ok = false;
TEST_ASSERT(false);
}
}
*gen_len = n;
ds4_session_free(session);
return ok;
}
static bool test_mpp_capture_logits_only(ds4_engine *engine,
const test_mpp_eq_case *tc,
float *logits) {
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, tc->ctx) == 0);
if (!session) return false;
char err[160];
bool ok = ds4_session_sync(session, &tc->prompt, err, sizeof(err)) == 0;
TEST_ASSERT(ok);
if (ok) {
ok = ds4_session_copy_logits(session, logits, tc->vocab_size) == tc->vocab_size;
TEST_ASSERT(ok);
}
ds4_session_free(session);
return ok;
}
static bool test_mpp_eq_case_selected(const char *id) {
const char *filter = getenv("DS4_TEST_MPP_EQ_CASE");
if (!filter || !filter[0]) return true;
char buf[256];
snprintf(buf, sizeof(buf), "%s", filter);
for (char *tok = strtok(buf, ","); tok; tok = strtok(NULL, ",")) {
tok = test_trim_line(tok);
if (tok[0] && strstr(id, tok)) return true;
}
return false;
}
static int test_load_mpp_cases(ds4_engine *engine, test_mpp_eq_case *cases, int cap) {
const char *path = getenv("DS4_TEST_VECTOR_FILE");
if (!path || !path[0]) {
path = "tests/test-vectors/flash-0731/official.vec";
}
FILE *fp = fopen(path, "rb");
TEST_ASSERT(fp != NULL);
if (!fp) return 0;
int ncase = 0;
test_vec_case vc;
while (ncase < cap && test_read_vector_case(fp, &vc)) {
if (!test_fill_vector_case(fp, &vc)) break;
if (!test_mpp_eq_case_selected(vc.id)) continue;
char *prompt_text = test_read_file(vc.prompt_path);
TEST_ASSERT(prompt_text != NULL);
if (!prompt_text) continue;
test_mpp_eq_case *tc = &cases[ncase++];
snprintf(tc->id, sizeof(tc->id), "%s", vc.id);
tc->ctx = vc.ctx;
tc->vocab_size = ds4_engine_vocab_size(engine);
tc->gen_steps = vc.nsteps < TEST_VEC_MAX_STEPS ? vc.nsteps : TEST_VEC_MAX_STEPS;
ds4_encode_chat_prompt(engine, "", prompt_text, DS4_THINK_NONE, &tc->prompt);
free(prompt_text);
TEST_ASSERT(tc->prompt.len > 0);
}
fclose(fp);
return ncase;
}
static void test_mpp_summary_init(test_mpp_eq_summary *summary, const char *label) {
memset(summary, 0, sizeof(*summary));
summary->label = label;
summary->min_overlap = TEST_MPP_EQ_TOPK;
summary->min_top5_overlap = TEST_MPP_EQ_TOP5;
}
static void test_mpp_summary_note_logits(test_mpp_eq_summary *summary,
const test_mpp_eq_result *result) {
if (!result->pass) summary->logits_failures++;
if (!result->same_top1) summary->top1_mismatches++;
if (result->overlap < summary->min_overlap) summary->min_overlap = result->overlap;
if (result->top5_overlap < summary->min_top5_overlap) {
summary->min_top5_overlap = result->top5_overlap;
}
if (result->max_rank_delta > summary->worst_rank_delta) {
summary->worst_rank_delta = result->max_rank_delta;
}
if (result->rms > summary->worst_rms) summary->worst_rms = result->rms;
if (result->max_abs > summary->worst_max_abs) summary->worst_max_abs = result->max_abs;
if (result->top20_max_abs > summary->worst_top20_max_abs) {
summary->worst_top20_max_abs = result->top20_max_abs;
}
}
static void test_mpp_summary_print(const test_mpp_eq_summary *summary) {
fprintf(stderr,
"ds4-test: Tensor summary route=%s cases=%d capture_fail=%d logits_fail=%d greedy_fail=%d top1_mismatch=%d min_top5_overlap=%d/%d min_overlap=%d/%d worst_rank_delta=%d worst_rms=%g worst_max_abs=%g worst_top20_max_abs=%g\n",
summary->label,
summary->cases,
summary->capture_failures,
summary->logits_failures,
summary->greedy_failures,
summary->top1_mismatches,
summary->min_top5_overlap,
TEST_MPP_EQ_TOP5,
summary->min_overlap,
TEST_MPP_EQ_TOPK,
summary->worst_rank_delta,
summary->worst_rms,
summary->worst_max_abs,
summary->worst_top20_max_abs);
}
static void test_run_mpp_candidate(const char *label,
test_mpp_eq_case *cases,
int ncase) {
fprintf(stderr, "ds4-test: Tensor equivalence candidate route=%s\n", label);
test_mpp_eq_summary summary;
test_mpp_summary_init(&summary, label);
ds4_engine *cand_engine = test_open_engine(false);
if (cand_engine) {
const int vocab_size = ncase > 0 ? cases[0].vocab_size : 0;
float *cand_logits = malloc((size_t)vocab_size * sizeof(cand_logits[0]));
TEST_ASSERT(cand_logits != NULL);
if (cand_logits) {
for (int i = 0; i < ncase; i++) {
test_mpp_eq_case *tc = &cases[i];
if (!tc->ref_logits) continue;
int cand_gen[TEST_VEC_MAX_STEPS] = {0};
int cand_gen_len = 0;
if (!test_mpp_capture(cand_engine, tc, cand_logits, cand_gen, &cand_gen_len)) {
summary.capture_failures++;
continue;
}
summary.cases++;
test_mpp_eq_result result = test_compare_mpp_logits(tc, cand_logits, true);
test_mpp_summary_note_logits(&summary, &result);
TEST_ASSERT(cand_gen_len == tc->ref_gen_len);
if (cand_gen_len != tc->ref_gen_len) summary.greedy_failures++;
for (int j = 0; j < tc->ref_gen_len && j < cand_gen_len; j++) {
if (cand_gen[j] != tc->ref_gen[j]) {
fprintf(stderr,
"ds4-test: Tensor equivalence %s greedy token mismatch step=%d ref=%d cand=%d\n",
tc->id, j, tc->ref_gen[j], cand_gen[j]);
summary.greedy_failures++;
}
TEST_ASSERT(cand_gen[j] == tc->ref_gen[j]);
}
}
free(cand_logits);
}
ds4_engine_close(cand_engine);
}
test_mpp_summary_print(&summary);
}
static void test_metal_mpp_equivalence(void) {
test_close_engines();
test_mpp_eq_case cases[TEST_MPP_EQ_MAX_CASES];
memset(cases, 0, sizeof(cases));
char *saved_disable_metal4 = test_save_env("DS4_METAL_DISABLE_METAL4");
setenv("DS4_METAL_DISABLE_METAL4", "1", 1);
ds4_engine *ref_engine = test_open_engine(false);
if (!ref_engine) {
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
return;
}
const int ncase = test_load_mpp_cases(ref_engine, cases, TEST_MPP_EQ_MAX_CASES);
TEST_ASSERT(ncase > 0);
for (int i = 0; i < ncase; i++) {
test_mpp_eq_case *tc = &cases[i];
tc->ref_logits = malloc((size_t)tc->vocab_size * sizeof(tc->ref_logits[0]));
TEST_ASSERT(tc->ref_logits != NULL);
if (!tc->ref_logits) continue;
TEST_ASSERT(test_mpp_capture(ref_engine, tc,
tc->ref_logits,
tc->ref_gen,
&tc->ref_gen_len));
}
ds4_engine_close(ref_engine);
test_restore_env("DS4_METAL_DISABLE_METAL4", saved_disable_metal4);
test_run_mpp_candidate("auto", cases, ncase);
for (int i = 0; i < ncase; i++) test_mpp_eq_case_free(&cases[i]);
}
static void test_streaming_decode_prefill_correctness(void) {
test_close_engines();
if (!test_env_bool("DS4_TEST_SSD_STREAMING")) {
fprintf(stderr,
"ds4-test: streaming decode-prefill correctness skipped "
"(set DS4_TEST_SSD_STREAMING=1 to enable)\n");
return;
}
test_mpp_eq_case cases[TEST_MPP_EQ_MAX_CASES];
memset(cases, 0, sizeof(cases));
test_streaming_prefill_env saved_canonical_streaming_prefill =
test_force_canonical_streaming_prefill();
ds4_engine *ref_engine = test_open_engine(false);
if (!ref_engine) {
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
return;
}
const int ncase = test_load_mpp_cases(ref_engine, cases, TEST_MPP_EQ_MAX_CASES);
TEST_ASSERT(ncase > 0);
for (int i = 0; i < ncase; i++) {
test_mpp_eq_case *tc = &cases[i];
tc->ref_logits = malloc((size_t)tc->vocab_size * sizeof(tc->ref_logits[0]));
TEST_ASSERT(tc->ref_logits != NULL);
if (!tc->ref_logits) continue;
TEST_ASSERT(test_mpp_capture(ref_engine, tc,
tc->ref_logits,
tc->ref_gen,
&tc->ref_gen_len));
}
ds4_engine_close(ref_engine);
unsetenv("DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL");
unsetenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR");
ds4_engine *cand_engine = test_open_engine(false);
if (cand_engine) {
for (int i = 0; i < ncase; i++) {
test_mpp_eq_case *tc = &cases[i];
if (!tc->ref_logits) continue;
float *cand_cold = malloc((size_t)tc->vocab_size * sizeof(cand_cold[0]));
float *cand_warm_a = malloc((size_t)tc->vocab_size * sizeof(cand_warm_a[0]));
float *cand_warm_b = malloc((size_t)tc->vocab_size * sizeof(cand_warm_b[0]));
TEST_ASSERT(cand_cold != NULL);
TEST_ASSERT(cand_warm_a != NULL);
TEST_ASSERT(cand_warm_b != NULL);
if (!cand_cold || !cand_warm_a || !cand_warm_b) {
free(cand_cold);
free(cand_warm_a);
free(cand_warm_b);
continue;
}
TEST_ASSERT(test_mpp_capture_logits_only(cand_engine, tc, cand_cold));
TEST_ASSERT(test_mpp_capture_logits_only(cand_engine, tc, cand_warm_a));
TEST_ASSERT(test_mpp_capture_logits_only(cand_engine, tc, cand_warm_b));
test_mpp_eq_result result = test_compare_mpp_logits(tc, cand_cold, false);
TEST_ASSERT(result.nonfinite == 0);
TEST_ASSERT(result.top5_overlap >= 2);
TEST_ASSERT(result.overlap >= 10);
TEST_ASSERT(result.rms <= 4.0f);
TEST_ASSERT(result.top20_max_abs <= 12.0f);
int cold_warm_neq = 0;
int warm_repeat_neq = 0;
int repeat_nonfinite = 0;
float cold_warm_max_abs = 0.0f;
float warm_repeat_max_abs = 0.0f;
for (int j = 0; j < tc->vocab_size; j++) {
if (!isfinite(cand_cold[j]) ||
!isfinite(cand_warm_a[j]) ||
!isfinite(cand_warm_b[j])) {
repeat_nonfinite++;
continue;
}
const float cold_warm_d = fabsf(cand_cold[j] - cand_warm_a[j]);
if (cold_warm_d != 0.0f) cold_warm_neq++;
if (cold_warm_d > cold_warm_max_abs) cold_warm_max_abs = cold_warm_d;
const float warm_repeat_d = fabsf(cand_warm_a[j] - cand_warm_b[j]);
if (warm_repeat_d != 0.0f) warm_repeat_neq++;
if (warm_repeat_d > warm_repeat_max_abs) {
warm_repeat_max_abs = warm_repeat_d;
}
}
TEST_ASSERT(repeat_nonfinite == 0);
TEST_ASSERT(cold_warm_neq == 0);
TEST_ASSERT(warm_repeat_neq == 0);
fprintf(stderr,
"ds4-test: streaming decode-prefill %s cold_warm_neq=%d "
"cold_warm_max_abs=%g warm_repeat_neq=%d "
"warm_repeat_max_abs=%g top1 canonical=%d decode=%d\n",
tc->id,
cold_warm_neq,
cold_warm_max_abs,
warm_repeat_neq,
warm_repeat_max_abs,
result.ref_top1,
result.cand_top1);
free(cand_cold);
free(cand_warm_a);
free(cand_warm_b);
}
ds4_engine_close(cand_engine);
}
test_restore_canonical_streaming_prefill(saved_canonical_streaming_prefill);
for (int i = 0; i < ncase; i++) test_mpp_eq_case_free(&cases[i]);
}
#define TEST_LIST_FILES_USER_PROMPT \
"Use the list_files tool to list the current directory exactly once, " \
"then report the listed files and stop."
#define TEST_LIST_FILES_TOOL_JSON \
"{\"type\":\"function\",\"function\":{" \
"\"name\":\"list_files\"," \
"\"description\":\"List files in a directory.\"," \
"\"parameters\":{\"type\":\"object\",\"properties\":{" \
"\"path\":{\"type\":\"string\",\"description\":\"Directory path to list.\"}" \
"},\"required\":[\"path\"]}" \
"}}"
#define TEST_LIST_FILES_RESULT "[\"README.md\",\"Makefile\",\"ds4.c\",\"metal\"]"
static const char *test_tool_call_request_json(void) {
return
"{"
"\"model\":\"deepseek-v4-flash\","
"\"messages\":[{\"role\":\"user\",\"content\":\""
TEST_LIST_FILES_USER_PROMPT
"\"}],"
"\"tools\":[" TEST_LIST_FILES_TOOL_JSON "],"
"\"tool_choice\":\"auto\","
"\"think\":false,"
"\"temperature\":0,"
"\"max_tokens\":256,"
"\"stream\":false"
"}";
}
static char *test_tool_result_request_json(const char *assistant_content,
const tool_call *call) {
if (!call || !call->id || !call->id[0]) return NULL;
buf b = {0};
buf_puts(&b,
"{\"model\":\"deepseek-v4-flash\",\"messages\":["
"{\"role\":\"user\",\"content\":");
json_escape(&b, TEST_LIST_FILES_USER_PROMPT);
buf_puts(&b, "},{\"role\":\"assistant\",\"content\":");
json_escape(&b, assistant_content ? assistant_content : "");
buf_puts(&b, ",\"tool_calls\":[{\"id\":");
json_escape(&b, call->id);
buf_puts(&b, ",\"type\":\"function\",\"function\":{\"name\":");
json_escape(&b, call->name ? call->name : "");
buf_puts(&b, ",\"arguments\":");
json_escape(&b, call->arguments ? call->arguments : "{}");
buf_puts(&b, "}}]},{\"role\":\"tool\",\"tool_call_id\":");
json_escape(&b, call->id);
buf_puts(&b, ",\"content\":");
json_escape(&b, TEST_LIST_FILES_RESULT);
buf_puts(&b,
"}],\"tools\":[" TEST_LIST_FILES_TOOL_JSON "],"
"\"tool_choice\":\"auto\",\"think\":false,"
"\"temperature\":0,\"max_tokens\":256,\"stream\":false}");
return buf_take(&b);
}
/* A complete tool call inside unclosed reasoning is recovered directly. The
* detector must wait for the complete block, and the parser must keep only the
* preceding prose in reasoning_content. */
static void test_think_tool_recovery(void) {
const char *generated =
"The user wants a directory listing.\n\n"
DS4_TOOL_CALLS_START "\n"
DS4_INVOKE_START " name=\"list_files\">\n"
DS4_PARAM_START " name=\"path\" string=\"true\">." DS4_PARAM_END "\n"
DS4_INVOKE_END "\n"
DS4_TOOL_CALLS_END;
buf text = {0};
size_t scan_from = 0;
bool complete = false;
for (size_t i = 0; generated[i]; i++) {
buf_append(&text, generated + i, 1);
complete = complete_tool_call_inside_thinking(text.ptr, text.len,
&scan_from);
TEST_ASSERT(complete == (generated[i + 1] == '\0'));
}
TEST_ASSERT(complete);
char *content = NULL;
char *reasoning = NULL;
tool_calls calls = {0};
bool parsed = parse_generated_message_ex(text.ptr, true,
&content, &reasoning, &calls);
TEST_ASSERT(parsed);
TEST_ASSERT(calls.len > 0 && !strcmp(calls.v[0].name, "list_files"));
TEST_ASSERT(calls.v[0].arguments && strstr(calls.v[0].arguments, "\"path\": \".\""));
TEST_ASSERT(content && content[0] == '\0');
TEST_ASSERT(reasoning && !strcmp(reasoning, "The user wants a directory listing."));
fprintf(stderr,
"ds4-test: think-tool-recovery complete=%d calls=%d name=%s\n",
complete ? 1 : 0, calls.len, calls.len ? calls.v[0].name : "-");
free(content);
free(reasoning);
tool_calls_free(&calls);
buf_free(&text);
}
typedef struct {
char *raw;
char *content;
char *reasoning;
tool_calls calls;
const char *finish;
} test_chat_turn;
static void test_chat_turn_free(test_chat_turn *turn) {
if (!turn) return;
free(turn->raw);
free(turn->content);
free(turn->reasoning);
tool_calls_free(&turn->calls);
memset(turn, 0, sizeof(*turn));
}
/* Run the same greedy stop/tool-marker/response parse path needed by the server,
* while keeping one session alive so the next request genuinely continues from
* this sampled turn. */
static bool test_generate_chat_turn(ds4_engine *engine, ds4_session *session,
const request *r, test_chat_turn *turn) {
memset(turn, 0, sizeof(*turn));
char err[160] = {0};
if (ds4_session_sync(session, &r->prompt, err, sizeof(err)) != 0) {
fprintf(stderr, "ds4-test: tool-call sync failed: %s\n", err);
turn->finish = "error";
return false;
}
buf text = {0};
uint64_t rng = 123;
const char *finish = "length";
bool saw_tool_start = false;
bool saw_tool_end = false;
bool decode_ok = true;
for (int i = 0; i < r->max_tokens; i++) {
int token = ds4_session_sample(session, r->temperature, r->top_k,
r->top_p, r->min_p, &rng);
if (ds4_token_is_stop_for_think_mode(engine, token, r->think_mode)) {
finish = "stop";
break;
}
if (ds4_session_eval(session, token, err, sizeof(err)) != 0) {
finish = "error";
decode_ok = false;
break;
}
size_t piece_len = 0;
char *piece = ds4_token_text(engine, token, &piece_len);
buf_append(&text, piece, piece_len);
free(piece);
if (r->has_tools) {
observe_tool_markers(text.ptr ? text.ptr : "",
&saw_tool_start, &saw_tool_end, NULL);
if (saw_tool_end) {
finish = "tool_calls";
break;
}
}
}
turn->raw = buf_take(&text);
turn->finish = finish;
if (!decode_ok) {
fprintf(stderr, "ds4-test: tool-call decode failed: %s\n", err);
return false;
}
bool recovered = false;
bool parsed = parse_generated_message_for_response_for_syntax(
r->model_syntax,
turn->raw ? turn->raw : "",
r->has_tools,
saw_tool_start,
ds4_think_mode_enabled(r->think_mode),
&turn->finish,
err,
sizeof(err),
&turn->content,
&turn->reasoning,
&turn->calls,
&recovered);
if (turn->calls.len > 0) turn->finish = "tool_calls";
if (!parsed) {
fprintf(stderr,
"ds4-test: generated message parse failed: %s recovered=%d raw=%s\n",
err, recovered ? 1 : 0, turn->raw ? turn->raw : "");
}
return parsed;
}
static void test_tool_call_quality_one(bool quality) {
ds4_engine *engine = test_get_engine(quality);
if (!engine) return;
server s = {0};
s.engine = engine;
pthread_mutex_init(&s.tool_mu, NULL);
request first_request = {0};
request second_request = {0};
ds4_session *session = NULL;
test_chat_turn first = {0};
test_chat_turn second = {0};
char *second_body = NULL;
char err[160] = {0};
bool request_ok = parse_chat_request(engine, &s,
test_tool_call_request_json(),
512, 32768, &first_request,
err, sizeof(err));
TEST_ASSERT(request_ok);
if (!request_ok) goto done;
bool session_ok = ds4_session_create(&session, engine, 32768) == 0;
TEST_ASSERT(session_ok);
if (!session_ok) goto done;
bool first_ok = test_generate_chat_turn(engine, session,
&first_request, &first);
TEST_ASSERT(first_ok);
TEST_ASSERT(first.finish && !strcmp(first.finish, "tool_calls"));
TEST_ASSERT(first.calls.len == 1);
TEST_ASSERT(first.calls.len == 1 && first.calls.v[0].name &&
!strcmp(first.calls.v[0].name, "list_files"));
TEST_ASSERT(first.calls.raw_tool_text && first.calls.raw_tool_text[0]);
if (!first_ok || first.calls.len != 1 ||
!first.calls.v[0].name ||
strcmp(first.calls.v[0].name, "list_files") ||
!first.calls.raw_tool_text || !first.calls.raw_tool_text[0]) {
goto done;
}
/* Use the real response-side id assignment and exact sampled-DSML memory.
* The same id is serialized on both the assistant call and tool result. */
assign_tool_call_ids(&s, &first.calls, API_OPENAI);
TEST_ASSERT(first.calls.v[0].id && first.calls.v[0].id[0]);
if (!first.calls.v[0].id || !first.calls.v[0].id[0]) goto done;
tool_memory_remember(&s, &first.calls);
second_body = test_tool_result_request_json(first.content,
&first.calls.v[0]);
TEST_ASSERT(second_body != NULL);
if (!second_body) goto done;
err[0] = '\0';
request_ok = parse_chat_request(engine, &s, second_body,
512, 32768, &second_request,
err, sizeof(err));
TEST_ASSERT(request_ok);
if (!request_ok) goto done;
TEST_ASSERT(second_request.tool_replay.mem == 1);
TEST_ASSERT(second_request.tool_replay.disk == 0);
TEST_ASSERT(second_request.tool_replay.canonical == 0);
TEST_ASSERT(second_request.tool_replay.missing_ids == 0);
bool second_ok = test_generate_chat_turn(engine, session,
&second_request, &second);
TEST_ASSERT(second_ok);
TEST_ASSERT(second.finish && !strcmp(second.finish, "stop"));
TEST_ASSERT(second.calls.len == 0);
TEST_ASSERT(second.content && second.content[0]);
fprintf(stderr,
"ds4-test: post-tool-result turn1 finish_reason=%s tool_calls=%d "
"turn2 finish_reason=%s tool_calls=%d replay_mem=%d\n",
first.finish ? first.finish : "-", first.calls.len,
second.finish ? second.finish : "-", second.calls.len,
second_request.tool_replay.mem);
done:
free(second_body);
test_chat_turn_free(&second);
test_chat_turn_free(&first);
ds4_session_free(session);
request_free(&second_request);
request_free(&first_request);
tool_memory_free(&s.tool_mem);
pthread_mutex_destroy(&s.tool_mu);
}
static void test_tool_call_quality(void) {
fprintf(stderr, "ds4-test: tool-call quality fast path\n");
test_tool_call_quality_one(false);
test_close_engine(false);
fprintf(stderr, "ds4-test: tool-call quality exact path\n");
test_tool_call_quality_one(true);
test_close_engine(true);
}
/* Greedy speculative decode: capture committed tokens and the largest accepted
* chunk, so the caller can confirm the multi-row verify path actually ran. */
static bool test_mtp_capture_speculative(ds4_engine *engine, const ds4_tokens *prompt,
int max_tokens, int *out, int *out_len,
int *max_chunk) {
*out_len = 0;
*max_chunk = 0;
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, 32768) == 0);
if (!session) return false;
char err[160];
bool ok = ds4_session_sync(session, prompt, err, sizeof(err)) == 0;
TEST_ASSERT(ok);
const int eos = ds4_token_eos(engine);
int n = 0;
bool stop = false;
while (ok && !stop && n < max_tokens) {
const int token = ds4_session_argmax(session);
if (token == eos) break;
int toks[17]; /* base token + draft depth, which the engine clamps to 16 */
const int ntok = ds4_session_eval_speculative_argmax(
session, token, max_tokens - n, eos, toks,
(int)(sizeof(toks) / sizeof(toks[0])), err, sizeof(err));
if (ntok < 0) { ok = false; TEST_ASSERT(false); break; }
if (ntok > *max_chunk) *max_chunk = ntok;
for (int j = 0; j < ntok; j++) {
if (toks[j] == eos) { stop = true; break; }
out[n++] = toks[j];
if (n >= max_tokens) { stop = true; break; }
}
}
*out_len = n;
ds4_session_free(session);
return ok;
}
/* Replay toks[] through plain decode and return the largest gap between a
* position's argmax logit and the committed token's logit. Correct speculation
* commits (near-)argmax tokens (gap ~0); a mis-committed token gives a big gap. */
static bool test_mtp_worst_argmax_gap(ds4_engine *engine, const ds4_tokens *prompt,
const int *toks, int n,
float *worst_gap, int *worst_at) {
*worst_gap = 0.0f;
*worst_at = -1;
ds4_session *session = NULL;
TEST_ASSERT(ds4_session_create(&session, engine, 32768) == 0);
if (!session) return false;
char err[160];
bool ok = ds4_session_sync(session, prompt, err, sizeof(err)) == 0;
TEST_ASSERT(ok);
for (int i = 0; ok && i < n; i++) {
ds4_token_score best, cur;
ok = ds4_session_top_logprobs(session, &best, 1) >= 1 &&
ds4_session_token_logprob(session, toks[i], &cur) == 1;
TEST_ASSERT(ok);
if (!ok) break;
const float gap = best.logit - cur.logit;
if (gap > *worst_gap) { *worst_gap = gap; *worst_at = i; }
if (ds4_session_eval(session, toks[i], err, sizeof(err)) != 0) { ok = false; TEST_ASSERT(false); break; }
}
ds4_session_free(session);
return ok;
}
/* Verbatim-copy task: keeps the model confident (a mis-committed token shows as
* a large argmax gap) and draft acceptance high (so the multi-row verify path is
* exercised across the generation). */
static const char *test_mtp_copy_prompt(void) {
return
"Reproduce the following C code EXACTLY, character for character, "
"inside a single code block and output nothing else:\n\n"
"```c\n"
"static uint32_t clamp_u32(uint32_t v, uint32_t lo, uint32_t hi) {\n"
" if (v < lo) return lo;\n"
" if (v > hi) return hi;\n"
" return v;\n"
"}\n"
"\n"
"static uint32_t ring_advance(uint32_t pos, uint32_t cap) {\n"
" uint32_t next = pos + 1u;\n"
" return next >= cap ? 0u : next;\n"
"}\n"
"\n"
"static int scratch_init(scratch *s, uint32_t ctx_size) {\n"
" if (ctx_size == 0u) ctx_size = 1u;\n"
" s->ctx_size = ctx_size;\n"
" s->comp_cap = ctx_size / 4u + 2u;\n"
" s->rows = clamp_u32(s->comp_cap, 1u, 4096u);\n"
" s->head = 0u;\n"
" return s->rows > 0u ? 0 : -1;\n"
"}\n"
"```\n";
}
#define TEST_MTP_MAXGEN 256
#define TEST_DSPARK_MAXGEN 128
static ds4_engine *test_open_dspark_engine(const char *support_path) {
ds4_engine *engine = NULL;
ds4_engine_options opt = {
.model_path = test_model_path(),
#ifdef __APPLE__
.backend = DS4_BACKEND_METAL,
#else
.backend = DS4_BACKEND_CUDA,
#endif
.quality = false,
.ssd_streaming = test_env_bool("DS4_TEST_SSD_STREAMING"),
.ssd_streaming_cold = test_env_bool("DS4_TEST_SSD_STREAMING_COLD"),
.ssd_streaming_cache_experts =
test_env_u32("DS4_TEST_SSD_STREAMING_CACHE_EXPERTS"),
.ssd_streaming_cache_bytes =
test_env_gib("DS4_TEST_SSD_STREAMING_CACHE_GB"),
.ssd_streaming_preload_experts =
test_env_u32("DS4_TEST_SSD_STREAMING_PRELOAD_EXPERTS"),
.mtp_path = support_path,
.mtp_draft_tokens = 0,
.dspark = true,
.dspark_confidence_threshold = 0.9f,
.dspark_confidence_threshold_set = true,
};
const int rc = ds4_engine_open(&engine, &opt);
TEST_ASSERT(rc == 0);
return rc == 0 ? engine : NULL;
}
/* Regression for the swapped top-k arguments in metal_graph_verify_suffix_tops
* at draft depth > 2. Replays the committed speculative tokens through plain
* decode and requires each to be a (near-)argmax: that is the verify invariant,
* and unlike comparing token streams it tolerates the near-greedy tie
* divergences. Needs an MTP head, so it self-skips without DS4_TEST_MTP. */
static void test_mtp_verify_depth(void) {
ds4_engine *engine = test_get_engine(false);
if (!engine || !ds4_engine_has_mtp(engine)) {
fprintf(stderr, "ds4-test: mtp-verify-depth skipped (set DS4_TEST_MTP to an MTP GGUF)\n");
return;
}
TEST_ASSERT(ds4_engine_mtp_draft_tokens(engine) > 2);
ds4_tokens prompt = {0};
ds4_chat_begin(engine, &prompt);
ds4_chat_append_message(engine, &prompt, "user", test_mtp_copy_prompt());
ds4_chat_append_assistant_prefix(engine, &prompt, DS4_THINK_NONE);
TEST_ASSERT(prompt.len > 0);
int *spec = malloc((size_t)TEST_MTP_MAXGEN * sizeof(*spec));
TEST_ASSERT(spec != NULL);
if (spec && prompt.len > 0) {
int nspec = 0, max_chunk = 0;
const bool ok_spec = test_mtp_capture_speculative(engine, &prompt, TEST_MTP_MAXGEN,
spec, &nspec, &max_chunk);
TEST_ASSERT(ok_spec);
TEST_ASSERT(max_chunk > 1); /* multi-token chunks committed: the multi-row path ran */
TEST_ASSERT(nspec > 128); /* enough output to surface the bug, incl. a spurious-EOS truncation */
float worst_gap = 0.0f;
int worst_at = -1;
const bool ok_check = test_mtp_worst_argmax_gap(engine, &prompt, spec, nspec,
&worst_gap, &worst_at);
TEST_ASSERT(ok_check);
fprintf(stderr, "ds4-test: mtp-verify-depth nspec=%d max_chunk=%d worst_argmax_gap=%.3f at=%d\n",
nspec, max_chunk, worst_gap, worst_at);
TEST_ASSERT(worst_gap <= 2.0f); /* correct: ~0; bug: ~21 on the reference model */
}
free(spec);
ds4_tokens_free(&prompt);
}
/* Same invariant as the MTP depth smoke, but for the DSpark support model. This
* is separate from the fixture because it teacher-forces every committed token
* through normal decode and directly checks that DSpark never commits a token
* that was not near the target argmax. */
static void test_dspark_verify_depth(void) {
const char *support = getenv("DS4_TEST_DSPARK");
if (!support || !support[0]) {
fprintf(stderr, "ds4-test: dspark-verify-depth skipped (set DS4_TEST_DSPARK to a DSpark support GGUF)\n");
return;
}
char *saved_scheduler = test_save_env("DS4_DSPARK_SCHEDULER");
setenv("DS4_DSPARK_SCHEDULER", "0", 1);
ds4_engine *engine = test_open_dspark_engine(support);
ds4_tokens prompt = {0};
int *spec = NULL;
if (engine) {
const int draft_depth = ds4_engine_mtp_draft_tokens(engine);
TEST_ASSERT(draft_depth > 2);
ds4_chat_begin(engine, &prompt);
ds4_chat_append_message(engine, &prompt, "user", test_mtp_copy_prompt());
ds4_chat_append_assistant_prefix(engine, &prompt, DS4_THINK_NONE);
TEST_ASSERT(prompt.len > 0);
spec = malloc((size_t)TEST_DSPARK_MAXGEN * sizeof(*spec));
TEST_ASSERT(spec != NULL);
if (draft_depth > 2 && spec && prompt.len > 0) {
int nspec = 0, max_chunk = 0;
const bool ok_spec = test_mtp_capture_speculative(engine, &prompt,
TEST_DSPARK_MAXGEN,
spec, &nspec,
&max_chunk);
TEST_ASSERT(ok_spec);
TEST_ASSERT(max_chunk > 1);
TEST_ASSERT(nspec > 64);
float worst_gap = 0.0f;
int worst_at = -1;
const bool ok_check = test_mtp_worst_argmax_gap(engine, &prompt,
spec, nspec,
&worst_gap,
&worst_at);
TEST_ASSERT(ok_check);
fprintf(stderr,
"ds4-test: dspark-verify-depth nspec=%d max_chunk=%d draft_depth=%d worst_argmax_gap=%.3f at=%d\n",
nspec, max_chunk, draft_depth, worst_gap, worst_at);
TEST_ASSERT(worst_gap <= 2.0f);
}
}
free(spec);
ds4_tokens_free(&prompt);
ds4_engine_close(engine);
test_restore_env("DS4_DSPARK_SCHEDULER", saved_scheduler);
}
#endif
static void test_server_unit_group(void) {
ds4_server_unit_tests_run();
}
typedef void (*test_fn)(void);
typedef struct {
const char *flag;
const char *name;
const char *desc;
test_fn fn;
} ds4_test_entry;
static const ds4_test_entry test_entries[] = {
#ifndef DS4_NO_GPU
{"--long-context", "long-context", "long-context story fact-recall regression", test_long_story_fact_recall},
{"--tool-call-quality", "tool-call-quality", "model tool call and post-result stop regression", test_tool_call_quality},
{"--think-tool-recovery", "think-tool-recovery", "recover a complete tool call emitted inside unclosed reasoning", test_think_tool_recovery},
{"--logprob-vectors", "logprob-vectors", "official API top-logprob vector comparison on the standard Metal path", test_official_logprob_vectors},
{"--metal-ssd-streaming-cache-pressure", "metal-ssd-streaming-cache-pressure", "Metal SSD-streaming layer-batched decode cache-pressure repro for issue #384", test_metal_ssd_streaming_cache_pressure},
{"--local-golden-vectors", "local-golden-vectors", "local top-k/logit drift regression for long Metal prefill", test_local_golden_vectors},
{"--metal-short-prefill", "metal-short-prefill", "Metal ratio-4 short prefill regression", test_metal_short_prefill_ratio4},
{"--metal-kernels", "metal-kernels", "isolated Metal kernel numeric regressions", test_metal_kernel_group},
{"--metal-tensor-equivalence", "metal-tensor-equivalence", "fast/quality Metal prompt-logit and greedy equivalence", test_metal_mpp_equivalence},
{"--streaming-decode-prefill-correctness", "streaming-decode-prefill-correctness", "streaming decode-style cold prefill drift and repeatability", test_streaming_decode_prefill_correctness},
{"--mtp-verify-depth", "mtp-verify-depth", "MTP speculative verify commits autoregressive-identical tokens at draft depth > 2", test_mtp_verify_depth},
{"--dspark-verify-depth", "dspark-verify-depth", "DSpark speculative verify commits autoregressive-identical tokens at draft depth > 2", test_dspark_verify_depth},
#endif
{"--server", "server", "server parser/rendering/cache unit tests", test_server_unit_group},
};
static void test_print_help(const char *prog) {
printf("Usage: %s [--all | TEST...]\n\n", prog);
puts("Tests:");
puts(" --all");
puts(" Run every test. This is the default, ordered from slower to faster.");
for (size_t i = 0; i < sizeof(test_entries) / sizeof(test_entries[0]); i++) {
printf(" %-20s %s\n", test_entries[i].flag, test_entries[i].desc);
}
puts(" --list");
puts(" Print test names only.");
#ifndef DS4_NO_GPU
puts(" --metal-mpp-equivalence");
puts(" Compatibility alias for --metal-tensor-equivalence.");
#endif
puts(" -h, --help");
puts(" Show this help.");
puts("\nEnvironment:");
puts(" DS4_TEST_MODEL=FILE Model path. Default: ds4flash.gguf");
puts(" DS4_TEST_BACKEND=cpu Run model tests on CPU instead of Metal/CUDA.");
puts(" DS4_TEST_SSD_STREAMING=1 Run model tests through Metal SSD streaming.");
puts(" DS4_TEST_SSD_STREAMING_CACHE_GB=N Streaming routed expert cache in GiB.");
puts(" DS4_TEST_SSD_STREAMING_CACHE_EXPERTS=N Streaming routed expert cache count.");
puts(" DS4_TEST_SSD_STREAMING_COLD=1 Skip streaming hot expert preload.");
puts(" DS4_METAL_DISABLE_STREAMING_COLD_DECODE_PREFILL=1 Force canonical streamed cold prefill.");
puts(" DS4_TEST_LONG_PROMPT=FILE Rendered long-context story fact prompt.");
puts(" DS4_TEST_VECTOR_FILE=FILE Official fixture. Default: flash-0731/official.vec.");
puts(" DS4_TEST_LOCAL_GOLDEN_FILE=FILE Local fixture. Default: flash-0731/local-golden.vec.");
puts(" DS4_TEST_MPP_EQ_CASE=NAME Run only Tensor equivalence cases whose id contains NAME.");
puts(" DS4_TEST_MTP=FILE Legacy MTP support GGUF for --mtp-verify-depth.");
puts(" DS4_TEST_DSPARK=FILE DSpark support GGUF for --dspark-verify-depth.");
}
static const ds4_test_entry *test_find_entry(const char *arg) {
#ifndef DS4_NO_GPU
if (!strcmp(arg, "--metal-mpp-equivalence")) {
arg = "--metal-tensor-equivalence";
}
#endif
for (size_t i = 0; i < sizeof(test_entries) / sizeof(test_entries[0]); i++) {
if (!strcmp(arg, test_entries[i].flag)) return &test_entries[i];
}
return NULL;
}
static void test_run_entry(const ds4_test_entry *entry) {
int before = test_failures;
fprintf(stderr, "%s:\n", entry->name);
entry->fn();
fprintf(stderr, "%s: ", entry->name);
ds4_log(stderr,
test_failures == before ? DS4_LOG_OK : DS4_LOG_ERROR,
"%s",
test_failures == before ? "OK" : "ERR");
fputc('\n', stderr);
}
int main(int argc, char **argv) {
bool run_all = argc == 1;
bool selected[sizeof(test_entries) / sizeof(test_entries[0])] = {0};
for (int i = 1; i < argc; i++) {
if (!strcmp(argv[i], "--all")) {
run_all = true;
} else if (!strcmp(argv[i], "--list")) {
for (size_t j = 0; j < sizeof(test_entries) / sizeof(test_entries[0]); j++) {
puts(test_entries[j].flag);
}
return 0;
} else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) {
test_print_help(argv[0]);
return 0;
} else {
const ds4_test_entry *entry = test_find_entry(argv[i]);
if (!entry) {
fprintf(stderr, "ds4-test: unknown test switch: %s\n", argv[i]);
test_print_help(argv[0]);
return 2;
}
selected[(size_t)(entry - test_entries)] = true;
}
}
if (run_all) {
for (size_t i = 0; i < sizeof(test_entries) / sizeof(test_entries[0]); i++) {
test_run_entry(&test_entries[i]);
}
} else {
for (size_t i = 0; i < sizeof(test_entries) / sizeof(test_entries[0]); i++) {
if (selected[i]) test_run_entry(&test_entries[i]);
}
}
#ifndef DS4_NO_GPU
test_close_engines();
#endif
if (test_failures) {
fprintf(stderr, "ds4 tests: %d failure(s)\n", test_failures);
return 1;
}
puts("ds4 tests: ok");
return 0;
}
Reference in New Issue View Git Blame Copy Permalink