mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-09-09 06:19:14 +02:00
qwen4exp: fix seq_cp, block position keying, mtmd input, cuda abort, add tests (#27941)
* qwen4exp: follow up fixes * -kvu NaN collapse fix Assisted-by: Claude * indexer cache ext.x/ext.y restore fix Assisted-by: Claude * kv-cells: rename seq_set to seq_get_all seq_get is already taken by the single-id getter, so the suggested name cannot be overloaded on return type alone. Assisted-by: Claude * memory-hybrid-idx: implement set_input_qsa on the memory class The context held the whole implementation, where the pattern elsewhere is a thin context forwarding to the memory class, as llama_kv_cache_context does for set_input_kq_mask. The body reads no context state, so it moves unchanged and the context keeps a forwarder. Also shortens the seq_get_all comment as suggested. * tests: check that a sequence state survives a save/restore round-trip Saves seq 0, erases it, restores the blob and saves again, requiring the two blobs to match. Compares blobs rather than generated text, which cannot see a field dropped on the way back in. Note this passes on master for qwen4exp, so it does not demonstrate the ext.x/ext.y drop this PR fixes; reaching that needs 2D mrope content. * tests: give the synthetic qwen4exp a PLE so the state test bites has_cell_ext() is n_pos_per_embd() > 1 || ple_n_heads > 0, and the indexer cache sets rope_type = NONE, so without a PLE it serializes no cell ext at all and the round-trip test cannot see a dropped ext.x/ext.y. With one, removing the ext_set restore in state_read_meta fails the test: 198 of 335692 bytes differ, first at offset 282092. Loading such a model needed two fixes: - the row count of per_layer_token_embd came from require_weight(), which a model synthesised from metadata alone has no file to answer. Derive it from the head ranges and prefer the file's padded count where there is one. - the PLE conv history is a row of the recurrent cache, so a PLE on a full attention layer dereferenced a null p_l. Reject it at load time instead. The meta mirror is skipped for qwen4exp. It returned NaN logits before this fixture carried a PLE, which the nmse check passes since a NaN comparison is false, and aborts with one. -sm tensor on real devices works. Assisted-by: Claude * llama: disable -sm tensor for qwen4exp test-llama-archs skipped the tensor split for this arch from inside the test, so the arch still advertised support it does not have. Declare it in llm_arch_supports_sm_tensor instead and drop the test-side exception; the existing llm_arch_supports_sm_tensor branch then does the skipping. Assisted-by: Claude
This commit is contained in:
Daniel Han
GitHub
@@ -1143,6 +1143,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
|
||||
case LLM_ARCH_BAILINGMOE3:
|
||||
case LLM_ARCH_KIMI_K3:
|
||||
case LLM_ARCH_QWEN3TTS:
|
||||
case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs
|
||||
return false;
|
||||
default:
|
||||
return true;
|
||||
|
||||
@@ -2387,6 +2387,12 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
|
||||
|
||||
ubatch.seq_id_unq[0] = dest_seq_id;
|
||||
|
||||
// the ext as it was saved, to put back after apply_ubatch()
|
||||
std::vector<llama_kv_cell_ext> exts;
|
||||
if (has_cell_ext()) {
|
||||
exts.resize(cell_count);
|
||||
}
|
||||
|
||||
for (uint32_t i = 0; i < cell_count; ++i) {
|
||||
llama_pos pos;
|
||||
uint32_t n_seq_id;
|
||||
@@ -2410,6 +2416,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
|
||||
|
||||
// apply_ubatch() below restores ext.tok from the ubatch tokens
|
||||
ubatch.token[i] = ext.tok;
|
||||
|
||||
exts[i] = ext;
|
||||
}
|
||||
|
||||
// read the sequence id, but directly discard it - we will use dest_seq_id instead
|
||||
@@ -2461,6 +2469,14 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
|
||||
// see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350
|
||||
apply_ubatch(sinfo, ubatch);
|
||||
|
||||
// apply_ubatch() takes the 2D position from the ubatch, and that ubatch is built with this
|
||||
// cache's own n_pos_per_embd. a cache that does not use M-RoPE itself but mirrors one that
|
||||
// does (the qwen4exp QSA indexer) would drop x and y. put the saved ext back instead, which
|
||||
// is what the whole-context path below already does.
|
||||
for (uint32_t i = 0; i < (uint32_t) exts.size(); ++i) {
|
||||
cells.ext_set(sinfo.idxs[0][i], exts[i]);
|
||||
}
|
||||
|
||||
LLAMA_LOG_DEBUG("%s: cell_count = %d, dest_seq_id = %d\n", __func__, cell_count, dest_seq_id);
|
||||
|
||||
// DEBUG CHECK: verify that all cells were allocated and have correct seq_id and pos values
|
||||
|
||||
@@ -35,6 +35,8 @@ struct llama_kv_cell_ext {
|
||||
// TODO: add unit tests
|
||||
class llama_kv_cells {
|
||||
public:
|
||||
using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;
|
||||
|
||||
void reset() {
|
||||
for (uint32_t i = 0; i < pos.size(); ++i) {
|
||||
pos[i] = -1;
|
||||
@@ -301,6 +303,13 @@ public:
|
||||
return seq[i].count();
|
||||
}
|
||||
|
||||
// the full set of sequences this cell is visible to
|
||||
const seq_set_t & seq_get_all(uint32_t i) const {
|
||||
assert(i < pos.size());
|
||||
|
||||
return seq[i];
|
||||
}
|
||||
|
||||
// check if the cell contains seq_id
|
||||
bool seq_has(uint32_t i, llama_seq_id seq_id) const {
|
||||
assert(i < pos.size());
|
||||
@@ -511,8 +520,6 @@ private:
|
||||
//
|
||||
std::vector<llama_pos> shift;
|
||||
|
||||
using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;
|
||||
|
||||
// the bitset seq[i] tells us which sequences are currently occupying the i-th cell
|
||||
std::vector<seq_set_t> seq;
|
||||
|
||||
|
||||
+329
-115
@@ -5,6 +5,7 @@
|
||||
#include "llama-io.h"
|
||||
#include "llama-model.h"
|
||||
|
||||
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <cmath>
|
||||
@@ -50,6 +51,10 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
|
||||
std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
|
||||
hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
|
||||
|
||||
// the cached indexer keys are raw, rotation happens after pooling at read time, so a
|
||||
// K-shift must not rotate them while the stream copies in the same update still apply
|
||||
hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE;
|
||||
|
||||
LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
|
||||
|
||||
return new llama_kv_cache(
|
||||
@@ -261,6 +266,324 @@ llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const {
|
||||
return mem_idx.get();
|
||||
}
|
||||
|
||||
void llama_memory_hybrid_idx::set_input_qsa(
|
||||
ggml_tensor * cell_blk,
|
||||
ggml_tensor * blk_cells,
|
||||
ggml_tensor * blk_pos,
|
||||
ggml_tensor * bias,
|
||||
const llama_ubatch * ubatch,
|
||||
uint32_t ratio,
|
||||
bool blk_bias) const {
|
||||
GGML_ASSERT(ratio > 0);
|
||||
GGML_ASSERT(get_mem_idx() != nullptr);
|
||||
|
||||
GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
|
||||
|
||||
const int64_t n_kv = cell_blk->ne[0];
|
||||
const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch
|
||||
const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
|
||||
const int64_t n_tokens = ubatch->n_tokens;
|
||||
const int64_t r = ratio;
|
||||
|
||||
GGML_ASSERT(n_tokens % n_ns == 0);
|
||||
const int64_t n_tps = n_tokens/n_ns; // tokens per stream
|
||||
|
||||
int32_t * dst_cell_blk = (int32_t *) cell_blk->data;
|
||||
int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
|
||||
int32_t * dst_blk_pos = (int32_t *) blk_pos->data;
|
||||
float * dst_bias = (float *) bias->data;
|
||||
|
||||
// a block is keyed on (sequence set, index bucket): a unified cache counts every sequence
|
||||
// from zero, so the bucket alone would pool two sequences into one block
|
||||
GGML_ASSERT(r <= 64);
|
||||
const uint64_t slots_full = r == 64 ? ~uint64_t(0) : ((uint64_t(1) << r) - 1);
|
||||
|
||||
// TODO: this runs per ubatch and is O(n_kv) per stream, about 865 us at 33k context. the cost
|
||||
// is the per-cell scan rather than these allocations, so hoisting them buys nothing
|
||||
std::vector<int32_t> blk_of(n_kv);
|
||||
std::vector<int32_t> cell_grp(n_kv);
|
||||
std::vector<int32_t> grp_head(n_blocks);
|
||||
std::vector<int32_t> grp_next;
|
||||
std::vector<int32_t> grp_first;
|
||||
std::vector<int32_t> grp_slot0;
|
||||
std::vector<uint64_t> grp_slots;
|
||||
std::vector<int32_t> grp_bid;
|
||||
std::vector<int32_t> bid_idx;
|
||||
std::vector<int32_t> bid_cell;
|
||||
std::vector<int32_t> bid_slot0;
|
||||
|
||||
std::vector<int32_t> order;
|
||||
std::vector<int32_t> rank;
|
||||
|
||||
std::fill(dst_blk_pos, dst_blk_pos + 4*n_blocks*n_ns, 0);
|
||||
|
||||
for (int64_t s = 0; s < n_ns; ++s) {
|
||||
// ubatch index s*n_tps belongs to this stream; ask which cells array it uses
|
||||
const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];
|
||||
const auto & cells = get_mem_idx()->get_cells(seq_of_stream);
|
||||
|
||||
int32_t * cur_cell_blk = dst_cell_blk + s*n_kv;
|
||||
int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
|
||||
|
||||
std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0);
|
||||
|
||||
bid_idx .clear();
|
||||
bid_cell .clear();
|
||||
bid_slot0.clear();
|
||||
|
||||
int n_seq_present = 0;
|
||||
|
||||
for (int sq = 0; sq < LLAMA_MAX_SEQ && n_seq_present < 2; ++sq) {
|
||||
if (cells.seq_pos_min(sq) >= 0) {
|
||||
n_seq_present++;
|
||||
}
|
||||
}
|
||||
|
||||
const bool one_seq = n_seq_present <= 1;
|
||||
|
||||
// a cell no block covers needs its own -inf, which a per-block bias cannot carry
|
||||
// every cache path keeps the position below the cell window, so this stays false
|
||||
bool oor = false;
|
||||
|
||||
bool dup = false;
|
||||
|
||||
bool ranked = false;
|
||||
|
||||
auto group_cells = [&]() {
|
||||
// -1 means no usable block: an incomplete or short group cannot be pooled
|
||||
std::fill(blk_of.begin(), blk_of.end(), -1);
|
||||
std::fill(cell_grp.begin(), cell_grp.end(), -1);
|
||||
std::fill(grp_head.begin(), grp_head.end(), -1);
|
||||
|
||||
grp_next .clear();
|
||||
grp_first.clear();
|
||||
grp_slot0.clear();
|
||||
grp_slots.clear();
|
||||
grp_bid .clear();
|
||||
|
||||
oor = false;
|
||||
dup = false;
|
||||
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
if (cells.is_empty(j)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t idx = ranked ? rank[j] : cells.pos_get(j);
|
||||
const int64_t pb = idx/r;
|
||||
|
||||
if (pb >= n_blocks) {
|
||||
oor = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
int32_t g = -1;
|
||||
|
||||
for (int32_t c = grp_head[pb]; c >= 0; c = grp_next[c]) {
|
||||
if (one_seq || cells.seq_get_all((uint32_t) grp_first[c]) == cells.seq_get_all((uint32_t) j)) {
|
||||
g = c;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (g < 0) {
|
||||
g = (int32_t) grp_first.size();
|
||||
|
||||
grp_next .push_back(grp_head[pb]);
|
||||
grp_first.push_back((int32_t) j);
|
||||
grp_slot0.push_back(-1);
|
||||
grp_slots.push_back(0);
|
||||
grp_bid .push_back(-1);
|
||||
|
||||
grp_head[pb] = g;
|
||||
}
|
||||
|
||||
const uint64_t bit = uint64_t(1) << (idx%r);
|
||||
|
||||
dup |= (grp_slots[g] & bit) != 0;
|
||||
|
||||
cell_grp[j] = g;
|
||||
grp_slots[g] |= bit;
|
||||
|
||||
if (idx%r == 0) {
|
||||
grp_slot0[g] = (int32_t) j;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
group_cells();
|
||||
|
||||
// mrope repeats one position across an image, so rank cells instead of using the position
|
||||
if (dup && ubatch->is_pos_2d() && one_seq) {
|
||||
order.clear();
|
||||
order.reserve(n_kv);
|
||||
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
if (!cells.is_empty(j)) {
|
||||
order.push_back((int32_t) j);
|
||||
}
|
||||
}
|
||||
|
||||
// same total order the mrope causal mask uses: pos, then ext.y, then ext.x
|
||||
std::sort(order.begin(), order.end(), [&cells](int32_t a, int32_t b) {
|
||||
const llama_pos pa = cells.pos_get(a);
|
||||
const llama_pos pb = cells.pos_get(b);
|
||||
|
||||
if (pa != pb) {
|
||||
return pa < pb;
|
||||
}
|
||||
|
||||
const auto & ea = cells.ext_get(a);
|
||||
|
||||
return cells.ext_get(b).is_2d_gt(ea.x, ea.y);
|
||||
});
|
||||
|
||||
rank.assign(n_kv, -1);
|
||||
|
||||
for (int64_t k = 0; k < (int64_t) order.size(); ++k) {
|
||||
rank[order[k]] = (int32_t) k;
|
||||
}
|
||||
|
||||
ranked = true;
|
||||
|
||||
group_cells();
|
||||
}
|
||||
|
||||
GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window");
|
||||
|
||||
int32_t n_bid = 0;
|
||||
|
||||
for (int64_t pb = 0; pb < n_blocks; ++pb) {
|
||||
for (int32_t g = grp_head[pb]; g >= 0; g = grp_next[g]) {
|
||||
if (grp_slots[g] != slots_full) {
|
||||
continue;
|
||||
}
|
||||
|
||||
grp_bid[g] = n_bid++;
|
||||
|
||||
bid_idx .push_back((int32_t) (pb*r));
|
||||
bid_cell .push_back(grp_first[g]);
|
||||
bid_slot0.push_back(grp_slot0[g]);
|
||||
}
|
||||
}
|
||||
|
||||
GGML_ASSERT(n_bid <= n_blocks);
|
||||
|
||||
for (int32_t b = 0; b < n_bid; ++b) {
|
||||
int32_t sec_pos[4] = { bid_idx[b], bid_idx[b], bid_idx[b], bid_idx[b] };
|
||||
|
||||
if (ranked) {
|
||||
const int32_t c = bid_slot0[b];
|
||||
const llama_pos p = cells.pos_get(c);
|
||||
const auto & e = cells.ext_get(c);
|
||||
|
||||
sec_pos[0] = p;
|
||||
sec_pos[1] = e.y;
|
||||
sec_pos[2] = e.x;
|
||||
sec_pos[3] = p;
|
||||
}
|
||||
|
||||
for (int64_t sec = 0; sec < 4; ++sec) {
|
||||
dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = sec_pos[sec];
|
||||
}
|
||||
}
|
||||
|
||||
// unpooled cells all point at one spare block. a spare block exists only when some
|
||||
// cell is unpooled: n_bid == n_blocks means every cell sits in a full block.
|
||||
const bool have_dead = n_bid < n_blocks;
|
||||
const int32_t dead_bid = have_dead ? n_bid : n_blocks - 1;
|
||||
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
const int32_t g = cell_grp[j];
|
||||
|
||||
blk_of[j] = g < 0 ? -1 : grp_bid[g];
|
||||
|
||||
if (blk_of[j] >= 0) {
|
||||
const int64_t idx = ranked ? rank[j] : cells.pos_get(j);
|
||||
|
||||
cur_blk_cells[blk_of[j]*r + (idx%r)] = (int32_t) j;
|
||||
}
|
||||
|
||||
cur_cell_blk[j] = blk_of[j] < 0 ? dead_bid : blk_of[j];
|
||||
}
|
||||
|
||||
for (int64_t ii = 0; ii < n_tps; ++ii) {
|
||||
const int64_t i = s*n_tps + ii;
|
||||
const llama_seq_id seq_id = ubatch->seq_id[i][0];
|
||||
|
||||
int64_t q = ubatch->pos[i];
|
||||
|
||||
if (ranked) {
|
||||
const llama_pos qt = ubatch->pos[i];
|
||||
const llama_pos qy = ubatch->pos[i + n_tokens];
|
||||
const llama_pos qx = ubatch->pos[i + n_tokens*2];
|
||||
|
||||
int64_t lo = 0;
|
||||
int64_t hi = (int64_t) order.size();
|
||||
|
||||
while (lo < hi) {
|
||||
const int64_t mid = (lo + hi)/2;
|
||||
const int32_t c = order[mid];
|
||||
const llama_pos pc = cells.pos_get(c);
|
||||
|
||||
if (pc < qt || (pc == qt && !cells.ext_get(c).is_2d_gt(qx, qy))) {
|
||||
lo = mid + 1;
|
||||
} else {
|
||||
hi = mid;
|
||||
}
|
||||
}
|
||||
|
||||
q = lo - 1;
|
||||
}
|
||||
|
||||
// the tail is an incomplete block and is always visible, as in the reference
|
||||
const int64_t tail_start = (q + 1)/r*r;
|
||||
|
||||
if (blk_bias) {
|
||||
// a block sits wholly inside or outside the tail, so one value covers it
|
||||
// the caller adds the attention mask, which drops empty, foreign and future cells
|
||||
float * cur_blk_bias = dst_bias + i*n_blocks;
|
||||
|
||||
for (int64_t b = 0; b < n_blocks; ++b) {
|
||||
if (b >= n_bid || !cells.seq_has((uint32_t) bid_cell[b], seq_id)) {
|
||||
cur_blk_bias[b] = -INFINITY;
|
||||
continue;
|
||||
}
|
||||
|
||||
// finite, so it can never meet a -inf and produce a nan
|
||||
cur_blk_bias[b] = bid_idx[b] >= tail_start ? 1e9f : 0.0f;
|
||||
}
|
||||
|
||||
// the spare block holds the unpooled cells, which are the incomplete tail, so
|
||||
// it gets the tail value. it must stay finite: a sequence with fewer than
|
||||
// `ratio` cells owns no full block, and a row of -inf only gives a nan.
|
||||
if (have_dead) {
|
||||
cur_blk_bias[dead_bid] = 1e9f;
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
float * cur_bias = dst_bias + i*n_kv;
|
||||
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
float v = -INFINITY;
|
||||
|
||||
if (!cells.is_empty(j) && cells.seq_has(j, seq_id)) {
|
||||
const int64_t idx = ranked ? rank[j] : cells.pos_get(j);
|
||||
|
||||
if (idx <= q) {
|
||||
// finite, so it can never meet a -inf and produce a nan
|
||||
v = idx >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
cur_bias[j] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// llama_memory_hybrid_idx_context
|
||||
//
|
||||
@@ -295,7 +618,10 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
|
||||
llama_context * lctx,
|
||||
bool optimize) :
|
||||
llama_memory_hybrid_context(mem, lctx, optimize),
|
||||
mem(mem) {}
|
||||
mem(mem),
|
||||
// update() applies a pending cross-stream seq_cp, else the copy keeps stale indexer keys
|
||||
ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
|
||||
mem->get_mem_idx()->init_update(lctx, optimize)) {}
|
||||
|
||||
llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
|
||||
llama_memory_hybrid_idx * mem,
|
||||
@@ -347,119 +673,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
|
||||
const llama_ubatch * ubatch,
|
||||
uint32_t ratio,
|
||||
bool blk_bias) const {
|
||||
GGML_ASSERT(ratio > 0);
|
||||
GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
|
||||
GGML_ASSERT(mem != nullptr);
|
||||
|
||||
GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
|
||||
|
||||
const int64_t n_kv = cell_blk->ne[0];
|
||||
const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch
|
||||
const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
|
||||
const int64_t n_tokens = ubatch->n_tokens;
|
||||
const int64_t r = ratio;
|
||||
|
||||
GGML_ASSERT(n_tokens % n_ns == 0);
|
||||
const int64_t n_tps = n_tokens/n_ns; // tokens per stream
|
||||
|
||||
int32_t * dst_cell_blk = (int32_t *) cell_blk->data;
|
||||
int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
|
||||
int32_t * dst_blk_pos = (int32_t *) blk_pos->data;
|
||||
float * dst_bias = (float *) bias->data;
|
||||
|
||||
// block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio
|
||||
// all mrope sections carry it: exact for text, approximate for images
|
||||
for (int64_t sec = 0; sec < 4; ++sec) {
|
||||
for (int64_t s = 0; s < n_ns; ++s) {
|
||||
for (int64_t b = 0; b < n_blocks; ++b) {
|
||||
dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// one pass per stream: cell j is a different token in each, so no mapping is shared
|
||||
std::vector<int32_t> blk_of(n_kv);
|
||||
std::vector<int32_t> filled(n_blocks);
|
||||
|
||||
for (int64_t s = 0; s < n_ns; ++s) {
|
||||
// ubatch index s*n_tps belongs to this stream; ask which cells array it uses
|
||||
const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];
|
||||
const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream);
|
||||
|
||||
int32_t * cur_cell_blk = dst_cell_blk + s*n_kv;
|
||||
int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
|
||||
|
||||
// an incomplete block cannot be pooled; the bias below forces those tail cells in
|
||||
// -1 means no usable block, and block 0 only keeps the gather in range
|
||||
std::fill(blk_of.begin(), blk_of.end(), -1);
|
||||
std::fill(filled.begin(), filled.end(), 0);
|
||||
std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0);
|
||||
|
||||
// a cell no block covers needs its own -inf, which a per-block bias cannot carry
|
||||
// every cache path keeps the position below the cell window, so this stays false
|
||||
bool oor = false;
|
||||
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
if (cells.is_empty(j)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const llama_pos p = cells.pos_get(j);
|
||||
const int64_t b = p/r;
|
||||
|
||||
if (b >= n_blocks) {
|
||||
oor = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
blk_of[j] = (int32_t) b;
|
||||
cur_blk_cells[b*r + (p%r)] = (int32_t) j;
|
||||
filled[b]++;
|
||||
}
|
||||
|
||||
GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window");
|
||||
|
||||
// per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it
|
||||
// per-cell mode carries that -inf itself and only needs the gather in range
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) {
|
||||
blk_of[j] = -1;
|
||||
}
|
||||
cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j];
|
||||
}
|
||||
|
||||
for (int64_t ii = 0; ii < n_tps; ++ii) {
|
||||
const int64_t i = s*n_tps + ii;
|
||||
const llama_seq_id seq_id = ubatch->seq_id[i][0];
|
||||
const llama_pos q = ubatch->pos[i];
|
||||
|
||||
// the tail is an incomplete block and is always visible, as in the reference
|
||||
const llama_pos tail_start = (q + 1)/r*r;
|
||||
|
||||
if (blk_bias) {
|
||||
// a block sits wholly inside or outside the tail, so one value covers it
|
||||
// the caller adds the attention mask, which drops empty, foreign and future cells
|
||||
float * cur_blk_bias = dst_bias + i*n_blocks;
|
||||
|
||||
for (int64_t b = 0; b < n_blocks; ++b) {
|
||||
// finite, so it can never meet a -inf and produce a nan
|
||||
cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f);
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
float * cur_bias = dst_bias + i*n_kv;
|
||||
|
||||
for (int64_t j = 0; j < n_kv; ++j) {
|
||||
float v = -INFINITY;
|
||||
|
||||
if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) {
|
||||
// finite, so it can never meet a -inf and produce a nan
|
||||
v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f);
|
||||
}
|
||||
|
||||
cur_bias[j] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias);
|
||||
}
|
||||
|
||||
@@ -75,6 +75,18 @@ public:
|
||||
|
||||
llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer
|
||||
|
||||
// block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache.
|
||||
// Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout:
|
||||
// cell_blk I32 [n_kv, ns] block each cell belongs to
|
||||
// blk_cells I32 [ratio*n_blocks, ns] cells making up each block
|
||||
// blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token
|
||||
// bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible
|
||||
// blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns]
|
||||
// the caller then adds the attention mask, the only part of the bias that varies within a block
|
||||
void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
|
||||
ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
|
||||
bool blk_bias) const;
|
||||
|
||||
private:
|
||||
// forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step
|
||||
// seq_id < 0 drops the whole context, as the caches themselves do on a failed restore
|
||||
@@ -123,20 +135,12 @@ public:
|
||||
// llama_memory_hybrid_idx_context specific API
|
||||
//
|
||||
|
||||
// nullptr with no indexer, and for the update context, which builds no sparse graph
|
||||
// nullptr with no indexer
|
||||
const llama_kv_cache_context * get_idx() const;
|
||||
|
||||
// streams in the current slot info, the `ns` of get_k/get_v; 1 if unified
|
||||
uint32_t get_n_stream() const;
|
||||
|
||||
// block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache.
|
||||
// Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout:
|
||||
// cell_blk I32 [n_kv, ns] block each cell belongs to
|
||||
// blk_cells I32 [ratio*n_blocks, ns] cells making up each block
|
||||
// blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token
|
||||
// bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible
|
||||
// blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns]
|
||||
// the caller then adds the attention mask, the only part of the bias that varies within a block
|
||||
void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
|
||||
ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
|
||||
bool blk_bias) const;
|
||||
@@ -148,7 +152,7 @@ private:
|
||||
// declared first, so it is initialised while sinfos_idx is still intact
|
||||
const std::vector<uint32_t> ns_ubatch;
|
||||
|
||||
// null unless the model has an indexer and this is a batch or full context
|
||||
// null unless the model has an indexer
|
||||
const llama_memory_context_ptr ctx_idx;
|
||||
|
||||
// mirrors the base class's ubatch cursor, which is private there
|
||||
|
||||
+70
-19
@@ -6,6 +6,23 @@
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
|
||||
// bad metadata must be catchable: GGML_ASSERT aborts the whole process
|
||||
static void qwen4exp_require_nonzero(const llama_model_loader & ml, llm_kv kid, uint32_t value) {
|
||||
if (value == 0) {
|
||||
throw std::runtime_error(format("%s must be greater than zero, got %u", ml.llm_kv(kid).c_str(), value));
|
||||
}
|
||||
}
|
||||
|
||||
// get_arr() copies a short array as-is, leaving a zero tail the n-gram hash silently drops
|
||||
static void qwen4exp_require_arr_len(llama_model_loader & ml, llm_kv kid, uint32_t n_min) {
|
||||
uint32_t n_arr = 0;
|
||||
ml.get_arr_n(kid, n_arr, true);
|
||||
if (n_arr < n_min) {
|
||||
throw std::runtime_error(format("%s has %u entries, but at least %u are required",
|
||||
ml.llm_kv(kid).c_str(), n_arr, n_min));
|
||||
}
|
||||
}
|
||||
|
||||
void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
|
||||
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
|
||||
@@ -18,21 +35,30 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);
|
||||
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
|
||||
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
|
||||
GGML_ASSERT(hparams.ssm_d_conv > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 &&
|
||||
hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
|
||||
|
||||
// HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank)
|
||||
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult);
|
||||
ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
|
||||
GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0);
|
||||
// a count of 1 has nothing to mix: transformers configuration_qwen4_exp.py:196, vLLM
|
||||
// config.py:49 and SGLang configs/qwen4_exp.py:38 all raise on hc_count <= 1
|
||||
if (hparams.dsv4_hc_mult <= 1) {
|
||||
throw std::runtime_error(format("%s must be greater than one, got %u",
|
||||
ml.llm_kv(LLM_KV_HYPER_CONNECTION_COUNT).c_str(), hparams.dsv4_hc_mult));
|
||||
}
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
|
||||
hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
|
||||
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
|
||||
ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
|
||||
GGML_ASSERT(hparams.indexer_n_head > 0
|
||||
&& hparams.indexer_head_size > 0
|
||||
&& hparams.indexer_top_k > 0);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
|
||||
ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false);
|
||||
|
||||
// PLE n-gram hash embeddings; if the key group is absent every field stays zero
|
||||
@@ -44,7 +70,11 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
if (n_ple > 0) {
|
||||
std::vector<uint32_t> ple_layers;
|
||||
ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers);
|
||||
GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer");
|
||||
if (n_ple != 1) {
|
||||
// hparams holds one set of hash constants, so several PLE modules cannot be represented
|
||||
throw std::runtime_error(format("%s lists %u layers, but only one PLE layer is supported",
|
||||
ml.llm_kv(LLM_KV_PLE_LAYERS).c_str(), n_ple));
|
||||
}
|
||||
for (uint32_t il : ple_layers) {
|
||||
if (il >= hparams.n_layer_all) {
|
||||
throw std::runtime_error(format("PLE layer %u is out of range", il));
|
||||
@@ -59,7 +89,8 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
// optional: files written before this key fall back to the EOS token
|
||||
ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false);
|
||||
ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);
|
||||
GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);
|
||||
|
||||
hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
|
||||
hparams.ple_head_dim = hparams.n_embd_per_layer;
|
||||
@@ -70,6 +101,10 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads));
|
||||
}
|
||||
|
||||
qwen4exp_require_arr_len(ml, LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_ngram_size);
|
||||
qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_OFFSETS, hparams.ple_n_heads);
|
||||
qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_VOCAB_SIZES, hparams.ple_n_heads);
|
||||
|
||||
ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers);
|
||||
|
||||
// the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses
|
||||
@@ -93,12 +128,19 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
|
||||
if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
|
||||
uint32_t full_attn_interval = 4;
|
||||
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
|
||||
GGML_ASSERT(full_attn_interval > 0);
|
||||
qwen4exp_require_nonzero(ml, LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval);
|
||||
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
|
||||
hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
|
||||
}
|
||||
}
|
||||
|
||||
// the PLE conv history is a row of the recurrent cache, which linear layers alone have
|
||||
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
|
||||
if (hparams.is_ple(i) && !hparams.is_recr(i)) {
|
||||
throw std::runtime_error(format("PLE layer %u is not a linear attention layer", i));
|
||||
}
|
||||
}
|
||||
|
||||
switch (hparams.n_layer()) {
|
||||
case 48: type = LLM_TYPE_A3B; break;
|
||||
default: type = LLM_TYPE_UNKNOWN;
|
||||
@@ -124,18 +166,24 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
|
||||
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
|
||||
}
|
||||
|
||||
// flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back
|
||||
// flat [ple_head_dim, n_rows] gather target
|
||||
if (hparams.ple_n_heads > 0) {
|
||||
const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
|
||||
const auto & ple_w = ml.require_weight(ple_name.c_str());
|
||||
const int64_t ple_rows = ple_w.tensor->ne[1];
|
||||
|
||||
// sanity check
|
||||
// the head ranges are what the gather indexes, so they set the minimum row count
|
||||
int64_t ple_rows = 0;
|
||||
for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {
|
||||
if ((int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h] > ple_rows) {
|
||||
throw std::runtime_error(format("PLE head %u range exceeds the %" PRId64 " table rows", h, ple_rows));
|
||||
}
|
||||
ple_rows = std::max(ple_rows, (int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h]);
|
||||
}
|
||||
|
||||
// the converter pads the table; a model synthesised from metadata has no tensor to ask
|
||||
const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
|
||||
if (const auto * ple_w = ml.get_weight(ple_name.c_str())) {
|
||||
if (ple_w->tensor->ne[1] < ple_rows) {
|
||||
throw std::runtime_error(format("%s has %" PRId64 " rows, too few for the PLE head ranges (%" PRId64 ")",
|
||||
ple_name.c_str(), ple_w->tensor->ne[1], ple_rows));
|
||||
}
|
||||
ple_rows = ple_w->tensor->ne[1];
|
||||
}
|
||||
|
||||
per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"),
|
||||
{ hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY);
|
||||
}
|
||||
@@ -556,9 +604,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
|
||||
pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r);
|
||||
cb(pooled, "indexer_k_pooled", il);
|
||||
|
||||
// count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536
|
||||
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks*n_stream, 1);
|
||||
pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
|
||||
|
||||
// rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after.
|
||||
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream);
|
||||
pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
|
||||
pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr,
|
||||
n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
|
||||
@@ -254,6 +254,30 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8));
|
||||
// without this the QSA layers fall back to dense and go uncovered
|
||||
ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>(n_layer, 4));
|
||||
|
||||
// has_cell_ext() needs ple_n_heads here: the indexer cache serializes no ext without it
|
||||
const uint32_t ple_ngram_size = 3;
|
||||
const uint32_t ple_heads_per_ngram = 2;
|
||||
const uint32_t ple_n_heads = (ple_ngram_size - 1)*ple_heads_per_ngram;
|
||||
GGML_ASSERT(n_embd % ple_n_heads == 0);
|
||||
const uint32_t ple_head_dim = n_embd/ple_n_heads;
|
||||
|
||||
std::vector<uint64_t> ple_head_offsets(ple_n_heads);
|
||||
std::vector<uint64_t> ple_head_vocab_sizes(ple_n_heads, n_vocab);
|
||||
for (uint32_t h = 0; h < ple_n_heads; h++) {
|
||||
ple_head_offsets[h] = uint64_t(h)*n_vocab;
|
||||
}
|
||||
|
||||
// the PLE history lives in the recurrent cache, so it must sit on a linear attention layer
|
||||
ms.add_kv(LLM_KV_PLE_LAYERS, std::vector<uint32_t>({ 0 }));
|
||||
ms.add_kv(LLM_KV_PLE_NGRAM_SIZE, ple_ngram_size);
|
||||
ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, ple_heads_per_ngram);
|
||||
ms.add_kv(LLM_KV_PLE_CONV_KERNEL, uint32_t(4));
|
||||
ms.add_kv(LLM_KV_PLE_EOS_TOKEN_ID, uint32_t(0));
|
||||
ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, ple_head_dim);
|
||||
ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector<uint64_t>({ 1, 3, 5 }));
|
||||
ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS, ple_head_offsets);
|
||||
ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, ple_head_vocab_sizes);
|
||||
}
|
||||
|
||||
// minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused
|
||||
@@ -709,6 +733,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in
|
||||
std::string status_nmse = "\033[1;33mSKIP\033[0m";
|
||||
std::string status_roundtrip = "\033[1;33mSKIP\033[0m";
|
||||
char nmse_str[12] = {0};
|
||||
|
||||
bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty());
|
||||
if (!skip) {
|
||||
if (logits_cpu.empty()) {
|
||||
|
||||
@@ -449,7 +449,66 @@ static bool test_seq_cp_scatter(struct llama_model * model, const struct common_
|
||||
}
|
||||
|
||||
|
||||
// Run the full save/load test suite (tests 1-7) for a single model.
|
||||
// Test 8: state blob round-trip
|
||||
// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens
|
||||
static bool test_state_roundtrip(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
|
||||
auto params_ctx = common_context_params_to_llama(params);
|
||||
auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
|
||||
|
||||
LOG("\n=== Test 8: state blob round-trip ===\n");
|
||||
|
||||
if (llama_decode(ctx.get(), llama_batch_get_one(const_cast<llama_token *>(tokens.data()), (int32_t) tokens.size()))) {
|
||||
LOG_ERR("\n%s: failed to decode prompt\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> blob_a(llama_state_seq_get_size(ctx.get(), 0));
|
||||
const size_t n_a = llama_state_seq_get_data(ctx.get(), blob_a.data(), blob_a.size(), 0);
|
||||
if (n_a != blob_a.size()) {
|
||||
LOG_ERR("\n%s: saved %zu bytes, expected %zu\n", __func__, n_a, blob_a.size());
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
|
||||
LOG_ERR("\n%s: failed to erase seq 0\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (llama_state_seq_set_data(ctx.get(), blob_a.data(), blob_a.size(), 0) != blob_a.size()) {
|
||||
LOG_ERR("\n%s: failed to restore seq 0\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> blob_b(llama_state_seq_get_size(ctx.get(), 0));
|
||||
const size_t n_b = llama_state_seq_get_data(ctx.get(), blob_b.data(), blob_b.size(), 0);
|
||||
if (n_b != n_a) {
|
||||
LOG_ERR("\n%s: re-saved %zu bytes, expected %zu\n", __func__, n_b, n_a);
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t n_diff = 0;
|
||||
size_t i_diff = 0;
|
||||
for (size_t i = 0; i < n_a; i++) {
|
||||
if (blob_a[i] != blob_b[i]) {
|
||||
if (n_diff == 0) {
|
||||
i_diff = i;
|
||||
}
|
||||
n_diff++;
|
||||
}
|
||||
}
|
||||
|
||||
if (n_diff > 0) {
|
||||
LOG_ERR("\n%s: state changed across a restore: %zu of %zu bytes differ, first at offset %zu\n",
|
||||
__func__, n_diff, n_a, i_diff);
|
||||
return false;
|
||||
}
|
||||
|
||||
LOG("\nPASS\n");
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
// Run the full save/load test suite (tests 1-8) for a single model.
|
||||
// Returns true if all tests pass, false otherwise.
|
||||
static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
|
||||
struct common_params params = base_params;
|
||||
@@ -526,6 +585,11 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const
|
||||
return false;
|
||||
}
|
||||
|
||||
// Test 8: state blob round-trip
|
||||
if (!test_state_roundtrip(model, params, tokens)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
LOG("\nAll tests passed.\n");
|
||||
|
||||
return true;
|
||||
|
||||
Reference in New Issue
Block a user