317 lines
11 KiB
C
317 lines
11 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_metal_norm.c — RMS normalization kernels (fm_rms_norm,
|
|
* fm_rms_norm_nx, fm_qk_rms_norm)
|
|
*
|
|
* fm_qk_rms_norm is checked against the CPU-reference golden set in
|
|
* metal/tests/golden (provenance in fm_golden_vectors.h). The batched and
|
|
* single-vector RMS forms are checked against a CPU reference computed here
|
|
* — never against each other, and never against another GPU result.
|
|
*
|
|
* Two properties get their own tests because the obvious test would not
|
|
* notice them:
|
|
* - the QK-norm weight is SHARED by every head, so a kernel that indexes
|
|
* it per head is wrong; a uniform weight vector hides that completely,
|
|
* so the trap test uses a weight buffer whose per-head-indexed region
|
|
* holds deliberately different values,
|
|
* - fm_rms_norm_nx's row stride is n, which only shows up when the rows
|
|
* hold different data.
|
|
*
|
|
* Helper functions are named case_* so lib/test.h's failure line, which
|
|
* prints __func__ + 5, still reads as the case that failed.
|
|
*
|
|
* Exits 77 when no Metal device is present so GPU-less CI skips cleanly.
|
|
*/
|
|
|
|
#include "lib/test.h"
|
|
#include "metal/fm_metal.h"
|
|
#include "metal/tests/fm_golden_vectors.h"
|
|
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
/* qwen35moe attention shape the golden set was generated for. */
|
|
#define QKNORM_HEAD_DIM 256
|
|
#define QKNORM_HEADS 16
|
|
#define NORM_EPS 1e-6f
|
|
|
|
/* The GPU sums squares in float32 where the reference uses double, so some
|
|
* drift is expected; measured on M4 / macOS 26 it stays sub-ULP at these
|
|
* widths, well inside this band and two orders tighter than the 1e-4..1e-3
|
|
* the golden set allows for Kepler. */
|
|
#define NORM_TOL_ABS 1e-6f
|
|
#define NORM_TOL_REL 1e-6f
|
|
|
|
static uint32_t norm_rng = 0x9e3779b9u;
|
|
|
|
static float next_random_float(void) {
|
|
norm_rng ^= norm_rng << 13;
|
|
norm_rng ^= norm_rng >> 17;
|
|
norm_rng ^= norm_rng << 5;
|
|
return (float)((int32_t)norm_rng) / 2147483648.0f;
|
|
}
|
|
|
|
/* CPU reference: out = x * rsqrt(mean(x^2) + eps) * weight, sum of squares
|
|
* accumulated in double exactly as pn_vrms_norm_f32 does. weight NULL means
|
|
* no gain. */
|
|
static void reference_rms_norm(const float *x, const float *weight, float *out,
|
|
int n, float eps) {
|
|
double sum_sq = 0.0;
|
|
for (int i = 0; i < n; i++)
|
|
sum_sq += (double)x[i] * (double)x[i];
|
|
float inv_rms = (float)(1.0 / sqrt(sum_sq / (double)n + (double)eps));
|
|
for (int i = 0; i < n; i++)
|
|
out[i] = x[i] * inv_rms * (weight ? weight[i] : 1.0f);
|
|
}
|
|
|
|
static fm_buffer *upload_floats(const float *src, size_t n) {
|
|
fm_buffer *buf = fm_buffer_alloc(n * sizeof(float));
|
|
if (buf) memcpy(fm_buffer_data(buf), src, n * sizeof(float));
|
|
return buf;
|
|
}
|
|
|
|
/* Diffs against the reference. `label` is printed just before the assertion
|
|
* fires, because lib/test.h names the failing function and every case here
|
|
* shares this one — without it a failure would not say which. */
|
|
static void case_compare(const char *label, const float *expected,
|
|
const float *got, size_t n) {
|
|
for (size_t i = 0; i < n; i++) {
|
|
float tol = NORM_TOL_ABS + NORM_TOL_REL * fabsf(expected[i]);
|
|
if (fabsf(expected[i] - got[i]) > tol)
|
|
printf(" %s: element %zu of %zu\n", label, i, n);
|
|
ASSERT_FLT_NEAR(expected[i], got[i], tol);
|
|
}
|
|
}
|
|
|
|
/* ================================================================
|
|
* fm_qk_rms_norm
|
|
* ================================================================ */
|
|
|
|
static void test_qk_rms_norm_matches_golden(void) {
|
|
size_t count = (size_t)QKNORM_HEADS * QKNORM_HEAD_DIM;
|
|
float *input = fm_golden_load_f32("qknorm_input.bin", count);
|
|
float *weight = fm_golden_load_f32("qknorm_weight.bin", QKNORM_HEAD_DIM);
|
|
float *golden = fm_golden_load_f32("qknorm_output.bin", count);
|
|
ASSERT_NOT_NULL(input);
|
|
ASSERT_NOT_NULL(weight);
|
|
ASSERT_NOT_NULL(golden);
|
|
|
|
fm_buffer *bx = upload_floats(input, count);
|
|
fm_buffer *bw = upload_floats(weight, QKNORM_HEAD_DIM);
|
|
fm_buffer *bo = fm_buffer_alloc(count * sizeof(float));
|
|
ASSERT_NOT_NULL(bx);
|
|
ASSERT_NOT_NULL(bw);
|
|
ASSERT_NOT_NULL(bo);
|
|
|
|
fm_qk_rms_norm(bx, bw, bo, QKNORM_HEADS, QKNORM_HEAD_DIM, NORM_EPS);
|
|
fm_sync();
|
|
|
|
case_compare("qk_rms_norm vs qknorm_output.bin",
|
|
golden, (const float *)fm_buffer_data(bo), count);
|
|
|
|
fm_buffer_free(bx);
|
|
fm_buffer_free(bw);
|
|
fm_buffer_free(bo);
|
|
free(input);
|
|
free(weight);
|
|
free(golden);
|
|
}
|
|
|
|
/* THE SHARED-WEIGHT TRAP. The weight buffer is n_heads * head_dim long: the
|
|
* first head_dim floats are the real shared weight, and the rest — exactly
|
|
* the region a per-head-indexing kernel would read for heads 1..3 — hold
|
|
* very different values. So `weight[head * head_dim + i]` produces a
|
|
* specific, deterministic wrong answer instead of reading past the buffer
|
|
* and maybe getting away with it. A uniform weight would make both
|
|
* indexings agree and the test would prove nothing. */
|
|
static void test_qk_rms_norm_weight_is_shared_across_heads(void) {
|
|
const int heads = 4, head_dim = 64;
|
|
const size_t count = (size_t)heads * head_dim;
|
|
float input[4 * 64], weight[4 * 64], expected[4 * 64];
|
|
|
|
for (size_t i = 0; i < count; i++)
|
|
input[i] = next_random_float() * 3.0f;
|
|
for (int i = 0; i < head_dim; i++)
|
|
weight[i] = 0.5f + 0.01f * (float)i; /* the shared weight */
|
|
for (size_t i = head_dim; i < count; i++)
|
|
weight[i] = -4.0f - 0.02f * (float)i; /* never to be read */
|
|
|
|
for (int h = 0; h < heads; h++)
|
|
reference_rms_norm(input + (size_t)h * head_dim, weight,
|
|
expected + (size_t)h * head_dim, head_dim, NORM_EPS);
|
|
|
|
fm_buffer *bx = upload_floats(input, count);
|
|
fm_buffer *bw = upload_floats(weight, count);
|
|
fm_buffer *bo = fm_buffer_alloc(count * sizeof(float));
|
|
ASSERT_NOT_NULL(bx);
|
|
ASSERT_NOT_NULL(bw);
|
|
ASSERT_NOT_NULL(bo);
|
|
|
|
fm_qk_rms_norm(bx, bw, bo, heads, head_dim, NORM_EPS);
|
|
fm_sync();
|
|
|
|
case_compare("qk_rms_norm with a shared non-uniform weight",
|
|
expected, (const float *)fm_buffer_data(bo), count);
|
|
|
|
fm_buffer_free(bx);
|
|
fm_buffer_free(bw);
|
|
fm_buffer_free(bo);
|
|
}
|
|
|
|
static void test_qk_rms_norm_honors_null_weight(void) {
|
|
const int heads = 3, head_dim = 128;
|
|
const size_t count = (size_t)heads * head_dim;
|
|
float input[3 * 128], expected[3 * 128];
|
|
|
|
for (size_t i = 0; i < count; i++)
|
|
input[i] = next_random_float() * 2.0f;
|
|
for (int h = 0; h < heads; h++)
|
|
reference_rms_norm(input + (size_t)h * head_dim, NULL,
|
|
expected + (size_t)h * head_dim, head_dim, NORM_EPS);
|
|
|
|
fm_buffer *bx = upload_floats(input, count);
|
|
fm_buffer *bo = fm_buffer_alloc(count * sizeof(float));
|
|
ASSERT_NOT_NULL(bx);
|
|
ASSERT_NOT_NULL(bo);
|
|
|
|
fm_qk_rms_norm(bx, NULL, bo, heads, head_dim, NORM_EPS);
|
|
fm_sync();
|
|
|
|
case_compare("qk_rms_norm with weight == NULL",
|
|
expected, (const float *)fm_buffer_data(bo), count);
|
|
|
|
fm_buffer_free(bx);
|
|
fm_buffer_free(bo);
|
|
}
|
|
|
|
/* ================================================================
|
|
* fm_rms_norm_nx / fm_rms_norm
|
|
* ================================================================ */
|
|
|
|
/* Rows hold independent random data and n is neither a multiple of 4 nor of
|
|
* the 32-lane SIMD width, so a wrong row stride, a dropped tail element or a
|
|
* reduction that spans rows all change the result. */
|
|
static void test_rms_norm_nx_matches_cpu_reference(void) {
|
|
const int rows = 5, n = 300;
|
|
const size_t count = (size_t)rows * n;
|
|
float input[5 * 300], weight[300], expected[5 * 300];
|
|
|
|
for (size_t i = 0; i < count; i++)
|
|
input[i] = next_random_float() * (1.0f + (float)(i % 7));
|
|
for (int i = 0; i < n; i++)
|
|
weight[i] = 0.75f + 0.002f * (float)i;
|
|
|
|
for (int r = 0; r < rows; r++)
|
|
reference_rms_norm(input + (size_t)r * n, weight,
|
|
expected + (size_t)r * n, n, NORM_EPS);
|
|
|
|
fm_buffer *bx = upload_floats(input, count);
|
|
fm_buffer *bw = upload_floats(weight, n);
|
|
fm_buffer *bo = fm_buffer_alloc(count * sizeof(float));
|
|
ASSERT_NOT_NULL(bx);
|
|
ASSERT_NOT_NULL(bw);
|
|
ASSERT_NOT_NULL(bo);
|
|
|
|
fm_rms_norm_nx(bx, bw, bo, rows, n, NORM_EPS);
|
|
fm_sync();
|
|
|
|
case_compare("rms_norm_nx rows=5 n=300",
|
|
expected, (const float *)fm_buffer_data(bo), count);
|
|
|
|
fm_buffer_free(bx);
|
|
fm_buffer_free(bw);
|
|
fm_buffer_free(bo);
|
|
}
|
|
|
|
static void test_rms_norm_nx_honors_null_weight(void) {
|
|
const int rows = 3, n = 129;
|
|
const size_t count = (size_t)rows * n;
|
|
float input[3 * 129], expected[3 * 129];
|
|
|
|
for (size_t i = 0; i < count; i++)
|
|
input[i] = next_random_float() * 5.0f;
|
|
for (int r = 0; r < rows; r++)
|
|
reference_rms_norm(input + (size_t)r * n, NULL,
|
|
expected + (size_t)r * n, n, NORM_EPS);
|
|
|
|
fm_buffer *bx = upload_floats(input, count);
|
|
fm_buffer *bo = fm_buffer_alloc(count * sizeof(float));
|
|
ASSERT_NOT_NULL(bx);
|
|
ASSERT_NOT_NULL(bo);
|
|
|
|
fm_rms_norm_nx(bx, NULL, bo, rows, n, NORM_EPS);
|
|
fm_sync();
|
|
|
|
case_compare("rms_norm_nx with weight == NULL",
|
|
expected, (const float *)fm_buffer_data(bo), count);
|
|
|
|
fm_buffer_free(bx);
|
|
fm_buffer_free(bo);
|
|
}
|
|
|
|
/* The single-vector entry point is the rows == 1 case of the same kernel;
|
|
* this pins its behaviour to the CPU reference independently. */
|
|
static void test_rms_norm_single_vector_matches_cpu_reference(void) {
|
|
const int n = 2048;
|
|
float *input = (float *)malloc(n * sizeof(float));
|
|
float *weight = (float *)malloc(n * sizeof(float));
|
|
float *expected = (float *)malloc(n * sizeof(float));
|
|
ASSERT_NOT_NULL(input);
|
|
ASSERT_NOT_NULL(weight);
|
|
ASSERT_NOT_NULL(expected);
|
|
|
|
for (int i = 0; i < n; i++) {
|
|
input[i] = next_random_float() * 4.0f;
|
|
weight[i] = 1.0f + 0.1f * next_random_float();
|
|
}
|
|
reference_rms_norm(input, weight, expected, n, NORM_EPS);
|
|
|
|
fm_buffer *bx = upload_floats(input, n);
|
|
fm_buffer *bw = upload_floats(weight, n);
|
|
fm_buffer *bo = fm_buffer_alloc(n * sizeof(float));
|
|
ASSERT_NOT_NULL(bx);
|
|
ASSERT_NOT_NULL(bw);
|
|
ASSERT_NOT_NULL(bo);
|
|
|
|
fm_rms_norm(bx, bw, bo, n, NORM_EPS);
|
|
fm_sync();
|
|
|
|
case_compare("rms_norm single vector n=2048",
|
|
expected, (const float *)fm_buffer_data(bo), (size_t)n);
|
|
|
|
fm_buffer_free(bx);
|
|
fm_buffer_free(bw);
|
|
fm_buffer_free(bo);
|
|
free(input);
|
|
free(weight);
|
|
free(expected);
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
(void)argc;
|
|
fm_golden_locate(argv[0]);
|
|
|
|
printf("\n Metal Normalization Kernel Tests\n");
|
|
printf(" ================================\n\n");
|
|
|
|
if (fm_init() != 0 || !fm_available()) {
|
|
fprintf(stderr, " Metal unavailable — skipping.\n\n");
|
|
return 77;
|
|
}
|
|
|
|
RUN_TEST(qk_rms_norm_matches_golden);
|
|
RUN_TEST(qk_rms_norm_weight_is_shared_across_heads);
|
|
RUN_TEST(qk_rms_norm_honors_null_weight);
|
|
RUN_TEST(rms_norm_nx_matches_cpu_reference);
|
|
RUN_TEST(rms_norm_nx_honors_null_weight);
|
|
RUN_TEST(rms_norm_single_vector_matches_cpu_reference);
|
|
|
|
fm_free();
|
|
|
|
TEST_REPORT();
|
|
return __ft_failed > 0 ? 1 : 0;
|
|
}
|