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

764 lines
28 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file test_metal_moe.c
* @brief Absolute-correctness tests for the Metal MoE router, the
* sort-by-expert prepass, the fused silu*mul and the shared-expert gate.
*
* The bar is ABSOLUTE, never GPU-against-GPU. Every assertion compares the
* kernel either to a Pantheon golden vector (produced by a third
* implementation from a fixed seed) or to the CPU reference written below,
* which is a term-for-term transcription of `qwen35moe_moe.c`'s router and a
* from-scratch CSR inversion of the routing. A GPU-against-GPU arm would
* stay green through a uniform scale, which is exactly the break Phase 1
* proved slips past self-comparison.
*
* The tie-break gets its own fixture. A test whose logits happen to be
* distinct does not test a tie-break at all, so `router_tie_break_*` makes
* sixteen experts share the top probability bit-for-bit and requires the
* eight LOWEST of them, in ascending order.
*
* The sort prepass is exercised at 512 tokens, not at 1: its entire value is
* at prefill, and a single-token run cannot expose a scan or a per-expert
* offset that is off by one.
*
* @status All tests pass on M4 / macOS 26. Each was proven able to fail
* against a deliberately broken kernel before being trusted — a 0.01% scaled
* weight, a tie broken toward the higher index, an off-by-one per-expert
* count and a top-k returning the lowest probabilities instead of the
* highest were all induced and all caught.
* @issues None
* @todo None
*/
#include "lib/test.h"
/* The family header directly, not the fm_metal.h umbrella: the umbrella lists
* one include per family and is owned by the integration step, so a test that
* went through it could not build until that line lands. */
#include "metal/kernels/fm_kernel_moe.h"
#include "metal/tests/fm_golden_vectors.h"
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
/* ================================================================
* Fixtures
* ================================================================ */
#define MOE_N_EXPERT 256
#define MOE_EXPERT_USED 8
#define MOE_N_EMBD 2048
/* Tokens for the sort prepass. The prepass exists for prefill, so it is
* measured at a prefill batch rather than at the decode batch of 1. */
#define MOE_SORT_TOKENS 512
/* Experts sharing the top probability in the tie fixture. Deliberately more
* than expert_used, so which ones are picked is decided by the tie-break. */
#define MOE_TIED_EXPERTS 16
static const int kTiedExperts[MOE_TIED_EXPERTS] = {
5, 9, 17, 33, 41, 60, 77, 88, 101, 130, 150, 171, 190, 201, 222, 250
};
/* Deterministic PRNG so a failure is always reproducible. */
static uint32_t moe_rand_state = 0x12345678u;
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_unit(void)
{
return (float)(moe_rand_u32() >> 8) / (float)(1u << 24);
}
/**
* Reads a PNGOLDEN1 i32 vector. fm_golden_vectors.h reads f32 only, and the
* router's selected indices — the one thing that must match exactly — are i32.
*/
static int *moe_golden_load_i32(const char *name, size_t expect_count)
{
char path[1152];
char header[256];
char magic[16] = {0}, vec_name[64] = {0}, dtype[8] = {0};
unsigned long count = 0;
int *values = NULL;
FILE *f = NULL;
snprintf(path, sizeof(path), "%s/%s", fm_golden_dir, name);
f = fopen(path, "rb");
if (!f) {
fprintf(stderr, " golden: cannot open %s\n", path);
return NULL;
}
if (!fgets(header, sizeof(header), f)) {
fprintf(stderr, " golden: %s has no header line\n", path);
goto cleanup;
}
if (sscanf(header, "%15s %63s %7s %lu", magic, vec_name, dtype, &count) != 4 ||
strcmp(magic, "PNGOLDEN1") != 0 || strcmp(dtype, "i32") != 0 ||
count != expect_count) {
fprintf(stderr, " golden: %s has a bad header: %s", path, header);
goto cleanup;
}
values = (int *)malloc(expect_count * sizeof(int));
if (values && fread(values, sizeof(int), expect_count, f) != expect_count) {
fprintf(stderr, " golden: %s payload is truncated\n", path);
free(values);
values = NULL;
}
cleanup:
fclose(f);
return values;
}
/* ================================================================
* CPU reference — transcribed from Pantheon's qwen35moe_moe.c
* ================================================================ */
/**
* Softmax over the router logits, top-k by probability, norm_topk_prob and the
* routed-weight scale. Mirrors qwen35moe_router_softmax + qwen35moe_select_topk
* + the renormalization block: selection scans experts in ascending order with
* strict `>`, so a tie resolves to the lower index.
*/
static void cpu_router_topk(const float *logits, int n_expert, int k, float scale,
int *top_idx, float *top_wt)
{
float *probs = (float *)malloc((size_t)n_expert * sizeof(float));
float max_logit = logits[0];
for (int e = 1; e < n_expert; e++)
if (logits[e] > max_logit)
max_logit = logits[e];
float sum_exp = 0.0f;
for (int e = 0; e < n_expert; e++) {
probs[e] = expf(logits[e] - max_logit);
sum_exp += probs[e];
}
float inv = 1.0f / sum_exp;
for (int e = 0; e < n_expert; e++)
probs[e] *= inv;
float wt_sum = 0.0f;
for (int s = 0; s < k; s++) {
int best = -1;
float best_val = -1.0f;
for (int e = 0; e < n_expert; e++) {
int skip = 0;
for (int j = 0; j < s; j++)
if (top_idx[j] == e) {
skip = 1;
break;
}
if (!skip && probs[e] > best_val) {
best = e;
best_val = probs[e];
}
}
top_idx[s] = best;
top_wt[s] = best_val;
wt_sum += best_val;
}
float inv_wt = (wt_sum > 0.0f) ? (1.0f / wt_sum) : 0.0f;
for (int s = 0; s < k; s++)
top_wt[s] *= scale * inv_wt;
free(probs);
}
/**
* Inverts a routing table into the CSR form the prepass must produce: per
* expert, the tokens that selected it in ascending token order, plus the slot
* each used. Written from the definition, not from the kernel's tiled scan.
*/
static void cpu_expert_sort(const int *top_idx, int n_tokens, int n_expert, int k,
int *counts, int *offsets, int *sorted_token,
int *sorted_slot, int *sorted_expert)
{
for (int e = 0; e < n_expert; e++)
counts[e] = 0;
for (int t = 0; t < n_tokens; t++)
for (int s = 0; s < k; s++) {
int e = top_idx[t * k + s];
if (e >= 0 && e < n_expert)
counts[e]++;
}
offsets[0] = 0;
for (int e = 0; e < n_expert; e++)
offsets[e + 1] = offsets[e] + counts[e];
for (int e = 0; e < n_expert; e++) {
int pos = offsets[e];
for (int t = 0; t < n_tokens; t++)
for (int s = 0; s < k; s++)
if (top_idx[t * k + s] == e) {
sorted_token[pos] = t;
sorted_slot[pos] = s;
sorted_expert[pos] = e;
pos++;
}
}
}
/* ================================================================
* Router
* ================================================================ */
/* Runs the router over one token's logits and hands back the device results. */
static int run_router(const float *logits, int n_tokens, int n_expert, int k,
float scale, int *out_idx, float *out_wt)
{
fm_buffer *lg = fm_buffer_alloc((size_t)n_tokens * n_expert * sizeof(float));
fm_buffer *ti = fm_buffer_alloc((size_t)n_tokens * k * sizeof(int));
fm_buffer *tw = fm_buffer_alloc((size_t)n_tokens * k * sizeof(float));
if (!lg || !ti || !tw)
return -1;
memcpy(fm_buffer_data(lg), logits,
(size_t)n_tokens * n_expert * sizeof(float));
memset(fm_buffer_data(ti), 0xFF, (size_t)n_tokens * k * sizeof(int));
memset(fm_buffer_data(tw), 0, (size_t)n_tokens * k * sizeof(float));
fm_moe_router_topk(lg, ti, tw, n_tokens, n_expert, k, scale);
fm_sync();
memcpy(out_idx, fm_buffer_data(ti), (size_t)n_tokens * k * sizeof(int));
memcpy(out_wt, fm_buffer_data(tw), (size_t)n_tokens * k * sizeof(float));
fm_buffer_free(lg);
fm_buffer_free(ti);
fm_buffer_free(tw);
return 0;
}
static void test_router_topk_matches_golden(void)
{
float *logits = fm_golden_load_f32("moe_router_logits.bin", MOE_N_EXPERT);
float *golden_wt = fm_golden_load_f32("moe_topk_weights.bin", MOE_EXPERT_USED);
int *golden_idx = moe_golden_load_i32("moe_topk_indices.bin", MOE_EXPERT_USED);
ASSERT_NOT_NULL(logits);
ASSERT_NOT_NULL(golden_wt);
ASSERT_NOT_NULL(golden_idx);
int gpu_idx[MOE_EXPERT_USED];
float gpu_wt[MOE_EXPERT_USED];
ASSERT_INT_EQ(0, run_router(logits, 1, MOE_N_EXPERT, MOE_EXPERT_USED, 1.0f,
gpu_idx, gpu_wt));
float weight_sum = 0.0f;
for (int s = 0; s < MOE_EXPERT_USED; s++) {
ASSERT_INT_EQ(golden_idx[s], gpu_idx[s]);
ASSERT_FLT_NEAR(golden_wt[s], gpu_wt[s], 2e-6f);
weight_sum += gpu_wt[s];
}
ASSERT_FLT_NEAR(1.0f, weight_sum, 2e-6f);
free(logits);
free(golden_wt);
free(golden_idx);
}
static void test_router_tie_break_prefers_lower_index(void)
{
/* Sixteen experts share the maximum logit exactly, so their softmax
* probabilities are bit-identical and the SELECTION ORDER is decided
* entirely by the tie-break rule. */
float logits[MOE_N_EXPERT];
for (int e = 0; e < MOE_N_EXPERT; e++)
logits[e] = 0.0f;
for (int i = 0; i < MOE_TIED_EXPERTS; i++)
logits[kTiedExperts[i]] = 1.0f;
int gpu_idx[MOE_EXPERT_USED];
float gpu_wt[MOE_EXPERT_USED];
ASSERT_INT_EQ(0, run_router(logits, 1, MOE_N_EXPERT, MOE_EXPERT_USED, 1.0f,
gpu_idx, gpu_wt));
/* Ties resolve to the LOWER index, so the eight lowest of the tied set are
* selected, in ascending order. A rule that preferred the higher index
* would return kTiedExperts[15..8] descending. */
for (int s = 0; s < MOE_EXPERT_USED; s++) {
ASSERT_INT_EQ(kTiedExperts[s], gpu_idx[s]);
ASSERT_FLT_NEAR(1.0f / (float)MOE_EXPERT_USED, gpu_wt[s], 1e-6f);
}
}
static void test_router_tie_break_matches_cpu_selection_sort(void)
{
/* Same fixture through the CPU selection sort, which resolves ties by
* first-seen-highest. The two must agree index for index. */
float logits[MOE_N_EXPERT];
for (int e = 0; e < MOE_N_EXPERT; e++)
logits[e] = 0.0f;
for (int i = 0; i < MOE_TIED_EXPERTS; i++)
logits[kTiedExperts[i]] = 1.0f;
int cpu_idx[MOE_EXPERT_USED];
float cpu_wt[MOE_EXPERT_USED];
cpu_router_topk(logits, MOE_N_EXPERT, MOE_EXPERT_USED, 1.0f, cpu_idx, cpu_wt);
int gpu_idx[MOE_EXPERT_USED];
float gpu_wt[MOE_EXPERT_USED];
ASSERT_INT_EQ(0, run_router(logits, 1, MOE_N_EXPERT, MOE_EXPERT_USED, 1.0f,
gpu_idx, gpu_wt));
for (int s = 0; s < MOE_EXPERT_USED; s++) {
ASSERT_INT_EQ(cpu_idx[s], gpu_idx[s]);
ASSERT_FLT_NEAR(cpu_wt[s], gpu_wt[s], 1e-6f);
}
}
static void test_router_applies_expert_weights_scale(void)
{
float *logits = fm_golden_load_f32("moe_router_logits.bin", MOE_N_EXPERT);
ASSERT_NOT_NULL(logits);
const float scale = 2.5f;
int cpu_idx[MOE_EXPERT_USED];
float cpu_wt[MOE_EXPERT_USED];
cpu_router_topk(logits, MOE_N_EXPERT, MOE_EXPERT_USED, scale, cpu_idx, cpu_wt);
int gpu_idx[MOE_EXPERT_USED];
float gpu_wt[MOE_EXPERT_USED];
ASSERT_INT_EQ(0, run_router(logits, 1, MOE_N_EXPERT, MOE_EXPERT_USED, scale,
gpu_idx, gpu_wt));
float weight_sum = 0.0f;
for (int s = 0; s < MOE_EXPERT_USED; s++) {
ASSERT_INT_EQ(cpu_idx[s], gpu_idx[s]);
ASSERT_FLT_NEAR(cpu_wt[s], gpu_wt[s], 2e-6f);
weight_sum += gpu_wt[s];
}
ASSERT_FLT_NEAR(scale, weight_sum, 5e-6f);
free(logits);
}
/* Per-token logits that are distinct by construction (a permutation on a
* 0.03-wide grid), so index selection cannot be flipped by the last ulp of a
* differently-associated softmax while still giving every token its own
* expert set. */
static void fill_permuted_logits(float *logits, int n_tokens, int n_expert)
{
for (int t = 0; t < n_tokens; t++) {
int *perm = (int *)malloc((size_t)n_expert * sizeof(int));
for (int e = 0; e < n_expert; e++)
perm[e] = e;
for (int e = n_expert - 1; e > 0; e--) {
int j = (int)(moe_rand_u32() % (uint32_t)(e + 1));
int tmp = perm[e];
perm[e] = perm[j];
perm[j] = tmp;
}
for (int e = 0; e < n_expert; e++)
logits[t * n_expert + e] = (float)perm[e] * 0.03f - 4.0f;
free(perm);
}
}
static void test_router_batch_matches_cpu_reference(void)
{
const int n_tokens = 64;
float *logits = (float *)malloc((size_t)n_tokens * MOE_N_EXPERT * sizeof(float));
int *gpu_idx = (int *)malloc((size_t)n_tokens * MOE_EXPERT_USED * sizeof(int));
float *gpu_wt = (float *)malloc((size_t)n_tokens * MOE_EXPERT_USED * sizeof(float));
ASSERT_NOT_NULL(logits);
ASSERT_NOT_NULL(gpu_idx);
ASSERT_NOT_NULL(gpu_wt);
moe_rand_state = 0x2468ACE0u;
fill_permuted_logits(logits, n_tokens, MOE_N_EXPERT);
ASSERT_INT_EQ(0, run_router(logits, n_tokens, MOE_N_EXPERT, MOE_EXPERT_USED,
1.0f, gpu_idx, gpu_wt));
for (int t = 0; t < n_tokens; t++) {
int cpu_idx[MOE_EXPERT_USED];
float cpu_wt[MOE_EXPERT_USED];
cpu_router_topk(logits + (size_t)t * MOE_N_EXPERT, MOE_N_EXPERT,
MOE_EXPERT_USED, 1.0f, cpu_idx, cpu_wt);
for (int s = 0; s < MOE_EXPERT_USED; s++) {
ASSERT_INT_EQ(cpu_idx[s], gpu_idx[t * MOE_EXPERT_USED + s]);
ASSERT_FLT_NEAR(cpu_wt[s], gpu_wt[t * MOE_EXPERT_USED + s], 2e-6f);
}
}
free(logits);
free(gpu_idx);
free(gpu_wt);
}
/* ================================================================
* Sort-by-expert prepass
* ================================================================ */
/* Routing for the sort tests: real router output over per-token permuted
* logits, so the token->expert map is as irregular as a real batch's. */
static int build_routing(int n_tokens, int *top_idx)
{
float *logits = (float *)malloc((size_t)n_tokens * MOE_N_EXPERT * sizeof(float));
float *wt = (float *)malloc((size_t)n_tokens * MOE_EXPERT_USED * sizeof(float));
if (!logits || !wt)
return -1;
moe_rand_state = 0x0BADC0DEu;
fill_permuted_logits(logits, n_tokens, MOE_N_EXPERT);
int rc = run_router(logits, n_tokens, MOE_N_EXPERT, MOE_EXPERT_USED, 1.0f,
top_idx, wt);
free(logits);
free(wt);
return rc;
}
/* Runs the prepass and copies its five outputs back. */
static int run_expert_sort(const int *top_idx, int n_tokens, int *counts,
int *offsets, int *sorted_token, int *sorted_slot,
int *sorted_expert)
{
const size_t requests = (size_t)n_tokens * MOE_EXPERT_USED;
fm_buffer *ti = fm_buffer_alloc(requests * sizeof(int));
fm_buffer *bc = fm_buffer_alloc(MOE_N_EXPERT * sizeof(int));
fm_buffer *bo = fm_buffer_alloc((MOE_N_EXPERT + 1) * sizeof(int));
fm_buffer *bt = fm_buffer_alloc(requests * sizeof(int));
fm_buffer *bs = fm_buffer_alloc(requests * sizeof(int));
fm_buffer *be = fm_buffer_alloc(requests * sizeof(int));
if (!ti || !bc || !bo || !bt || !bs || !be)
return -1;
memcpy(fm_buffer_data(ti), top_idx, requests * sizeof(int));
memset(fm_buffer_data(bc), 0xFF, MOE_N_EXPERT * sizeof(int));
memset(fm_buffer_data(bo), 0xFF, (MOE_N_EXPERT + 1) * sizeof(int));
memset(fm_buffer_data(bt), 0xFF, requests * sizeof(int));
memset(fm_buffer_data(bs), 0xFF, requests * sizeof(int));
memset(fm_buffer_data(be), 0xFF, requests * sizeof(int));
fm_moe_expert_sort(ti, bc, bo, bt, bs, be, n_tokens, MOE_N_EXPERT,
MOE_EXPERT_USED);
fm_sync();
memcpy(counts, fm_buffer_data(bc), MOE_N_EXPERT * sizeof(int));
memcpy(offsets, fm_buffer_data(bo), (MOE_N_EXPERT + 1) * sizeof(int));
memcpy(sorted_token, fm_buffer_data(bt), requests * sizeof(int));
memcpy(sorted_slot, fm_buffer_data(bs), requests * sizeof(int));
memcpy(sorted_expert, fm_buffer_data(be), requests * sizeof(int));
fm_buffer_free(ti);
fm_buffer_free(bc);
fm_buffer_free(bo);
fm_buffer_free(bt);
fm_buffer_free(bs);
fm_buffer_free(be);
return 0;
}
static void test_expert_sort_matches_cpu_reference_at_prefill(void)
{
const int n_tokens = MOE_SORT_TOKENS;
const size_t requests = (size_t)n_tokens * MOE_EXPERT_USED;
int *top_idx = (int *)malloc(requests * sizeof(int));
int *counts = (int *)malloc(MOE_N_EXPERT * sizeof(int));
int *offsets = (int *)malloc((MOE_N_EXPERT + 1) * sizeof(int));
int *sorted_token = (int *)malloc(requests * sizeof(int));
int *sorted_slot = (int *)malloc(requests * sizeof(int));
int *sorted_expert = (int *)malloc(requests * sizeof(int));
int *ref_counts = (int *)malloc(MOE_N_EXPERT * sizeof(int));
int *ref_offsets = (int *)malloc((MOE_N_EXPERT + 1) * sizeof(int));
int *ref_token = (int *)malloc(requests * sizeof(int));
int *ref_slot = (int *)malloc(requests * sizeof(int));
int *ref_expert = (int *)malloc(requests * sizeof(int));
ASSERT_NOT_NULL(top_idx);
ASSERT_NOT_NULL(ref_expert);
ASSERT_INT_EQ(0, build_routing(n_tokens, top_idx));
ASSERT_INT_EQ(0, run_expert_sort(top_idx, n_tokens, counts, offsets,
sorted_token, sorted_slot, sorted_expert));
cpu_expert_sort(top_idx, n_tokens, MOE_N_EXPERT, MOE_EXPERT_USED, ref_counts,
ref_offsets, ref_token, ref_slot, ref_expert);
for (int e = 0; e < MOE_N_EXPERT; e++)
ASSERT_INT_EQ(ref_counts[e], counts[e]);
for (int e = 0; e <= MOE_N_EXPERT; e++)
ASSERT_INT_EQ(ref_offsets[e], offsets[e]);
ASSERT_INT_EQ((int)requests, offsets[MOE_N_EXPERT]);
for (size_t r = 0; r < requests; r++) {
ASSERT_INT_EQ(ref_token[r], sorted_token[r]);
ASSERT_INT_EQ(ref_slot[r], sorted_slot[r]);
ASSERT_INT_EQ(ref_expert[r], sorted_expert[r]);
}
free(top_idx);
free(counts);
free(offsets);
free(sorted_token);
free(sorted_slot);
free(sorted_expert);
free(ref_counts);
free(ref_offsets);
free(ref_token);
free(ref_slot);
free(ref_expert);
}
static void test_expert_sort_is_a_permutation_of_the_routing(void)
{
/* Independent of the CPU reference: every (token, slot) request must appear
* exactly once, at a position inside its own expert's CSR range, with the
* expert the router actually chose. This catches a scan that drops or
* duplicates an entry even if the reference made the same mistake. */
const int n_tokens = MOE_SORT_TOKENS;
const size_t requests = (size_t)n_tokens * MOE_EXPERT_USED;
int *top_idx = (int *)malloc(requests * sizeof(int));
int *counts = (int *)malloc(MOE_N_EXPERT * sizeof(int));
int *offsets = (int *)malloc((MOE_N_EXPERT + 1) * sizeof(int));
int *sorted_token = (int *)malloc(requests * sizeof(int));
int *sorted_slot = (int *)malloc(requests * sizeof(int));
int *sorted_expert = (int *)malloc(requests * sizeof(int));
char *seen = (char *)calloc(requests, 1);
ASSERT_NOT_NULL(top_idx);
ASSERT_NOT_NULL(seen);
ASSERT_INT_EQ(0, build_routing(n_tokens, top_idx));
ASSERT_INT_EQ(0, run_expert_sort(top_idx, n_tokens, counts, offsets,
sorted_token, sorted_slot, sorted_expert));
for (int e = 0; e < MOE_N_EXPERT; e++) {
ASSERT_INT_EQ(counts[e], offsets[e + 1] - offsets[e]);
for (int r = offsets[e]; r < offsets[e + 1]; r++) {
ASSERT_INT_EQ(e, sorted_expert[r]);
int t = sorted_token[r];
int s = sorted_slot[r];
ASSERT_TRUE(t >= 0 && t < n_tokens);
ASSERT_TRUE(s >= 0 && s < MOE_EXPERT_USED);
ASSERT_INT_EQ(e, top_idx[t * MOE_EXPERT_USED + s]);
ASSERT_INT_EQ(0, seen[(size_t)t * MOE_EXPERT_USED + s]);
seen[(size_t)t * MOE_EXPERT_USED + s] = 1;
/* Ascending token order within an expert is what makes the output
* deterministic; an atomic-ordered scatter would break here. */
if (r > offsets[e])
ASSERT_TRUE(sorted_token[r] > sorted_token[r - 1]);
}
}
for (size_t r = 0; r < requests; r++)
ASSERT_INT_EQ(1, seen[r]);
free(top_idx);
free(counts);
free(offsets);
free(sorted_token);
free(sorted_slot);
free(sorted_expert);
free(seen);
}
static void test_expert_sort_single_token(void)
{
int top_idx[MOE_EXPERT_USED];
int counts[MOE_N_EXPERT];
int offsets[MOE_N_EXPERT + 1];
int sorted_token[MOE_EXPERT_USED];
int sorted_slot[MOE_EXPERT_USED];
int sorted_expert[MOE_EXPERT_USED];
ASSERT_INT_EQ(0, build_routing(1, top_idx));
ASSERT_INT_EQ(0, run_expert_sort(top_idx, 1, counts, offsets, sorted_token,
sorted_slot, sorted_expert));
int total = 0;
for (int e = 0; e < MOE_N_EXPERT; e++) {
int expected = 0;
for (int s = 0; s < MOE_EXPERT_USED; s++)
if (top_idx[s] == e)
expected = 1;
ASSERT_INT_EQ(expected, counts[e]);
total += counts[e];
}
ASSERT_INT_EQ(MOE_EXPERT_USED, total);
for (int r = 0; r < MOE_EXPERT_USED; r++) {
ASSERT_INT_EQ(0, sorted_token[r]);
ASSERT_INT_EQ(sorted_expert[r], top_idx[sorted_slot[r]]);
}
}
/* ================================================================
* Shared expert
* ================================================================ */
static void test_silu_mul_matches_cpu_reference(void)
{
const int n = 4096;
float *a = (float *)malloc((size_t)n * sizeof(float));
float *b = (float *)malloc((size_t)n * sizeof(float));
ASSERT_NOT_NULL(a);
ASSERT_NOT_NULL(b);
moe_rand_state = 0x51A7ABCDu;
for (int i = 0; i < n; i++) {
a[i] = moe_rand_unit() * 8.0f - 4.0f;
b[i] = moe_rand_unit() * 4.0f - 2.0f;
}
fm_buffer *ba = fm_buffer_alloc((size_t)n * sizeof(float));
fm_buffer *bb = fm_buffer_alloc((size_t)n * sizeof(float));
fm_buffer *bo = fm_buffer_alloc((size_t)n * sizeof(float));
ASSERT_NOT_NULL(ba);
ASSERT_NOT_NULL(bb);
ASSERT_NOT_NULL(bo);
memcpy(fm_buffer_data(ba), a, (size_t)n * sizeof(float));
memcpy(fm_buffer_data(bb), b, (size_t)n * sizeof(float));
memset(fm_buffer_data(bo), 0, (size_t)n * sizeof(float));
fm_moe_silu_mul(ba, bb, bo, n);
fm_sync();
const float *got = (const float *)fm_buffer_data(bo);
for (int i = 0; i < n; i++) {
float expected = (a[i] / (1.0f + expf(-a[i]))) * b[i];
ASSERT_FLT_NEAR(expected, got[i], 1e-5f);
}
fm_buffer_free(ba);
fm_buffer_free(bb);
fm_buffer_free(bo);
free(a);
free(b);
}
static void test_shexp_dot_gate_matches_golden_scalar(void)
{
float *gate_w = fm_golden_load_f32("moe_shared_gate_weight.bin", MOE_N_EMBD);
float *x = fm_golden_load_f32("moe_shared_gate_input.bin", MOE_N_EMBD);
float *golden = fm_golden_load_f32("moe_shared_gate_scalar.bin", 1);
ASSERT_NOT_NULL(gate_w);
ASSERT_NOT_NULL(x);
ASSERT_NOT_NULL(golden);
fm_buffer *bw = fm_buffer_alloc(MOE_N_EMBD * sizeof(float));
fm_buffer *bx = fm_buffer_alloc(MOE_N_EMBD * sizeof(float));
fm_buffer *bs = fm_buffer_alloc(MOE_N_EMBD * sizeof(float));
ASSERT_NOT_NULL(bw);
ASSERT_NOT_NULL(bx);
ASSERT_NOT_NULL(bs);
memcpy(fm_buffer_data(bw), gate_w, MOE_N_EMBD * sizeof(float));
memcpy(fm_buffer_data(bx), x, MOE_N_EMBD * sizeof(float));
/* Seeding the branch with 1.0 makes the in-place kernel report the scalar
* gate itself, which is what the golden vector holds. */
float *shexp = (float *)fm_buffer_data(bs);
for (int i = 0; i < MOE_N_EMBD; i++)
shexp[i] = 1.0f;
fm_moe_shexp_dot_gate(bw, bx, bs, 1, MOE_N_EMBD);
fm_sync();
for (int i = 0; i < MOE_N_EMBD; i++)
ASSERT_FLT_NEAR(golden[0], shexp[i], 5e-7f);
fm_buffer_free(bw);
fm_buffer_free(bx);
fm_buffer_free(bs);
free(gate_w);
free(x);
free(golden);
}
static void test_shexp_dot_gate_scales_the_branch_per_token(void)
{
const int n_tokens = 5;
float *gate_w = fm_golden_load_f32("moe_shared_gate_weight.bin", MOE_N_EMBD);
float *x0 = fm_golden_load_f32("moe_shared_gate_input.bin", MOE_N_EMBD);
ASSERT_NOT_NULL(gate_w);
ASSERT_NOT_NULL(x0);
fm_buffer *bw = fm_buffer_alloc(MOE_N_EMBD * sizeof(float));
fm_buffer *bx = fm_buffer_alloc((size_t)n_tokens * MOE_N_EMBD * sizeof(float));
fm_buffer *bs = fm_buffer_alloc((size_t)n_tokens * MOE_N_EMBD * sizeof(float));
ASSERT_NOT_NULL(bw);
ASSERT_NOT_NULL(bx);
ASSERT_NOT_NULL(bs);
memcpy(fm_buffer_data(bw), gate_w, MOE_N_EMBD * sizeof(float));
/* Each token scales the golden input differently, so every token must end
* up with its own gate — a kernel that reduced token 0 for everyone passes
* the single-token test and fails here. */
float *x = (float *)fm_buffer_data(bx);
float *shexp = (float *)fm_buffer_data(bs);
for (int t = 0; t < n_tokens; t++)
for (int i = 0; i < MOE_N_EMBD; i++) {
x[t * MOE_N_EMBD + i] = x0[i] * (0.5f + 0.25f * (float)t);
shexp[t * MOE_N_EMBD + i] = 2.0f;
}
fm_moe_shexp_dot_gate(bw, bx, bs, n_tokens, MOE_N_EMBD);
fm_sync();
for (int t = 0; t < n_tokens; t++) {
double dot = 0.0;
for (int i = 0; i < MOE_N_EMBD; i++)
dot += (double)gate_w[i] * (double)x[t * MOE_N_EMBD + i];
float expected = 2.0f * (float)(1.0 / (1.0 + exp(-dot)));
for (int i = 0; i < MOE_N_EMBD; i += 257)
ASSERT_FLT_NEAR(expected, shexp[t * MOE_N_EMBD + i], 1e-5f);
}
fm_buffer_free(bw);
fm_buffer_free(bx);
fm_buffer_free(bs);
free(gate_w);
free(x0);
}
/* ================================================================
* main
* ================================================================ */
int main(int argc, char **argv)
{
(void)argc;
printf("\n Metal MoE — Router, Sort Prepass, Shared Expert\n");
printf(" ==============================================\n\n");
if (fm_init() != 0 || !fm_available()) {
fprintf(stderr, " Metal unavailable — no device to test against.\n");
return 77;
}
fm_golden_locate(argv[0]);
float *probe = fm_golden_load_f32("moe_router_logits.bin", MOE_N_EXPERT);
if (!probe) {
fprintf(stderr, " Golden vectors missing — every assertion below is "
"absolute and needs them.\n");
fm_free();
return 77;
}
free(probe);
RUN_TEST(router_topk_matches_golden);
RUN_TEST(router_tie_break_prefers_lower_index);
RUN_TEST(router_tie_break_matches_cpu_selection_sort);
RUN_TEST(router_applies_expert_weights_scale);
RUN_TEST(router_batch_matches_cpu_reference);
RUN_TEST(expert_sort_matches_cpu_reference_at_prefill);
RUN_TEST(expert_sort_is_a_permutation_of_the_routing);
RUN_TEST(expert_sort_single_token);
RUN_TEST(silu_mul_matches_cpu_reference);
RUN_TEST(shexp_dot_gate_matches_golden_scalar);
RUN_TEST(shexp_dot_gate_scales_the_branch_per_token);
fm_free();
TEST_REPORT();
return __ft_failed > 0 ? 1 : 0;
}
Reference in New Issue View Git Blame Copy Permalink