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
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200/* qxmx_batch: multi-slot batched-decode benchmark (M4.1 spike). * * Loads the real model, opens N slots, prefills each with the same prompt * (so all reach a common depth), then drives N steps of batched decode via * engine::forward_batch (FFN run as one M=N gemm_gpu_v4 call). Reports * per-step wall-clock and aggregate tok/s, A/B vs N separate single-slot * forward() calls (the serialized baseline). * * ./qxmx_batch model.gguf -f prompt.txt -n 64 -N 4 * * -N: number of slots (default 4). Padded to a multiple of 16 internally. */#include "qxmx_engine.h"#include "qxmx_gpu.h"#include "qxmx_model.h"#include "qxmx_tokenizer.h"#include "qxmx_device.h"
#include <chrono>#include <cstdio>#include <cstring>#include <fstream>#include <memory>#include <sstream>#include <string>#include <vector>
namespace qx { extern qx_model* g_model; }
int main(int argc, char** argv) { int n_gen = 64; int N_slots = 4; const char* prompt_file = nullptr; const char* model_path = nullptr; qx::engine::KCT k_ctk = qx::engine::KCT::q8_0;
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 if (a == "-N" && i + 1 < argc) N_slots = std::atoi(argv[++i]); else if (a == "-f" && i + 1 < argc) prompt_file = argv[++i]; else if (a == "-ctk" && i + 1 < argc) { std::string v = argv[++i]; if (v == "fp16") k_ctk = qx::engine::KCT::fp16; else if (v == "q8_0") k_ctk = qx::engine::KCT::q8_0; else if (v == "fp8") k_ctk = qx::engine::KCT::fp8; else { std::fprintf(stderr, "-ctk: unknown '%s'\n", v.c_str()); return 1; } } else if (a[0] != '-') model_path = argv[i]; } if (!model_path || !prompt_file) { std::fprintf(stderr, "usage: %s model.gguf -f prompt.txt -n N [-N slots] [-ctk fp16|q8_0|fp8]\n", argv[0]); 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::ifstream pf(prompt_file); std::stringstream ss; ss << pf.rdbuf(); std::string prompt_str = ss.str(); std::vector<int32_t> ids; tok.encode(prompt_str, &ids); std::printf("prompt: %zu tokens, N=%d slots, %d decode steps\n", ids.size(), N_slots, n_gen);
int max_ctx = (int)ids.size() + n_gen + 16; if (max_ctx < (1 << 14)) max_ctx = 1 << 14; qx::g_model = &m; auto eng = std::make_unique<qx::engine>(); eng->k_ctk = k_ctk; if (!eng->init(&m, max_ctx)) { std::fprintf(stderr, "engine init failed\n"); return 1; } std::printf("device: %s\n", qx::device_name(eng->q).c_str());
/* Open N slots, prefill each with the same prompt. */ std::vector<int> slots(N_slots); for (int i = 0; i < N_slots; i++) { slots[i] = eng->open_slot(); eng->prefill(slots[i], ids.data(), (int)ids.size()); } std::printf("prefilled %d slots to depth %zu\n", N_slots, ids.size());
/* Sample the first generated token per slot (greedy). */ qx::sample_params sp; /* greedy */ std::vector<int32_t> next_tok(N_slots); for (int i = 0; i < N_slots; i++) next_tok[i] = eng->sample(slots[i], sp);
/* forward_batch writes per-slot logits to eng->logits_batch_buf[i*VOCAB] * (USM shared). sample_slot copies row i into eng->logits before sampling, * since sample() reads eng->logits. */ auto sample_slot = [&](int i) -> int32_t { std::memcpy(eng->logits, eng->logits_batch_buf + (size_t)i * QX_N_VOCAB, (size_t)QX_N_VOCAB * sizeof(float)); return eng->sample(slots[i], sp); };
/* Correctness check: one batched step vs one serialized step per slot. * Both start from the same post-prefill state, run one forward, and we * compare the argmax token per slot. Reassociation drift may flip a * near-tie argmax occasionally; we report mismatches but don't fail. */ { std::vector<int32_t> batch_tok(N_slots), ser_tok(N_slots); eng->forward_batch(slots.data(), next_tok.data(), N_slots); for (int i = 0; i < N_slots; i++) batch_tok[i] = sample_slot(i); /* Reset and run serialized. */ for (int i = 0; i < N_slots; i++) { eng->reset(slots[i]); eng->prefill(slots[i], ids.data(), (int)ids.size()); next_tok[i] = eng->sample(slots[i], sp); eng->forward(slots[i], next_tok[i]); ser_tok[i] = eng->sample(slots[i], sp); } int mism = 0; for (int i = 0; i < N_slots; i++) if (batch_tok[i] != ser_tok[i]) mism++; std::printf("correctness: %d/%d slots match serialized argmax (%d mismatch)\n", N_slots - mism, N_slots, mism); /* Re-prepare for the timed batched run: reset + prefill all slots. */ for (int i = 0; i < N_slots; i++) { eng->reset(slots[i]); eng->prefill(slots[i], ids.data(), (int)ids.size()); next_tok[i] = eng->sample(slots[i], sp); } }
/* Warm up (2 steps). */ for (int w = 0; w < 2; w++) { eng->forward_batch(slots.data(), next_tok.data(), N_slots); for (int i = 0; i < N_slots; i++) next_tok[i] = sample_slot(i); }
/* Timed batched run. forward_batch submits async; we drain via a final * reset-prefill-free dummy: simplest is to read a USM value the last * kernel wrote. Easier: forward_batch internally does q.wait() when * QXMX_PROFILE is unset (non-profiling mode), so each call already * drained. No explicit wait needed. */ auto t0 = std::chrono::steady_clock::now(); for (int s = 0; s < n_gen; s++) eng->forward_batch(slots.data(), next_tok.data(), N_slots); auto t1 = std::chrono::steady_clock::now(); double batch_ms = std::chrono::duration<double, std::milli>(t1 - t0).count(); double batch_per_step = batch_ms / n_gen; double batch_agg_tps = N_slots / (batch_per_step / 1000.0);
std::printf("\n=== batched (FFN as one M=%d gemm_gpu_v4) ===\n", N_slots); std::printf(" total: %.1f ms over %d steps\n", batch_ms, n_gen); std::printf(" per-step: %.1f ms\n", batch_per_step); std::printf(" aggregate: %.1f tok/s (per-slot %.1f tok/s)\n", batch_agg_tps, batch_agg_tps / N_slots); std::fflush(stdout);
/* === Serialized baseline: N separate forward() calls per step === * Reset slots and re-prefill. Wrapped in try/catch so a DEVICE_LOST * here doesn't lose the batched result above. */ double ser_ms = 0, ser_per_step = 0, ser_agg_tps = 0; try { for (int i = 0; i < N_slots; i++) { eng->reset(slots[i]); eng->prefill(slots[i], ids.data(), (int)ids.size()); next_tok[i] = eng->sample(slots[i], sp); } auto t2 = std::chrono::steady_clock::now(); for (int s = 0; s < n_gen; s++) { for (int i = 0; i < N_slots; i++) { eng->forward(slots[i], next_tok[i]); next_tok[i] = eng->sample(slots[i], sp); } } auto t3 = std::chrono::steady_clock::now(); ser_ms = std::chrono::duration<double, std::milli>(t3 - t2).count(); ser_per_step = ser_ms / n_gen; ser_agg_tps = N_slots / (ser_per_step / 1000.0);
std::printf("\n=== serialized (N=%d separate forward() calls per step) ===\n", N_slots); std::printf(" total: %.1f ms over %d steps\n", ser_ms, n_gen); std::printf(" per-step: %.1f ms\n", ser_per_step); std::printf(" aggregate: %.1f tok/s (per-slot %.1f tok/s)\n", ser_agg_tps, ser_agg_tps / N_slots);
std::printf("\n=== speedup ===\n"); std::printf(" batched / serialized: %.2fx aggregate tok/s\n", batch_agg_tps / ser_agg_tps); std::printf(" per-step: %.1f ms -> %.1f ms (-%.0f%%)\n", ser_per_step, batch_per_step, (ser_per_step - batch_per_step) * 100.0 / ser_per_step); } catch (const std::exception& e) { std::fprintf(stderr, "\nserialized baseline crashed: %s\n", e.what()); std::fprintf(stderr, "(batched result above is still valid)\n"); } std::fflush(stdout);
for (int i = 0; i < N_slots; i++) eng->close_slot(slots[i]); return 0;}