1193 lines
63 KiB
C
1193 lines
63 KiB
C
#include "gem/metal.h"
|
|
#include "gem/ops.h"
|
|
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <sys/mman.h>
|
|
#include <unistd.h>
|
|
|
|
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;
|
|
}
|