From 3dd5505795a5cfc3d3055621cf8dc29c010a4e38 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 29 Aug 2026 09:02:11 +0000 Subject: [PATCH 1/8] qwen4exp: follow up fixes --- src/llama-kv-cells.h | 11 +- src/llama-memory-hybrid-idx.cpp | 267 +++++++++++++++++++++++++++----- src/llama-memory-hybrid-idx.h | 4 +- src/models/qwen4exp.cpp | 58 +++++-- 4 files changed, 287 insertions(+), 53 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index a4292c79e443..22d8c068354e 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -35,6 +35,8 @@ struct llama_kv_cell_ext { // TODO: add unit tests class llama_kv_cells { public: + using seq_set_t = std::bitset; + void reset() { for (uint32_t i = 0; i < pos.size(); ++i) { pos[i] = -1; @@ -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()); @@ -511,8 +520,6 @@ class llama_kv_cells { // std::vector shift; - using seq_set_t = std::bitset; - // the bitset seq[i] tells us which sequences are currently occupying the i-th cell std::vector seq; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d4e59d77e570..04f6ad6d4559 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -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( @@ -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, @@ -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); - // one pass per stream: cell j is a different token in each, so no mapping is shared - std::vector blk_of(n_kv); - std::vector filled(n_blocks); + std::vector blk_of(n_kv); + std::vector cell_grp(n_kv); + std::vector grp_head(n_blocks); + std::vector grp_next; + std::vector grp_first; + std::vector grp_slot0; + std::vector grp_slots; + std::vector grp_bid; + std::vector bid_idx; + std::vector bid_cell; + std::vector bid_slot0; + + std::vector order; + std::vector 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 @@ -388,52 +403,217 @@ 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) { + 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 @@ -441,8 +621,13 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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; @@ -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; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e3472646d0f6..e205d4afe484 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -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 @@ -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 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 diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index abf6a0502fbf..f2a189ef844a 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -6,6 +6,23 @@ #include #include +// 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 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,7 +128,7 @@ 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); } @@ -556,9 +591,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); From 4cbf055d07fad25ffd6b8c837458893ae2e73140 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 30 Aug 2026 10:25:22 +0000 Subject: [PATCH 2/8] -kvu NaN collapse fix Assisted-by: Claude --- src/llama-memory-hybrid-idx.cpp | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 04f6ad6d4559..0511aef46d81 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -566,8 +566,10 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } - // 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; + // 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]; @@ -630,6 +632,13 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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; } From 9bfa91ce27fd74d73adc5463b95e4acc3ad8a86d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 30 Aug 2026 10:47:31 +0000 Subject: [PATCH 3/8] indexer cache ext.x/ext.y restore fix Assisted-by: Claude --- src/llama-kv-cache.cpp | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 0e095f8b6b3d..fd7ce0bb6e4e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -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 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 From c76ea01d12bb181b7d65974b5a2aa5437c00056d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 30 Aug 2026 21:10:21 +0000 Subject: [PATCH 4/8] 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 --- src/llama-kv-cells.h | 6 ++++-- src/llama-memory-hybrid-idx.cpp | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 22d8c068354e..9a06997c37a1 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -303,8 +303,10 @@ 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 { + // the full set of sequences this cell is visible to. two cells with the same + // set are visible to exactly the same sequences. + // note: seq_get() above returns the single id and requires a one-sequence cell + const seq_set_t & seq_get_all(uint32_t i) const { assert(i < pos.size()); return seq[i]; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 0511aef46d81..751a0afc4e06 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -458,7 +458,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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)) { + if (one_seq || cells.seq_get_all((uint32_t) grp_first[c]) == cells.seq_get_all((uint32_t) j)) { g = c; break; } From 4ebbc37da03721fc9ad152b5e7d852722188ced1 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 31 Aug 2026 08:42:11 +0000 Subject: [PATCH 5/8] 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. --- src/llama-kv-cells.h | 4 +- src/llama-memory-hybrid-idx.cpp | 184 +++++++++++++++++--------------- src/llama-memory-hybrid-idx.h | 20 ++-- 3 files changed, 113 insertions(+), 95 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 9a06997c37a1..e9adffc09891 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -303,9 +303,7 @@ class llama_kv_cells { return seq[i].count(); } - // the full set of sequences this cell is visible to. two cells with the same - // set are visible to exactly the same sequences. - // note: seq_get() above returns the single id and requires a one-sequence cell + // 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()); diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 751a0afc4e06..93b468784a33 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -5,6 +5,7 @@ #include "llama-io.h" #include "llama-model.h" + #include #include #include @@ -265,88 +266,7 @@ llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const { return mem_idx.get(); } -// -// llama_memory_hybrid_idx_context -// - -// streams in each ubatch's slot info, matching get_k/get_v's `ns` -static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::slot_info_vec_t & sinfos) { - std::vector res; - res.reserve(sinfos.size()); - - for (const auto & sinfo : sinfos) { - res.push_back(sinfo.s1 - sinfo.s0 + 1); - } - - return res; -} - -llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) : - llama_memory_hybrid_context(status) {} - -llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem) : - llama_memory_hybrid_context(mem), - mem(mem), - // graph reservation walks a full context, and qwen4exp builds the sparse attention only when this is set - // without it the reserved worst case is the dense graph, so ggml-alloc must grow the buffer on the first decode - ns_ubatch(mem->get_mem_idx() == nullptr ? - std::vector() : std::vector{ mem->get_mem_idx()->get_n_stream() }), - ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : - new llama_kv_cache_context(mem->get_mem_idx())) {} - -llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( - llama_memory_hybrid_idx * mem, - llama_context * lctx, - bool optimize) : - llama_memory_hybrid_context(mem, lctx, optimize), - 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, - slot_info_vec_t sinfos_attn, - slot_info_vec_t sinfos_idx, - std::vector ubatches) : - // note: the base copies the ubatches; ctx_idx gets a copy of its own - llama_memory_hybrid_context(mem, std::move(sinfos_attn), ubatches), - mem(mem), - ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)), - ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : - new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {} - -bool llama_memory_hybrid_idx_context::next() { - if (ctx_idx) { - ctx_idx->next(); - } - - ++i_cur; - - return llama_memory_hybrid_context::next(); -} - -bool llama_memory_hybrid_idx_context::apply() { - bool res = llama_memory_hybrid_context::apply(); - - if (ctx_idx) { - res = res & ctx_idx->apply(); - } - - return res; -} - -const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const { - return static_cast(ctx_idx.get()); -} - -uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { - GGML_ASSERT(i_cur < ns_ubatch.size()); - - return ns_ubatch[i_cur]; -} - -void llama_memory_hybrid_idx_context::set_input_qsa( +void llama_memory_hybrid_idx::set_input_qsa( ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, @@ -355,7 +275,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( uint32_t ratio, bool blk_bias) const { GGML_ASSERT(ratio > 0); - GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + GGML_ASSERT(get_mem_idx() != nullptr); GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); @@ -378,6 +298,8 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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 blk_of(n_kv); std::vector cell_grp(n_kv); std::vector grp_head(n_blocks); @@ -398,7 +320,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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); + 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); @@ -661,3 +583,97 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } } + +// +// llama_memory_hybrid_idx_context +// + +// streams in each ubatch's slot info, matching get_k/get_v's `ns` +static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::slot_info_vec_t & sinfos) { + std::vector res; + res.reserve(sinfos.size()); + + for (const auto & sinfo : sinfos) { + res.push_back(sinfo.s1 - sinfo.s0 + 1); + } + + return res; +} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) : + llama_memory_hybrid_context(status) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem) : + llama_memory_hybrid_context(mem), + mem(mem), + // graph reservation walks a full context, and qwen4exp builds the sparse attention only when this is set + // without it the reserved worst case is the dense graph, so ggml-alloc must grow the buffer on the first decode + ns_ubatch(mem->get_mem_idx() == nullptr ? + std::vector() : std::vector{ mem->get_mem_idx()->get_n_stream() }), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx())) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + llama_context * lctx, + bool optimize) : + llama_memory_hybrid_context(mem, lctx, optimize), + 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, + slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, + std::vector ubatches) : + // note: the base copies the ubatches; ctx_idx gets a copy of its own + llama_memory_hybrid_context(mem, std::move(sinfos_attn), ubatches), + mem(mem), + ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {} + +bool llama_memory_hybrid_idx_context::next() { + if (ctx_idx) { + ctx_idx->next(); + } + + ++i_cur; + + return llama_memory_hybrid_context::next(); +} + +bool llama_memory_hybrid_idx_context::apply() { + bool res = llama_memory_hybrid_context::apply(); + + if (ctx_idx) { + res = res & ctx_idx->apply(); + } + + return res; +} + +const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const { + return static_cast(ctx_idx.get()); +} + +uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { + GGML_ASSERT(i_cur < ns_ubatch.size()); + + return ns_ubatch[i_cur]; +} + +void llama_memory_hybrid_idx_context::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(mem != nullptr); + + mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e205d4afe484..705189e7eb58 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -75,6 +75,18 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { 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 @@ -129,14 +141,6 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // 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; From 52bc9ff157fffd10b79e73b08f6661c9871f87e0 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 31 Aug 2026 08:42:28 +0000 Subject: [PATCH 6/8] 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/test-save-load-state.cpp | 66 +++++++++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 35c76905870d..6179e6c10848 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -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(tokens.data()), (int32_t) tokens.size()))) { + LOG_ERR("\n%s: failed to decode prompt\n", __func__); + return false; + } + + std::vector 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 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; From 21c13c48b265e9a76854666b45dba99cdc4e1efa Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 31 Aug 2026 23:15:48 +0000 Subject: [PATCH 7/8] 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 --- src/models/qwen4exp.cpp | 29 +++++++++++++++++++++-------- tests/test-llama-archs.cpp | 29 ++++++++++++++++++++++++++++- 2 files changed, 49 insertions(+), 9 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index f2a189ef844a..4d356fa5ad1e 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -134,6 +134,13 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { } } + // 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; @@ -159,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); } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 35a3286e4a1b..2836ae75dc62 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -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(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 ple_head_offsets(ple_n_heads); + std::vector 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({ 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({ 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 @@ -699,7 +723,10 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg 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()); + // qwen4exp has never worked on the meta mirror; -sm tensor on real devices does + const bool skip_meta = dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && arch == LLM_ARCH_QWEN4EXP; + + bool skip = !arch_supported(arch) || skip_meta || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); if (!skip) { if (logits_cpu.empty()) { model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode); From c961cd319dc13960c4da07e79bb592d13dddf516 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 1 Sep 2026 08:40:13 +0000 Subject: [PATCH 8/8] 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 --- src/llama-arch.cpp | 1 + tests/test-llama-archs.cpp | 4 +--- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5e61f61f7f0d..52124dc73995 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1142,6 +1142,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; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 2836ae75dc62..958142c313f7 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -723,10 +723,8 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg std::string status_nmse = "\033[1;33mSKIP\033[0m"; std::string status_roundtrip = "\033[1;33mSKIP\033[0m"; char nmse_str[12] = {0}; - // qwen4exp has never worked on the meta mirror; -sm tensor on real devices does - const bool skip_meta = dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && arch == LLM_ARCH_QWEN4EXP; - bool skip = !arch_supported(arch) || skip_meta || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); + bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); if (!skip) { if (logits_cpu.empty()) { model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode);