283 lines
11 KiB
C
283 lines
11 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_eviction_boundary.c — HDC Memory Viability Test 11
|
|
*
|
|
* Eviction Boundary Analysis
|
|
*
|
|
* Method:
|
|
* 1. Run 3.6B (or 0.2B) needle-in-haystack with HybridMem
|
|
* 2. After prefill, query the HDC store for the fact embedding
|
|
* 3. Report:
|
|
* - Total entries stored to disk
|
|
* - Per-layer entry count
|
|
* - Rank of fact entry per layer (how well can we retrieve it?)
|
|
* - Score distribution of retrieved entries
|
|
*
|
|
* Usage:
|
|
* ./examples/test_eviction_boundary \
|
|
* --model ../models/base/RWKV-X-0.2B-64k-Base.safetensors \
|
|
* --vocab rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
* --filler ../docs/architecture/rwkv-x-deep-dive.md
|
|
*/
|
|
|
|
#include "rwkvx/rwkvx.h"
|
|
#include "hybridmem/hybridmem.h"
|
|
#include "tensor/tensor.h"
|
|
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
|
|
/* Read entire file into malloc'd string. */
|
|
static char* read_file(const char *path) {
|
|
FILE *fp = fopen(path, "r");
|
|
if (!fp) return NULL;
|
|
fseek(fp, 0, SEEK_END);
|
|
long len = ftell(fp);
|
|
fseek(fp, 0, SEEK_SET);
|
|
char *buf = (char*)malloc(len + 1);
|
|
size_t nread = fread(buf, 1, len, fp);
|
|
buf[nread] = '\0';
|
|
fclose(fp);
|
|
return buf;
|
|
}
|
|
|
|
static void rmdir_recursive(const char *path) {
|
|
char cmd[512];
|
|
snprintf(cmd, sizeof(cmd), "rm -rf %s", path);
|
|
system(cmd);
|
|
}
|
|
|
|
/* Cosine similarity between two float vectors */
|
|
static float cosine_sim(const float *a, const float *b, int dim) {
|
|
float dot = 0, na = 0, nb = 0;
|
|
for (int i = 0; i < dim; i++) {
|
|
dot += a[i] * b[i];
|
|
na += a[i] * a[i];
|
|
nb += b[i] * b[i];
|
|
}
|
|
float d = sqrtf(na) * sqrtf(nb);
|
|
return d > 1e-12f ? dot / d : 0;
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *model_path = NULL;
|
|
const char *vocab_path = NULL;
|
|
const char *filler_path = NULL;
|
|
int cache_size = 8;
|
|
int hm_top_k = 10; /* retrieve more for rank analysis */
|
|
int max_filler_chars = 2000;
|
|
int quantize = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (strcmp(argv[i], "--model") == 0 && i + 1 < argc)
|
|
model_path = argv[++i];
|
|
else if (strcmp(argv[i], "--vocab") == 0 && i + 1 < argc)
|
|
vocab_path = argv[++i];
|
|
else if (strcmp(argv[i], "--filler") == 0 && i + 1 < argc)
|
|
filler_path = argv[++i];
|
|
else if (strcmp(argv[i], "--cache-size") == 0 && i + 1 < argc)
|
|
cache_size = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--hm-top-k") == 0 && i + 1 < argc)
|
|
hm_top_k = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--max-filler") == 0 && i + 1 < argc)
|
|
max_filler_chars = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--quantize") == 0)
|
|
quantize = 1;
|
|
else if (strcmp(argv[i], "--help") == 0) {
|
|
fprintf(stderr,
|
|
"Usage: %s --model PATH --vocab PATH --filler PATH [options]\n"
|
|
" --cache-size N KV cache slots per attn layer (default: 8)\n"
|
|
" --hm-top-k N Retrieval count for rank analysis (default: 10)\n"
|
|
" --max-filler N Max chars of filler text (default: 2000)\n"
|
|
" --quantize Quantize weights to INT8\n",
|
|
argv[0]);
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
if (!model_path || !vocab_path || !filler_path) {
|
|
fprintf(stderr, "Error: --model, --vocab, and --filler are required\n");
|
|
return 1;
|
|
}
|
|
|
|
/* Load model + tokenizer */
|
|
fprintf(stderr, "Loading tokenizer...\n");
|
|
rx_tokenizer *tok = rx_tokenizer_new(vocab_path);
|
|
if (!tok) { fprintf(stderr, "Failed to load tokenizer\n"); return 1; }
|
|
|
|
fprintf(stderr, "Loading model%s...\n", quantize ? " (INT8)" : "");
|
|
rx_model *model = quantize ? rx_model_load_q8(model_path)
|
|
: rx_model_load(model_path);
|
|
if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; }
|
|
rx_config cfg = rx_model_config(model);
|
|
int c = cfg.n_embd;
|
|
|
|
/* Build prompt */
|
|
const char *fact = "The secret project codename is AURORA-7 and the lead researcher is Dr. Kimura.";
|
|
const char *question = "\n\nQuestion: What is the secret project codename and who leads it?\nAnswer:";
|
|
|
|
char *filler_raw = read_file(filler_path);
|
|
if (!filler_raw) {
|
|
fprintf(stderr, "Failed to read filler file: %s\n", filler_path);
|
|
return 1;
|
|
}
|
|
if ((int)strlen(filler_raw) > max_filler_chars)
|
|
filler_raw[max_filler_chars] = '\0';
|
|
|
|
int prompt_len = (int)(strlen(fact) + 2 + strlen(filler_raw) + strlen(question) + 1);
|
|
char *prompt = (char*)malloc(prompt_len);
|
|
snprintf(prompt, prompt_len, "%s\n\n%s%s", fact, filler_raw, question);
|
|
free(filler_raw);
|
|
|
|
/* Tokenize fact separately (to get fact token range) */
|
|
uint32_t *fact_tokens = (uint32_t*)malloc(1024 * sizeof(uint32_t));
|
|
int n_fact_tokens = rx_tokenizer_encode(tok, fact, fact_tokens, 1024);
|
|
|
|
uint32_t *tokens = (uint32_t*)malloc(32768 * sizeof(uint32_t));
|
|
int n_tokens = rx_tokenizer_encode(tok, prompt, tokens, 32768);
|
|
|
|
fprintf(stderr, "\n==========================================\n");
|
|
fprintf(stderr, "HDC Memory Viability — Test 11\n");
|
|
fprintf(stderr, "Eviction Boundary Analysis\n");
|
|
fprintf(stderr, "==========================================\n");
|
|
fprintf(stderr, "Prompt tokens: %d\n", n_tokens);
|
|
fprintf(stderr, "Fact tokens: %d (first %d tokens of prompt)\n",
|
|
n_fact_tokens, n_fact_tokens);
|
|
fprintf(stderr, "Cache size: %d\n", cache_size);
|
|
fprintf(stderr, "HM top-k: %d\n", hm_top_k);
|
|
fprintf(stderr, "Attn layers: %d\n\n", cfg.n_attn_layer);
|
|
|
|
/* Run prefill with HybridMem */
|
|
const char *hm_dir = "/tmp/hm_eviction_test";
|
|
rmdir_recursive(hm_dir);
|
|
|
|
rx_hybridmem_config hm_cfg = rx_hybridmem_config_defaults();
|
|
hm_cfg.store_dir = hm_dir;
|
|
hm_cfg.max_cache_len = cache_size;
|
|
hm_cfg.top_k = hm_top_k;
|
|
|
|
rx_state *state = rx_state_new_with_hybridmem(model, &hm_cfg);
|
|
if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; }
|
|
|
|
/* Also capture fact embeddings for querying */
|
|
int max_capture = n_fact_tokens + 10;
|
|
float *fact_embeddings = (float*)malloc((size_t)max_capture * c * sizeof(float));
|
|
rx_state_set_embed_capture(state, fact_embeddings, max_capture, -1);
|
|
|
|
fprintf(stderr, "Running prefill...\n");
|
|
struct timespec t0, t1;
|
|
clock_gettime(CLOCK_MONOTONIC, &t0);
|
|
|
|
tm_tensor *logits = NULL;
|
|
for (int i = 0; i < n_tokens; i++) {
|
|
if (logits) tm_free(logits);
|
|
logits = rx_forward(model, tokens[i], state);
|
|
|
|
/* Stop capturing after fact tokens */
|
|
if (i == n_fact_tokens - 1) {
|
|
rx_state_set_embed_capture(state, NULL, 0, -1);
|
|
}
|
|
}
|
|
if (logits) tm_free(logits);
|
|
|
|
clock_gettime(CLOCK_MONOTONIC, &t1);
|
|
double ms = (t1.tv_sec - t0.tv_sec) * 1000.0 + (t1.tv_nsec - t0.tv_nsec) / 1e6;
|
|
|
|
int n_captured = rx_state_embed_capture_count(state);
|
|
fprintf(stderr, "Prefill: %d tokens in %.1fms (%.1f tok/s)\n",
|
|
n_tokens, ms, n_tokens * 1000.0 / ms);
|
|
fprintf(stderr, "Captured %d fact embeddings\n\n", n_captured);
|
|
|
|
/* Store stats */
|
|
fm_stats stats = fm_stats_get(rx_state_get_hm_store(state));
|
|
fprintf(stderr, "HybridMem store stats:\n");
|
|
fprintf(stderr, " Total stores: %zu\n", stats.total_stores);
|
|
fprintf(stderr, " Total entries: %zu\n", stats.num_entries);
|
|
fprintf(stderr, " Disk bytes: %zu\n", stats.disk_bytes_used);
|
|
fprintf(stderr, "\n");
|
|
|
|
/* Flush store before querying */
|
|
rx_state_flush(state);
|
|
|
|
/* Query the HDC store for each captured fact embedding */
|
|
fprintf(stderr, "--- Fact Embedding Retrieval Analysis ---\n\n");
|
|
|
|
if (n_captured == 0) {
|
|
fprintf(stderr, "No fact embeddings captured (model may have 0 attn layers in fact range).\n");
|
|
fprintf(stderr, "With %d fact tokens and cache_size=%d, the fact may still be in cache.\n",
|
|
n_fact_tokens, cache_size);
|
|
}
|
|
|
|
/* For each captured embedding, query the store */
|
|
for (int e = 0; e < n_captured && e < 4; e++) {
|
|
float *query_emb = &fact_embeddings[e * c];
|
|
|
|
fprintf(stderr, "Fact embedding %d:\n", e);
|
|
|
|
/* Query each layer */
|
|
for (int l = 0; l < cfg.n_attn_layer; l++) {
|
|
int64_t shape_q[] = {c};
|
|
tm_tensor *query_t = tm_from_data(query_emb, shape_q, 1, TM_F32);
|
|
fm_results *results = fm_query(rx_state_get_hm_store(state), query_t, l, hm_top_k);
|
|
tm_free(query_t);
|
|
|
|
if (results && results->count > 0) {
|
|
fprintf(stderr, " Layer %d: %d results\n", l, results->count);
|
|
for (int r = 0; r < results->count && r < 5; r++) {
|
|
int n_tok = (int)tm_size(results->items[r].keys, 0);
|
|
|
|
/* Compute cosine similarity between query and first key in result */
|
|
float *rk = tm_data_f32(results->items[r].keys);
|
|
float cos = cosine_sim(query_emb, rk, c);
|
|
|
|
fprintf(stderr, " Rank %d: score=%.4f, cos_sim=%.4f, n_tokens=%d\n",
|
|
r, results->items[r].score, cos, n_tok);
|
|
}
|
|
fm_results_free(results);
|
|
} else {
|
|
fprintf(stderr, " Layer %d: no results\n", l);
|
|
if (results) fm_results_free(results);
|
|
}
|
|
}
|
|
fprintf(stderr, "\n");
|
|
}
|
|
|
|
/* Eviction count analysis */
|
|
int eviction_count = n_tokens / cache_size; /* approximate */
|
|
fprintf(stderr, "--- Eviction Analysis ---\n");
|
|
fprintf(stderr, "Approximate evictions: %d (tokens=%d / cache=%d)\n",
|
|
eviction_count, n_tokens, cache_size);
|
|
fprintf(stderr, "Entries in store: %zu\n", stats.num_entries);
|
|
fprintf(stderr, "Tokens per eviction batch: %d (cache_size * evict_fraction)\n",
|
|
(int)(cache_size * hm_cfg.evict_fraction));
|
|
|
|
fprintf(stderr, "\nVERDICT:\n");
|
|
if (stats.num_entries > 0 && n_captured > 0) {
|
|
fprintf(stderr, " Store populated with %zu entries — eviction pipeline works.\n",
|
|
stats.num_entries);
|
|
fprintf(stderr, " Check retrieval scores above: high scores indicate fact is retrievable.\n");
|
|
} else if (stats.num_entries > 0) {
|
|
fprintf(stderr, " Store populated but no fact embeddings captured.\n");
|
|
fprintf(stderr, " The fact may be in cache (within first %d tokens).\n", cache_size);
|
|
} else {
|
|
fprintf(stderr, " FAIL — No entries stored. Check eviction path.\n");
|
|
}
|
|
|
|
/* Cleanup */
|
|
free(fact_embeddings);
|
|
free(fact_tokens);
|
|
free(tokens);
|
|
free(prompt);
|
|
rx_state_free(state);
|
|
rx_model_free(model);
|
|
rx_tokenizer_free(tok);
|
|
rmdir_recursive(hm_dir);
|
|
|
|
fprintf(stderr, "\n==========================================\n");
|
|
return 0;
|
|
}
|