David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/metal/tests/test_metal_moe_experts.c

660 lines
26 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file test_metal_moe_experts.c
* @brief Absolute-correctness tests for the grouped MoE expert kernels: the
* Q3_K gate+up SwiGLU (slot-major and sort-prepass-driven) and the Q3_K /
* Q4_K down projection with its routed-weight sum.
*
* Every assertion is against the CPU reference in this file, never against
* another GPU arm. The reference dequantizes each super-block with a
* CLOSED-FORM per-weight index (and, for Q3_K, a per-scale bit-field unpack),
* while the kernels walk the block by half, shift level and 16-weight group —
* the same mapping derived two different ways, so a transcription slip in one
* cannot hide in the other.
*
* The weights are REAL: blk.0.ffn_{gate,up}_exps (Q3_K), blk.5.ffn_down_exps
* (Q3_K) and blk.0.ffn_down_exps (Q4_K) are read out of the shipped GGUF at
* their true file offsets. A synthetic block cannot catch a misread of the
* merged 3-D expert layout, which is the thing most likely to be wrong here.
*
* Half the model's layers ship ffn_down_exps as Q4_K, not Q3_K — the plan's
* "all Q3_K in the MVP" is wrong on this point — so both decoders are tested.
*
* @status All tests pass on M4 / macOS 26. Proven able to fail: a 0.01% scaled
* row sum, a swapped gate/up half assignment, an off-by-one super-block
* stride and a dropped routed weight were induced and all caught.
* @issues None
* @todo None
*/
#include "lib/test.h"
/* The family header directly, not the fm_metal.h umbrella — see the note in
* test_metal_moe.c. */
#include "metal/kernels/fm_kernel_moe.h"
#include "metal/tests/gguf_test_reader.h"
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
/* ================================================================
* Fixtures
* ================================================================ */
#define Q3_K_BLOCK_BYTES 110
#define Q4_K_BLOCK_BYTES 144
#define BLOCK_WEIGHTS 256
#define GGUF_TYPE_Q3_K 11
#define GGUF_TYPE_Q4_K 12
#define MOE_N_EMBD 2048
#define MOE_EXPERT_FF 512
#define MOE_EXPERT_USED 8
/* Experts lifted out of the real merged tensors, spread across the 256 so the
* test exercises far-apart expert strides rather than the first few. */
#define MOE_TEST_EXPERTS 8
static const int kSourceExperts[MOE_TEST_EXPERTS] = { 0, 3, 17, 42, 91, 128, 200, 255 };
/* A byte offset that is a multiple of 32 (GGUF's tensor alignment) but not of
* the 16384-byte page — the real placement a weight gets inside a region. */
#define WEIGHT_BIND_OFFSET (1024u * 1024u + 32u)
static const char *kDefaultModelPath =
"models/Nex-N2-mini/nex-agi_Nex-N2-mini-Q3_K_M.gguf";
static gguf_test_file g_model;
static int g_model_is_open = 0;
static uint32_t moe_rand_state = 0xC0FFEE11u;
static uint32_t moe_rand_u32(void)
{
moe_rand_state ^= moe_rand_state << 13;
moe_rand_state ^= moe_rand_state >> 17;
moe_rand_state ^= moe_rand_state << 5;
return moe_rand_state;
}
static float moe_rand_signed(float magnitude)
{
float unit = (float)(moe_rand_u32() >> 8) / (float)(1u << 24);
return (unit * 2.0f - 1.0f) * magnitude;
}
/* ================================================================
* CPU dequant reference — closed form, from the GGUF block layouts
* ================================================================ */
/* IEEE-754 binary16 -> binary32, including subnormals, Inf and NaN. */
static float half_bits_to_float(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 {
/* Subnormal: value = mant * 2^-24. Shifting left until bit 10 is
* set leaves 1.f * 2^(-14 - shift), so the binary32 exponent field
* is 127 - 14 - shift. Getting this off by one halves every
* subnormal scale, and a real Q4_K `d` IS subnormal — block 1 of
* blk.0.ffn_down_exps row 0 is 4.55e-05, below half's 6.10e-05
* smallest normal. */
int shift = 0;
while ((mant & 0x0400u) == 0) { mant <<= 1; shift++; }
mant &= 0x03FFu;
bits = sign | ((uint32_t)(127 - 14 - shift) << 23) | (mant << 13);
}
} else if (exp == 31) {
bits = sign | 0x7F800000u | (mant << 13);
} else {
bits = sign | ((exp + 127 - 15) << 23) | (mant << 13);
}
float f;
memcpy(&f, &bits, sizeof(f));
return f;
}
static uint16_t read_u16(const uint8_t *p)
{
return (uint16_t)((uint16_t)p[0] | ((uint16_t)p[1] << 8));
}
/**
* Dequantizes one 110-byte Q3_K super-block: hmask[32] @0, qs[64] @32,
* scales[12] @96, d (fp16) @108. Weight i takes its 2 low bits from
* qs[(i/128)*32 + i%32] at shift 2*((i%128)/32), its high bit from hmask[i%32]
* bit (i/128)*4 + (i%128)/32 — CLEARED means subtract 4 — and its scale from
* sub-block i/16. Six-bit scale j takes its low 4 bits from packed byte j%8
* (low nibble for j < 8, high nibble after) and its high 2 bits from packed
* byte 8 + j%4 at bit 2*(j/4).
*/
static void cpu_q3k_dequant_block(const uint8_t *blk, float *out)
{
const uint8_t *hmask = blk;
const uint8_t *qs = blk + 32;
float d = half_bits_to_float(read_u16(blk + 108));
const uint8_t *packed = blk + 96;
uint8_t scales[16];
for (int j = 0; j < 16; j++) {
int low = (j < 8) ? (packed[j] & 0x0f) : (packed[j - 8] >> 4);
int high = (packed[8 + j % 4] >> (2 * (j / 4))) & 3;
scales[j] = (uint8_t)(low | (high << 4));
}
for (int i = 0; i < BLOCK_WEIGHTS; i++) {
int qs_index = (i / 128) * 32 + (i % 32);
int shift = 2 * ((i % 128) / 32);
int high_bit = (i / 128) * 4 + (i % 128) / 32;
int quant = (qs[qs_index] >> shift) & 3;
int bias = (hmask[i % 32] & (1 << high_bit)) ? 0 : 4;
out[i] = d * (float)((int)scales[i / 16] - 32) * (float)(quant - bias);
}
}
/**
* Dequantizes one 144-byte Q4_K super-block: d @0, dmin @2 (both fp16),
* scales[12] @4, qs[128] @16. Weight i sits in 64-weight group i/64: the
* first 32 take the low nibble of qs[32*(i/64) + i%32] with sub-block scale
* 2*(i/64), the second 32 the high nibble with sub-block 2*(i/64)+1.
* value = d*scale*q - dmin*min.
*/
static void cpu_q4k_dequant_block(const uint8_t *blk, float *out)
{
float d = half_bits_to_float(read_u16(blk));
float dmin = half_bits_to_float(read_u16(blk + 2));
const uint8_t *sr = blk + 4;
const uint8_t *qs = blk + 16;
for (int i = 0; i < BLOCK_WEIGHTS; i++) {
int group = i / 64;
int within = i % 64;
int sub = 2 * group + (within >= 32 ? 1 : 0);
uint8_t packed = qs[32 * group + (within % 32)];
int quant = (within < 32) ? (packed & 0x0F) : (packed >> 4);
int scale, minimum;
if (sub < 4) {
scale = sr[sub] & 63;
minimum = sr[sub + 4] & 63;
} else {
scale = (sr[sub + 4] & 0x0F) | ((sr[sub - 4] >> 6) << 4);
minimum = (sr[sub + 4] >> 4) | ((sr[sub] >> 6) << 4);
}
out[i] = d * (float)scale * (float)quant - dmin * (float)minimum;
}
}
static float cpu_row_dot(const uint8_t *row, const float *x, int cols,
int block_bytes)
{
float block[BLOCK_WEIGHTS];
float acc = 0.0f;
for (int b = 0; b < cols / BLOCK_WEIGHTS; b++) {
if (block_bytes == Q3_K_BLOCK_BYTES)
cpu_q3k_dequant_block(row + (size_t)b * block_bytes, block);
else
cpu_q4k_dequant_block(row + (size_t)b * block_bytes, block);
for (int l = 0; l < BLOCK_WEIGHTS; l++)
acc += block[l] * x[b * BLOCK_WEIGHTS + l];
}
return acc;
}
/* ================================================================
* CPU MoE reference
* ================================================================ */
static void cpu_gateup(const uint8_t *gate_all, const uint8_t *up_all,
const float *x, const int *top_idx, int n_requests,
int expert_used, int ff, int cols, int row_stride,
float *hid)
{
size_t expert_stride = (size_t)ff * (size_t)row_stride;
for (int r = 0; r < n_requests; r++) {
int expert = top_idx[r];
int token = r / expert_used;
const float *xt = x + (size_t)token * cols;
for (int row = 0; row < ff; row++) {
size_t base = (size_t)expert * expert_stride + (size_t)row * row_stride;
float g = cpu_row_dot(gate_all + base, xt, cols, Q3_K_BLOCK_BYTES);
float u = cpu_row_dot(up_all + base, xt, cols, Q3_K_BLOCK_BYTES);
hid[(size_t)r * ff + row] = (g / (1.0f + expf(-g))) * u;
}
}
}
static void cpu_down(const uint8_t *down_all, const float *hid, const int *top_idx,
const float *top_wt, int n_tokens, int expert_used, int rows,
int ff, int row_stride, int block_bytes, float *out)
{
size_t expert_stride = (size_t)rows * (size_t)row_stride;
for (int t = 0; t < n_tokens; t++)
for (int row = 0; row < rows; row++) {
float acc = 0.0f;
for (int s = 0; s < expert_used; s++) {
int request = t * expert_used + s;
int expert = top_idx[request];
const uint8_t *w = down_all + (size_t)expert * expert_stride +
(size_t)row * row_stride;
acc += top_wt[request] *
cpu_row_dot(w, hid + (size_t)request * ff, ff, block_bytes);
}
out[(size_t)t * rows + row] = acc;
}
}
/* ================================================================
* Comparison
*
* The acceptance bar from the spec is max_abs_err / RMS(reference), which
* stays meaningful where individual reference values pass through zero.
* ================================================================ */
typedef struct error_report {
double max_abs;
double rms;
double ratio;
size_t worst_index;
} error_report;
static error_report compare_vectors(const float *got, const float *ref, size_t n)
{
error_report report = {0.0, 0.0, 0.0, 0};
double sum_sq = 0.0;
for (size_t i = 0; i < n; i++) {
double diff = fabs((double)got[i] - (double)ref[i]);
if (diff > report.max_abs) {
report.max_abs = diff;
report.worst_index = i;
}
sum_sq += (double)ref[i] * (double)ref[i];
}
report.rms = sqrt(sum_sq / (double)n);
report.ratio = (report.rms > 0.0) ? report.max_abs / report.rms : report.max_abs;
return report;
}
/* ================================================================
* Loading real expert weights
* ================================================================ */
/* Copies MOE_TEST_EXPERTS real experts out of a merged 3-D expert tensor into a
* contiguous table, so a routing over slots 0..7 addresses real bytes from
* experts scattered across the full 256. */
static uint8_t *load_expert_table(const char *tensor_name, int expected_type,
int expected_cols, int expected_rows,
size_t *expert_stride_out)
{
gguf_tensor_location location;
if (gguf_test_find_tensor(&g_model, tensor_name, &location) != 0) {
fprintf(stderr, " missing tensor %s\n", tensor_name);
return NULL;
}
if ((int)location.ggml_type != expected_type || location.n_dims != 3 ||
(int)location.dims[0] != expected_cols ||
(int)location.dims[1] != expected_rows) {
fprintf(stderr, " %s has an unexpected type/shape (%u, %llu x %llu x %llu)\n",
tensor_name, location.ggml_type,
(unsigned long long)location.dims[0],
(unsigned long long)location.dims[1],
(unsigned long long)location.dims[2]);
return NULL;
}
int block_bytes = (expected_type == GGUF_TYPE_Q3_K) ? Q3_K_BLOCK_BYTES
: Q4_K_BLOCK_BYTES;
size_t row_stride = (size_t)(expected_cols / BLOCK_WEIGHTS) * block_bytes;
size_t expert_stride = (size_t)expected_rows * row_stride;
uint8_t *table = (uint8_t *)malloc(expert_stride * MOE_TEST_EXPERTS);
if (!table)
return NULL;
for (int i = 0; i < MOE_TEST_EXPERTS; i++) {
uint64_t offset = location.file_offset +
(uint64_t)kSourceExperts[i] * expert_stride;
if (gguf_test_read_at(&g_model, offset, table + (size_t)i * expert_stride,
expert_stride) != 0) {
fprintf(stderr, " short read of %s expert %d\n", tensor_name,
kSourceExperts[i]);
free(table);
return NULL;
}
}
*expert_stride_out = expert_stride;
return table;
}
/* Slot-major routing over the local expert table: every token uses all
* MOE_TEST_EXPERTS slots, in a token-dependent rotation so no two tokens share
* a slot->expert map. */
static void build_local_routing(int n_tokens, int *top_idx, float *top_wt)
{
for (int t = 0; t < n_tokens; t++) {
float sum = 0.0f;
for (int s = 0; s < MOE_EXPERT_USED; s++) {
top_idx[t * MOE_EXPERT_USED + s] =
(s + t) % MOE_TEST_EXPERTS;
float w = 0.1f + (float)(moe_rand_u32() % 1000u) / 1000.0f;
top_wt[t * MOE_EXPERT_USED + s] = w;
sum += w;
}
for (int s = 0; s < MOE_EXPERT_USED; s++)
top_wt[t * MOE_EXPERT_USED + s] /= sum;
}
}
/* ================================================================
* Gate + up
* ================================================================ */
/* Shared fixture for the gate+up tests: real weights, random activations, and
* both the GPU and CPU results. `weight_bind_offset` places the weight table at
* a non-page-aligned byte offset inside its buffer when nonzero. */
static int run_gateup_case(int n_tokens, int use_sorted, size_t weight_bind_offset,
error_report *report)
{
const int cols = MOE_N_EMBD;
const int ff = MOE_EXPERT_FF;
const int n_requests = n_tokens * MOE_EXPERT_USED;
const int row_stride = (cols / BLOCK_WEIGHTS) * Q3_K_BLOCK_BYTES;
size_t gate_stride = 0, up_stride = 0;
int rc = -1;
uint8_t *gate_table = load_expert_table("blk.0.ffn_gate_exps.weight",
GGUF_TYPE_Q3_K, cols, ff, &gate_stride);
uint8_t *up_table = load_expert_table("blk.0.ffn_up_exps.weight",
GGUF_TYPE_Q3_K, cols, ff, &up_stride);
float *x = (float *)malloc((size_t)n_tokens * cols * sizeof(float));
int *top_idx = (int *)malloc((size_t)n_requests * sizeof(int));
float *top_wt = (float *)malloc((size_t)n_requests * sizeof(float));
float *ref_hid = (float *)malloc((size_t)n_requests * ff * sizeof(float));
fm_buffer *gate_buf = NULL, *up_buf = NULL, *x_buf = NULL, *hid_buf = NULL;
fm_buffer *idx_buf = NULL, *st_buf = NULL, *ss_buf = NULL, *se_buf = NULL;
fm_buffer *count_buf = NULL, *offset_buf = NULL;
if (!gate_table || !up_table || !x || !top_idx || !top_wt || !ref_hid)
goto cleanup;
moe_rand_state = 0x1337BEEFu;
for (int i = 0; i < n_tokens * cols; i++)
x[i] = moe_rand_signed(1.0f);
build_local_routing(n_tokens, top_idx, top_wt);
gate_buf = fm_buffer_alloc(weight_bind_offset + gate_stride * MOE_TEST_EXPERTS);
up_buf = fm_buffer_alloc(weight_bind_offset + up_stride * MOE_TEST_EXPERTS);
x_buf = fm_buffer_alloc((size_t)n_tokens * cols * sizeof(float));
idx_buf = fm_buffer_alloc((size_t)n_requests * sizeof(int));
hid_buf = fm_buffer_alloc((size_t)n_requests * ff * sizeof(float));
if (!gate_buf || !up_buf || !x_buf || !idx_buf || !hid_buf)
goto cleanup;
memcpy((uint8_t *)fm_buffer_data(gate_buf) + weight_bind_offset, gate_table,
gate_stride * MOE_TEST_EXPERTS);
memcpy((uint8_t *)fm_buffer_data(up_buf) + weight_bind_offset, up_table,
up_stride * MOE_TEST_EXPERTS);
memcpy(fm_buffer_data(x_buf), x, (size_t)n_tokens * cols * sizeof(float));
memcpy(fm_buffer_data(idx_buf), top_idx, (size_t)n_requests * sizeof(int));
memset(fm_buffer_data(hid_buf), 0x7F, (size_t)n_requests * ff * sizeof(float));
if (use_sorted) {
count_buf = fm_buffer_alloc(MOE_TEST_EXPERTS * sizeof(int));
offset_buf = fm_buffer_alloc((MOE_TEST_EXPERTS + 1) * sizeof(int));
st_buf = fm_buffer_alloc((size_t)n_requests * sizeof(int));
ss_buf = fm_buffer_alloc((size_t)n_requests * sizeof(int));
se_buf = fm_buffer_alloc((size_t)n_requests * sizeof(int));
if (!count_buf || !offset_buf || !st_buf || !ss_buf || !se_buf)
goto cleanup;
fm_moe_expert_sort(idx_buf, count_buf, offset_buf, st_buf, ss_buf, se_buf,
n_tokens, MOE_TEST_EXPERTS, MOE_EXPERT_USED);
fm_moe_gateup_q3k_sorted_at(fm_bufref_at(gate_buf, weight_bind_offset),
fm_bufref_at(up_buf, weight_bind_offset),
fm_bufref_at(x_buf, 0), fm_bufref_at(st_buf, 0),
fm_bufref_at(ss_buf, 0), fm_bufref_at(se_buf, 0),
fm_bufref_at(hid_buf, 0), n_requests,
MOE_EXPERT_USED, ff, cols, row_stride);
} else {
fm_moe_gateup_q3k_at(fm_bufref_at(gate_buf, weight_bind_offset),
fm_bufref_at(up_buf, weight_bind_offset),
fm_bufref_at(x_buf, 0), fm_bufref_at(idx_buf, 0),
fm_bufref_at(hid_buf, 0), n_tokens, MOE_EXPERT_USED,
ff, cols, row_stride);
}
fm_sync();
cpu_gateup(gate_table, up_table, x, top_idx, n_requests, MOE_EXPERT_USED, ff,
cols, row_stride, ref_hid);
*report = compare_vectors((const float *)fm_buffer_data(hid_buf), ref_hid,
(size_t)n_requests * ff);
rc = 0;
cleanup:
fm_buffer_free(gate_buf);
fm_buffer_free(up_buf);
fm_buffer_free(x_buf);
fm_buffer_free(idx_buf);
fm_buffer_free(hid_buf);
fm_buffer_free(count_buf);
fm_buffer_free(offset_buf);
fm_buffer_free(st_buf);
fm_buffer_free(ss_buf);
fm_buffer_free(se_buf);
free(gate_table);
free(up_table);
free(x);
free(top_idx);
free(top_wt);
free(ref_hid);
return rc;
}
static void test_gateup_q3k_matches_cpu_reference(void)
{
error_report report;
ASSERT_INT_EQ(0, run_gateup_case(1, 0, 0, &report));
printf(" gate+up max_abs=%.3e rms=%.3e ratio=%.2e (worst %zu)\n",
report.max_abs, report.rms, report.ratio, report.worst_index);
ASSERT_TRUE(report.rms > 0.0);
ASSERT_TRUE(report.ratio < 5e-5);
}
static void test_gateup_q3k_batched_matches_cpu_reference(void)
{
error_report report;
ASSERT_INT_EQ(0, run_gateup_case(3, 0, 0, &report));
printf(" gate+up x3 max_abs=%.3e rms=%.3e ratio=%.2e\n",
report.max_abs, report.rms, report.ratio);
ASSERT_TRUE(report.rms > 0.0);
ASSERT_TRUE(report.ratio < 5e-5);
}
static void test_gateup_q3k_reads_weights_at_a_nonzero_offset(void)
{
error_report report;
ASSERT_INT_EQ(0, run_gateup_case(1, 0, WEIGHT_BIND_OFFSET, &report));
printf(" gate+up @%u max_abs=%.3e ratio=%.2e\n", WEIGHT_BIND_OFFSET,
report.max_abs, report.ratio);
ASSERT_TRUE(report.rms > 0.0);
ASSERT_TRUE(report.ratio < 5e-5);
}
static void test_gateup_q3k_sorted_matches_cpu_reference(void)
{
error_report report;
ASSERT_INT_EQ(0, run_gateup_case(3, 1, 0, &report));
printf(" gate+up sorted max_abs=%.3e rms=%.3e ratio=%.2e\n",
report.max_abs, report.rms, report.ratio);
ASSERT_TRUE(report.rms > 0.0);
ASSERT_TRUE(report.ratio < 5e-5);
}
/* ================================================================
* Down projection
* ================================================================ */
static int run_down_case(const char *tensor_name, int expected_type, int n_tokens,
error_report *report, float *garbage_survivor)
{
const int rows = MOE_N_EMBD;
const int ff = MOE_EXPERT_FF;
const int n_requests = n_tokens * MOE_EXPERT_USED;
const int block_bytes = (expected_type == GGUF_TYPE_Q3_K) ? Q3_K_BLOCK_BYTES
: Q4_K_BLOCK_BYTES;
const int row_stride = (ff / BLOCK_WEIGHTS) * block_bytes;
size_t expert_stride = 0;
int rc = -1;
uint8_t *down_table = load_expert_table(tensor_name, expected_type, ff, rows,
&expert_stride);
float *hid = (float *)malloc((size_t)n_requests * ff * sizeof(float));
int *top_idx = (int *)malloc((size_t)n_requests * sizeof(int));
float *top_wt = (float *)malloc((size_t)n_requests * sizeof(float));
float *ref_out = (float *)malloc((size_t)n_tokens * rows * sizeof(float));
fm_buffer *down_buf = NULL, *hid_buf = NULL, *idx_buf = NULL;
fm_buffer *wt_buf = NULL, *out_buf = NULL;
if (!down_table || !hid || !top_idx || !top_wt || !ref_out)
goto cleanup;
moe_rand_state = 0x5EED1234u;
for (int i = 0; i < n_requests * ff; i++)
hid[i] = moe_rand_signed(2.0f);
build_local_routing(n_tokens, top_idx, top_wt);
down_buf = fm_buffer_alloc(expert_stride * MOE_TEST_EXPERTS);
hid_buf = fm_buffer_alloc((size_t)n_requests * ff * sizeof(float));
idx_buf = fm_buffer_alloc((size_t)n_requests * sizeof(int));
wt_buf = fm_buffer_alloc((size_t)n_requests * sizeof(float));
out_buf = fm_buffer_alloc((size_t)n_tokens * rows * sizeof(float));
if (!down_buf || !hid_buf || !idx_buf || !wt_buf || !out_buf)
goto cleanup;
memcpy(fm_buffer_data(down_buf), down_table, expert_stride * MOE_TEST_EXPERTS);
memcpy(fm_buffer_data(hid_buf), hid, (size_t)n_requests * ff * sizeof(float));
memcpy(fm_buffer_data(idx_buf), top_idx, (size_t)n_requests * sizeof(int));
memcpy(fm_buffer_data(wt_buf), top_wt, (size_t)n_requests * sizeof(float));
/* The kernel writes `out` rather than accumulating, so a pre-existing
* value must not survive. */
float *out_seed = (float *)fm_buffer_data(out_buf);
for (int i = 0; i < n_tokens * rows; i++)
out_seed[i] = 12345.0f;
if (expected_type == GGUF_TYPE_Q3_K)
fm_moe_down_q3k(down_buf, hid_buf, idx_buf, wt_buf, out_buf, n_tokens,
MOE_EXPERT_USED, rows, ff, row_stride);
else
fm_moe_down_q4k(down_buf, hid_buf, idx_buf, wt_buf, out_buf, n_tokens,
MOE_EXPERT_USED, rows, ff, row_stride);
fm_sync();
cpu_down(down_table, hid, top_idx, top_wt, n_tokens, MOE_EXPERT_USED, rows, ff,
row_stride, block_bytes, ref_out);
*report = compare_vectors((const float *)fm_buffer_data(out_buf), ref_out,
(size_t)n_tokens * rows);
if (garbage_survivor)
*garbage_survivor = ((const float *)fm_buffer_data(out_buf))[0];
rc = 0;
cleanup:
fm_buffer_free(down_buf);
fm_buffer_free(hid_buf);
fm_buffer_free(idx_buf);
fm_buffer_free(wt_buf);
fm_buffer_free(out_buf);
free(down_table);
free(hid);
free(top_idx);
free(top_wt);
free(ref_out);
return rc;
}
static void test_down_q3k_matches_cpu_reference(void)
{
error_report report;
float first = 0.0f;
ASSERT_INT_EQ(0, run_down_case("blk.5.ffn_down_exps.weight", GGUF_TYPE_Q3_K, 1,
&report, &first));
printf(" down Q3_K max_abs=%.3e rms=%.3e ratio=%.2e (worst %zu)\n",
report.max_abs, report.rms, report.ratio, report.worst_index);
ASSERT_TRUE(report.rms > 0.0);
/* Bounds the one reassociation in this family: the routed-weight sum is
* split across lanes by (slot, half super-block) and reduced by simd_sum,
* where the CPU finishes an expert dot before weighting it. */
ASSERT_TRUE(report.ratio < 1e-5);
ASSERT_TRUE(first != 12345.0f);
}
static void test_down_q3k_batched_matches_cpu_reference(void)
{
error_report report;
ASSERT_INT_EQ(0, run_down_case("blk.5.ffn_down_exps.weight", GGUF_TYPE_Q3_K, 3,
&report, NULL));
printf(" down Q3_K x3 max_abs=%.3e rms=%.3e ratio=%.2e\n",
report.max_abs, report.rms, report.ratio);
ASSERT_TRUE(report.rms > 0.0);
ASSERT_TRUE(report.ratio < 1e-5);
}
static void test_down_q4k_matches_cpu_reference(void)
{
error_report report;
ASSERT_INT_EQ(0, run_down_case("blk.0.ffn_down_exps.weight", GGUF_TYPE_Q4_K, 1,
&report, NULL));
printf(" down Q4_K max_abs=%.3e rms=%.3e ratio=%.2e (worst %zu)\n",
report.max_abs, report.rms, report.ratio, report.worst_index);
ASSERT_TRUE(report.rms > 0.0);
ASSERT_TRUE(report.ratio < 1e-5);
}
/* ================================================================
* main
* ================================================================ */
int main(int argc, char **argv)
{
printf("\n Metal MoE — Grouped Experts (real GGUF weights)\n");
printf(" ==============================================\n\n");
if (fm_init() != 0 || !fm_available()) {
fprintf(stderr, " Metal unavailable — no device to test against.\n");
return 77;
}
const char *model_path = (argc > 1) ? argv[1] : getenv("FOUNDRY_TEST_GGUF");
if (!model_path) model_path = kDefaultModelPath;
g_model_is_open = gguf_test_open(model_path, &g_model) == 0;
if (!g_model_is_open) {
fprintf(stderr, " Model file missing — every test here is driven by real\n"
" expert blocks and cannot run without it:\n %s\n",
model_path);
fm_free();
return 77;
}
printf(" model = %s\n\n", model_path);
RUN_TEST(gateup_q3k_matches_cpu_reference);
RUN_TEST(gateup_q3k_batched_matches_cpu_reference);
RUN_TEST(gateup_q3k_reads_weights_at_a_nonzero_offset);
RUN_TEST(gateup_q3k_sorted_matches_cpu_reference);
RUN_TEST(down_q3k_matches_cpu_reference);
RUN_TEST(down_q3k_batched_matches_cpu_reference);
RUN_TEST(down_q4k_matches_cpu_reference);
gguf_test_close(&g_model);
fm_free();
TEST_REPORT();
return __ft_failed > 0 ? 1 : 0;
}
Reference in New Issue View Git Blame Copy Permalink