315 lines
11 KiB
C
315 lines
11 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_context_scaling.c — HDC Memory Viability Test 7
|
|
*
|
|
* Context Length Scaling
|
|
*
|
|
* Method:
|
|
* Vary prompt length (200, 500, 1000, 2000, 4000 tokens via filler)
|
|
* Vary cache_size (8, 64, 256)
|
|
* Same AURORA-7 fact + question.
|
|
* Find crossover point where HybridMem starts helping.
|
|
*
|
|
* Usage:
|
|
* ./examples/test_context_scaling \
|
|
* --model ../models/base/RWKV-X-3.6B-64k-Base.safetensors \
|
|
* --vocab rwkvx/data/rwkv_vocab_v20230424.txt \
|
|
* --filler ../docs/architecture/rwkv-x-deep-dive.md
|
|
*/
|
|
|
|
#include "rwkvx/rwkvx.h"
|
|
#include "tensor/tensor.h"
|
|
|
|
#include <ctype.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
|
|
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);
|
|
}
|
|
|
|
/* Case-insensitive substring search */
|
|
static int contains_ci(const char *haystack, const char *needle) {
|
|
if (!haystack || !needle) return 0;
|
|
int hlen = (int)strlen(haystack);
|
|
int nlen = (int)strlen(needle);
|
|
for (int i = 0; i <= hlen - nlen; i++) {
|
|
int match = 1;
|
|
for (int j = 0; j < nlen; j++) {
|
|
if (tolower((unsigned char)haystack[i + j]) !=
|
|
tolower((unsigned char)needle[j])) {
|
|
match = 0;
|
|
break;
|
|
}
|
|
}
|
|
if (match) return 1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
/* Generate tokens greedily from a prefilled state. */
|
|
static char* generate(rx_model *model, rx_tokenizer *tok, rx_state *state,
|
|
tm_tensor *logits, int max_tokens, rx_config cfg) {
|
|
rx_sampling_params params = rx_sampling_defaults();
|
|
params.temperature = 0.0f;
|
|
|
|
uint32_t *gen_tokens = (uint32_t*)calloc(max_tokens, sizeof(uint32_t));
|
|
int n_gen = 0;
|
|
tm_tensor *cur_logits = logits;
|
|
|
|
for (int i = 0; i < max_tokens; i++) {
|
|
if (!cur_logits) break;
|
|
uint32_t token = rx_sample(cur_logits, ¶ms, NULL, cfg.vocab_size);
|
|
if (i > 0) tm_free(cur_logits);
|
|
if (token == 0) { cur_logits = NULL; break; }
|
|
gen_tokens[n_gen++] = token;
|
|
cur_logits = rx_forward(model, token, state);
|
|
}
|
|
if (cur_logits && cur_logits != logits) tm_free(cur_logits);
|
|
|
|
char *text = rx_tokenizer_decode(tok, gen_tokens, n_gen);
|
|
free(gen_tokens);
|
|
return text;
|
|
}
|
|
|
|
/* Extend filler to target character count by repeating */
|
|
static char* extend_filler(const char *base, int target_chars) {
|
|
int base_len = (int)strlen(base);
|
|
if (base_len == 0) return strdup("");
|
|
char *buf = (char*)malloc(target_chars + base_len + 1);
|
|
int pos = 0;
|
|
while (pos < target_chars) {
|
|
int chunk = base_len;
|
|
if (pos + chunk > target_chars) chunk = target_chars - pos;
|
|
memcpy(buf + pos, base, chunk);
|
|
pos += chunk;
|
|
}
|
|
buf[target_chars] = '\0';
|
|
return buf;
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *model_path = NULL;
|
|
const char *vocab_path = NULL;
|
|
const char *filler_path = NULL;
|
|
int hm_top_k = 5;
|
|
int max_tokens = 48;
|
|
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], "--hm-top-k") == 0 && i + 1 < argc) hm_top_k = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--max-tokens") == 0 && i + 1 < argc) max_tokens = 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", argv[0]);
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
if (!model_path || !vocab_path || !filler_path) {
|
|
fprintf(stderr, "Error: --model, --vocab, and --filler are required\n");
|
|
return 1;
|
|
}
|
|
|
|
/* Load shared resources */
|
|
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);
|
|
|
|
/* Load filler */
|
|
char *filler_raw = read_file(filler_path);
|
|
if (!filler_raw) { fprintf(stderr, "Failed to read filler\n"); return 1; }
|
|
|
|
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:";
|
|
const char *key_detail = "AURORA";
|
|
|
|
/* Target token counts and cache sizes to test */
|
|
int target_tokens[] = {200, 500, 1000, 2000, 4000};
|
|
int n_targets = 5;
|
|
int cache_sizes[] = {8, 64, 256};
|
|
int n_caches = 3;
|
|
|
|
fprintf(stderr, "\n==========================================\n");
|
|
fprintf(stderr, "HDC Memory Viability — Test 7\n");
|
|
fprintf(stderr, "Context Length Scaling\n");
|
|
fprintf(stderr, "==========================================\n\n");
|
|
|
|
/* Results: [target_idx][cache_idx] => 0=both_miss, 1=A_only, 2=B_only, 3=both_hit */
|
|
int results[5][3];
|
|
memset(results, 0, sizeof(results));
|
|
|
|
/* Header */
|
|
fprintf(stderr, "%-10s", "tokens");
|
|
for (int c = 0; c < n_caches; c++)
|
|
fprintf(stderr, " cache=%-6d", cache_sizes[c]);
|
|
fprintf(stderr, "\n");
|
|
for (int i = 0; i < 10 + n_caches * 14; i++) fprintf(stderr, "-");
|
|
fprintf(stderr, "\n");
|
|
|
|
for (int ti = 0; ti < n_targets; ti++) {
|
|
int target = target_tokens[ti];
|
|
|
|
/* Estimate chars needed for target token count (~3.5 chars/token) */
|
|
int filler_chars = (int)(target * 3.5);
|
|
char *filler = extend_filler(filler_raw, filler_chars);
|
|
|
|
/* Build prompt */
|
|
int pl = (int)(strlen(fact) + 2 + strlen(filler) + strlen(question) + 1);
|
|
char *prompt = (char*)malloc(pl);
|
|
snprintf(prompt, pl, "%s\n\n%s%s", fact, filler, question);
|
|
|
|
uint32_t *tokens = (uint32_t*)malloc(65536 * sizeof(uint32_t));
|
|
int n_tokens = rx_tokenizer_encode(tok, prompt, tokens, 65536);
|
|
|
|
fprintf(stderr, "%-10d", n_tokens);
|
|
|
|
for (int ci = 0; ci < n_caches; ci++) {
|
|
int cs = cache_sizes[ci];
|
|
|
|
/* Run A: no HM */
|
|
rx_hybridmem_config hm_cfg = rx_hybridmem_config_defaults();
|
|
hm_cfg.store_dir = NULL;
|
|
hm_cfg.max_cache_len = cs;
|
|
rx_state *state_a = rx_state_new_with_hybridmem(model, &hm_cfg);
|
|
|
|
tm_tensor *logits = NULL;
|
|
for (int i = 0; i < n_tokens; i++) {
|
|
if (logits) tm_free(logits);
|
|
logits = rx_forward(model, tokens[i], state_a);
|
|
}
|
|
char *out_a = generate(model, tok, state_a, logits, max_tokens, cfg);
|
|
tm_free(logits);
|
|
rx_state_free(state_a);
|
|
int hit_a = contains_ci(out_a, key_detail);
|
|
|
|
/* Run B: with HM */
|
|
const char *hm_dir = "/tmp/hm_context_scaling";
|
|
rmdir_recursive(hm_dir);
|
|
hm_cfg.store_dir = hm_dir;
|
|
hm_cfg.max_cache_len = cs;
|
|
hm_cfg.top_k = hm_top_k;
|
|
rx_state *state_b = rx_state_new_with_hybridmem(model, &hm_cfg);
|
|
|
|
logits = NULL;
|
|
for (int i = 0; i < n_tokens; i++) {
|
|
if (logits) tm_free(logits);
|
|
logits = rx_forward(model, tokens[i], state_b);
|
|
}
|
|
char *out_b = generate(model, tok, state_b, logits, max_tokens, cfg);
|
|
tm_free(logits);
|
|
rx_state_free(state_b);
|
|
int hit_b = contains_ci(out_b, key_detail);
|
|
rmdir_recursive(hm_dir);
|
|
|
|
/* Encode result */
|
|
if (hit_a && hit_b) results[ti][ci] = 3;
|
|
else if (hit_a) results[ti][ci] = 1;
|
|
else if (hit_b) results[ti][ci] = 2;
|
|
else results[ti][ci] = 0;
|
|
|
|
const char *labels[] = {"--/--", "A/--", "--/B", "A/B"};
|
|
fprintf(stderr, " %-12s", labels[results[ti][ci]]);
|
|
|
|
free(out_a);
|
|
free(out_b);
|
|
}
|
|
fprintf(stderr, "\n");
|
|
|
|
free(prompt);
|
|
free(tokens);
|
|
free(filler);
|
|
}
|
|
|
|
/* Analysis */
|
|
fprintf(stderr, "\nLegend: A=no_HM hit, B=HM hit, --=miss\n\n");
|
|
|
|
/* Find crossover points */
|
|
fprintf(stderr, "--- Crossover Analysis ---\n");
|
|
for (int ci = 0; ci < n_caches; ci++) {
|
|
int crossover = -1;
|
|
for (int ti = 0; ti < n_targets; ti++) {
|
|
if (results[ti][ci] == 2) { /* B hits but A doesn't */
|
|
crossover = target_tokens[ti];
|
|
break;
|
|
}
|
|
}
|
|
if (crossover > 0) {
|
|
fprintf(stderr, " cache=%d: HybridMem starts helping at ~%d tokens\n",
|
|
cache_sizes[ci], crossover);
|
|
} else {
|
|
/* Check if both always hit or both always miss */
|
|
int all_both = 1, all_miss = 1;
|
|
for (int ti = 0; ti < n_targets; ti++) {
|
|
if (results[ti][ci] != 3) all_both = 0;
|
|
if (results[ti][ci] != 0) all_miss = 0;
|
|
}
|
|
if (all_both) {
|
|
fprintf(stderr, " cache=%d: Both methods always succeed (cache large enough)\n",
|
|
cache_sizes[ci]);
|
|
} else if (all_miss) {
|
|
fprintf(stderr, " cache=%d: Both methods always fail\n", cache_sizes[ci]);
|
|
} else {
|
|
fprintf(stderr, " cache=%d: No clear crossover pattern\n", cache_sizes[ci]);
|
|
}
|
|
}
|
|
}
|
|
|
|
fprintf(stderr, "\nVERDICT:\n");
|
|
/* Count how many conditions HybridMem helps (result==2) */
|
|
int hm_helps = 0, total_cells = n_targets * n_caches;
|
|
for (int ti = 0; ti < n_targets; ti++)
|
|
for (int ci = 0; ci < n_caches; ci++)
|
|
if (results[ti][ci] == 2) hm_helps++;
|
|
|
|
if (hm_helps > 0) {
|
|
fprintf(stderr, " PASS — HybridMem helps in %d/%d conditions\n", hm_helps, total_cells);
|
|
} else {
|
|
/* Check if cache always sufficient */
|
|
int both_hit = 0;
|
|
for (int ti = 0; ti < n_targets; ti++)
|
|
for (int ci = 0; ci < n_caches; ci++)
|
|
if (results[ti][ci] == 3) both_hit++;
|
|
if (both_hit == total_cells) {
|
|
fprintf(stderr, " N/A — Cache always sufficient at tested sizes\n");
|
|
} else {
|
|
fprintf(stderr, " FAIL — HybridMem doesn't help at any tested condition\n");
|
|
}
|
|
}
|
|
|
|
/* Cleanup */
|
|
free(filler_raw);
|
|
rx_model_free(model);
|
|
rx_tokenizer_free(tok);
|
|
|
|
fprintf(stderr, "\n==========================================\n");
|
|
return 0;
|
|
}
|