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
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107/* qxmx_engine.cpp: base-class sampling + RNG plumbing for engine_i. * * M2.2: pure-host sampler over `logits` (already synced by forward/prefill). * Greedy when temperature==0; otherwise temperature scaling + top-k + top-p + * seeded std::mt19937_64. Implemented once here so both backends share it. * * Algorithm: * - greedy: single argmax pass (O(VOCAB), ~0.5 ms on 248K). * - temp>0: scale logits by 1/temperature, then top-k (partial sort down to * k candidates), then top-p (cumulative-sum filter on softmax of the k), * then multinomial draw. When both filters are disabled this degenerates * to a full-vocab softmax + multinomial — the ~1-2 ms the ROADMAP budgets. * * Scratch (idx_buf, prob_buf) is hoisted on engine_i; the only per-sample * allocation is the partial_sort lambda's swaps, all in-place. */#include "qxmx_engine.h"
#include <algorithm>#include <cmath>#include <cstdint>
namespace qx {
void engine_i::seed(int slot, uint64_t s) { if (slot < 0 || slot >= (int)rngs.size()) return; rngs[slot].seed(s);}
int engine_i::open_slot() { int slot = (int)rngs.size(); rngs.emplace_back(); /* default-constructed mt19937_64 (seed 0) */ on_open_slot(slot); return slot;}
void engine_i::close_slot(int slot) { if (slot < 0 || slot >= (int)rngs.size()) return; on_close_slot(slot); if (slot == (int)rngs.size() - 1) rngs.pop_back(); else rngs.erase(rngs.begin() + slot);}
int engine_i::sample(int slot, const sample_params& p) { /* Greedy: single argmax pass. No scratch, no RNG. */ if (p.temperature == 0.0f) { int best = 0; float bv = -1e30f; for (int v = 0; v < QX_N_VOCAB; v++) if (logits[v] > bv) { bv = logits[v]; best = v; } return best; }
/* Non-greedy: temperature scale + top-k + top-p + multinomial. */ std::mt19937_64& rng = rngs.at(slot); const int V = QX_N_VOCAB; idx_buf.assign(V, 0); for (int i = 0; i < V; i++) idx_buf[i] = i; prob_buf.assign(V, 0.0f);
/* Scale logits by 1/temperature into prob_buf. */ const float inv_t = 1.0f / p.temperature; for (int i = 0; i < V; i++) prob_buf[i] = logits[i] * inv_t;
/* Top-k: keep the k largest logits. partial_sort is O(V log k) — far * cheaper than a full sort, and the ROADMAP's "~1-2 ms" budget holds. */ int k = V; if (p.top_k > 0 && p.top_k < V) k = p.top_k; if (k < V) { std::partial_sort(idx_buf.begin(), idx_buf.begin() + k, idx_buf.end(), [&](int a, int b){ return prob_buf[a] > prob_buf[b]; }); } else { std::sort(idx_buf.begin(), idx_buf.end(), [&](int a, int b){ return prob_buf[a] > prob_buf[b]; }); }
/* Softmax over the k surviving candidates (numerically stable: subtract * the max, which is idx_buf[0] after the sort above). */ const float max_l = prob_buf[idx_buf[0]]; double z = 0.0; for (int i = 0; i < k; i++) { float e = std::exp(prob_buf[idx_buf[i]] - max_l); prob_buf[idx_buf[i]] = e; z += e; } if (z <= 0.0) return idx_buf[0]; /* degenerate — return the argmax */
/* Top-p: accumulate over the (already prob-sorted) candidates until the * cumulative probability reaches top_p; cut everything after. */ int keep = k; if (p.top_p < 1.0f) { double c = 0.0; for (int i = 0; i < k; i++) { c += prob_buf[idx_buf[i]] / z; if (c >= p.top_p) { keep = i + 1; break; } } }
/* Multinomial draw over the kept candidates. */ double r = std::uniform_real_distribution<double>(0.0, 1.0)(rng); double c = 0.0; for (int i = 0; i < keep; i++) { c += prob_buf[idx_buf[i]] / z; if (r < c) return idx_buf[i]; } return idx_buf[keep - 1]; /* FP roundoff fallback */}
} // namespace qx