288 lines
14 KiB
C
288 lines
14 KiB
C
#include "gem/cpu.h"
|
|
#include "gem/metal.h"
|
|
#include "gem/metal_layer.h"
|
|
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <string.h>
|
|
#include <sys/mman.h>
|
|
#include <unistd.h>
|
|
|
|
static size_t append_region(size_t *cursor, size_t bytes)
|
|
{
|
|
*cursor = (*cursor + 255) & ~(size_t)255;
|
|
size_t result = *cursor;
|
|
*cursor += bytes;
|
|
return result;
|
|
}
|
|
|
|
static size_t q4_size(size_t input, size_t output)
|
|
{
|
|
return input / 32 * 18 * output;
|
|
}
|
|
|
|
static void fill_q4(uint8_t *matrix, size_t input, size_t output, unsigned seed)
|
|
{
|
|
size_t row_bytes = input / 32 * 18;
|
|
for (size_t row = 0; row < output; ++row)
|
|
for (size_t block = 0; block < input / 32; ++block) {
|
|
uint8_t *destination = matrix + row * row_bytes + block * 18;
|
|
destination[0] = 0x00;
|
|
destination[1] = 0x24; /* binary16 0.015625 */
|
|
for (size_t i = 0; i < 16; ++i) {
|
|
unsigned low = (unsigned)(row * 3 + block * 5 + i * 7 + seed) & 15;
|
|
unsigned high = (unsigned)(row * 11 + block * 3 + i * 5 + seed * 2) & 15;
|
|
destination[2 + i] = (uint8_t)(low | high << 4);
|
|
}
|
|
}
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
enum { embedding = 64, query_heads = 2, kv_heads = 1, head_dimension = 32,
|
|
query_size = 64, kv_size = 32, feed_forward = 64, cache_capacity = 2 };
|
|
long page_size = sysconf(_SC_PAGESIZE);
|
|
if (page_size <= 0)
|
|
return 2;
|
|
size_t cursor = 0;
|
|
size_t q = append_region(&cursor, q4_size(embedding, query_size));
|
|
size_t k = append_region(&cursor, q4_size(embedding, kv_size));
|
|
size_t v = append_region(&cursor, q4_size(embedding, kv_size));
|
|
size_t o = append_region(&cursor, q4_size(query_size, embedding));
|
|
size_t gate = append_region(&cursor, q4_size(embedding, feed_forward));
|
|
size_t up = append_region(&cursor, q4_size(embedding, feed_forward));
|
|
size_t down = append_region(&cursor, q4_size(feed_forward, embedding));
|
|
size_t attn_norm = append_region(&cursor, embedding * sizeof(float));
|
|
size_t q_norm = append_region(&cursor, head_dimension * sizeof(float));
|
|
size_t k_norm = append_region(&cursor, head_dimension * sizeof(float));
|
|
size_t post_attn_norm = append_region(&cursor, embedding * sizeof(float));
|
|
size_t ffn_norm = append_region(&cursor, embedding * sizeof(float));
|
|
size_t post_ffn_norm = append_region(&cursor, embedding * sizeof(float));
|
|
size_t mapping_size = (cursor + (size_t)page_size - 1) / (size_t)page_size *
|
|
(size_t)page_size;
|
|
uint8_t *mapping = mmap(NULL, mapping_size, PROT_READ | PROT_WRITE,
|
|
MAP_PRIVATE | MAP_ANON, -1, 0);
|
|
if (mapping == MAP_FAILED)
|
|
return 2;
|
|
fill_q4(mapping + q, embedding, query_size, 1);
|
|
fill_q4(mapping + k, embedding, kv_size, 2);
|
|
fill_q4(mapping + v, embedding, kv_size, 3);
|
|
fill_q4(mapping + o, query_size, embedding, 4);
|
|
fill_q4(mapping + gate, embedding, feed_forward, 5);
|
|
fill_q4(mapping + up, embedding, feed_forward, 6);
|
|
fill_q4(mapping + down, feed_forward, embedding, 7);
|
|
size_t norm_offsets[] = {attn_norm, q_norm, k_norm, post_attn_norm, ffn_norm,
|
|
post_ffn_norm};
|
|
size_t norm_counts[] = {embedding, head_dimension, head_dimension, embedding, embedding,
|
|
embedding};
|
|
for (size_t n = 0; n < sizeof(norm_offsets) / sizeof(norm_offsets[0]); ++n) {
|
|
float *weight = (float *)(mapping + norm_offsets[n]);
|
|
for (size_t i = 0; i < norm_counts[n]; ++i)
|
|
weight[i] = 0.75f + (float)((i + n) % 9) * 0.03125f;
|
|
}
|
|
|
|
struct gem_cpu_layer layer = {
|
|
.embedding_size = embedding,
|
|
.query_heads = query_heads,
|
|
.kv_heads = kv_heads,
|
|
.head_dimension = head_dimension,
|
|
.rotary_dimension = head_dimension,
|
|
.feed_forward_size = feed_forward,
|
|
.rms_epsilon = 1e-6f,
|
|
.rope_frequency_base = 10000.0f,
|
|
.rope_frequency_scale = 1.0f,
|
|
.attention_scale = 1.0f,
|
|
.attention_query = mapping + q,
|
|
.attention_key = mapping + k,
|
|
.attention_value = mapping + v,
|
|
.attention_output = mapping + o,
|
|
.ffn_gate = mapping + gate,
|
|
.ffn_up = mapping + up,
|
|
.ffn_down = mapping + down,
|
|
.attention_norm = (float *)(mapping + attn_norm),
|
|
.query_norm = (float *)(mapping + q_norm),
|
|
.key_norm = (float *)(mapping + k_norm),
|
|
.post_attention_norm = (float *)(mapping + post_attn_norm),
|
|
.ffn_norm = (float *)(mapping + ffn_norm),
|
|
.post_ffn_norm = (float *)(mapping + post_ffn_norm),
|
|
.layer_output_scale = 0.8125f,
|
|
};
|
|
float input[embedding];
|
|
float expected[embedding];
|
|
float actual[embedding];
|
|
float initial_key_cache[cache_capacity * kv_size];
|
|
float initial_value_cache[cache_capacity * kv_size];
|
|
float cpu_key_cache[cache_capacity * kv_size];
|
|
float cpu_value_cache[cache_capacity * kv_size];
|
|
for (size_t i = 0; i < embedding; ++i)
|
|
input[i] = ((int)(i % 19) - 9) * 0.035f;
|
|
for (size_t i = 0; i < cache_capacity * kv_size; ++i) {
|
|
initial_key_cache[i] = ((int)(i % 13) - 6) * 0.025f;
|
|
initial_value_cache[i] = ((int)(i % 11) - 5) * 0.04f;
|
|
}
|
|
memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
|
|
memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
|
|
char error[512];
|
|
if (gem_cpu_layer_decode(&layer, input, 2, cpu_key_cache, cpu_value_cache,
|
|
cache_capacity, 2, cache_capacity, expected,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "CPU layer: %s\n", error);
|
|
return 1;
|
|
}
|
|
|
|
struct gem_metal *metal = NULL;
|
|
if (gem_metal_open(&metal, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "SKIP: %s\n", error);
|
|
return 77;
|
|
}
|
|
struct gem_metal_layer_workspace *overflow_workspace = NULL;
|
|
size_t overflowing_query_heads = SIZE_MAX / (32 * sizeof(float)) + 1;
|
|
if (gem_metal_layer_workspace_open(&overflow_workspace, metal, 1, 1, 1, 1,
|
|
overflowing_query_heads, 1,
|
|
error, sizeof(error)) == 0 || overflow_workspace) {
|
|
fprintf(stderr, "Metal score scratch byte overflow was accepted\n");
|
|
gem_metal_layer_workspace_close(overflow_workspace);
|
|
return 1;
|
|
}
|
|
struct gem_metal_buffer *model_buffer = NULL;
|
|
struct gem_metal_buffer *cache_buffer = NULL;
|
|
struct gem_metal_layer_workspace *workspace = NULL;
|
|
size_t one_cache_bytes = sizeof(initial_key_cache);
|
|
if (gem_metal_buffer_wrap(metal, mapping, mapping_size, &model_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_buffer_new(metal, one_cache_bytes * 4, &cache_buffer,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_workspace_open(&workspace, metal, embedding, query_size, kv_size,
|
|
feed_forward, query_heads, cache_capacity,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "Metal setup: %s\n", error);
|
|
return 1;
|
|
}
|
|
size_t initial_score_bytes = query_heads * cache_capacity * 32 * sizeof(float);
|
|
size_t grown_score_bytes = query_heads * cache_capacity * 64 * sizeof(float);
|
|
if (gem_metal_layer_workspace_score_bytes(workspace) != initial_score_bytes) {
|
|
fprintf(stderr, "Metal initial score scratch is %zu, expected %zu\n",
|
|
gem_metal_layer_workspace_score_bytes(workspace), initial_score_bytes);
|
|
return 1;
|
|
}
|
|
if (gem_metal_layer_workspace_reserve_score_rows(workspace, 64,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "Metal score scratch reserve: %s\n", error);
|
|
return 1;
|
|
}
|
|
if (gem_metal_layer_workspace_score_bytes(workspace) != grown_score_bytes) {
|
|
fprintf(stderr, "Metal grown score scratch is %zu, expected %zu\n",
|
|
gem_metal_layer_workspace_score_bytes(workspace), grown_score_bytes);
|
|
return 1;
|
|
}
|
|
if (gem_metal_layer_workspace_reserve_score_rows(workspace, 1,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "Metal score scratch no-shrink reserve: %s\n", error);
|
|
return 1;
|
|
}
|
|
if (gem_metal_layer_workspace_score_bytes(workspace) != grown_score_bytes) {
|
|
fprintf(stderr, "Metal score scratch shrank to %zu from %zu\n",
|
|
gem_metal_layer_workspace_score_bytes(workspace), grown_score_bytes);
|
|
return 1;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), initial_key_cache, one_cache_bytes);
|
|
memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + one_cache_bytes,
|
|
initial_value_cache, one_cache_bytes);
|
|
if (gem_metal_layer_upload(workspace, input, embedding, error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_decode(workspace, &layer, mapping, mapping_size, model_buffer,
|
|
cache_buffer, 0, one_cache_bytes, cache_capacity, 2,
|
|
cache_capacity, 2, error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_download(workspace, actual, embedding, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "Metal layer: %s\n", error);
|
|
return 1;
|
|
}
|
|
for (size_t i = 0; i < embedding; ++i)
|
|
if (fabsf(actual[i] - expected[i]) > 2e-3f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "layer %zu: Metal %.9g CPU %.9g\n", i, actual[i], expected[i]);
|
|
return 1;
|
|
}
|
|
/* Encode two complete layers into one command buffer. This checks the candidate token-span
|
|
* scheduler boundary without changing the promoted layer-span default. */
|
|
layer.attention_value = mapping + v;
|
|
float middle[embedding];
|
|
memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
|
|
memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
|
|
if (gem_cpu_layer_decode(&layer, input, 0, cpu_key_cache, cpu_value_cache,
|
|
cache_capacity, 0, cache_capacity, middle,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "CPU spanned layer one: %s\n", error);
|
|
return 1;
|
|
}
|
|
memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
|
|
memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
|
|
if (gem_cpu_layer_decode(&layer, middle, 0, cpu_key_cache, cpu_value_cache,
|
|
cache_capacity, 0, cache_capacity, expected,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "CPU spanned layer two: %s\n", error);
|
|
return 1;
|
|
}
|
|
for (size_t layer_index = 0; layer_index < 2; ++layer_index) {
|
|
size_t base = layer_index * one_cache_bytes * 2;
|
|
memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + base,
|
|
initial_key_cache, one_cache_bytes);
|
|
memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + base + one_cache_bytes,
|
|
initial_value_cache, one_cache_bytes);
|
|
}
|
|
struct gem_metal_batch *batch = NULL;
|
|
if (gem_metal_layer_upload(workspace, input, embedding, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_begin(metal, &batch, error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_encode_decode(workspace, batch, &layer, mapping, mapping_size,
|
|
model_buffer, cache_buffer, 0, one_cache_bytes,
|
|
cache_capacity, 0, cache_capacity, 0,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_encode_decode(workspace, batch, &layer, mapping, mapping_size,
|
|
model_buffer, cache_buffer, one_cache_bytes * 2,
|
|
one_cache_bytes * 3, cache_capacity, 0,
|
|
cache_capacity, 0, error, sizeof(error)) != 0 ||
|
|
gem_metal_batch_submit(batch, error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_download(workspace, actual, embedding, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "Metal spanned layers: %s\n", error);
|
|
gem_metal_batch_close(batch);
|
|
return 1;
|
|
}
|
|
gem_metal_batch_close(batch);
|
|
for (size_t i = 0; i < embedding; ++i)
|
|
if (fabsf(actual[i] - expected[i]) > 4e-3f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "spanned layer %zu: Metal %.9g CPU %.9g\n",
|
|
i, actual[i], expected[i]);
|
|
return 1;
|
|
}
|
|
layer.attention_value = NULL;
|
|
memcpy(cpu_key_cache, initial_key_cache, sizeof(cpu_key_cache));
|
|
memcpy(cpu_value_cache, initial_value_cache, sizeof(cpu_value_cache));
|
|
if (gem_cpu_layer_decode(&layer, input, 1, cpu_key_cache, cpu_value_cache,
|
|
cache_capacity, 1, 0, expected,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "CPU K-as-V layer: %s\n", error);
|
|
return 1;
|
|
}
|
|
memcpy(gem_metal_buffer_contents(cache_buffer), initial_key_cache, one_cache_bytes);
|
|
memcpy((uint8_t *)gem_metal_buffer_contents(cache_buffer) + one_cache_bytes,
|
|
initial_value_cache, one_cache_bytes);
|
|
if (gem_metal_layer_upload(workspace, input, embedding, error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_decode(workspace, &layer, mapping, mapping_size, model_buffer,
|
|
cache_buffer, 0, one_cache_bytes, cache_capacity, 1,
|
|
0, 1, error, sizeof(error)) != 0 ||
|
|
gem_metal_layer_download(workspace, actual, embedding, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "Metal K-as-V layer: %s\n", error);
|
|
return 1;
|
|
}
|
|
for (size_t i = 0; i < embedding; ++i)
|
|
if (fabsf(actual[i] - expected[i]) > 2e-3f * (1.0f + fabsf(expected[i]))) {
|
|
fprintf(stderr, "K-as-V layer %zu: Metal %.9g CPU %.9g\n",
|
|
i, actual[i], expected[i]);
|
|
return 1;
|
|
}
|
|
gem_metal_layer_workspace_close(workspace);
|
|
gem_metal_buffer_close(cache_buffer);
|
|
gem_metal_buffer_close(model_buffer);
|
|
gem_metal_close(metal);
|
|
munmap(mapping, mapping_size);
|
|
return 0;
|
|
}
|