100 lines
4 KiB
C
100 lines
4 KiB
C
#include "gem/runtime.h"
|
|
|
|
#include <stdio.h>
|
|
#include <string.h>
|
|
|
|
static int check_plan(enum gem_model_size size, uint32_t layers, uint64_t embedding,
|
|
uint64_t query_heads, uint64_t swa_kv, uint64_t full_kv,
|
|
uint64_t ff, uint64_t expected_bytes, uint64_t expected_mixed_bytes,
|
|
uint64_t expected_mixed_128k_bytes)
|
|
{
|
|
struct gem_model model;
|
|
memset(&model, 0, sizeof(model));
|
|
model.size = size;
|
|
model.block_count = layers;
|
|
model.context_length = 262144;
|
|
model.sliding_window = 1024;
|
|
model.embedding_length = embedding;
|
|
model.head_count = query_heads;
|
|
model.head_count_kv_swa = swa_kv;
|
|
model.head_count_kv_full = full_kv;
|
|
model.head_dimension_swa = 256;
|
|
model.head_dimension_full = 512;
|
|
model.feed_forward_length = ff;
|
|
|
|
struct gem_runtime_plan plan;
|
|
char error[256];
|
|
if (gem_runtime_plan_init(&plan, &model, 32768, 2, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s\n", error);
|
|
return 1;
|
|
}
|
|
if (plan.kv_bytes != expected_bytes || plan.layer_count != layers ||
|
|
plan.swa_kv_element_size != 2 || plan.full_kv_element_size != 2 ||
|
|
!plan.layers[0].sliding || plan.layers[5].sliding ||
|
|
plan.layers[0].cache_capacity != 1024 || plan.layers[5].cache_capacity != 32768 ||
|
|
plan.layers[0].kv_heads != swa_kv || plan.layers[5].kv_heads != full_kv ||
|
|
plan.layers[0].head_dimension != 256 || plan.layers[5].head_dimension != 512 ||
|
|
plan.layers[0].kv_element_size != 2 || plan.layers[5].kv_element_size != 2) {
|
|
fprintf(stderr, "runtime plan mismatch for %uB\n", (unsigned)size);
|
|
gem_runtime_plan_destroy(&plan);
|
|
return 1;
|
|
}
|
|
for (uint32_t i = 0; i < plan.layer_count; ++i)
|
|
if (plan.layers[i].key_offset % 256 || plan.layers[i].value_offset % 256) {
|
|
fprintf(stderr, "unaligned layer %u\n", i);
|
|
gem_runtime_plan_destroy(&plan);
|
|
return 1;
|
|
}
|
|
gem_runtime_plan_destroy(&plan);
|
|
if (gem_runtime_plan_init_types(&plan, &model, 32768, 4, 2,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s\n", error);
|
|
return 1;
|
|
}
|
|
if (plan.kv_bytes != expected_mixed_bytes || plan.swa_kv_element_size != 4 ||
|
|
plan.full_kv_element_size != 2 || plan.layers[0].kv_element_size != 4 ||
|
|
plan.layers[5].kv_element_size != 2) {
|
|
fprintf(stderr, "mixed runtime plan mismatch for %uB\n", (unsigned)size);
|
|
gem_runtime_plan_destroy(&plan);
|
|
return 1;
|
|
}
|
|
gem_runtime_plan_destroy(&plan);
|
|
if (gem_runtime_plan_init_types(&plan, &model, 131072, 4, 2,
|
|
error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s\n", error);
|
|
return 1;
|
|
}
|
|
if (plan.kv_bytes != expected_mixed_128k_bytes ||
|
|
plan.layers[0].cache_capacity != 1024 ||
|
|
plan.layers[5].cache_capacity != 131072) {
|
|
fprintf(stderr, "128K mixed runtime plan mismatch for %uB\n", (unsigned)size);
|
|
gem_runtime_plan_destroy(&plan);
|
|
return 1;
|
|
}
|
|
gem_runtime_plan_destroy(&plan);
|
|
if (gem_runtime_plan_init_types(&plan, &model, 32768, 4, 3,
|
|
error, sizeof(error)) == 0) {
|
|
gem_runtime_plan_destroy(&plan);
|
|
fprintf(stderr, "accepted invalid mixed KV element size\n");
|
|
return 1;
|
|
}
|
|
if (gem_runtime_plan_init(&plan, &model, 262145, 2, error, sizeof(error)) == 0) {
|
|
gem_runtime_plan_destroy(&plan);
|
|
fprintf(stderr, "accepted oversized context\n");
|
|
return 1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
if (check_plan(GEM_MODEL_12B, 48, 3840, 16, 8, 1, 15360,
|
|
UINT64_C(872415232), UINT64_C(1207959552),
|
|
UINT64_C(2818572288)) != 0)
|
|
return 1;
|
|
if (check_plan(GEM_MODEL_31B, 60, 5376, 32, 16, 4, 21504,
|
|
UINT64_C(3523215360), UINT64_C(4362076160),
|
|
UINT64_C(12415139840)) != 0)
|
|
return 1;
|
|
return 0;
|
|
}
|