#include "gem/metal.h" #include "gem/ops.h" #include #include #include #include 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; }