81 lines
2.4 KiB
C
81 lines
2.4 KiB
C
#include "gem/cpu.h"
|
|
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
static uint8_t *zero_q4_matrix(size_t input_size, size_t output_size)
|
|
{
|
|
size_t bytes = input_size / 32 * 18 * output_size;
|
|
uint8_t *matrix = calloc(bytes, 1);
|
|
if (!matrix)
|
|
return NULL;
|
|
for (size_t offset = 0; offset < bytes; offset += 18) {
|
|
matrix[offset] = 0x00;
|
|
matrix[offset + 1] = 0x3c; /* F16 1.0 */
|
|
memset(matrix + offset + 2, 0x88, 16);
|
|
}
|
|
return matrix;
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
uint8_t *embedding_matrix = zero_q4_matrix(32, 32);
|
|
uint8_t *kv_matrix = zero_q4_matrix(32, 16);
|
|
if (!embedding_matrix || !kv_matrix)
|
|
return 2;
|
|
float norm32[32];
|
|
float norm16[16];
|
|
float input[32];
|
|
float output[32];
|
|
for (size_t i = 0; i < 32; ++i) {
|
|
norm32[i] = 1.0f;
|
|
input[i] = (float)i - 16.0f;
|
|
}
|
|
for (size_t i = 0; i < 16; ++i)
|
|
norm16[i] = 1.0f;
|
|
struct gem_cpu_layer layer = {
|
|
.embedding_size = 32,
|
|
.query_heads = 2,
|
|
.kv_heads = 1,
|
|
.head_dimension = 16,
|
|
.rotary_dimension = 16,
|
|
.feed_forward_size = 32,
|
|
.rms_epsilon = 1e-6f,
|
|
.rope_frequency_base = 10000.0f,
|
|
.rope_frequency_scale = 1.0f,
|
|
.attention_scale = 1.0f,
|
|
.attention_query = embedding_matrix,
|
|
.attention_key = kv_matrix,
|
|
.attention_value = NULL, /* exercise K-as-V */
|
|
.attention_output = embedding_matrix,
|
|
.ffn_gate = embedding_matrix,
|
|
.ffn_up = embedding_matrix,
|
|
.ffn_down = embedding_matrix,
|
|
.attention_norm = norm32,
|
|
.query_norm = norm16,
|
|
.key_norm = norm16,
|
|
.post_attention_norm = norm32,
|
|
.ffn_norm = norm32,
|
|
.post_ffn_norm = norm32,
|
|
.layer_output_scale = 0.5f,
|
|
};
|
|
float key_cache[4 * 16] = {0};
|
|
float value_cache[4 * 16] = {0};
|
|
char error[256];
|
|
if (gem_cpu_layer_decode(&layer, input, 0, key_cache, value_cache, 4, 0, 0,
|
|
output, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s\n", error);
|
|
return 1;
|
|
}
|
|
for (size_t i = 0; i < 32; ++i)
|
|
if (fabsf(output[i] - input[i] * 0.5f) > 1e-6f) {
|
|
fprintf(stderr, "layer output mismatch at %zu: %.8f\n", i, output[i]);
|
|
return 1;
|
|
}
|
|
free(kv_matrix);
|
|
free(embedding_matrix);
|
|
return 0;
|
|
}
|