1594 lines
86 KiB
C
1594 lines
86 KiB
C
#include "gem/metal.h"
|
|
#include "gem/ops.h"
|
|
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
static uint64_t hash_f32(const float *values, size_t count)
|
|
{
|
|
uint64_t hash = UINT64_C(1469598103934665603);
|
|
for (size_t i = 0; i < count; ++i) {
|
|
uint32_t bits;
|
|
memcpy(&bits, values + i, sizeof(bits));
|
|
hash ^= bits;
|
|
hash *= UINT64_C(1099511628211);
|
|
}
|
|
return hash;
|
|
}
|
|
|
|
static int dispatch_attention(
|
|
struct gem_metal *metal, int split,
|
|
struct gem_metal_buffer *query_buffer,
|
|
struct gem_metal_buffer *cache_buffer, size_t value_cache_offset,
|
|
size_t query_heads, size_t kv_heads, size_t head_dimension,
|
|
size_t total_tokens, size_t cache_capacity, size_t window,
|
|
struct gem_metal_buffer *score_buffer,
|
|
struct gem_metal_buffer *output_buffer,
|
|
char *error, size_t error_size)
|
|
{
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, error_size);
|
|
if (status == 0) {
|
|
if (split == 3)
|
|
status = gem_metal_batch_attention_split_gqa_scores_dim32_f32(
|
|
batch, query_buffer, 0, cache_buffer, 0, value_cache_offset,
|
|
query_heads, kv_heads, head_dimension, total_tokens, cache_capacity,
|
|
window, 1.0f / sqrtf((float)head_dimension), score_buffer, 0,
|
|
output_buffer, 0, error, error_size);
|
|
else if (split == 2)
|
|
status = gem_metal_batch_attention_split_gqa_scores_f32(
|
|
batch, query_buffer, 0, cache_buffer, 0, value_cache_offset,
|
|
query_heads, kv_heads, head_dimension, total_tokens, cache_capacity,
|
|
window, 1.0f / sqrtf((float)head_dimension), score_buffer, 0,
|
|
output_buffer, 0, error, error_size);
|
|
else if (split)
|
|
status = gem_metal_batch_attention_split_gqa_f32(
|
|
batch, query_buffer, 0, cache_buffer, 0, value_cache_offset,
|
|
query_heads, kv_heads, head_dimension, total_tokens, cache_capacity,
|
|
window, 1.0f / sqrtf((float)head_dimension), score_buffer, 0,
|
|
output_buffer, 0, error, error_size);
|
|
else
|
|
status = gem_metal_batch_attention_parallel_f32(
|
|
batch, query_buffer, 0, cache_buffer, 0, value_cache_offset,
|
|
query_heads, kv_heads, head_dimension, total_tokens, cache_capacity,
|
|
window, 1.0f / sqrtf((float)head_dimension), score_buffer, 0,
|
|
output_buffer, 0, error, error_size);
|
|
}
|
|
if (status == 0)
|
|
status = gem_metal_batch_submit(batch, error, error_size);
|
|
gem_metal_batch_close(batch);
|
|
return status;
|
|
}
|
|
|
|
static int dispatch_partitioned_attention(
|
|
struct gem_metal *metal, size_t partitions, int score_tile8,
|
|
struct gem_metal_buffer *query_buffer,
|
|
struct gem_metal_buffer *cache_buffer, size_t value_cache_offset,
|
|
size_t query_heads, size_t kv_heads, size_t head_dimension,
|
|
size_t total_tokens, size_t cache_capacity,
|
|
struct gem_metal_buffer *score_buffer,
|
|
struct gem_metal_buffer *value_scratch_buffer,
|
|
struct gem_metal_buffer *output_buffer,
|
|
char *error, size_t error_size)
|
|
{
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, error_size);
|
|
if (status == 0)
|
|
status = (score_tile8 ?
|
|
gem_metal_batch_attention_split_gqa_scores_tile8_dim32_partitioned_f32 :
|
|
gem_metal_batch_attention_split_gqa_scores_dim32_partitioned_f32)(
|
|
batch, query_buffer, 0, cache_buffer, 0, value_cache_offset,
|
|
query_heads, kv_heads, head_dimension, total_tokens, cache_capacity, 0,
|
|
1.0f / sqrtf((float)head_dimension), score_buffer, 0,
|
|
value_scratch_buffer, 0, partitions, output_buffer, 0, error, error_size);
|
|
if (status == 0)
|
|
status = gem_metal_batch_submit(batch, error, error_size);
|
|
gem_metal_batch_close(batch);
|
|
return status;
|
|
}
|
|
|
|
static int run_case(struct gem_metal *metal, const char *name,
|
|
size_t query_heads, size_t kv_heads, size_t head_dimension,
|
|
size_t total_tokens, size_t cache_capacity, size_t window)
|
|
{
|
|
size_t live = window && total_tokens > window ? window : total_tokens;
|
|
size_t query_elements = query_heads * head_dimension;
|
|
size_t cache_elements = cache_capacity * kv_heads * head_dimension;
|
|
size_t score_elements = query_heads * live;
|
|
float *query = calloc(query_elements, sizeof(*query));
|
|
float *cache = calloc(cache_elements * 2, sizeof(*cache));
|
|
float *cpu_scores = calloc(score_elements, sizeof(*cpu_scores));
|
|
float *expected = calloc(query_elements, sizeof(*expected));
|
|
float *baseline = calloc(query_elements, sizeof(*baseline));
|
|
float *baseline_repeat = calloc(query_elements, sizeof(*baseline_repeat));
|
|
float *actual = calloc(query_elements, sizeof(*actual));
|
|
float *actual_repeat = calloc(query_elements, sizeof(*actual_repeat));
|
|
float *candidate = calloc(query_elements, sizeof(*candidate));
|
|
float *candidate_repeat = calloc(query_elements, sizeof(*candidate_repeat));
|
|
float *split_scores = calloc(score_elements, sizeof(*split_scores));
|
|
float *split_scores_repeat = calloc(score_elements, sizeof(*split_scores_repeat));
|
|
struct gem_metal_buffer *query_buffer = NULL;
|
|
struct gem_metal_buffer *cache_buffer = NULL;
|
|
struct gem_metal_buffer *score_buffer = NULL;
|
|
struct gem_metal_buffer *value_scratch_buffer = NULL;
|
|
struct gem_metal_buffer *baseline_output_buffer = NULL;
|
|
struct gem_metal_buffer *output_buffer = NULL;
|
|
char error[512];
|
|
int result = 1;
|
|
if (!query || !cache || !cpu_scores || !expected || !baseline || !baseline_repeat ||
|
|
!actual || !actual_repeat || !candidate || !candidate_repeat)
|
|
goto cleanup;
|
|
if (!split_scores || !split_scores_repeat)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < cache_elements * 2; ++i)
|
|
cache[i] = ((int)((i * 13 + i / 17) % 251) - 125) * (1.0f / 1024.0f);
|
|
if (gem_metal_buffer_new(metal, query_elements * sizeof(float), &query_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, cache_elements * 2 * sizeof(float), &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, score_elements * sizeof(float), &score_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * 16 * sizeof(float),
|
|
&value_scratch_buffer, error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * sizeof(float), &baseline_output_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * sizeof(float), &output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s allocation: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache,
|
|
cache_elements * 2 * sizeof(float));
|
|
for (size_t repetition = 0; repetition < 3; ++repetition) {
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 17 + repetition * 31) % 257) - 128) * (1.0f / 512.0f);
|
|
if (gem_attention_decode_ring_f32(
|
|
query, cache, cache + cache_elements, query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
1.0f / sqrtf((float)head_dimension), cpu_scores, expected) != 0) {
|
|
fprintf(stderr, "%s CPU oracle failed\n", name);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query,
|
|
query_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(baseline_output_buffer), 0xa5,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_attention(metal, 0, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, baseline_output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s baseline dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(baseline, gem_metal_buffer_contents(baseline_output_buffer),
|
|
query_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(baseline_output_buffer), 0x5a,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_attention(metal, 0, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, baseline_output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s baseline repeat dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(baseline_repeat, gem_metal_buffer_contents(baseline_output_buffer),
|
|
query_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(output_buffer), 0xa5,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_attention(metal, 2, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(actual, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
memcpy(split_scores, gem_metal_buffer_contents(score_buffer),
|
|
score_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_attention(metal, 2, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s repeat dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(actual_repeat, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
memcpy(split_scores_repeat, gem_metal_buffer_contents(score_buffer),
|
|
score_elements * sizeof(float));
|
|
int test_dim32 = query_heads == 32 && kv_heads == 4 &&
|
|
head_dimension == 512 && window == 0 && total_tokens >= 8192;
|
|
if (test_dim32) {
|
|
memset(gem_metal_buffer_contents(output_buffer), 0xa5,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_attention(metal, 3, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s dim32 dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(candidate, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_attention(metal, 3, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s dim32 repeat dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(candidate_repeat, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
static const size_t partitions[] = {2, 4, 8, 16};
|
|
for (size_t p = 0; p < sizeof(partitions) / sizeof(partitions[0]); ++p) {
|
|
memset(gem_metal_buffer_contents(output_buffer), 0xa5,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_partitioned_attention(
|
|
metal, partitions[p], 0, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, score_buffer,
|
|
value_scratch_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s split%zu dispatch: %s\n",
|
|
name, partitions[p], error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(candidate, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_partitioned_attention(
|
|
metal, partitions[p], 0, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, score_buffer,
|
|
value_scratch_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s split%zu repeat dispatch: %s\n",
|
|
name, partitions[p], error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(candidate_repeat, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
float difference = fabsf(candidate[i] - expected[i]);
|
|
if (candidate[i] != candidate_repeat[i] || !isfinite(candidate[i]) ||
|
|
difference > 1e-3f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "%s split%zu repetition %zu element %zu: "
|
|
"candidate %.9g / %.9g CPU %.9g difference %.9g\n",
|
|
name, partitions[p], repetition, i, candidate[i],
|
|
candidate_repeat[i], expected[i], difference);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
}
|
|
for (size_t attempt = 0; attempt < 2; ++attempt) {
|
|
memset(gem_metal_buffer_contents(output_buffer), attempt ? 0x5a : 0xa5,
|
|
query_elements * sizeof(float));
|
|
if (dispatch_partitioned_attention(
|
|
metal, 16, 1, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, score_buffer,
|
|
value_scratch_buffer, output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s tile8 score dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
const float *tile8 = gem_metal_buffer_contents(output_buffer);
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
if (tile8[i] != candidate[i]) {
|
|
fprintf(stderr, "%s repetition %zu element %zu tile8 score mismatch: "
|
|
"candidate %.9g split16 %.9g\n",
|
|
name, repetition, i, tile8[i], candidate[i]);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
}
|
|
/* Restore the unsplit dim32 candidate for the exact parity check below. */
|
|
if (dispatch_attention(metal, 3, query_buffer, cache_buffer,
|
|
cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, total_tokens, cache_capacity, window,
|
|
score_buffer, output_buffer, error, sizeof(error)) != 0)
|
|
goto cleanup;
|
|
memcpy(candidate, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
memcpy(candidate_repeat, candidate, query_elements * sizeof(float));
|
|
}
|
|
double square_error = 0.0;
|
|
double square_reference = 0.0;
|
|
double parity_square_error = 0.0;
|
|
double parity_square_reference = 0.0;
|
|
float maximum_error = 0.0f;
|
|
for (size_t i = 0; i < score_elements; ++i) {
|
|
if (split_scores[i] != split_scores_repeat[i]) {
|
|
fprintf(stderr, "%s repetition %zu score %zu nondeterministic: %.9g / %.9g\n",
|
|
name, repetition, i, split_scores[i], split_scores_repeat[i]);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
float difference = fabsf(actual[i] - expected[i]);
|
|
float baseline_difference = fabsf(baseline[i] - expected[i]);
|
|
float parity_difference = fabsf(actual[i] - baseline[i]);
|
|
if (actual[i] != actual_repeat[i] || baseline[i] != baseline_repeat[i]) {
|
|
fprintf(stderr, "%s repetition %zu element %zu nondeterministic: "
|
|
"split %.9g / %.9g parallel %.9g / %.9g CPU %.9g\n",
|
|
name, repetition, i, actual[i], actual_repeat[i], baseline[i],
|
|
baseline_repeat[i], expected[i]);
|
|
goto cleanup;
|
|
}
|
|
if (test_dim32 &&
|
|
(candidate[i] != candidate_repeat[i] || candidate[i] != actual[i])) {
|
|
fprintf(stderr, "%s repetition %zu element %zu dim32 mismatch: "
|
|
"candidate %.9g / %.9g baseline-value %.9g CPU %.9g\n",
|
|
name, repetition, i, candidate[i], candidate_repeat[i],
|
|
actual[i], expected[i]);
|
|
goto cleanup;
|
|
}
|
|
if (!isfinite(actual[i]) || !isfinite(baseline[i]) ||
|
|
difference > 1e-3f * (1.0f + fabsf(expected[i])) ||
|
|
baseline_difference > 2e-3f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "%s repetition %zu element %zu: split %.9g "
|
|
"parallel %.9g CPU %.9g split-CPU %.9g "
|
|
"parallel-CPU %.9g parity %.9g\n",
|
|
name, repetition, i, actual[i], baseline[i], expected[i],
|
|
difference, baseline_difference, parity_difference);
|
|
goto cleanup;
|
|
}
|
|
if (difference > maximum_error)
|
|
maximum_error = difference;
|
|
square_error += (double)difference * difference;
|
|
square_reference += (double)expected[i] * expected[i];
|
|
parity_square_error += (double)parity_difference * parity_difference;
|
|
parity_square_reference += (double)baseline[i] * baseline[i];
|
|
}
|
|
double nmse = square_reference > 0.0 ? square_error / square_reference : square_error;
|
|
double parity_nmse = parity_square_reference > 0.0 ?
|
|
parity_square_error / parity_square_reference : parity_square_error;
|
|
if (nmse > 5e-6 || parity_nmse > 1e-8) {
|
|
fprintf(stderr, "%s repetition %zu NMSE %.9g parity NMSE %.9g max error %.9g\n",
|
|
name, repetition, nmse, parity_nmse, maximum_error);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
result = 0;
|
|
|
|
cleanup:
|
|
gem_metal_buffer_close(value_scratch_buffer);
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(baseline_output_buffer);
|
|
gem_metal_buffer_close(score_buffer);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(split_scores_repeat);
|
|
free(split_scores);
|
|
free(actual_repeat);
|
|
free(actual);
|
|
free(candidate_repeat);
|
|
free(candidate);
|
|
free(baseline_repeat);
|
|
free(baseline);
|
|
free(expected);
|
|
free(cpu_scores);
|
|
free(cache);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
static size_t tree_slot(size_t local, size_t committed_count, size_t committed_start,
|
|
size_t cache_tokens, size_t row, size_t depth,
|
|
size_t speculative_count, const uint8_t parents[8])
|
|
{
|
|
if (local < committed_count)
|
|
return committed_start + local;
|
|
size_t target_depth = depth + 1 - speculative_count + local - committed_count;
|
|
size_t node = row;
|
|
for (size_t current_depth = depth; current_depth > target_depth; --current_depth)
|
|
node = parents[node];
|
|
return cache_tokens + node;
|
|
}
|
|
|
|
static int run_tree_case(struct gem_metal *metal, size_t window)
|
|
{
|
|
const size_t rows = 8, query_heads = 4, kv_heads = 2, head_dimension = 64;
|
|
const size_t cache_tokens = 5, cache_capacity = 16;
|
|
const uint8_t parents[8] = {UINT8_MAX, 0, 0, 1, 1, 2, 2, 4};
|
|
const uint8_t depths[8] = {0, 1, 1, 2, 2, 2, 2, 3};
|
|
const float scale = 1.0f / 8.0f;
|
|
size_t query_elements = rows * query_heads * head_dimension;
|
|
size_t cache_elements = cache_capacity * kv_heads * head_dimension;
|
|
float *query = calloc(query_elements, sizeof(*query));
|
|
float *cache = calloc(2 * cache_elements, sizeof(*cache));
|
|
float *expected = calloc(query_elements, sizeof(*expected));
|
|
float *actual = calloc(query_elements, sizeof(*actual));
|
|
struct gem_metal_buffer *query_buffer = NULL;
|
|
struct gem_metal_buffer *cache_buffer = NULL;
|
|
struct gem_metal_buffer *output_buffer = NULL;
|
|
char error[512];
|
|
int result = 1;
|
|
if (!query || !cache || !expected || !actual)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 17 + 3) % 113) - 56) * (1.0f / 256.0f);
|
|
for (size_t i = 0; i < 2 * cache_elements; ++i)
|
|
cache[i] = ((int)((i * 29 + i / 31 + 7) % 127) - 63) * (1.0f / 256.0f);
|
|
for (size_t row = 0; row < rows; ++row) {
|
|
size_t depth = depths[row];
|
|
size_t speculative_count = window && depth + 1 > window ? window : depth + 1;
|
|
size_t committed_budget = window ? window - speculative_count : cache_tokens;
|
|
size_t committed_count = cache_tokens < committed_budget ?
|
|
cache_tokens : committed_budget;
|
|
size_t committed_start = cache_tokens - committed_count;
|
|
size_t live = committed_count + speculative_count;
|
|
for (size_t head = 0; head < query_heads; ++head) {
|
|
size_t kv_head = head / (query_heads / kv_heads);
|
|
const float *q = query + (row * query_heads + head) * head_dimension;
|
|
float scores[13];
|
|
float maximum = -INFINITY;
|
|
for (size_t local = 0; local < live; ++local) {
|
|
size_t slot = tree_slot(local, committed_count, committed_start,
|
|
cache_tokens, row, depth, speculative_count, parents);
|
|
const float *key = cache +
|
|
(slot * kv_heads + kv_head) * head_dimension;
|
|
float score = 0.0f;
|
|
for (size_t i = 0; i < head_dimension; ++i) score += q[i] * key[i];
|
|
scores[local] = score * scale;
|
|
if (scores[local] > maximum) maximum = scores[local];
|
|
}
|
|
float denominator = 0.0f;
|
|
for (size_t local = 0; local < live; ++local) {
|
|
scores[local] = expf(scores[local] - maximum);
|
|
denominator += scores[local];
|
|
}
|
|
float *destination = expected +
|
|
(row * query_heads + head) * head_dimension;
|
|
for (size_t dimension = 0; dimension < head_dimension; ++dimension)
|
|
for (size_t local = 0; local < live; ++local) {
|
|
size_t slot = tree_slot(local, committed_count, committed_start,
|
|
cache_tokens, row, depth, speculative_count, parents);
|
|
size_t value = cache_elements +
|
|
(slot * kv_heads + kv_head) * head_dimension + dimension;
|
|
destination[dimension] += scores[local] / denominator * cache[value];
|
|
}
|
|
}
|
|
}
|
|
if (gem_metal_buffer_new(metal, query_elements * sizeof(float), &query_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * cache_elements * sizeof(float), &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * sizeof(float), &output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "tree attention allocation: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query, query_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache, 2 * cache_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
if (gem_metal_batch_begin(metal, &batch, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_attention_tree8_f32(
|
|
batch, query_buffer, 0, cache_buffer, 0, cache_elements * sizeof(float),
|
|
query_heads, kv_heads, head_dimension, cache_tokens, cache_capacity,
|
|
window, scale, parents, depths, output_buffer, 0, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_submit(batch, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "tree attention dispatch: %s\n", error);
|
|
gem_metal_batch_close(batch);
|
|
goto cleanup;
|
|
}
|
|
gem_metal_batch_close(batch);
|
|
memcpy(actual, gem_metal_buffer_contents(output_buffer), query_elements * sizeof(float));
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
float difference = fabsf(actual[i] - expected[i]);
|
|
if (!isfinite(actual[i]) || difference > 2e-5f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "tree attention window %zu element %zu: Metal %.9g CPU %.9g\n",
|
|
window, i, actual[i], expected[i]);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
uint8_t invalid_parents[8];
|
|
memcpy(invalid_parents, parents, sizeof(invalid_parents));
|
|
invalid_parents[1] = 1;
|
|
batch = NULL;
|
|
if (gem_metal_batch_begin(metal, &batch, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_attention_tree8_f32(
|
|
batch, query_buffer, 0, cache_buffer, 0, cache_elements * sizeof(float),
|
|
query_heads, kv_heads, head_dimension, cache_tokens, cache_capacity,
|
|
window, scale, invalid_parents, depths, output_buffer, 0,
|
|
error, sizeof(error)) == 0) {
|
|
fprintf(stderr, "tree attention accepted a non-topological parent table\n");
|
|
gem_metal_batch_close(batch);
|
|
goto cleanup;
|
|
}
|
|
gem_metal_batch_close(batch);
|
|
result = 0;
|
|
|
|
cleanup:
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(actual);
|
|
free(expected);
|
|
free(cache);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
static int run_batched_wrapped_prefill_case(struct gem_metal *metal)
|
|
{
|
|
enum { rows = 8, query_heads = 4, kv_heads = 2, head_dimension = 64,
|
|
cache_capacity = 16, cache_tokens = 32 };
|
|
const float scale = 1.0f / 8.0f;
|
|
size_t query_elements = (size_t)rows * query_heads * head_dimension;
|
|
size_t pending_elements = (size_t)rows * kv_heads * head_dimension;
|
|
size_t one_cache_elements = (size_t)cache_capacity * kv_heads * head_dimension;
|
|
size_t score_elements = (size_t)rows * query_heads * cache_capacity;
|
|
float *query = calloc(query_elements, sizeof(*query));
|
|
float *pending_key = calloc(pending_elements, sizeof(*pending_key));
|
|
float *pending_value = calloc(pending_elements, sizeof(*pending_value));
|
|
float *cache = calloc(2 * one_cache_elements, sizeof(*cache));
|
|
float *expected = calloc(query_elements, sizeof(*expected));
|
|
float *actual = calloc(query_elements, sizeof(*actual));
|
|
struct gem_metal_buffer *query_buffer = NULL, *pending_key_buffer = NULL;
|
|
struct gem_metal_buffer *pending_value_buffer = NULL, *cache_buffer = NULL;
|
|
struct gem_metal_buffer *score_buffer = NULL, *output_buffer = NULL;
|
|
char error[512];
|
|
int result = 1;
|
|
if (!query || !pending_key || !pending_value || !cache || !expected || !actual)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 17 + 3) % 113) - 56) * (1.0f / 256.0f);
|
|
for (size_t i = 0; i < pending_elements; ++i) {
|
|
pending_key[i] = ((int)((i * 19 + 5) % 127) - 63) * (1.0f / 256.0f);
|
|
pending_value[i] = ((int)((i * 23 + 7) % 131) - 65) * (1.0f / 256.0f);
|
|
}
|
|
for (size_t i = 0; i < 2 * one_cache_elements; ++i)
|
|
cache[i] = ((int)((i * 29 + i / 31 + 11) % 137) - 68) * (1.0f / 256.0f);
|
|
|
|
for (size_t row = 0; row < rows; ++row) {
|
|
size_t total = cache_tokens + row + 1;
|
|
size_t start = total - cache_capacity;
|
|
for (size_t head = 0; head < query_heads; ++head) {
|
|
size_t kv_head = head / (query_heads / kv_heads);
|
|
const float *q = query + (row * query_heads + head) * head_dimension;
|
|
float scores[cache_capacity];
|
|
float maximum = -INFINITY;
|
|
for (size_t local = 0; local < cache_capacity; ++local) {
|
|
size_t logical = start + local;
|
|
const float *key = logical < cache_tokens ?
|
|
cache + ((logical % cache_capacity) * kv_heads + kv_head) * head_dimension :
|
|
pending_key + ((logical - cache_tokens) * kv_heads + kv_head) * head_dimension;
|
|
float score = 0.0f;
|
|
for (size_t dimension = 0; dimension < head_dimension; ++dimension)
|
|
score += q[dimension] * key[dimension];
|
|
scores[local] = score * scale;
|
|
if (scores[local] > maximum) maximum = scores[local];
|
|
}
|
|
float denominator = 0.0f;
|
|
for (size_t local = 0; local < cache_capacity; ++local) {
|
|
scores[local] = expf(scores[local] - maximum);
|
|
denominator += scores[local];
|
|
}
|
|
float *destination = expected +
|
|
(row * query_heads + head) * head_dimension;
|
|
for (size_t local = 0; local < cache_capacity; ++local) {
|
|
size_t logical = start + local;
|
|
const float *value = logical < cache_tokens ?
|
|
cache + one_cache_elements +
|
|
((logical % cache_capacity) * kv_heads + kv_head) * head_dimension :
|
|
pending_value +
|
|
((logical - cache_tokens) * kv_heads + kv_head) * head_dimension;
|
|
float weight = scores[local] / denominator;
|
|
for (size_t dimension = 0; dimension < head_dimension; ++dimension)
|
|
destination[dimension] += weight * value[dimension];
|
|
}
|
|
}
|
|
}
|
|
|
|
if (gem_metal_buffer_new(metal, query_elements * sizeof(float), &query_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_elements * sizeof(float), &pending_key_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_elements * sizeof(float), &pending_value_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * one_cache_elements * sizeof(float), &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, score_elements * sizeof(float), &score_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * sizeof(float), &output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "wrapped batched prefill allocation: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query, query_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(pending_key_buffer), pending_key,
|
|
pending_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(pending_value_buffer), pending_value,
|
|
pending_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache,
|
|
2 * one_cache_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
if (gem_metal_batch_begin(metal, &batch, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_attention_prefill_split_gqa_scores_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, rows, cache_tokens, cache_capacity, cache_capacity, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_submit(batch, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "wrapped batched prefill dispatch: %s\n", error);
|
|
gem_metal_batch_close(batch);
|
|
goto cleanup;
|
|
}
|
|
gem_metal_batch_close(batch);
|
|
memcpy(actual, gem_metal_buffer_contents(output_buffer), query_elements * sizeof(float));
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
float difference = fabsf(actual[i] - expected[i]);
|
|
if (!isfinite(actual[i]) || difference > 3e-5f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "wrapped batched prefill element %zu: Metal %.9g CPU %.9g\n",
|
|
i, actual[i], expected[i]);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
result = 0;
|
|
|
|
cleanup:
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(score_buffer);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(pending_value_buffer);
|
|
gem_metal_buffer_close(pending_key_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(actual);
|
|
free(expected);
|
|
free(cache);
|
|
free(pending_value);
|
|
free(pending_key);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
static int run_q8k8_prefill_case(struct gem_metal *metal, const char *name,
|
|
size_t query_heads, size_t kv_heads,
|
|
size_t head_dimension, size_t cache_capacity,
|
|
size_t cache_tokens, size_t window)
|
|
{
|
|
const size_t rows = 32;
|
|
size_t maximum_total = cache_tokens + rows;
|
|
size_t maximum_start = window && maximum_total > window ? maximum_total - window : 0;
|
|
size_t score_stride = maximum_total - maximum_start;
|
|
const float scale = 1.0f / sqrtf((float)head_dimension);
|
|
size_t query_elements = (size_t)rows * query_heads * head_dimension;
|
|
size_t pending_elements = (size_t)rows * kv_heads * head_dimension;
|
|
size_t one_cache_elements = (size_t)cache_capacity * kv_heads * head_dimension;
|
|
size_t score_elements = (size_t)rows * query_heads * score_stride;
|
|
float *query = calloc(query_elements, sizeof(*query));
|
|
float *pending_key = calloc(pending_elements, sizeof(*pending_key));
|
|
float *pending_value = calloc(pending_elements, sizeof(*pending_value));
|
|
float *cache = calloc(2 * one_cache_elements, sizeof(*cache));
|
|
float *baseline = calloc(query_elements, sizeof(*baseline));
|
|
float *candidate = calloc(query_elements, sizeof(*candidate));
|
|
float *mma_candidate = calloc(query_elements, sizeof(*mma_candidate));
|
|
float *online_candidate = calloc(query_elements, sizeof(*online_candidate));
|
|
float *online_repeat = calloc(query_elements, sizeof(*online_repeat));
|
|
float *row2_candidate = calloc(query_elements, sizeof(*row2_candidate));
|
|
float *row2_repeat = calloc(query_elements, sizeof(*row2_repeat));
|
|
float *row4_candidate = calloc(query_elements, sizeof(*row4_candidate));
|
|
float *row4_repeat = calloc(query_elements, sizeof(*row4_repeat));
|
|
float *bq8_candidate = calloc(query_elements, sizeof(*bq8_candidate));
|
|
float *bq8_repeat = calloc(query_elements, sizeof(*bq8_repeat));
|
|
float *bq8_bk128_candidate = calloc(query_elements, sizeof(*bq8_bk128_candidate));
|
|
float *bq8_bk128_repeat = calloc(query_elements, sizeof(*bq8_bk128_repeat));
|
|
float *chunked_candidate = calloc(query_elements, sizeof(*chunked_candidate));
|
|
float *baseline_scores = calloc(score_elements, sizeof(*baseline_scores));
|
|
float *q8k8_scores = calloc(score_elements, sizeof(*q8k8_scores));
|
|
float *mma_scores = calloc(score_elements, sizeof(*mma_scores));
|
|
float *swa_row2_scores = calloc(score_elements, sizeof(*swa_row2_scores));
|
|
struct gem_metal_buffer *query_buffer = NULL, *pending_key_buffer = NULL;
|
|
struct gem_metal_buffer *pending_value_buffer = NULL, *cache_buffer = NULL;
|
|
struct gem_metal_buffer *score_buffer = NULL, *output_buffer = NULL;
|
|
char error[512];
|
|
int result = 1;
|
|
if (!query || !pending_key || !pending_value || !cache || !baseline || !candidate ||
|
|
!mma_candidate || !online_candidate || !online_repeat ||
|
|
!row2_candidate || !row2_repeat ||
|
|
!row4_candidate || !row4_repeat ||
|
|
!bq8_candidate || !bq8_repeat || !bq8_bk128_candidate || !bq8_bk128_repeat ||
|
|
!chunked_candidate ||
|
|
!baseline_scores || !q8k8_scores || !mma_scores || !swa_row2_scores)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 17 + i / 29 + 3) % 251) - 125) * (1.0f / 128.0f);
|
|
for (size_t i = 0; i < pending_elements; ++i) {
|
|
pending_key[i] = ((int)((i * 19 + i / 31 + 5) % 241) - 120) * (1.0f / 128.0f);
|
|
pending_value[i] = ((int)((i * 23 + i / 37 + 7) % 239) - 119) * (1.0f / 128.0f);
|
|
}
|
|
for (size_t i = 0; i < 2 * one_cache_elements; ++i)
|
|
cache[i] = ((int)((i * 29 + i / 41 + 11) % 233) - 116) * (1.0f / 128.0f);
|
|
if (gem_metal_buffer_new(metal, query_elements * sizeof(float), &query_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_elements * sizeof(float), &pending_key_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_elements * sizeof(float), &pending_value_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * one_cache_elements * sizeof(float), &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, score_elements * sizeof(float), &score_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * sizeof(float), &output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s allocation: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query, query_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(pending_key_buffer), pending_key,
|
|
pending_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(pending_value_buffer), pending_value,
|
|
pending_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache,
|
|
2 * one_cache_elements * sizeof(float));
|
|
if (!window && query_heads == 32 && kv_heads == 4 && head_dimension == 512) {
|
|
size_t kv_row_bytes = kv_heads * head_dimension * sizeof(float);
|
|
size_t overflow_capacity = SIZE_MAX / kv_row_bytes + 1;
|
|
struct gem_metal_batch *invalid_batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &invalid_batch, error, sizeof(error));
|
|
int row1_status = status == 0 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_f32(
|
|
invalid_batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error)) : status;
|
|
int row2_status = status == 0 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_row2_f32(
|
|
invalid_batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error)) : status;
|
|
int row4_status = status == 0 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_row4_f32(
|
|
invalid_batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error)) : status;
|
|
int bq8_status = status == 0 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_bq8_f32(
|
|
invalid_batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error)) : status;
|
|
int bq8_bk128_status = status == 0 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_bq8_bk128_f32(
|
|
invalid_batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error)) : status;
|
|
gem_metal_batch_close(invalid_batch);
|
|
if (status != 0 || row1_status == 0 || row2_status == 0 || row4_status == 0 ||
|
|
bq8_status == 0 || bq8_bk128_status == 0) {
|
|
fprintf(stderr, "%s online capacity overflow was accepted\n", name);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
if (!window) {
|
|
float *outputs[] = {bq8_bk128_candidate, bq8_bk128_repeat};
|
|
for (size_t pass = 0; pass < 2; ++pass) {
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a,
|
|
query_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_bq8_bk128_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "%s BQ8/BK128 dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(outputs[pass], gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
}
|
|
if (memcmp(bq8_bk128_candidate, bq8_bk128_repeat,
|
|
query_elements * sizeof(float)) != 0) {
|
|
fprintf(stderr, "%s BQ8/BK128 repeat changed output hash %016llx/%016llx\n",
|
|
name,
|
|
(unsigned long long)hash_f32(bq8_bk128_candidate, query_elements),
|
|
(unsigned long long)hash_f32(bq8_bk128_repeat, query_elements));
|
|
goto cleanup;
|
|
}
|
|
}
|
|
int candidate_count = window ? 4 : 7;
|
|
for (int candidate_kernel = 0; candidate_kernel < candidate_count; ++candidate_kernel) {
|
|
if (candidate_kernel < 4)
|
|
memset(gem_metal_buffer_contents(score_buffer), 0,
|
|
score_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(output_buffer), 0,
|
|
query_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0) {
|
|
if (candidate_kernel == 6)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_bq8_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else if (candidate_kernel == 5)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_row4_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else if (candidate_kernel == 4)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_row2_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else if (candidate_kernel == 3)
|
|
status = window ?
|
|
gem_metal_batch_attention_prefill_split_gqa_scores_swa_row2_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, rows, cache_tokens, cache_capacity, window, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error)) :
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else if (candidate_kernel == 2)
|
|
status = gem_metal_batch_attention_prefill_split_gqa_q8k8_mma_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), query_heads,
|
|
kv_heads, head_dimension, rows, cache_tokens, cache_capacity, window, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error));
|
|
else if (candidate_kernel == 1)
|
|
status = gem_metal_batch_attention_prefill_split_gqa_scores_q8k8_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), query_heads,
|
|
kv_heads, head_dimension, rows, cache_tokens, cache_capacity, window, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error));
|
|
else
|
|
status = gem_metal_batch_attention_prefill_split_gqa_scores_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), query_heads,
|
|
kv_heads, head_dimension, rows, cache_tokens, cache_capacity, window, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error));
|
|
}
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "%s dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(candidate_kernel == 6 ? bq8_candidate :
|
|
(candidate_kernel == 5 ? row4_candidate :
|
|
(candidate_kernel == 4 ? row2_candidate :
|
|
(candidate_kernel == 3 ? online_candidate :
|
|
(candidate_kernel == 2 ? mma_candidate :
|
|
(candidate_kernel == 1 ? candidate : baseline))))),
|
|
gem_metal_buffer_contents(output_buffer), query_elements * sizeof(float));
|
|
if (candidate_kernel == 0)
|
|
memcpy(baseline_scores, gem_metal_buffer_contents(score_buffer),
|
|
score_elements * sizeof(float));
|
|
else if (candidate_kernel == 1)
|
|
memcpy(q8k8_scores, gem_metal_buffer_contents(score_buffer),
|
|
score_elements * sizeof(float));
|
|
else if (candidate_kernel == 2)
|
|
memcpy(mma_scores, gem_metal_buffer_contents(score_buffer),
|
|
score_elements * sizeof(float));
|
|
else if (candidate_kernel == 3 && window)
|
|
memcpy(swa_row2_scores, gem_metal_buffer_contents(score_buffer),
|
|
score_elements * sizeof(float));
|
|
}
|
|
if (window &&
|
|
(memcmp(baseline_scores, swa_row2_scores, score_elements * sizeof(float)) != 0 ||
|
|
memcmp(baseline, online_candidate, query_elements * sizeof(float)) != 0)) {
|
|
fprintf(stderr,
|
|
"%s exact-order SWA row2 changed score/output hashes %016llx/%016llx "
|
|
"%016llx/%016llx\n",
|
|
name,
|
|
(unsigned long long)hash_f32(baseline_scores, score_elements),
|
|
(unsigned long long)hash_f32(swa_row2_scores, score_elements),
|
|
(unsigned long long)hash_f32(baseline, query_elements),
|
|
(unsigned long long)hash_f32(online_candidate, query_elements));
|
|
goto cleanup;
|
|
}
|
|
if (!window) {
|
|
float *repeats[] = { online_repeat, row2_repeat, row4_repeat, bq8_repeat };
|
|
const float *firsts[] = { online_candidate, row2_candidate, row4_candidate,
|
|
bq8_candidate };
|
|
for (size_t repeat_kernel = 0; repeat_kernel < 4; ++repeat_kernel) {
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a,
|
|
query_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0) {
|
|
if (repeat_kernel == 0)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else if (repeat_kernel == 1)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_row2_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else if (repeat_kernel == 2)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_row4_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
else
|
|
status = gem_metal_batch_attention_prefill_online_gqa_31b_full_bq8_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), rows, cache_tokens,
|
|
cache_capacity, scale, output_buffer, 0, error, sizeof(error));
|
|
}
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "%s online repeat dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(repeats[repeat_kernel], gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
if (memcmp(firsts[repeat_kernel], repeats[repeat_kernel],
|
|
query_elements * sizeof(float)) != 0) {
|
|
fprintf(stderr, "%s online %s repeat changed output hash %016llx/%016llx\n",
|
|
name, repeat_kernel == 0 ? "row1" :
|
|
(repeat_kernel == 1 ? "row2" :
|
|
(repeat_kernel == 2 ? "row4" : "bq8")),
|
|
(unsigned long long)hash_f32(firsts[repeat_kernel], query_elements),
|
|
(unsigned long long)hash_f32(repeats[repeat_kernel], query_elements));
|
|
goto cleanup;
|
|
}
|
|
}
|
|
}
|
|
if (!window) {
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache,
|
|
2 * one_cache_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
size_t query_row_bytes = query_heads * head_dimension * sizeof(float);
|
|
size_t kv_row_bytes = kv_heads * head_dimension * sizeof(float);
|
|
for (size_t row_base = 0; status == 0 && row_base < rows; row_base += 8) {
|
|
status = gem_metal_batch_attention_prefill_split_gqa_q8k8_mma_f32(
|
|
batch, query_buffer, row_base * query_row_bytes,
|
|
pending_key_buffer, row_base * kv_row_bytes,
|
|
pending_value_buffer, row_base * kv_row_bytes,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), query_heads,
|
|
kv_heads, head_dimension, 8, cache_tokens + row_base, cache_capacity,
|
|
0, scale, score_buffer, 0, output_buffer,
|
|
row_base * query_row_bytes, error, sizeof(error));
|
|
for (size_t row = 0; status == 0 && row < 8; ++row)
|
|
status = gem_metal_batch_kv_store_f32(
|
|
batch, pending_key_buffer, (row_base + row) * kv_row_bytes,
|
|
pending_value_buffer, (row_base + row) * kv_row_bytes,
|
|
kv_heads * head_dimension, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), cache_capacity,
|
|
cache_tokens + row_base + row, error, sizeof(error));
|
|
}
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "%s chunked dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(chunked_candidate, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
}
|
|
double output_square_error = 0.0, output_square_reference = 0.0;
|
|
double mma_square_error = 0.0, mma_square_reference = 0.0;
|
|
double online_square_error = 0.0, online_square_reference = 0.0;
|
|
double row2_square_error = 0.0, row2_square_reference = 0.0;
|
|
double row4_square_error = 0.0, row4_square_reference = 0.0;
|
|
double bq8_square_error = 0.0, bq8_square_reference = 0.0;
|
|
double bq8_bk128_square_error = 0.0, bq8_bk128_square_reference = 0.0;
|
|
double chunked_square_error = 0.0, chunked_square_reference = 0.0;
|
|
double score_square_error = 0.0, score_square_reference = 0.0;
|
|
double mma_score_square_error = 0.0;
|
|
float maximum_score_error = 0.0f;
|
|
size_t maximum_score_index = 0;
|
|
float *candidate_scores = mma_scores;
|
|
if (getenv("GEM_ATTENTION_HASHES"))
|
|
fprintf(stderr, "%s score hashes scalar=%016llx q8=%016llx mma=%016llx\n", name,
|
|
(unsigned long long)hash_f32(baseline_scores, score_elements),
|
|
(unsigned long long)hash_f32(q8k8_scores, score_elements),
|
|
(unsigned long long)hash_f32(candidate_scores, score_elements));
|
|
for (size_t row = 0; row < rows; ++row) {
|
|
size_t total = cache_tokens + row + 1;
|
|
size_t live = window && total > window ? window : total;
|
|
for (size_t head = 0; head < query_heads; ++head)
|
|
for (size_t token = 0; token < live; ++token) {
|
|
size_t index = (row * query_heads + head) * score_stride + token;
|
|
float difference = q8k8_scores[index] - baseline_scores[index];
|
|
float mma_difference = candidate_scores[index] - q8k8_scores[index];
|
|
if (fabsf(difference) > maximum_score_error) {
|
|
maximum_score_error = fabsf(difference);
|
|
maximum_score_index = index;
|
|
}
|
|
if (!isfinite(candidate_scores[index])) {
|
|
fprintf(stderr, "%s non-finite score %zu\n", name, index);
|
|
goto cleanup;
|
|
}
|
|
score_square_error += (double)difference * difference;
|
|
mma_score_square_error += (double)mma_difference * mma_difference;
|
|
score_square_reference +=
|
|
(double)baseline_scores[index] * baseline_scores[index];
|
|
}
|
|
}
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
float difference = candidate[i] - baseline[i];
|
|
float mma_difference = mma_candidate[i] - candidate[i];
|
|
float online_difference = online_candidate[i] - baseline[i];
|
|
float row2_difference = row2_candidate[i] - baseline[i];
|
|
float row4_difference = row4_candidate[i] - baseline[i];
|
|
float bq8_difference = bq8_candidate[i] - baseline[i];
|
|
float bq8_bk128_difference = bq8_bk128_candidate[i] - baseline[i];
|
|
float chunked_difference = chunked_candidate[i] - mma_candidate[i];
|
|
if (!isfinite(candidate[i]) || !isfinite(mma_candidate[i])) {
|
|
fprintf(stderr, "%s non-finite output %zu\n", name, i);
|
|
goto cleanup;
|
|
}
|
|
output_square_error += (double)difference * difference;
|
|
output_square_reference += (double)baseline[i] * baseline[i];
|
|
mma_square_error += (double)mma_difference * mma_difference;
|
|
mma_square_reference += (double)candidate[i] * candidate[i];
|
|
if (!window) {
|
|
if (!isfinite(online_candidate[i]) || !isfinite(row2_candidate[i]) ||
|
|
!isfinite(row4_candidate[i]) || !isfinite(bq8_candidate[i]) ||
|
|
!isfinite(bq8_bk128_candidate[i])) {
|
|
fprintf(stderr, "%s non-finite online output %zu\n", name, i);
|
|
goto cleanup;
|
|
}
|
|
online_square_error += (double)online_difference * online_difference;
|
|
online_square_reference += (double)baseline[i] * baseline[i];
|
|
row2_square_error += (double)row2_difference * row2_difference;
|
|
row2_square_reference += (double)baseline[i] * baseline[i];
|
|
row4_square_error += (double)row4_difference * row4_difference;
|
|
row4_square_reference += (double)baseline[i] * baseline[i];
|
|
bq8_square_error += (double)bq8_difference * bq8_difference;
|
|
bq8_square_reference += (double)baseline[i] * baseline[i];
|
|
bq8_bk128_square_error +=
|
|
(double)bq8_bk128_difference * bq8_bk128_difference;
|
|
bq8_bk128_square_reference += (double)baseline[i] * baseline[i];
|
|
chunked_square_error += (double)chunked_difference * chunked_difference;
|
|
chunked_square_reference += (double)mma_candidate[i] * mma_candidate[i];
|
|
}
|
|
}
|
|
double score_nmse = score_square_error / score_square_reference;
|
|
double output_nmse = output_square_error / output_square_reference;
|
|
double mma_score_nmse = mma_score_square_error / score_square_reference;
|
|
double mma_output_nmse = mma_square_error / mma_square_reference;
|
|
double online_output_nmse = window ? 0.0 : online_square_error / online_square_reference;
|
|
double row2_output_nmse = window ? 0.0 : row2_square_error / row2_square_reference;
|
|
double row4_output_nmse = window ? 0.0 : row4_square_error / row4_square_reference;
|
|
double bq8_output_nmse = window ? 0.0 : bq8_square_error / bq8_square_reference;
|
|
double bq8_bk128_output_nmse =
|
|
window ? 0.0 : bq8_bk128_square_error / bq8_bk128_square_reference;
|
|
double chunked_output_nmse = window ? 0.0 : chunked_square_error / chunked_square_reference;
|
|
if (getenv("GEM_ATTENTION_HASHES"))
|
|
fprintf(stderr,
|
|
"%s output NMSE q8=%.12g mma=%.12g online=%.12g row2=%.12g row4=%.12g bq8=%.12g bq8-bk128=%.12g chunked=%.12g\n",
|
|
name, output_nmse, mma_output_nmse,
|
|
online_output_nmse, row2_output_nmse, row4_output_nmse,
|
|
bq8_output_nmse, bq8_bk128_output_nmse, chunked_output_nmse);
|
|
if (score_nmse > 1e-8 || output_nmse > 1e-8 ||
|
|
mma_score_nmse > 1e-8 || mma_output_nmse > 1e-8 || online_output_nmse > 1e-6 ||
|
|
row2_output_nmse > 1e-6 || row4_output_nmse > 1e-6 || bq8_output_nmse > 1e-6 ||
|
|
bq8_bk128_output_nmse > 1e-6 ||
|
|
chunked_output_nmse > 1e-8) {
|
|
size_t maximum_local = maximum_score_index % score_stride;
|
|
size_t maximum_head_row = maximum_score_index / score_stride;
|
|
fprintf(stderr, "%s q8 score/output NMSE %.9g/%.9g MMA score/output NMSE %.9g/%.9g "
|
|
"online/row2/row4/bq8/bq8-bk128/chunked output NMSE %.9g/%.9g/%.9g/%.9g/%.9g/%.9g max score row/head/token=%zu/%zu/%zu "
|
|
"values=%.9g/%.9g error=%.9g\n",
|
|
name, score_nmse, output_nmse, mma_score_nmse, mma_output_nmse,
|
|
online_output_nmse, row2_output_nmse, row4_output_nmse,
|
|
bq8_output_nmse, bq8_bk128_output_nmse, chunked_output_nmse,
|
|
maximum_head_row / query_heads, maximum_head_row % query_heads, maximum_local,
|
|
baseline_scores[maximum_score_index], q8k8_scores[maximum_score_index],
|
|
maximum_score_error);
|
|
goto cleanup;
|
|
}
|
|
result = 0;
|
|
|
|
cleanup:
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(score_buffer);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(pending_value_buffer);
|
|
gem_metal_buffer_close(pending_key_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(q8k8_scores);
|
|
free(mma_scores);
|
|
free(swa_row2_scores);
|
|
free(baseline_scores);
|
|
free(mma_candidate);
|
|
free(row2_repeat);
|
|
free(row2_candidate);
|
|
free(row4_repeat);
|
|
free(row4_candidate);
|
|
free(bq8_repeat);
|
|
free(bq8_candidate);
|
|
free(bq8_bk128_repeat);
|
|
free(bq8_bk128_candidate);
|
|
free(online_repeat);
|
|
free(online_candidate);
|
|
free(chunked_candidate);
|
|
free(candidate);
|
|
free(baseline);
|
|
free(cache);
|
|
free(pending_value);
|
|
free(pending_key);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
static int run_long_scalar_repeat_case(struct gem_metal *metal)
|
|
{
|
|
enum { rows = 8, query_heads = 32, kv_heads = 4, head_dimension = 512,
|
|
cache_tokens = 4096, cache_capacity = 4104, repetitions = 20 };
|
|
const float scale = 1.0f / sqrtf((float)head_dimension);
|
|
const size_t query_elements = (size_t)rows * query_heads * head_dimension;
|
|
const size_t pending_elements = (size_t)rows * kv_heads * head_dimension;
|
|
const size_t one_cache_elements = (size_t)cache_capacity * kv_heads * head_dimension;
|
|
const size_t score_elements = (size_t)rows * query_heads * cache_capacity;
|
|
float *query = malloc(query_elements * sizeof(*query));
|
|
float *pending_key = malloc(pending_elements * sizeof(*pending_key));
|
|
float *pending_value = malloc(pending_elements * sizeof(*pending_value));
|
|
float *cache = malloc(2 * one_cache_elements * sizeof(*cache));
|
|
float *reference_scores = malloc(score_elements * sizeof(*reference_scores));
|
|
float *reference_output = malloc(query_elements * sizeof(*reference_output));
|
|
float *rowwise_output = malloc(query_elements * sizeof(*rowwise_output));
|
|
struct gem_metal_buffer *query_buffer = NULL, *pending_key_buffer = NULL;
|
|
struct gem_metal_buffer *pending_value_buffer = NULL, *cache_buffer = NULL;
|
|
struct gem_metal_buffer *score_buffer = NULL, *output_buffer = NULL;
|
|
char error[512] = {0};
|
|
int result = 1;
|
|
if (!query || !pending_key || !pending_value || !cache ||
|
|
!reference_scores || !reference_output || !rowwise_output)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 17 + i / 29 + 3) % 251) - 125) * (1.0f / 128.0f);
|
|
for (size_t i = 0; i < pending_elements; ++i) {
|
|
pending_key[i] = ((int)((i * 19 + i / 31 + 5) % 241) - 120) * (1.0f / 128.0f);
|
|
pending_value[i] = ((int)((i * 23 + i / 37 + 7) % 239) - 119) * (1.0f / 128.0f);
|
|
}
|
|
for (size_t i = 0; i < 2 * one_cache_elements; ++i)
|
|
cache[i] = ((int)((i * 29 + i / 41 + 11) % 233) - 116) * (1.0f / 128.0f);
|
|
if (gem_metal_buffer_new(metal, query_elements * sizeof(float), &query_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_elements * sizeof(float), &pending_key_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_elements * sizeof(float), &pending_value_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * one_cache_elements * sizeof(float), &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, score_elements * sizeof(float), &score_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_elements * sizeof(float), &output_buffer,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "long scalar repeat allocation: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query, query_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(pending_key_buffer), pending_key,
|
|
pending_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(pending_value_buffer), pending_value,
|
|
pending_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache,
|
|
2 * one_cache_elements * sizeof(float));
|
|
{
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
size_t query_row_bytes = query_heads * head_dimension * sizeof(float);
|
|
size_t kv_row_bytes = kv_heads * head_dimension * sizeof(float);
|
|
size_t score_row_bytes = query_heads * cache_capacity * sizeof(float);
|
|
for (size_t row = 0; status == 0 && row < rows; ++row)
|
|
status = gem_metal_batch_kv_store_f32(
|
|
batch, pending_key_buffer, row * kv_row_bytes,
|
|
pending_value_buffer, row * kv_row_bytes, kv_heads * head_dimension,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), cache_capacity,
|
|
cache_tokens + row, error, sizeof(error));
|
|
for (size_t row = 0; status == 0 && row < rows; ++row)
|
|
status = gem_metal_batch_attention_split_gqa_scores_f32(
|
|
batch, query_buffer, row * query_row_bytes, cache_buffer, 0,
|
|
one_cache_elements * sizeof(float), query_heads, kv_heads,
|
|
head_dimension, cache_tokens + row + 1, cache_capacity, 0, scale,
|
|
score_buffer, row * score_row_bytes, output_buffer,
|
|
row * query_row_bytes, error, sizeof(error));
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "long scalar rowwise dispatch: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(rowwise_output, gem_metal_buffer_contents(output_buffer),
|
|
query_elements * sizeof(float));
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache,
|
|
2 * one_cache_elements * sizeof(float));
|
|
}
|
|
for (size_t repetition = 0; repetition < repetitions; ++repetition) {
|
|
memset(gem_metal_buffer_contents(score_buffer), 0,
|
|
score_elements * sizeof(float));
|
|
memset(gem_metal_buffer_contents(output_buffer), 0,
|
|
query_elements * sizeof(float));
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0)
|
|
status = gem_metal_batch_attention_prefill_split_gqa_scores_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_elements * sizeof(float), query_heads,
|
|
kv_heads, head_dimension, rows, cache_tokens, cache_capacity, 0, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error));
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "long scalar repeat dispatch: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
float *scores = gem_metal_buffer_contents(score_buffer);
|
|
float *output = gem_metal_buffer_contents(output_buffer);
|
|
if (repetition == 0) {
|
|
memcpy(reference_scores, scores, score_elements * sizeof(float));
|
|
memcpy(reference_output, output, query_elements * sizeof(float));
|
|
double square_error = 0.0, square_reference = 0.0;
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
double difference = (double)output[i] - rowwise_output[i];
|
|
square_error += difference * difference;
|
|
square_reference += (double)rowwise_output[i] * rowwise_output[i];
|
|
}
|
|
double nmse = square_error / square_reference;
|
|
if (getenv("GEM_ATTENTION_HASHES") || nmse > 1e-8)
|
|
fprintf(stderr, "long scalar batched-vs-rowwise output NMSE %.12g\n", nmse);
|
|
if (nmse > 1e-6) goto cleanup;
|
|
} else if (memcmp(reference_scores, scores, score_elements * sizeof(float)) != 0 ||
|
|
memcmp(reference_output, output, query_elements * sizeof(float)) != 0) {
|
|
fprintf(stderr,
|
|
"long scalar repeat %zu changed hashes scores=%016llx/%016llx "
|
|
"output=%016llx/%016llx\n", repetition,
|
|
(unsigned long long)hash_f32(reference_scores, score_elements),
|
|
(unsigned long long)hash_f32(scores, score_elements),
|
|
(unsigned long long)hash_f32(reference_output, query_elements),
|
|
(unsigned long long)hash_f32(output, query_elements));
|
|
goto cleanup;
|
|
}
|
|
}
|
|
result = 0;
|
|
|
|
cleanup:
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(score_buffer);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(pending_value_buffer);
|
|
gem_metal_buffer_close(pending_key_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(reference_output);
|
|
free(rowwise_output);
|
|
free(reference_scores);
|
|
free(cache);
|
|
free(pending_value);
|
|
free(pending_key);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
static int run_online_f16_case(struct gem_metal *metal, const char *name,
|
|
size_t query_heads, size_t kv_heads)
|
|
{
|
|
enum { rows = 8, head_dimension = 512, cache_tokens = 128, cache_capacity = 136 };
|
|
const float scale = 1.0f / sqrtf((float)head_dimension);
|
|
const size_t query_elements = (size_t)rows * query_heads * head_dimension;
|
|
const size_t pending_elements = (size_t)rows * kv_heads * head_dimension;
|
|
const size_t one_cache_elements = (size_t)cache_capacity * kv_heads * head_dimension;
|
|
const size_t score_elements = (size_t)rows * query_heads * cache_capacity;
|
|
const size_t query_bytes = query_elements * sizeof(float);
|
|
const size_t pending_bytes = pending_elements * sizeof(float);
|
|
const size_t kv_row_bytes = (size_t)kv_heads * head_dimension * sizeof(float);
|
|
const size_t one_cache_f32_bytes = one_cache_elements * sizeof(float);
|
|
const size_t one_cache_f16_bytes = one_cache_elements * sizeof(uint16_t);
|
|
float *query = malloc(query_bytes);
|
|
float *pending_key = malloc(pending_bytes);
|
|
float *pending_value = malloc(pending_bytes);
|
|
float *cache = malloc(2 * one_cache_f32_bytes);
|
|
float *baseline = malloc(query_bytes);
|
|
float *candidate = malloc(query_bytes);
|
|
float *repeat = malloc(query_bytes);
|
|
struct gem_metal_buffer *query_buffer = NULL, *pending_key_buffer = NULL;
|
|
struct gem_metal_buffer *pending_value_buffer = NULL, *cache_f32_buffer = NULL;
|
|
struct gem_metal_buffer *cache_f16_buffer = NULL, *score_buffer = NULL;
|
|
struct gem_metal_buffer *output_buffer = NULL;
|
|
char error[512] = {0};
|
|
int result = 1;
|
|
if (!query || !pending_key || !pending_value || !cache || !baseline ||
|
|
!candidate || !repeat)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 17 + i / 31 + 3) % 251) - 125) / 128.0f;
|
|
for (size_t i = 0; i < pending_elements; ++i) {
|
|
pending_key[i] = ((int)((i * 19 + i / 37 + 5) % 241) - 120) / 128.0f;
|
|
pending_value[i] = ((int)((i * 23 + i / 41 + 7) % 239) - 119) / 128.0f;
|
|
}
|
|
for (size_t i = 0; i < 2 * one_cache_elements; ++i)
|
|
cache[i] = ((int)((i * 29 + i / 43 + 11) % 233) - 116) / 127.0f;
|
|
if (gem_metal_buffer_new(metal, query_bytes, &query_buffer, error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_bytes, &pending_key_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_bytes, &pending_value_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * one_cache_f32_bytes, &cache_f32_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * one_cache_f16_bytes, &cache_f16_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, score_elements * sizeof(float), &score_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_bytes, &output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "online F16 %s allocation: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query, query_bytes);
|
|
memcpy(gem_metal_buffer_contents(pending_key_buffer), pending_key, pending_bytes);
|
|
memcpy(gem_metal_buffer_contents(pending_value_buffer), pending_value, pending_bytes);
|
|
memcpy(gem_metal_buffer_contents(cache_f32_buffer), cache, 2 * one_cache_f32_bytes);
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
for (size_t token = 0; status == 0 && token < cache_tokens; ++token)
|
|
status = gem_metal_batch_kv_store_f16(
|
|
batch, cache_f32_buffer, token * kv_row_bytes,
|
|
cache_f32_buffer, one_cache_f32_bytes + token * kv_row_bytes,
|
|
(size_t)kv_heads * head_dimension, cache_f16_buffer, 0,
|
|
one_cache_f16_bytes, cache_capacity, token, error, sizeof(error));
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "online F16 %s cache conversion: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
for (int pass = 0; pass < 3; ++pass) {
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a, query_bytes);
|
|
batch = NULL;
|
|
status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0 && pass == 0)
|
|
status = gem_metal_batch_attention_prefill_split_gqa_scores_f16(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_f16_buffer, 0, one_cache_f16_bytes, query_heads, kv_heads,
|
|
head_dimension, rows, cache_tokens, cache_capacity, 0, scale,
|
|
score_buffer, 0, output_buffer, 0, error, sizeof(error));
|
|
else if (status == 0)
|
|
status = (query_heads == 16 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_12b_full_f16 :
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_f16)(
|
|
batch, query_buffer, 0, pending_key_buffer, 0,
|
|
pending_value_buffer, 0, cache_f16_buffer, 0,
|
|
one_cache_f16_bytes, rows, cache_tokens, cache_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error));
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "online F16 %s dispatch: %s\n", name, error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(pass == 0 ? baseline : (pass == 1 ? candidate : repeat),
|
|
gem_metal_buffer_contents(output_buffer), query_bytes);
|
|
}
|
|
if (memcmp(candidate, repeat, query_bytes) != 0) {
|
|
fprintf(stderr, "online F16 %s repeat changed output hash %016llx/%016llx\n", name,
|
|
(unsigned long long)hash_f32(candidate, query_elements),
|
|
(unsigned long long)hash_f32(repeat, query_elements));
|
|
goto cleanup;
|
|
}
|
|
double square_error = 0.0, square_reference = 0.0;
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
if (!isfinite(candidate[i])) {
|
|
fprintf(stderr, "online F16 %s produced non-finite output %zu\n", name, i);
|
|
goto cleanup;
|
|
}
|
|
double difference = (double)candidate[i] - baseline[i];
|
|
square_error += difference * difference;
|
|
square_reference += (double)baseline[i] * baseline[i];
|
|
}
|
|
double nmse = square_error / square_reference;
|
|
if (nmse > 1e-6) {
|
|
fprintf(stderr, "online F16 %s output NMSE %.9g\n", name, nmse);
|
|
goto cleanup;
|
|
}
|
|
size_t kv_row_f16_bytes = (size_t)kv_heads * head_dimension * sizeof(uint16_t);
|
|
size_t overflow_capacity = SIZE_MAX / kv_row_f16_bytes + 1;
|
|
batch = NULL;
|
|
status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0)
|
|
status = (query_heads == 16 ?
|
|
gem_metal_batch_attention_prefill_online_gqa_12b_full_f16 :
|
|
gem_metal_batch_attention_prefill_online_gqa_31b_full_f16)(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_f16_buffer, 0, one_cache_f16_bytes, rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status == 0) {
|
|
fprintf(stderr, "online F16 %s capacity overflow was accepted\n", name);
|
|
goto cleanup;
|
|
}
|
|
result = 0;
|
|
cleanup:
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(score_buffer);
|
|
gem_metal_buffer_close(cache_f16_buffer);
|
|
gem_metal_buffer_close(cache_f32_buffer);
|
|
gem_metal_buffer_close(pending_value_buffer);
|
|
gem_metal_buffer_close(pending_key_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(repeat);
|
|
free(candidate);
|
|
free(baseline);
|
|
free(cache);
|
|
free(pending_value);
|
|
free(pending_key);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
static int run_online_12b_case(struct gem_metal *metal)
|
|
{
|
|
enum { rows = 8, query_heads = 16, kv_heads = 1, head_dimension = 512,
|
|
cache_tokens = 128, cache_capacity = 136 };
|
|
const float scale = 1.0f / sqrtf((float)head_dimension);
|
|
const size_t query_elements = (size_t)rows * query_heads * head_dimension;
|
|
const size_t pending_elements = (size_t)rows * kv_heads * head_dimension;
|
|
const size_t one_cache_elements = (size_t)cache_capacity * kv_heads * head_dimension;
|
|
const size_t score_elements = (size_t)rows * query_heads * cache_capacity;
|
|
float *query = malloc(query_elements * sizeof(float));
|
|
float *pending_key = malloc(pending_elements * sizeof(float));
|
|
float *pending_value = malloc(pending_elements * sizeof(float));
|
|
float *cache = malloc(2 * one_cache_elements * sizeof(float));
|
|
float *baseline = malloc(query_elements * sizeof(float));
|
|
float *candidate = malloc(query_elements * sizeof(float));
|
|
float *repeat = malloc(query_elements * sizeof(float));
|
|
struct gem_metal_buffer *query_buffer = NULL, *pending_key_buffer = NULL;
|
|
struct gem_metal_buffer *pending_value_buffer = NULL, *cache_buffer = NULL;
|
|
struct gem_metal_buffer *score_buffer = NULL, *output_buffer = NULL;
|
|
char error[512] = {0};
|
|
int result = 1;
|
|
if (!query || !pending_key || !pending_value || !cache || !baseline ||
|
|
!candidate || !repeat)
|
|
goto cleanup;
|
|
for (size_t i = 0; i < query_elements; ++i)
|
|
query[i] = ((int)((i * 31 + i / 17 + 3) % 251) - 125) / 128.0f;
|
|
for (size_t i = 0; i < pending_elements; ++i) {
|
|
pending_key[i] = ((int)((i * 37 + i / 19 + 5) % 241) - 120) / 128.0f;
|
|
pending_value[i] = ((int)((i * 41 + i / 23 + 7) % 239) - 119) / 128.0f;
|
|
}
|
|
for (size_t i = 0; i < 2 * one_cache_elements; ++i)
|
|
cache[i] = ((int)((i * 43 + i / 29 + 11) % 233) - 116) / 128.0f;
|
|
size_t query_bytes = query_elements * sizeof(float);
|
|
size_t pending_bytes = pending_elements * sizeof(float);
|
|
size_t one_cache_bytes = one_cache_elements * sizeof(float);
|
|
if (gem_metal_buffer_new(metal, query_bytes, &query_buffer, error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_bytes, &pending_key_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, pending_bytes, &pending_value_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, 2 * one_cache_bytes, &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, score_elements * sizeof(float), &score_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, query_bytes, &output_buffer, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "online 12B allocation: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(query_buffer), query, query_bytes);
|
|
memcpy(gem_metal_buffer_contents(pending_key_buffer), pending_key, pending_bytes);
|
|
memcpy(gem_metal_buffer_contents(pending_value_buffer), pending_value, pending_bytes);
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), cache, 2 * one_cache_bytes);
|
|
for (int pass = 0; pass < 3; ++pass) {
|
|
memset(gem_metal_buffer_contents(output_buffer), 0x5a, query_bytes);
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0 && pass == 0)
|
|
status = gem_metal_batch_attention_prefill_split_gqa_scores_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_bytes, query_heads, kv_heads, head_dimension,
|
|
rows, cache_tokens, cache_capacity, 0, scale, score_buffer, 0,
|
|
output_buffer, 0, error, sizeof(error));
|
|
else if (status == 0)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_12b_full_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_bytes, rows, cache_tokens, cache_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error));
|
|
if (status == 0) status = gem_metal_batch_submit(batch, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status != 0) {
|
|
fprintf(stderr, "online 12B dispatch: %s\n", error);
|
|
goto cleanup;
|
|
}
|
|
memcpy(pass == 0 ? baseline : (pass == 1 ? candidate : repeat),
|
|
gem_metal_buffer_contents(output_buffer), query_bytes);
|
|
}
|
|
if (memcmp(candidate, repeat, query_bytes) != 0) {
|
|
fprintf(stderr, "online 12B repeat changed output hash %016llx/%016llx\n",
|
|
(unsigned long long)hash_f32(candidate, query_elements),
|
|
(unsigned long long)hash_f32(repeat, query_elements));
|
|
goto cleanup;
|
|
}
|
|
double square_error = 0.0, square_reference = 0.0;
|
|
for (size_t i = 0; i < query_elements; ++i) {
|
|
if (!isfinite(candidate[i])) {
|
|
fprintf(stderr, "online 12B produced non-finite output %zu\n", i);
|
|
goto cleanup;
|
|
}
|
|
double difference = (double)candidate[i] - baseline[i];
|
|
square_error += difference * difference;
|
|
square_reference += (double)baseline[i] * baseline[i];
|
|
}
|
|
double nmse = square_error / square_reference;
|
|
if (nmse > 1e-6) {
|
|
fprintf(stderr, "online 12B output NMSE %.9g\n", nmse);
|
|
goto cleanup;
|
|
}
|
|
size_t kv_row_bytes = (size_t)kv_heads * head_dimension * sizeof(float);
|
|
size_t overflow_capacity = SIZE_MAX / kv_row_bytes + 1;
|
|
struct gem_metal_batch *batch = NULL;
|
|
int status = gem_metal_batch_begin(metal, &batch, error, sizeof(error));
|
|
if (status == 0)
|
|
status = gem_metal_batch_attention_prefill_online_gqa_12b_full_f32(
|
|
batch, query_buffer, 0, pending_key_buffer, 0, pending_value_buffer, 0,
|
|
cache_buffer, 0, one_cache_bytes, rows, 0, overflow_capacity,
|
|
scale, output_buffer, 0, error, sizeof(error));
|
|
gem_metal_batch_close(batch);
|
|
if (status == 0) {
|
|
fprintf(stderr, "online 12B capacity overflow was accepted\n");
|
|
goto cleanup;
|
|
}
|
|
result = 0;
|
|
cleanup:
|
|
gem_metal_buffer_close(output_buffer);
|
|
gem_metal_buffer_close(score_buffer);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(pending_value_buffer);
|
|
gem_metal_buffer_close(pending_key_buffer);
|
|
gem_metal_buffer_close(query_buffer);
|
|
free(repeat);
|
|
free(candidate);
|
|
free(baseline);
|
|
free(cache);
|
|
free(pending_value);
|
|
free(pending_key);
|
|
free(query);
|
|
return result;
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
struct gem_metal *metal = NULL;
|
|
char error[512];
|
|
if (gem_metal_open(&metal, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "SKIP: %s\n", error);
|
|
return 77;
|
|
}
|
|
int result =
|
|
run_case(metal, "31B full wrapped", 32, 4, 512, 769, 512, 512) ||
|
|
run_case(metal, "31B full dim32 threshold", 32, 4, 512, 8192, 8192, 0) ||
|
|
run_case(metal, "31B full dim32 nonmultiple", 32, 4, 512, 8193, 8193, 0) ||
|
|
run_case(metal, "31B SWA wrapped", 32, 16, 256, 1537, 1024, 1024) ||
|
|
run_case(metal, "12B full wrapped", 16, 1, 512, 769, 512, 512) ||
|
|
run_case(metal, "12B SWA wrapped", 16, 8, 256, 1537, 1024, 1024) ||
|
|
run_batched_wrapped_prefill_case(metal) ||
|
|
run_long_scalar_repeat_case(metal) ||
|
|
run_online_f16_case(metal, "31B", 32, 4) ||
|
|
run_online_f16_case(metal, "12B", 16, 1) ||
|
|
run_online_12b_case(metal) ||
|
|
run_q8k8_prefill_case(metal, "q8k8 31B full", 32, 4, 512, 4096, 2496, 0) ||
|
|
run_q8k8_prefill_case(metal, "q8k8 31B SWA", 32, 16, 256, 1024, 2496, 1024) ||
|
|
run_tree_case(metal, 0) ||
|
|
run_tree_case(metal, 4);
|
|
gem_metal_close(metal);
|
|
return result;
|
|
}
|