#include "gem/metal.h" #include "gem/ops.h" #include #include #include #include #include #include #include int main(void) { /* Exercise both a complete 32-row MMA tile and the tail path. */ enum { input_size = 64, output_size = 64, batch_rows = 33 }; uint8_t matrix[output_size * (input_size / 32) * 18]; float input[input_size]; float norm_weight[input_size]; float normalized[input_size]; float normalized_rows[input_size]; float unit_normalized_rows[input_size]; float norm_projected[output_size]; float rope_expected[input_size]; float rope_positions_expected[input_size]; float rope_explicit_expected[input_size]; float gelu_temporary[input_size]; float gelu_residual_expected[input_size]; uint8_t q8_matrix[output_size * (input_size / 32) * 34]; float q8_expected[output_size]; float q8_batch_expected[8 * output_size]; float q8_batch_actual[8 * output_size]; float embedding_expected[input_size]; float expected[output_size]; float second_expected[output_size]; float actual[output_size]; float batch_input[batch_rows * input_size]; float batch_expected[batch_rows * output_size]; float batch_actual[batch_rows * output_size]; for (size_t i = 0; i < input_size; ++i) { input[i] = ((int)(i % 13) - 6) * 0.125f; norm_weight[i] = 0.75f + (float)(i % 7) * 0.0625f; } for (size_t row = 0; row < batch_rows; ++row) for (size_t i = 0; i < input_size; ++i) batch_input[row * input_size + i] = input[i] + (float)row * 0.03125f; for (size_t i = 0; i < sizeof(matrix); i += 18) { matrix[i] = 0x00; matrix[i + 1] = 0x3c; /* binary16 1.0 */ for (size_t j = 0; j < 16; ++j) matrix[i + 2 + j] = (uint8_t)(((i + j * 7) & 15) | (((i / 3 + j * 5) & 15) << 4)); } for (size_t row = 0; row < output_size; ++row) for (size_t block = 0; block < input_size / 32; ++block) { uint8_t *destination = q8_matrix + (row * (input_size / 32) + block) * 34; destination[0] = 0x00; destination[1] = 0x30; /* binary16 0.125 */ for (size_t i = 0; i < 32; ++i) destination[2 + i] = (uint8_t)(int8_t)((int)((row * 3 + block * 5 + i) % 31) - 15); } if (gem_matvec_q4_0(matrix, input_size, output_size, input, expected) != 0) return 2; for (size_t row = 0; row < batch_rows; ++row) if (gem_matvec_q4_0(matrix, input_size, output_size, batch_input + row * input_size, batch_expected + row * output_size) != 0) return 2; if (gem_matvec_q4_0(matrix, input_size, output_size, expected, second_expected) != 0) return 2; gem_rms_norm_f32(input, norm_weight, 1e-6f, input_size, normalized); gem_rms_norm_f32(input, norm_weight, 1e-6f, input_size / 2, normalized_rows); gem_rms_norm_f32(input + input_size / 2, norm_weight, 1e-6f, input_size / 2, normalized_rows + input_size / 2); float unit_weight[input_size / 2]; for (size_t i = 0; i < input_size / 2; ++i) unit_weight[i] = 1.0f; gem_rms_norm_f32(input, unit_weight, 1e-6f, input_size / 2, unit_normalized_rows); gem_rms_norm_f32(input + input_size / 2, unit_weight, 1e-6f, input_size / 2, unit_normalized_rows + input_size / 2); if (gem_matvec_q4_0(matrix, input_size, output_size, normalized, norm_projected) != 0) return 2; if (gem_rope_neox_f32(input, input_size / 2, input_size / 2, 7, 10000.0f, 1.0f, norm_weight, rope_expected) != 0 || gem_rope_neox_f32(input + input_size / 2, input_size / 2, input_size / 2, 7, 10000.0f, 1.0f, norm_weight, rope_expected + input_size / 2) != 0) return 2; for (size_t token = 0; token < 2; ++token) for (size_t head = 0; head < 2; ++head) { size_t offset = (token * 2 + head) * (input_size / 4); if (gem_rope_neox_f32(input + offset, input_size / 4, input_size / 4, 7 + token, 10000.0f, 1.0f, norm_weight, rope_positions_expected + offset) != 0) return 2; } const uint8_t rope_position_offsets[2] = {1, 0}; for (size_t token = 0; token < 2; ++token) for (size_t head = 0; head < 2; ++head) { size_t offset = (token * 2 + head) * (input_size / 4); if (gem_rope_neox_f32(input + offset, input_size / 4, input_size / 4, 7 + rope_position_offsets[token], 10000.0f, 1.0f, norm_weight, rope_explicit_expected + offset) != 0) return 2; } gem_gelu_f32(input, input_size, gelu_temporary); for (size_t i = 0; i < input_size; ++i) gelu_residual_expected[i] = (gelu_temporary[i] * norm_weight[i] + input[i]) * 0.375f; if (gem_matvec_q8_0(q8_matrix, input_size, output_size, input, q8_expected) != 0 || gem_dequantize_q8_0(q8_matrix + 5 * (input_size / 32 * 34), input_size, embedding_expected) != 0) return 2; gem_softcap_f32(q8_expected, output_size, 3.0f); for (size_t row = 0; row < 8; ++row) { if (gem_matvec_q8_0(q8_matrix, input_size, output_size, batch_input + row * input_size, q8_batch_expected + row * output_size) != 0) return 2; gem_softcap_f32(q8_batch_expected + row * output_size, output_size, 3.0f); } for (size_t i = 0; i < input_size; ++i) embedding_expected[i] *= 8.0f; struct gem_metal *metal; char error[512]; if (gem_metal_open(&metal, error, sizeof(error)) != 0) { fprintf(stderr, "SKIP: %s\n", error); return 77; } struct gem_metal_init_profile init_profile; if (gem_metal_get_init_profile(metal, &init_profile, error, sizeof(error)) != 0 || init_profile.source_bytes == 0 || init_profile.device_seconds < 0.0 || init_profile.source_assembly_seconds < 0.0 || init_profile.library_compile_seconds <= 0.0 || init_profile.pipeline_create_seconds <= 0.0 || init_profile.queue_create_seconds < 0.0 || init_profile.total_seconds <= 0.0 || init_profile.total_seconds < init_profile.library_compile_seconds || init_profile.total_seconds < init_profile.pipeline_create_seconds) { fprintf(stderr, "invalid Metal initialization profile: %s\n", error); gem_metal_close(metal); return 1; } struct gem_metal_prefill_row2_resources row2_resources = {0}; struct gem_metal_prefill_bq8_resources bq8_resources = {0}; struct gem_metal_prefill_bq8_resources bq8_bk128_resources = {0}; if (gem_metal_prefill_row2_resources(metal, &row2_resources, error, sizeof(error)) != 0 || !row2_resources.thread_execution_width || row2_resources.max_total_threads_per_threadgroup < 256 || row2_resources.static_threadgroup_memory_bytes != 0 || row2_resources.dynamic_threadgroup_memory_bytes != 16384 || row2_resources.threads_per_threadgroup != 256) { fprintf(stderr, "invalid row2 pipeline resources: %s\n", error); return 1; } if (gem_metal_prefill_bq8_resources(metal, &bq8_resources, error, sizeof(error)) != 0 || !bq8_resources.thread_execution_width || bq8_resources.max_total_threads_per_threadgroup < 256 || bq8_resources.static_threadgroup_memory_bytes != 0 || bq8_resources.dynamic_threadgroup_memory_bytes != 18496 || bq8_resources.threads_per_threadgroup != 256) { fprintf(stderr, "invalid BQ8 pipeline resources: %s\n", error); return 1; } if (gem_metal_prefill_bq8_bk128_resources(metal, &bq8_bk128_resources, error, sizeof(error)) != 0 || !bq8_bk128_resources.thread_execution_width || bq8_bk128_resources.max_total_threads_per_threadgroup < 256 || bq8_bk128_resources.static_threadgroup_memory_bytes != 0 || bq8_bk128_resources.dynamic_threadgroup_memory_bytes != 20544 || bq8_bk128_resources.threads_per_threadgroup != 256) { fprintf(stderr, "invalid BQ8/BK128 pipeline resources: %s\n", error); return 1; } int result = gem_metal_matvec_q4_0(metal, matrix, input_size, output_size, input, actual, error, sizeof(error)); if (result != 0) { fprintf(stderr, "%s\n", error); return 1; } for (size_t i = 0; i < output_size; ++i) if (fabsf(actual[i] - expected[i]) > 1e-4f * (1.0f + fabsf(expected[i]))) { fprintf(stderr, "row %zu: Metal %.9g CPU %.9g\n", i, actual[i], expected[i]); return 1; } long page_size = sysconf(_SC_PAGESIZE); if (page_size <= 0) { fprintf(stderr, "could not determine page size\n"); return 1; } size_t matrix_bytes = sizeof(matrix); size_t mapped_bytes = (matrix_bytes * 2 + (size_t)page_size - 1) / (size_t)page_size * (size_t)page_size; void *mapping = mmap(NULL, mapped_bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0); struct gem_metal_buffer *matrix_buffer = NULL; struct gem_metal_buffer *input_buffer = NULL; struct gem_metal_buffer *output_buffer = NULL; struct gem_metal_buffer *weight_buffer = NULL; struct gem_metal_buffer *batch_input_buffer = NULL; struct gem_metal_buffer *batch_output_buffer = NULL; if (mapping == MAP_FAILED) { fprintf(stderr, "could not allocate no-copy test mapping\n"); return 1; } memcpy(mapping, matrix, matrix_bytes); memcpy((uint8_t *)mapping + matrix_bytes, matrix, matrix_bytes); if (gem_metal_buffer_wrap(metal, mapping, mapped_bytes, &matrix_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(input), &input_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(actual), &output_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(batch_input), &batch_input_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(batch_actual), &batch_output_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(norm_weight), &weight_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "%s\n", error); return 1; } memcpy(gem_metal_buffer_contents(weight_buffer), norm_weight, sizeof(norm_weight)); memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); memcpy(gem_metal_buffer_contents(batch_input_buffer), batch_input, sizeof(batch_input)); struct gem_metal_batch *matmul_batch = NULL; if (gem_metal_batch_begin(metal, &matmul_batch, error, sizeof(error)) != 0 || gem_metal_batch_matmul_q4_0(matmul_batch, matrix_buffer, 0, input_size, output_size, batch_input_buffer, 0, batch_rows, batch_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(matmul_batch, error, sizeof(error)) != 0) { fprintf(stderr, "Q4 matmul dispatch: %s\n", error); return 1; } gem_metal_batch_close(matmul_batch); memcpy(batch_actual, gem_metal_buffer_contents(batch_output_buffer), sizeof(batch_actual)); for (size_t i = 0; i < batch_rows * output_size; ++i) if (fabsf(batch_actual[i] - batch_expected[i]) > 2e-4f * (1.0f + fabsf(batch_expected[i]))) { fprintf(stderr, "Q4 matmul %zu: Metal %.9g CPU %.9g\n", i, batch_actual[i], batch_expected[i]); return 1; } struct gem_metal_batch *verify8_batch = NULL; if (gem_metal_batch_begin(metal, &verify8_batch, error, sizeof(error)) != 0 || gem_metal_batch_matmul_q4_0_64x8(verify8_batch, matrix_buffer, 0, input_size, output_size, batch_input_buffer, 0, batch_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(verify8_batch, error, sizeof(error)) != 0) { fprintf(stderr, "Q4 64x8 matmul dispatch: %s\n", error); return 1; } gem_metal_batch_close(verify8_batch); memcpy(batch_actual, gem_metal_buffer_contents(batch_output_buffer), 8 * output_size * sizeof(*batch_actual)); for (size_t i = 0; i < 8 * output_size; ++i) if (fabsf(batch_actual[i] - batch_expected[i]) > 2e-4f * (1.0f + fabsf(batch_expected[i]))) { fprintf(stderr, "Q4 64x8 matmul %zu: Metal %.9g CPU %.9g\n", i, batch_actual[i], batch_expected[i]); return 1; } verify8_batch = NULL; if (gem_metal_batch_begin(metal, &verify8_batch, error, sizeof(error)) != 0 || gem_metal_batch_matmul_q4_0_64x8_4sg(verify8_batch, matrix_buffer, 0, input_size, output_size, batch_input_buffer, 0, batch_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(verify8_batch, error, sizeof(error)) != 0) { fprintf(stderr, "Q4 64x8 4SG matmul dispatch: %s\n", error); return 1; } gem_metal_batch_close(verify8_batch); memcpy(batch_actual, gem_metal_buffer_contents(batch_output_buffer), 8 * output_size * sizeof(*batch_actual)); for (size_t i = 0; i < 8 * output_size; ++i) if (fabsf(batch_actual[i] - batch_expected[i]) > 2e-4f * (1.0f + fabsf(batch_expected[i]))) { fprintf(stderr, "Q4 64x8 4SG matmul %zu: Metal %.9g CPU %.9g\n", i, batch_actual[i], batch_expected[i]); return 1; } if (gem_metal_matvec_q4_0_buffer(metal, matrix_buffer, 0, input_size, output_size, input_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0) { fprintf(stderr, "%s\n", error); return 1; } memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < output_size; ++i) if (fabsf(actual[i] - expected[i]) > 1e-4f * (1.0f + fabsf(expected[i]))) { fprintf(stderr, "persistent row %zu: Metal %.9g CPU %.9g\n", i, actual[i], expected[i]); return 1; } struct gem_metal_batch *profile_batch = NULL; struct gem_metal_batch_profile profile_result; double stage_seconds[2]; if (gem_metal_batch_begin_stage_profiled(metal, 2, &profile_batch, error, sizeof(error)) != 0 || gem_metal_batch_matvec_q4_0(profile_batch, matrix_buffer, 0, input_size, output_size, input_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_next_profile_stage(profile_batch, error, sizeof(error)) != 0 || gem_metal_batch_matvec_q4_0(profile_batch, matrix_buffer, 0, input_size, output_size, output_buffer, 0, input_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit_stage_profiled(profile_batch, &profile_result, stage_seconds, 2, error, sizeof(error)) != 0) { fprintf(stderr, "stage-profiled Q4 matvec: %s\n", error); return 1; } gem_metal_batch_close(profile_batch); memcpy(actual, gem_metal_buffer_contents(input_buffer), sizeof(actual)); if (!(profile_result.gpu_seconds > 0.0) || !(stage_seconds[0] > 0.0) || !(stage_seconds[1] > 0.0)) { fprintf(stderr, "stage-profiled Q4 timestamps are not positive\n"); return 1; } for (size_t i = 0; i < output_size; ++i) if (fabsf(actual[i] - second_expected[i]) > 1e-4f * (1.0f + fabsf(second_expected[i]))) { fprintf(stderr, "stage-profiled row %zu: Metal %.9g CPU %.9g\n", i, actual[i], second_expected[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); if (gem_metal_rms_norm_f32_buffer(metal, input_buffer, 0, weight_buffer, 0, input_size, 1e-6f, output_buffer, 0, error, sizeof(error)) != 0) { fprintf(stderr, "RMSNorm dispatch: %s\n", error); return 1; } memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - normalized[i]) > 2e-5f * (1.0f + fabsf(normalized[i]))) { fprintf(stderr, "RMSNorm %zu: Metal %.9g CPU %.9g\n", i, actual[i], normalized[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); struct gem_metal_batch *rope_batch = NULL; if (gem_metal_batch_begin(metal, &rope_batch, error, sizeof(error)) != 0 || gem_metal_batch_rope_neox_f32(rope_batch, input_buffer, 0, input_size / 2, 2, 7, 10000.0f, 1.0f, weight_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(rope_batch, error, sizeof(error)) != 0) { fprintf(stderr, "RoPE dispatch: %s\n", error); return 1; } gem_metal_batch_close(rope_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - rope_expected[i]) > 3e-5f * (1.0f + fabsf(rope_expected[i]))) { fprintf(stderr, "RoPE %zu: Metal %.9g CPU %.9g\n", i, actual[i], rope_expected[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); rope_batch = NULL; if (gem_metal_batch_begin(metal, &rope_batch, error, sizeof(error)) != 0 || gem_metal_batch_rope_neox_positions_f32( rope_batch, input_buffer, 0, input_size / 4, 2, 2, 7, 10000.0f, 1.0f, weight_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(rope_batch, error, sizeof(error)) != 0) { fprintf(stderr, "position-batched RoPE dispatch: %s\n", error); return 1; } gem_metal_batch_close(rope_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - rope_positions_expected[i]) > 3e-5f * (1.0f + fabsf(rope_positions_expected[i]))) { fprintf(stderr, "position-batched RoPE %zu: Metal %.9g CPU %.9g\n", i, actual[i], rope_positions_expected[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); rope_batch = NULL; if (gem_metal_batch_begin(metal, &rope_batch, error, sizeof(error)) != 0 || gem_metal_batch_rope_neox_explicit_f32( rope_batch, input_buffer, 0, input_size / 4, 2, 2, 7, rope_position_offsets, 10000.0f, 1.0f, weight_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(rope_batch, error, sizeof(error)) != 0) { fprintf(stderr, "explicit-position RoPE dispatch: %s\n", error); return 1; } gem_metal_batch_close(rope_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - rope_explicit_expected[i]) > 3e-5f * (1.0f + fabsf(rope_explicit_expected[i]))) { fprintf(stderr, "explicit-position RoPE %zu: Metal %.9g CPU %.9g\n", i, actual[i], rope_explicit_expected[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); struct gem_metal_batch *rows_batch = NULL; if (gem_metal_batch_begin(metal, &rows_batch, error, sizeof(error)) != 0 || gem_metal_batch_rms_norm_rows_f32(rows_batch, input_buffer, 0, weight_buffer, 0, input_size / 2, 2, 1e-6f, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(rows_batch, error, sizeof(error)) != 0) { fprintf(stderr, "row RMSNorm dispatch: %s\n", error); return 1; } gem_metal_batch_close(rows_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - normalized_rows[i]) > 2e-5f * (1.0f + fabsf(normalized_rows[i]))) { fprintf(stderr, "row RMSNorm %zu: Metal %.9g CPU %.9g\n", i, actual[i], normalized_rows[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); struct gem_metal_batch *unit_batch = NULL; if (gem_metal_batch_begin(metal, &unit_batch, error, sizeof(error)) != 0 || gem_metal_batch_rms_norm_rows_f32(unit_batch, input_buffer, 0, NULL, 0, input_size / 2, 2, 1e-6f, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(unit_batch, error, sizeof(error)) != 0) { fprintf(stderr, "unit RMSNorm dispatch: %s\n", error); return 1; } gem_metal_batch_close(unit_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - unit_normalized_rows[i]) > 2e-5f * (1.0f + fabsf(unit_normalized_rows[i]))) { fprintf(stderr, "unit RMSNorm %zu: Metal %.9g CPU %.9g\n", i, actual[i], unit_normalized_rows[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); struct gem_metal_batch *norm_batch = NULL; if (gem_metal_batch_begin(metal, &norm_batch, error, sizeof(error)) != 0 || gem_metal_batch_rms_norm_f32(norm_batch, input_buffer, 0, weight_buffer, 0, input_size, 1e-6f, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_matvec_q4_0(norm_batch, matrix_buffer, 0, input_size, output_size, output_buffer, 0, input_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(norm_batch, error, sizeof(error)) != 0) { fprintf(stderr, "RMSNorm projection batch: %s\n", error); return 1; } gem_metal_batch_close(norm_batch); memcpy(actual, gem_metal_buffer_contents(input_buffer), sizeof(actual)); for (size_t i = 0; i < output_size; ++i) if (fabsf(actual[i] - norm_projected[i]) > 2e-4f * (1.0f + fabsf(norm_projected[i]))) { fprintf(stderr, "RMSNorm projection %zu: Metal %.9g CPU %.9g\n", i, actual[i], norm_projected[i]); return 1; } memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); struct gem_metal_batch *batch = NULL; if (gem_metal_batch_begin(metal, &batch, error, sizeof(error)) != 0 || gem_metal_batch_matvec_q4_0(batch, matrix_buffer, 0, input_size, output_size, input_buffer, 0, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_matvec_q4_0(batch, matrix_buffer, matrix_bytes, input_size, output_size, output_buffer, 0, input_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(batch, error, sizeof(error)) != 0) { fprintf(stderr, "batched dispatch: %s\n", error); return 1; } gem_metal_batch_close(batch); memcpy(actual, gem_metal_buffer_contents(input_buffer), sizeof(actual)); for (size_t i = 0; i < output_size; ++i) if (fabsf(actual[i] - second_expected[i]) > 2e-4f * (1.0f + fabsf(second_expected[i]))) { fprintf(stderr, "batched row %zu: Metal %.9g CPU %.9g\n", i, actual[i], second_expected[i]); return 1; } enum { attention_q_heads = 4, attention_kv_heads = 2, attention_dim = 64, attention_capacity = 2, attention_q_elements = attention_q_heads * attention_dim, attention_kv_elements = attention_kv_heads * attention_dim, attention_cache_elements = 2 * attention_capacity * attention_kv_elements }; float attention_query[attention_q_elements]; float current_key[attention_kv_elements]; float current_value[attention_kv_elements]; float attention_cache[attention_cache_elements]; float attention_expected[attention_q_elements]; float attention_actual[attention_q_elements]; float attention_scores[attention_q_heads * attention_capacity]; memset(attention_cache, 0, sizeof(attention_cache)); for (size_t i = 0; i < attention_q_elements; ++i) attention_query[i] = ((int)(i % 17) - 8) * 0.03125f; for (size_t i = 0; i < attention_kv_elements; ++i) { attention_cache[attention_kv_elements + i] = ((int)(i % 11) - 5) * 0.0625f; attention_cache[attention_capacity * attention_kv_elements + attention_kv_elements + i] = ((int)(i % 9) - 4) * 0.125f; current_key[i] = ((int)(i % 13) - 6) * 0.046875f; current_value[i] = ((int)(i % 7) - 3) * 0.15625f; } memcpy(attention_cache, current_key, sizeof(current_key)); memcpy(attention_cache + attention_capacity * attention_kv_elements, current_value, sizeof(current_value)); if (gem_attention_decode_ring_f32(attention_query, attention_cache, attention_cache + attention_capacity * attention_kv_elements, attention_q_heads, attention_kv_heads, attention_dim, 3, attention_capacity, attention_capacity, 1.0f, attention_scores, attention_expected) != 0) return 2; memset(attention_cache, 0, attention_kv_elements * sizeof(float)); memset(attention_cache + attention_capacity * attention_kv_elements, 0, attention_kv_elements * sizeof(float)); struct gem_metal_buffer *attention_query_buffer = NULL; struct gem_metal_buffer *current_key_buffer = NULL; struct gem_metal_buffer *current_value_buffer = NULL; struct gem_metal_buffer *attention_cache_buffer = NULL; struct gem_metal_buffer *attention_score_buffer = NULL; struct gem_metal_buffer *attention_output_buffer = NULL; if (gem_metal_buffer_new(metal, sizeof(attention_query), &attention_query_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(current_key), ¤t_key_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(current_value), ¤t_value_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(attention_cache), &attention_cache_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(attention_scores), &attention_score_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(attention_actual), &attention_output_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "attention allocation: %s\n", error); return 1; } memcpy(gem_metal_buffer_contents(attention_query_buffer), attention_query, sizeof(attention_query)); memcpy(gem_metal_buffer_contents(current_key_buffer), current_key, sizeof(current_key)); memcpy(gem_metal_buffer_contents(current_value_buffer), current_value, sizeof(current_value)); memcpy(gem_metal_buffer_contents(attention_cache_buffer), attention_cache, sizeof(attention_cache)); struct gem_metal_batch *attention_batch = NULL; if (gem_metal_batch_begin(metal, &attention_batch, error, sizeof(error)) != 0 || gem_metal_batch_kv_store_f32(attention_batch, current_key_buffer, 0, current_value_buffer, 0, attention_kv_elements, attention_cache_buffer, 0, attention_capacity * attention_kv_elements * sizeof(float), attention_capacity, 2, error, sizeof(error)) != 0 || gem_metal_batch_attention_decode_f32( attention_batch, attention_query_buffer, 0, attention_cache_buffer, 0, attention_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, 3, attention_capacity, attention_capacity, 1.0f, attention_score_buffer, 0, attention_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(attention_batch, error, sizeof(error)) != 0) { fprintf(stderr, "attention dispatch: %s\n", error); return 1; } gem_metal_batch_close(attention_batch); memcpy(attention_actual, gem_metal_buffer_contents(attention_output_buffer), sizeof(attention_actual)); for (size_t i = 0; i < attention_q_elements; ++i) if (fabsf(attention_actual[i] - attention_expected[i]) > 3e-5f * (1.0f + fabsf(attention_expected[i]))) { fprintf(stderr, "attention %zu: Metal %.9g CPU %.9g\n", i, attention_actual[i], attention_expected[i]); return 1; } struct gem_metal_batch *split_attention_batch = NULL; if (gem_metal_batch_begin(metal, &split_attention_batch, error, sizeof(error)) != 0 || gem_metal_batch_attention_split_gqa_f32( split_attention_batch, attention_query_buffer, 0, attention_cache_buffer, 0, attention_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, 3, attention_capacity, attention_capacity, 1.0f, attention_score_buffer, 0, attention_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(split_attention_batch, error, sizeof(error)) != 0) { fprintf(stderr, "split GQA attention dispatch: %s\n", error); return 1; } gem_metal_batch_close(split_attention_batch); memcpy(attention_actual, gem_metal_buffer_contents(attention_output_buffer), sizeof(attention_actual)); for (size_t i = 0; i < attention_q_elements; ++i) if (fabsf(attention_actual[i] - attention_expected[i]) > 3e-5f * (1.0f + fabsf(attention_expected[i]))) { fprintf(stderr, "split GQA attention %zu: Metal %.9g CPU %.9g\n", i, attention_actual[i], attention_expected[i]); return 1; } struct gem_metal_buffer *attention_cache_f16_buffer = NULL; if (gem_metal_buffer_new(metal, attention_cache_elements * sizeof(uint16_t), &attention_cache_f16_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "F16 attention allocation: %s\n", error); return 1; } float prior_key[attention_kv_elements]; float prior_value[attention_kv_elements]; for (size_t i = 0; i < attention_kv_elements; ++i) { prior_key[i] = ((int)(i % 11) - 5) * 0.0625f; prior_value[i] = ((int)(i % 9) - 4) * 0.125f; } memcpy(gem_metal_buffer_contents(current_key_buffer), prior_key, sizeof(prior_key)); memcpy(gem_metal_buffer_contents(current_value_buffer), prior_value, sizeof(prior_value)); struct gem_metal_batch *f16_attention_batch = NULL; size_t f16_one_cache_bytes = attention_capacity * attention_kv_elements * sizeof(uint16_t); if (gem_metal_batch_begin(metal, &f16_attention_batch, error, sizeof(error)) != 0 || gem_metal_batch_kv_store_f16(f16_attention_batch, current_key_buffer, 0, current_value_buffer, 0, attention_kv_elements, attention_cache_f16_buffer, 0, f16_one_cache_bytes, attention_capacity, 1, error, sizeof(error)) != 0 || gem_metal_batch_submit(f16_attention_batch, error, sizeof(error)) != 0) { fprintf(stderr, "F16 prior KV dispatch: %s\n", error); return 1; } gem_metal_batch_close(f16_attention_batch); f16_attention_batch = NULL; memcpy(gem_metal_buffer_contents(current_key_buffer), current_key, sizeof(current_key)); memcpy(gem_metal_buffer_contents(current_value_buffer), current_value, sizeof(current_value)); if (gem_metal_batch_begin(metal, &f16_attention_batch, error, sizeof(error)) != 0 || gem_metal_batch_kv_store_f16(f16_attention_batch, current_key_buffer, 0, current_value_buffer, 0, attention_kv_elements, attention_cache_f16_buffer, 0, f16_one_cache_bytes, attention_capacity, 2, error, sizeof(error)) != 0 || gem_metal_batch_attention_parallel_f16( f16_attention_batch, attention_query_buffer, 0, attention_cache_f16_buffer, 0, f16_one_cache_bytes, attention_q_heads, attention_kv_heads, attention_dim, 3, attention_capacity, attention_capacity, 1.0f, attention_score_buffer, 0, attention_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(f16_attention_batch, error, sizeof(error)) != 0) { fprintf(stderr, "F16 attention dispatch: %s\n", error); return 1; } gem_metal_batch_close(f16_attention_batch); memcpy(attention_actual, gem_metal_buffer_contents(attention_output_buffer), sizeof(attention_actual)); for (size_t i = 0; i < attention_q_elements; ++i) if (fabsf(attention_actual[i] - attention_expected[i]) > 6e-4f * (1.0f + fabsf(attention_expected[i]))) { fprintf(stderr, "F16 attention %zu: Metal %.9g CPU %.9g\n", i, attention_actual[i], attention_expected[i]); return 1; } gem_metal_buffer_close(attention_cache_f16_buffer); struct gem_metal_batch *parallel_attention_batch = NULL; struct gem_metal_batch_profile parallel_attention_profile; if (gem_metal_batch_begin(metal, ¶llel_attention_batch, error, sizeof(error)) != 0 || gem_metal_batch_attention_parallel_f32( parallel_attention_batch, attention_query_buffer, 0, attention_cache_buffer, 0, attention_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, 3, attention_capacity, attention_capacity, 1.0f, attention_score_buffer, 0, attention_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit_profiled(parallel_attention_batch, ¶llel_attention_profile, error, sizeof(error)) != 0) { fprintf(stderr, "parallel attention dispatch: %s\n", error); return 1; } gem_metal_batch_close(parallel_attention_batch); if (!(parallel_attention_profile.gpu_seconds > 0.0) || parallel_attention_profile.dispatches != 1) { fprintf(stderr, "parallel attention GPU profile is invalid\n"); return 1; } memcpy(attention_actual, gem_metal_buffer_contents(attention_output_buffer), sizeof(attention_actual)); for (size_t i = 0; i < attention_q_elements; ++i) if (fabsf(attention_actual[i] - attention_expected[i]) > 3e-5f * (1.0f + fabsf(attention_expected[i]))) { fprintf(stderr, "parallel attention %zu: Metal %.9g CPU %.9g\n", i, attention_actual[i], attention_expected[i]); return 1; } enum { prefill_rows = 8, prefill_prior = 300, prefill_capacity = 308 }; float prefill_query[prefill_rows * attention_q_elements]; float prefill_cache[2 * prefill_capacity * attention_kv_elements]; float prefill_expected[prefill_rows * attention_q_elements]; float prefill_actual[prefill_rows * attention_q_elements]; float prefill_oracle[prefill_rows * attention_q_elements]; float prefill_scores[prefill_rows * attention_q_heads * prefill_capacity]; for (size_t i = 0; i < prefill_rows * attention_q_elements; ++i) prefill_query[i] = ((int)(i % 29) - 14) * 0.015625f; for (size_t i = 0; i < prefill_capacity * attention_kv_elements; ++i) { prefill_cache[i] = ((int)(i % 23) - 11) * 0.0234375f; prefill_cache[prefill_capacity * attention_kv_elements + i] = ((int)(i % 19) - 9) * 0.03125f; } for (size_t row = 0; row < prefill_rows; ++row) if (gem_attention_decode_ring_f32( prefill_query + row * attention_q_elements, prefill_cache, prefill_cache + prefill_capacity * attention_kv_elements, attention_q_heads, attention_kv_heads, attention_dim, prefill_prior + row + 1, prefill_capacity, prefill_capacity, 1.0f, prefill_scores, prefill_expected + row * attention_q_elements) != 0) return 2; struct gem_metal_buffer *prefill_query_buffer = NULL; struct gem_metal_buffer *prefill_cache_buffer = NULL; struct gem_metal_buffer *prefill_output_buffer = NULL; struct gem_metal_buffer *prefill_score_buffer = NULL; if (gem_metal_buffer_new(metal, sizeof(prefill_query), &prefill_query_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(prefill_cache), &prefill_cache_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(prefill_scores), &prefill_score_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(prefill_actual), &prefill_output_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "causal prefill attention allocation: %s\n", error); return 1; } memcpy(gem_metal_buffer_contents(prefill_query_buffer), prefill_query, sizeof(prefill_query)); memcpy(gem_metal_buffer_contents(prefill_cache_buffer), prefill_cache, sizeof(prefill_cache)); struct gem_metal_batch *prefill_attention_batch = NULL; struct gem_metal_batch_profile prefill_attention_profile; if (gem_metal_batch_begin(metal, &prefill_attention_batch, error, sizeof(error)) != 0 || gem_metal_batch_attention_prefill_causal_f32( prefill_attention_batch, prefill_query_buffer, 0, prefill_cache_buffer, 0, prefill_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, prefill_rows, prefill_prior, prefill_capacity, prefill_capacity, 1.0f, prefill_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit_profiled(prefill_attention_batch, &prefill_attention_profile, error, sizeof(error)) != 0) { fprintf(stderr, "causal prefill attention dispatch: %s\n", error); return 1; } gem_metal_batch_close(prefill_attention_batch); if (!(prefill_attention_profile.gpu_seconds > 0.0) || prefill_attention_profile.dispatches != 1) { fprintf(stderr, "causal prefill attention GPU profile is invalid\n"); return 1; } memcpy(prefill_actual, gem_metal_buffer_contents(prefill_output_buffer), sizeof(prefill_actual)); for (size_t i = 0; i < prefill_rows * attention_q_elements; ++i) if (fabsf(prefill_actual[i] - prefill_expected[i]) > 3e-5f * (1.0f + fabsf(prefill_expected[i]))) { fprintf(stderr, "causal prefill attention %zu: Metal %.9g CPU %.9g\n", i, prefill_actual[i], prefill_expected[i]); return 1; } struct gem_metal_batch *prefill_oracle_batch = NULL; if (gem_metal_batch_begin(metal, &prefill_oracle_batch, error, sizeof(error)) != 0) { fprintf(stderr, "causal prefill oracle begin: %s\n", error); return 1; } for (size_t row = 0; row < prefill_rows; ++row) if (gem_metal_batch_attention_split_gqa_scores_f32( prefill_oracle_batch, prefill_query_buffer, row * attention_q_elements * sizeof(float), prefill_cache_buffer, 0, prefill_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, prefill_prior + row + 1, prefill_capacity, prefill_capacity, 1.0f, prefill_score_buffer, 0, prefill_output_buffer, row * attention_q_elements * sizeof(float), error, sizeof(error)) != 0) { fprintf(stderr, "causal prefill oracle encode: %s\n", error); return 1; } if (gem_metal_batch_submit(prefill_oracle_batch, error, sizeof(error)) != 0) { fprintf(stderr, "causal prefill oracle submit: %s\n", error); return 1; } gem_metal_batch_close(prefill_oracle_batch); memcpy(prefill_oracle, gem_metal_buffer_contents(prefill_output_buffer), sizeof(prefill_oracle)); for (size_t i = 0; i < prefill_rows * attention_q_elements; ++i) if (fabsf(prefill_actual[i] - prefill_oracle[i]) > 3e-5f * (1.0f + fabsf(prefill_oracle[i]))) { fprintf(stderr, "causal prefill oracle %zu: fused %.9g split %.9g\n", i, prefill_actual[i], prefill_oracle[i]); return 1; } struct gem_metal_batch *prefill_split_batch = NULL; struct gem_metal_batch_profile prefill_split_profile; if (gem_metal_batch_begin(metal, &prefill_split_batch, error, sizeof(error)) != 0 || gem_metal_batch_attention_prefill_split_gqa_scores_f32( prefill_split_batch, prefill_query_buffer, 0, prefill_cache_buffer, prefill_prior * attention_kv_elements * sizeof(float), prefill_cache_buffer, (prefill_capacity + prefill_prior) * attention_kv_elements * sizeof(float), prefill_cache_buffer, 0, prefill_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, prefill_rows, prefill_prior, prefill_capacity, prefill_capacity, 1.0f, prefill_score_buffer, 0, prefill_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit_profiled(prefill_split_batch, &prefill_split_profile, error, sizeof(error)) != 0) { fprintf(stderr, "batched split GQA prefill attention dispatch: %s\n", error); return 1; } gem_metal_batch_close(prefill_split_batch); if (!(prefill_split_profile.gpu_seconds > 0.0) || prefill_split_profile.dispatches != 3) { fprintf(stderr, "batched split GQA prefill GPU profile is invalid\n"); return 1; } memcpy(prefill_actual, gem_metal_buffer_contents(prefill_output_buffer), sizeof(prefill_actual)); if (memcmp(prefill_actual, prefill_oracle, sizeof(prefill_actual)) != 0) { for (size_t i = 0; i < prefill_rows * attention_q_elements; ++i) if (prefill_actual[i] != prefill_oracle[i]) { fprintf(stderr, "batched split GQA prefill %zu: batched %.9g rowwise %.9g\n", i, prefill_actual[i], prefill_oracle[i]); break; } return 1; } struct gem_metal_buffer *prefill_cache_f16_buffer = NULL; size_t prefill_f16_one_cache_bytes = prefill_capacity * attention_kv_elements * sizeof(uint16_t); if (gem_metal_buffer_new(metal, 2 * prefill_f16_one_cache_bytes, &prefill_cache_f16_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "F16 batched prefill allocation: %s\n", error); return 1; } struct gem_metal_batch *prefill_f16_store_batch = NULL; if (gem_metal_batch_begin(metal, &prefill_f16_store_batch, error, sizeof(error)) != 0) { fprintf(stderr, "F16 batched prefill store begin: %s\n", error); return 1; } for (size_t position = 0; position < prefill_capacity; ++position) if (gem_metal_batch_kv_store_f16( prefill_f16_store_batch, prefill_cache_buffer, position * attention_kv_elements * sizeof(float), prefill_cache_buffer, (prefill_capacity + position) * attention_kv_elements * sizeof(float), attention_kv_elements, prefill_cache_f16_buffer, 0, prefill_f16_one_cache_bytes, prefill_capacity, position, error, sizeof(error)) != 0) { fprintf(stderr, "F16 batched prefill store encode: %s\n", error); return 1; } if (gem_metal_batch_submit(prefill_f16_store_batch, error, sizeof(error)) != 0) { fprintf(stderr, "F16 batched prefill store submit: %s\n", error); return 1; } gem_metal_batch_close(prefill_f16_store_batch); struct gem_metal_batch *prefill_f16_batch = NULL; if (gem_metal_batch_begin(metal, &prefill_f16_batch, error, sizeof(error)) != 0 || gem_metal_batch_attention_prefill_split_gqa_scores_f16( prefill_f16_batch, prefill_query_buffer, 0, prefill_cache_buffer, prefill_prior * attention_kv_elements * sizeof(float), prefill_cache_buffer, (prefill_capacity + prefill_prior) * attention_kv_elements * sizeof(float), prefill_cache_f16_buffer, 0, prefill_f16_one_cache_bytes, attention_q_heads, attention_kv_heads, attention_dim, prefill_rows, prefill_prior, prefill_capacity, prefill_capacity, 1.0f, prefill_score_buffer, 0, prefill_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit_profiled(prefill_f16_batch, &prefill_split_profile, error, sizeof(error)) != 0) { fprintf(stderr, "F16 batched split GQA prefill dispatch: %s\n", error); return 1; } gem_metal_batch_close(prefill_f16_batch); if (prefill_split_profile.dispatches != 3) { fprintf(stderr, "F16 batched split GQA prefill profile is invalid\n"); return 1; } memcpy(prefill_actual, gem_metal_buffer_contents(prefill_output_buffer), sizeof(prefill_actual)); for (size_t i = 0; i < prefill_rows * attention_q_elements; ++i) if (fabsf(prefill_actual[i] - prefill_expected[i]) > 3e-5f * (1.0f + fabsf(prefill_expected[i]))) { fprintf(stderr, "F16 batched split GQA prefill %zu: Metal %.9g CPU %.9g\n", i, prefill_actual[i], prefill_expected[i]); return 1; } enum { extended_prefill_prior = 1024, extended_prefill_capacity = 1032 }; size_t extended_cache_elements = 2 * extended_prefill_capacity * attention_kv_elements; size_t extended_score_elements = prefill_rows * attention_q_heads * extended_prefill_capacity; struct gem_metal_buffer *extended_cache_buffer = NULL; struct gem_metal_buffer *extended_score_buffer = NULL; if (gem_metal_buffer_new(metal, extended_cache_elements * sizeof(float), &extended_cache_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, extended_score_elements * sizeof(float), &extended_score_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "extended prefill attention allocation: %s\n", error); return 1; } float *extended_cache = gem_metal_buffer_contents(extended_cache_buffer); for (size_t i = 0; i < extended_prefill_capacity * attention_kv_elements; ++i) { extended_cache[i] = ((int)(i % 31) - 15) * 0.01953125f; extended_cache[extended_prefill_capacity * attention_kv_elements + i] = ((int)(i % 27) - 13) * 0.02734375f; } prefill_oracle_batch = NULL; if (gem_metal_batch_begin(metal, &prefill_oracle_batch, error, sizeof(error)) != 0) { fprintf(stderr, "extended prefill oracle begin: %s\n", error); return 1; } for (size_t row = 0; row < prefill_rows; ++row) if (gem_metal_batch_attention_split_gqa_scores_f32( prefill_oracle_batch, prefill_query_buffer, row * attention_q_elements * sizeof(float), extended_cache_buffer, 0, extended_prefill_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, extended_prefill_prior + row + 1, extended_prefill_capacity, extended_prefill_capacity, 1.0f, extended_score_buffer, 0, prefill_output_buffer, row * attention_q_elements * sizeof(float), error, sizeof(error)) != 0) { fprintf(stderr, "extended prefill oracle encode: %s\n", error); return 1; } if (gem_metal_batch_submit(prefill_oracle_batch, error, sizeof(error)) != 0) { fprintf(stderr, "extended prefill oracle submit: %s\n", error); return 1; } gem_metal_batch_close(prefill_oracle_batch); memcpy(prefill_oracle, gem_metal_buffer_contents(prefill_output_buffer), sizeof(prefill_oracle)); prefill_split_batch = NULL; if (gem_metal_batch_begin(metal, &prefill_split_batch, error, sizeof(error)) != 0 || gem_metal_batch_attention_prefill_split_gqa_scores_f32( prefill_split_batch, prefill_query_buffer, 0, extended_cache_buffer, extended_prefill_prior * attention_kv_elements * sizeof(float), extended_cache_buffer, (extended_prefill_capacity + extended_prefill_prior) * attention_kv_elements * sizeof(float), extended_cache_buffer, 0, extended_prefill_capacity * attention_kv_elements * sizeof(float), attention_q_heads, attention_kv_heads, attention_dim, prefill_rows, extended_prefill_prior, extended_prefill_capacity, extended_prefill_capacity, 1.0f, extended_score_buffer, 0, prefill_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit_profiled(prefill_split_batch, &prefill_split_profile, error, sizeof(error)) != 0) { fprintf(stderr, "extended batched split GQA prefill dispatch: %s\n", error); return 1; } gem_metal_batch_close(prefill_split_batch); if (prefill_split_profile.dispatches != 3) { fprintf(stderr, "extended batched split GQA prefill profile is invalid\n"); return 1; } memcpy(prefill_actual, gem_metal_buffer_contents(prefill_output_buffer), sizeof(prefill_actual)); if (memcmp(prefill_actual, prefill_oracle, sizeof(prefill_actual)) != 0) { fprintf(stderr, "extended batched split GQA prefill differs from rowwise oracle\n"); return 1; } gem_metal_buffer_close(extended_score_buffer); gem_metal_buffer_close(extended_cache_buffer); gem_metal_buffer_close(prefill_cache_f16_buffer); gem_metal_buffer_close(prefill_output_buffer); gem_metal_buffer_close(prefill_score_buffer); gem_metal_buffer_close(prefill_cache_buffer); gem_metal_buffer_close(prefill_query_buffer); gem_metal_buffer_close(attention_output_buffer); gem_metal_buffer_close(attention_score_buffer); gem_metal_buffer_close(attention_cache_buffer); gem_metal_buffer_close(current_value_buffer); gem_metal_buffer_close(current_key_buffer); gem_metal_buffer_close(attention_query_buffer); enum { topk_vocabulary = 262144, topk_count = 40, topk_blocks = 256 }; float *topk_logits = malloc((size_t)topk_vocabulary * sizeof(*topk_logits)); struct gem_metal_buffer *topk_logits_buffer = NULL, *topk_scratch_a = NULL; struct gem_metal_buffer *topk_scratch_b = NULL, *topk_sample_buffer = NULL; struct gem_metal_buffer *topk_tokens_buffer = NULL, *topk_probabilities_buffer = NULL; if (!topk_logits) { fprintf(stderr, "top-k host allocation failed\n"); return 1; } for (size_t i = 0; i < topk_vocabulary; ++i) topk_logits[i] = -50.0f - (float)(i % 101) * 0.001f; for (size_t i = 0; i < topk_count; ++i) topk_logits[1000 + i] = 20.0f - (float)i * 0.125f; size_t topk_scratch_bytes = (size_t)topk_blocks * topk_count * sizeof(uint64_t); if (gem_metal_buffer_new(metal, (size_t)topk_vocabulary * sizeof(float), &topk_logits_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, topk_scratch_bytes, &topk_scratch_a, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, topk_scratch_bytes, &topk_scratch_b, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, sizeof(uint32_t), &topk_sample_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, topk_count * sizeof(uint32_t), &topk_tokens_buffer, error, sizeof(error)) != 0 || gem_metal_buffer_new(metal, topk_count * sizeof(float), &topk_probabilities_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "top-k Metal allocation: %s\n", error); return 1; } memcpy(gem_metal_buffer_contents(topk_logits_buffer), topk_logits, (size_t)topk_vocabulary * sizeof(float)); struct gem_metal_batch *topk_batch = NULL; if (gem_metal_batch_begin(metal, &topk_batch, error, sizeof(error)) != 0 || gem_metal_batch_topk40_sample_f32( topk_batch, topk_logits_buffer, 0, topk_vocabulary, 0.8f, 0.95f, 0.37f, topk_scratch_a, topk_scratch_b, topk_sample_buffer, 0, topk_tokens_buffer, 0, topk_probabilities_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(topk_batch, error, sizeof(error)) != 0) { fprintf(stderr, "top-k Metal dispatch: %s\n", error); return 1; } gem_metal_batch_close(topk_batch); const uint32_t *topk_tokens = gem_metal_buffer_contents(topk_tokens_buffer); const float *topk_probabilities = gem_metal_buffer_contents(topk_probabilities_buffer); float topk_weights[topk_count]; float topk_full_total = 0.0f; for (size_t i = 0; i < topk_count; ++i) { topk_weights[i] = expf((topk_logits[1000 + i] - topk_logits[1000]) / 0.8f); topk_full_total += topk_weights[i]; } float topk_cumulative = 0.0f; size_t topk_cutoff = topk_count; for (size_t i = 0; i < topk_count; ++i) { topk_cumulative += topk_weights[i] / topk_full_total; if (topk_cumulative >= 0.95f) { topk_cutoff = i + 1; break; } } float topk_cutoff_total = 0.0f; for (size_t i = 0; i < topk_cutoff; ++i) topk_cutoff_total += topk_weights[i]; float topk_target = 0.37f * topk_cutoff_total; float topk_running = 0.0f; uint32_t topk_expected_sample = 1000 + (uint32_t)(topk_cutoff - 1); for (size_t i = 0; i < topk_count; ++i) { float expected_probability = i < topk_cutoff ? topk_weights[i] / topk_cutoff_total : 0.0f; if (topk_tokens[i] != 1000 + i || fabsf(topk_probabilities[i] - expected_probability) > 2e-6f) { fprintf(stderr, "top-k rank %zu: token %u probability %.9g expected %zu %.9g\n", i, topk_tokens[i], topk_probabilities[i], 1000 + i, expected_probability); return 1; } if (i < topk_cutoff && topk_running < topk_target) { topk_running += topk_weights[i]; if (topk_running >= topk_target) topk_expected_sample = 1000 + (uint32_t)i; } } uint32_t topk_sample = *(const uint32_t *)gem_metal_buffer_contents(topk_sample_buffer); if (topk_sample != topk_expected_sample) { fprintf(stderr, "top-k sampled token %u expected %u\n", topk_sample, topk_expected_sample); return 1; } gem_metal_buffer_close(topk_probabilities_buffer); gem_metal_buffer_close(topk_tokens_buffer); gem_metal_buffer_close(topk_sample_buffer); gem_metal_buffer_close(topk_scratch_b); gem_metal_buffer_close(topk_scratch_a); gem_metal_buffer_close(topk_logits_buffer); free(topk_logits); struct gem_metal_buffer *q8_buffer = NULL; if (gem_metal_buffer_new(metal, sizeof(q8_matrix), &q8_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "Q8 allocation: %s\n", error); return 1; } memcpy(gem_metal_buffer_contents(q8_buffer), q8_matrix, sizeof(q8_matrix)); struct gem_metal_batch *q8_batch = NULL; if (gem_metal_batch_begin(metal, &q8_batch, error, sizeof(error)) != 0 || gem_metal_batch_embedding_q8_0(q8_batch, q8_buffer, 0, input_size, output_size, 5, 8.0f, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(q8_batch, error, sizeof(error)) != 0) { fprintf(stderr, "Q8 embedding dispatch: %s\n", error); return 1; } gem_metal_batch_close(q8_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - embedding_expected[i]) > 1e-5f) { fprintf(stderr, "Q8 embedding %zu: Metal %.9g CPU %.9g\n", i, actual[i], embedding_expected[i]); return 1; } struct gem_metal_buffer *indirect_token_buffer = NULL; if (gem_metal_buffer_new(metal, sizeof(uint32_t), &indirect_token_buffer, error, sizeof(error)) != 0) { fprintf(stderr, "indirect token allocation: %s\n", error); return 1; } *(uint32_t *)gem_metal_buffer_contents(indirect_token_buffer) = 5; q8_batch = NULL; if (gem_metal_batch_begin(metal, &q8_batch, error, sizeof(error)) != 0 || gem_metal_batch_embedding_q8_0_indirect( q8_batch, q8_buffer, 0, input_size, output_size, indirect_token_buffer, 0, 8.0f, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(q8_batch, error, sizeof(error)) != 0) { fprintf(stderr, "indirect Q8 embedding dispatch: %s\n", error); return 1; } gem_metal_batch_close(q8_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - embedding_expected[i]) > 1e-5f) { fprintf(stderr, "indirect Q8 embedding %zu: Metal %.9g CPU %.9g\n", i, actual[i], embedding_expected[i]); return 1; } gem_metal_buffer_close(indirect_token_buffer); memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); q8_batch = NULL; if (gem_metal_batch_begin(metal, &q8_batch, error, sizeof(error)) != 0 || gem_metal_batch_matvec_q8_0(q8_batch, q8_buffer, 0, input_size, output_size, input_buffer, 0, 3.0f, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(q8_batch, error, sizeof(error)) != 0) { fprintf(stderr, "Q8 matvec dispatch: %s\n", error); return 1; } gem_metal_batch_close(q8_batch); memcpy(actual, gem_metal_buffer_contents(output_buffer), sizeof(actual)); for (size_t i = 0; i < output_size; ++i) if (fabsf(actual[i] - q8_expected[i]) > 1e-4f * (1.0f + fabsf(q8_expected[i]))) { fprintf(stderr, "Q8 matvec %zu: Metal %.9g CPU %.9g\n", i, actual[i], q8_expected[i]); return 1; } q8_batch = NULL; if (gem_metal_batch_begin(metal, &q8_batch, error, sizeof(error)) != 0 || gem_metal_batch_matmul_q8_0_64x8(q8_batch, q8_buffer, 0, input_size, output_size, batch_input_buffer, 0, 3.0f, batch_output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(q8_batch, error, sizeof(error)) != 0) { fprintf(stderr, "Q8 64x8 matmul dispatch: %s\n", error); return 1; } gem_metal_batch_close(q8_batch); memcpy(q8_batch_actual, gem_metal_buffer_contents(batch_output_buffer), sizeof(q8_batch_actual)); for (size_t i = 0; i < 8 * output_size; ++i) if (fabsf(q8_batch_actual[i] - q8_batch_expected[i]) > 2e-4f * (1.0f + fabsf(q8_batch_expected[i]))) { fprintf(stderr, "Q8 64x8 matmul %zu: Metal %.9g CPU %.9g\n", i, q8_batch_actual[i], q8_batch_expected[i]); return 1; } gem_metal_buffer_close(q8_buffer); memcpy(gem_metal_buffer_contents(input_buffer), input, sizeof(input)); memcpy(gem_metal_buffer_contents(weight_buffer), norm_weight, sizeof(norm_weight)); struct gem_metal_batch *ffn_batch = NULL; if (gem_metal_batch_begin(metal, &ffn_batch, error, sizeof(error)) != 0 || gem_metal_batch_gelu_mul_f32(ffn_batch, input_buffer, 0, weight_buffer, 0, input_size, output_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_add_scaled_f32(ffn_batch, output_buffer, 0, input_buffer, 0, input_size, 0.375f, weight_buffer, 0, error, sizeof(error)) != 0 || gem_metal_batch_submit(ffn_batch, error, sizeof(error)) != 0) { fprintf(stderr, "FFN epilogue dispatch: %s\n", error); return 1; } gem_metal_batch_close(ffn_batch); memcpy(actual, gem_metal_buffer_contents(weight_buffer), sizeof(actual)); for (size_t i = 0; i < input_size; ++i) if (fabsf(actual[i] - gelu_residual_expected[i]) > 2e-5f * (1.0f + fabsf(gelu_residual_expected[i]))) { fprintf(stderr, "FFN epilogue %zu: Metal %.9g CPU %.9g\n", i, actual[i], gelu_residual_expected[i]); return 1; } if (gem_metal_matvec_q4_0_buffer(metal, matrix_buffer, mapped_bytes - 1, input_size, output_size, input_buffer, 0, output_buffer, 0, error, sizeof(error)) == 0) { fprintf(stderr, "out-of-bounds Metal range was accepted\n"); return 1; } gem_metal_buffer_close(output_buffer); gem_metal_buffer_close(batch_output_buffer); gem_metal_buffer_close(batch_input_buffer); gem_metal_buffer_close(input_buffer); gem_metal_buffer_close(weight_buffer); gem_metal_buffer_close(matrix_buffer); munmap(mapping, mapped_bytes); gem_metal_close(metal); return 0; }