Something went wrong. Try again.
A single-model (Bonsai-27B) inference engine for Intel Arc GPUs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120/* qxmx_prefix_cache: M2.6 implementation. See header for the design. */#include "qxmx_prefix_cache.h"
namespace qx {
bool prefix_cache::kv_has(block_hash h, const int32_t* ids, int n) { auto it = kv_.find(h.h); if (it == kv_.end()) return false; if (!ids_match(it->second, ids, n)) return false; /* collision -> miss */ it->second.lru_seq = tick(); return true;}
const kv_block* prefix_cache::kv_get(block_hash h, const int32_t* ids, int n) { auto it = kv_.find(h.h); if (it == kv_.end()) return nullptr; if (!ids_match(it->second, ids, n)) return nullptr; it->second.lru_seq = tick(); return &it->second;}
const kv_block* prefix_cache::kv_put(block_hash h, const int32_t* ids, int n, const uint8_t* k, const uint8_t* v) { auto& b = kv_[h.h]; bool was_present = !b.ids.empty(); b.ids.assign(ids, ids + n); b.k.assign(k, k + k_block_); b.v.assign(v, v + v_block_); b.lru_seq = tick(); if (!was_present) { kv_bytes_ += k_block_ + v_block_; /* Evict-from-anywhere LRU while over budget (KV blocks are ~14 MB; * the leaf-first refinement is M2.6.4). Never evict the block we just * put (it has the freshest lru_seq). */ while (kv_bytes_ > kv_budget_ && kv_.size() > 1) { if (!kv_evict_one()) break; } } else { /* replaced; size unchanged (same block granularity) */ } return &b;}
void prefix_cache::kv_ref(block_hash h) { auto it = kv_.find(h.h); if (it != kv_.end()) it->second.ref++;}
void prefix_cache::kv_deref(block_hash h) { auto it = kv_.find(h.h); if (it != kv_.end() && it->second.ref > 0) it->second.ref--;}
bool prefix_cache::ssm_has(block_hash h) const { return ssm_.find(h.h) != ssm_.end();}
const ssm_checkpoint* prefix_cache::ssm_get(block_hash h) { auto it = ssm_.find(h.h); if (it == ssm_.end()) return nullptr; it->second.lru_seq = tick(); return &it->second;}
const ssm_checkpoint* prefix_cache::ssm_put(block_hash h, int n_blocks, const uint8_t* state, const uint8_t* conv) { auto& c = ssm_[h.h]; bool was_present = !c.state.empty(); c.n_blocks = n_blocks; c.state.assign(state, state + (size_t)n_ssm_ * ssm_state_); c.conv.assign(conv, conv + (size_t)n_ssm_ * ssm_conv_); c.lru_seq = tick(); if (!was_present) { ssm_bytes_ += per_ckpt_; /* Evict-from-anywhere LRU while over budget. ~157 MB/checkpoint, so a * 16 GB budget holds ~100; the every-boundary placement policy relies * on this to keep host RAM bounded. Never evict the just-put one. */ while (ssm_bytes_ > ssm_budget_ && ssm_.size() > 1) { if (!ssm_evict_one()) break; } } return &c;}
bool prefix_cache::kv_evict_one() { /* Leaf-first: evict the LRU block with ref==0. A block is a "leaf" if no * other cached block has it as a parent in the hash chain -- but the chain * is implicit (parent hash mixed in), so we approximate: any ref==0 block * is evictable. True leaf-first (preserve prefix chains) would need the * parent->child map; deferred to M2.6.4's real LRU. For MVP this is * correct (never evicts a live block), just not optimal. */ uint64_t best_seq = UINT64_MAX; uint64_t best_key = 0; bool found = false; for (auto& [k, b] : kv_) { if (b.ref != 0) continue; if (b.lru_seq < best_seq) { best_seq = b.lru_seq; best_key = k; found = true; } } if (!found) return false; kv_.erase(best_key); kv_bytes_ -= k_block_ + v_block_; return true;}
bool prefix_cache::ssm_evict_one() { /* SSM checkpoints evict from ANYWHERE (SGLang pattern -- they're huge). */ uint64_t best_seq = UINT64_MAX; uint64_t best_key = 0; bool found = false; for (auto& [k, c] : ssm_) { if (c.lru_seq < best_seq) { best_seq = c.lru_seq; best_key = k; found = true; } } if (!found) return false; ssm_.erase(best_key); ssm_bytes_ -= per_ckpt_; return true;}
} // namespace qx