113 lines
4.3 KiB
C
113 lines
4.3 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* bench_attn_split.c — decode wall time vs split-K workgroup count.
|
|
*
|
|
* This is the measurement FM_ATTN_DECODE_TARGET_WORKGROUPS is set from, kept
|
|
* runnable so the default can be re-derived on another machine instead of
|
|
* trusted. It sweeps the split plan at three context lengths and prints
|
|
* milliseconds per decoded token for the decode + reduce pair.
|
|
*
|
|
* Not a correctness test — it asserts nothing. Correctness lives in
|
|
* test_metal_attn.c. Exits 77 with no Metal device.
|
|
*/
|
|
|
|
#include "metal/kernels/fm_kernel_attn.h"
|
|
|
|
#include <mach/mach_time.h>
|
|
#include <stdio.h>
|
|
|
|
#define BENCH_HEAD_DIM 256
|
|
#define BENCH_N_HEAD 16
|
|
#define BENCH_N_HEAD_KV 2
|
|
#define BENCH_KV_GROUPS (BENCH_N_HEAD / BENCH_N_HEAD_KV)
|
|
#define BENCH_MAX_SEQ 4096
|
|
#define BENCH_WARMUP 5
|
|
#define BENCH_ITERS 50
|
|
|
|
static double bench_now_ms(void) {
|
|
static mach_timebase_info_data_t timebase;
|
|
if (timebase.denom == 0) mach_timebase_info(&timebase);
|
|
return (double)mach_absolute_time() * timebase.numer / timebase.denom / 1e6;
|
|
}
|
|
|
|
/* Deterministic, non-degenerate fill. The values only have to be finite and
|
|
* varied — this measures dispatch shape, not numerics. */
|
|
static void bench_fill(fm_buffer *buf, size_t n, int period, float step) {
|
|
float *p = (float *)fm_buffer_data(buf);
|
|
for (size_t i = 0; i < n; i++) p[i] = (float)(i % (size_t)period) * step;
|
|
}
|
|
|
|
static double bench_one_plan(fm_buffer *q, fm_buffer *k, fm_buffer *v,
|
|
fm_buffer *out, int seq_len, int n_splits,
|
|
int split_len, float scale) {
|
|
size_t m_floats = 0, acc_floats = 0;
|
|
fm_attn_decode_scratch_floats(BENCH_N_HEAD, BENCH_HEAD_DIM, n_splits,
|
|
&m_floats, &acc_floats);
|
|
fm_buffer *m = fm_buffer_alloc(m_floats * sizeof(float));
|
|
fm_buffer *l = fm_buffer_alloc(m_floats * sizeof(float));
|
|
fm_buffer *acc = fm_buffer_alloc(acc_floats * sizeof(float));
|
|
if (!m || !l || !acc) return -1.0;
|
|
|
|
double elapsed = 0.0;
|
|
for (int i = 0; i < BENCH_WARMUP + BENCH_ITERS; i++) {
|
|
if (i == BENCH_WARMUP) elapsed = bench_now_ms();
|
|
fm_attn_flash_decode(q, k, v, m, l, acc, BENCH_KV_GROUPS,
|
|
BENCH_N_HEAD_KV, BENCH_HEAD_DIM, BENCH_MAX_SEQ,
|
|
seq_len, n_splits, split_len, scale);
|
|
fm_attn_flash_decode_reduce(m, l, acc, out, BENCH_N_HEAD,
|
|
BENCH_HEAD_DIM, n_splits);
|
|
fm_sync();
|
|
}
|
|
elapsed = (bench_now_ms() - elapsed) / BENCH_ITERS;
|
|
|
|
fm_buffer_free(m);
|
|
fm_buffer_free(l);
|
|
fm_buffer_free(acc);
|
|
return elapsed;
|
|
}
|
|
|
|
int main(void) {
|
|
if (fm_init() != 0 || !fm_available()) {
|
|
fprintf(stderr, "Metal unavailable — skipping.\n");
|
|
return 77;
|
|
}
|
|
|
|
size_t cache_n = (size_t)BENCH_N_HEAD_KV * BENCH_MAX_SEQ * BENCH_HEAD_DIM;
|
|
size_t row_n = (size_t)BENCH_N_HEAD * BENCH_HEAD_DIM;
|
|
fm_buffer *k = fm_buffer_alloc(cache_n * sizeof(float));
|
|
fm_buffer *v = fm_buffer_alloc(cache_n * sizeof(float));
|
|
fm_buffer *q = fm_buffer_alloc(row_n * sizeof(float));
|
|
fm_buffer *out = fm_buffer_alloc(row_n * sizeof(float));
|
|
if (!k || !v || !q || !out) return 1;
|
|
bench_fill(k, cache_n, 17, 0.01f);
|
|
bench_fill(v, cache_n, 13, 0.02f);
|
|
bench_fill(q, row_n, 7, 0.03f);
|
|
|
|
const int seq_lens[] = { 128, 512, 4096 };
|
|
const int workgroups[] = { 2, 4, 8, 16, 32, 64 };
|
|
const float scale = 1.0f / 16.0f; /* 1/sqrt(256) */
|
|
|
|
printf("decode n_head=%d n_head_kv=%d head_dim=%d (ms per token)\n",
|
|
BENCH_N_HEAD, BENCH_N_HEAD_KV, BENCH_HEAD_DIM);
|
|
for (int s = 0; s < 3; s++) {
|
|
printf(" seq %5d:\n", seq_lens[s]);
|
|
for (int w = 0; w < 6; w++) {
|
|
int split_len = 0;
|
|
int n_splits = fm_attn_decode_plan(seq_lens[s], BENCH_N_HEAD_KV,
|
|
workgroups[w], &split_len);
|
|
double ms = bench_one_plan(q, k, v, out, seq_lens[s], n_splits,
|
|
split_len, scale);
|
|
printf(" %2d workgroups (%2d splits x %5d keys) %7.3f\n",
|
|
n_splits * BENCH_N_HEAD_KV, n_splits, split_len, ms);
|
|
}
|
|
}
|
|
|
|
fm_buffer_free(k);
|
|
fm_buffer_free(v);
|
|
fm_buffer_free(q);
|
|
fm_buffer_free(out);
|
|
fm_free();
|
|
return 0;
|
|
}
|