Something went wrong. Try again.
A single-model (Bonsai-27B) inference engine for Intel Arc GPUs
Something went wrong. Try again.
C++
at extra-debug
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178/* prefix_cache_engine_test: M2.6.3 gate -- in-engine cache correctness. * * Gate (ROADMAP M2.6 / MVP done-criteria #5): a cached prefill must produce * bit-identical greedy output to a fresh full prefill. This is the * corruption guard for the device<->host DMA path: if the restored KV/SSM * bytes don't match what a fresh prefill would have written, the generated * tokens diverge. * * Test: * 1. Prefill prompt P (>= 1 block = 512 tok) on slot A, snapshotting to a * shared cache. Decode a few tokens to confirm the engine works post- * snapshot. * 2. Reset slot B, prefill_cached(P + delta) -> should restore P's prefix * (restored_to > 0), prefill only delta. Greedy-decode N tokens. * 3. Fresh slot C: prefill(P + delta) in full, greedy-decode N tokens. * 4. Slots B and C must produce IDENTICAL decoded text. * * ./prefix_cache_engine_test model.gguf prompt.txt delta.txt [-n N] * (prompt.txt should be >= 512 tokens; delta.txt any length.) */#include "qxmx_gpu.h"#include "qxmx_model.h"#include "qxmx_prefix_cache.h"#include "qxmx_tokenizer.h"
#include <cstdio>#include <fstream>#include <cstring>#include <sstream>#include <string>#include <vector>
namespace qx { extern qx_model* g_model; }
static std::string read_file(const char* path) { std::ifstream f(path); if (!f) { std::fprintf(stderr, "cannot open %s\n", path); return ""; } std::stringstream ss; ss << f.rdbuf(); return ss.str();}
/* Greedy decode N tokens after the slot is pre-filled (logits already set for * the last prompt token). Returns decoded text. */static std::string gen_decode(qx::engine& eng, int slot, int n_gen, const qx::tokenizer& tok, uint32_t eos_id) { qx::incremental_decoder dec(tok); std::string out; int next = eng.sample(slot, qx::sample_params{}); out += dec.push(next); for (int step = 1; step < n_gen; step++) { if (next == (int)eos_id) break; eng.forward(slot, next); next = eng.sample(slot, qx::sample_params{}); out += dec.push(next); } out += dec.finish(); return out;}
int main(int argc, char** argv) { std::vector<const char*> pos; int n_gen = 32; for (int i = 1; i < argc; i++) { std::string a = argv[i]; if (a == "-n" && i + 1 < argc) n_gen = std::atoi(argv[++i]); else pos.push_back(argv[i]); } if (pos.size() < 3) { std::fprintf(stderr, "usage: %s model.gguf prompt.txt delta.txt [-n N]\n", argv[0]); return 1; } const char* model_path = pos[0]; std::string prompt_str = read_file(pos[1]); std::string delta_str = read_file(pos[2]); if (prompt_str.empty() || delta_str.empty()) return 1;
qx_model m; char err[256]; if (qx_model_open(&m, model_path, err, sizeof err)) { std::fprintf(stderr, "load: %s\n", err); return 1; } qx::tokenizer tok; if (!tok.load(&m.gguf, err, sizeof err)) { std::fprintf(stderr, "tok: %s\n", err); return 1; }
std::vector<int32_t> P, D, PplusD; tok.encode(prompt_str, &P); tok.encode(delta_str, &D); PplusD = P; PplusD.insert(PplusD.end(), D.begin(), D.end());
if ((int)P.size() < qx::PC_BLOCK) { std::fprintf(stderr, "prompt too short: %zu tokens (need >= %d for one block)\n", P.size(), qx::PC_BLOCK); return 1; } std::printf("prompt: %zu tokens (%d full blocks), delta: %zu tokens\n", P.size(), (int)(P.size() / qx::PC_BLOCK), D.size());
int max_ctx = std::max((int)PplusD.size() + n_gen + 16, 1 << 14); qx::g_model = &m; qx::engine eng; eng.init(&m, max_ctx);
/* Cache geometry: 16 attn, 48 ssm, q8_0 K (1088 B/row), V (640 B/row), * fp32 SSM state/conv. Matches engine defaults. */ qx::prefix_cache cache; size_t s_state = (size_t)QX_SSM_DT_RANK * QX_SSM_STATE * QX_SSM_STATE * sizeof(float); size_t s_conv = (size_t)(QX_SSM_CONV_K - 1) * QX_SSM_CONVD * sizeof(float); cache.init(eng.n_attn, eng.n_ssm, 1088, 640, s_state, s_conv);
/* ---- 1. Populate the cache by prefilling P on slot A with snapshotting ---- */ int slotA = eng.open_slot(); int restored = eng.prefill_cached(slotA, P.data(), (int)P.size(), &cache); std::printf("[A] prefill_cached(P) restored_to=%d (expect 0, cold cache)\n", restored); if (restored != 0) { std::fprintf(stderr, "FAIL: cold cache should restore 0, got %d\n", restored); return 1; } std::printf("[A] cache after P: kv=%zu ssm=%zu\n", cache.kv_size(), cache.ssm_size()); if (cache.ssm_size() == 0) { std::fprintf(stderr, "FAIL: cache empty after prefilling P (snapshotting broken)\n"); return 1; } /* Sanity: with stride=1 every block is a group of 1, so kv==ssm==blocks. * With stride>1, KV is cached for EVERY block (the group snapshot DMAs * all stride blocks' KV) but SSM checkpoints only at group ends -- so * kv==full_blocks and ssm==full_blocks/stride. */ int snap_stride = []{ const char* e = std::getenv("QXMX_SNAP_STRIDE"); int v = e ? std::atoi(e) : 1; return v < 1 ? 1 : v; }(); size_t full_blocks = P.size() / qx::PC_BLOCK; size_t expect_kv = full_blocks; size_t expect_ssm = full_blocks / snap_stride; size_t expect_blocks = expect_ssm * snap_stride; /* whole groups, in blocks */ if (cache.kv_size() != expect_kv || cache.ssm_size() != expect_ssm) { std::fprintf(stderr, "FAIL: cache counts kv=%zu ssm=%zu, expected kv=%zu ssm=%zu (stride=%d)\n", cache.kv_size(), cache.ssm_size(), expect_kv, expect_ssm, snap_stride); return 1; } eng.close_slot(slotA);
/* ---- 2. Slot B: prefill_cached(P+D) -- should restore P, prefill only D ---- */ int slotB = eng.open_slot(); restored = eng.prefill_cached(slotB, PplusD.data(), (int)PplusD.size(), &cache); size_t restored_blocks = (size_t)restored / qx::PC_BLOCK; std::printf("[B] prefill_cached(P+D) restored_to=%d (%zu blocks of %zu)\n", restored, restored_blocks, expect_blocks); /* The cache may hold FEWER than expect_blocks if the budget evicted some; * the 16 GB default holds all of P's blocks (P ~10 blocks). Require at * least the majority restored. */ if (restored_blocks < expect_blocks) { std::fprintf(stderr, "FAIL: only %zu/%zu blocks restored\n", restored_blocks, expect_blocks); return 1; } if (restored_blocks == 0) { std::fprintf(stderr, "FAIL: cache hit but restored 0 blocks\n"); return 1; } std::string outB = gen_decode(eng, slotB, n_gen, tok, tok.eos_id); eng.close_slot(slotB);
/* ---- 3. Slot C: fresh full prefill of P+D, no cache ---- */ int slotC = eng.open_slot(); eng.prefill(slotC, PplusD.data(), (int)PplusD.size()); std::string outC = gen_decode(eng, slotC, n_gen, tok, tok.eos_id); eng.close_slot(slotC);
/* ---- 4. Compare ---- */ std::printf("[B] cached: \"%s\"\n", outB.c_str()); std::printf("[C] fresh: \"%s\"\n", outC.c_str()); if (outB != outC) { std::fprintf(stderr, "FAIL: cached vs fresh prefill diverge (%zu vs %zu bytes)\n", outB.size(), outC.size()); return 1; } std::printf("PASS: cached prefill produces bit-identical output to fresh prefill.\n"); std::printf(" restored %zu/%zu prefix blocks (%zu tokens), prefilled only the delta.\n", restored_blocks, expect_blocks, (size_t)restored); qx_model_close(&m); return 0;}