Skip to content

qwen4exp: fix QSA correctness defects and harden metadata loading #143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
danielhanchen wants to merge 2 commits into from
+287 −53
Closed

qwen4exp: fix QSA correctness defects and harden metadata loading #143

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
efaa1ed
qwen4exp: fix QSA correctness defects and harden metadata loading
shimmyshimmer Aug 29, 2026
70a88db
qwen4exp: trim comments
shimmyshimmer Aug 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions src/llama-kv-cells.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -301,6 +303,13 @@ class llama_kv_cells {
return seq[i].count();
}

// two cells with the same set are visible to exactly the same sequences
const seq_set_t & seq_set(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());
Expand Down Expand Up @@ -510,8 +519,6 @@ class llama_kv_cells {
//
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;

Expand Down
267 changes: 228 additions & 39 deletions src/llama-memory-hybrid-idx.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,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(
Expand Down Expand Up @@ -295,7 +299,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,
Expand Down Expand Up @@ -366,19 +373,27 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
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);
}
}
}
// 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);
Comment on lines +378 to +379

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Reject oversized compression ratios during loading

For a Qwen4Exp GGUF whose nonzero attention.compress_ratios entry exceeds 64, metadata loading still accepts the value, but the first sparse-attention graph input terminates the process at this GGML_ASSERT. This also undermines the new catchable-error handling for malformed metadata. Validate the upper bound in load_arch_hparams and throw a load error, or represent block occupancy without the 64-bit limit.

Useful? React with 👍 / 👎.


// 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);
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
Expand All @@ -388,61 +403,231 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
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);

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;

for (int64_t j = 0; j < n_kv; ++j) {
if (cells.is_empty(j)) {
continue;
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_set((uint32_t) grp_first[c]) == cells.seq_set((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;
}
}
};

const llama_pos p = cells.pos_get(j);
const int64_t b = p/r;
group_cells();

if (b >= n_blocks) {
oor = true;
continue;
// mrope repeats one position across an image, so rank cells instead of using the position
if (dup && ubatch->is_pos_2d() && one_seq) {

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Rank repeated positions after unified sequence copies

When a multimodal prompt is copied to another sequence in a unified cache, the shared prompt cells contain both sequence IDs, so n_seq_present is greater than one and one_seq is false even though all those cells have the same sequence set. This condition therefore skips ranking for the repeated M-RoPE image positions; group_cells() leaves the image cells competing for one slot, causing their block to be discarded as incomplete or pooled from only one image cell once QSA is active. The repeated-position ranking needs to handle copied/shared sequence sets, not only caches containing one sequence ID.

Useful? React with 👍 / 👎.

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);
}
}

blk_of[j] = (int32_t) b;
cur_blk_cells[b*r + (p%r)] = (int32_t) j;
filled[b]++;
// 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");

// 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
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 point at a dead block whose per-block bias is -inf
const int32_t dead_bid = n_bid < n_blocks ? n_bid : n_blocks - 1;

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;
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 ? 0 : blk_of[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];
const llama_pos q = ubatch->pos[i];

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 llama_pos tail_start = (q + 1)/r*r;
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] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f);
cur_blk_bias[b] = bid_idx[b] >= tail_start ? 1e9f : 0.0f;
}

continue;
Expand All @@ -453,9 +638,13 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
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);
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;
Expand Down
4 changes: 2 additions & 2 deletions src/llama-memory-hybrid-idx.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
// 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
Expand All @@ -148,7 +148,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
// 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
Expand Down
Loading
Loading