From efaa1ed73895eeb9169a07ae8ada05e5f8259756 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 29 Aug 2026 08:36:45 +0000 Subject: [PATCH 1/2] qwen4exp: fix QSA correctness defects and harden metadata loading Four runtime side fixes, no GGUF format change: - copy the indexer cache on a cross stream seq_cp. The update constructor never built ctx_idx, so a copied sequence kept the destination stream's old indexer keys while its cell metadata claimed the parent's positions. Reachable through the plain OpenAI n parameter with no flags. - key QSA blocks on the pair of sequence set and position bucket. Keying on position alone is only correct when the cells array holds one sequence, so under a unified cache the last writer won and a block's pooled key was built from another sequence's cells. - cut blocks on each cell's rank in sequence order under M-RoPE. Every token of an image shares one position, so all of them mapped to the same block slot and overwrote each other. - throw instead of aborting on malformed metadata. Eight GGML_ASSERT sites in load_arch_hparams were reachable from a hand edited GGUF, and a short fixed size array was accepted in silence. Also counts pooled blocks along ne1 for the norm, which avoids a pre existing CUDA abort at n_kv 262144 where the block count reached the 65535 gridDim.y limit. Perplexity is identical before and after over 1 and 16 chunks, and test-llama-archs -a qwen4exp passes on CPU and CUDA with roundtrip. --- src/llama-kv-cells.h | 13 +- src/llama-memory-hybrid-idx.cpp | 296 +++++++++++++++++++++++++++----- src/llama-memory-hybrid-idx.h | 4 +- src/models/qwen4exp.cpp | 74 ++++++-- 4 files changed, 333 insertions(+), 54 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 5167c037db3..9af6c7d21c4 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,15 @@ class llama_kv_cells { return seq[i].count(); } + // the whole set of sequences occupying the cell. Two cells that carry the + // same set are visible to exactly the same sequences, which is what lets a + // caller group cells without having to test one sequence at a time. + 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()); @@ -510,8 +521,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 d4e59d77e57..b272d86a2df 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,11 @@ 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), + // llama_kv_cache::update() is what applies a pending cross-stream seq_cp to the key buffer + // without this the copied sequence keeps the destination stream's old 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 +374,38 @@ 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); - } - } - } - - // 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); + // a block is r cells one sequence holds at consecutive indices, so it is keyed on + // (sequence set, index bucket). A unified cache shares one cells array between all + // sequences and every sequence counts positions from zero, so the bucket alone pools + // different sequences into one block. Cells with the same sequence set are visible to + // the same sequences, so that set is the coarsest key that never mixes two of them. + // A bucket whose cells disagree on the set splits into short groups, and a short group + // is dropped like the partial tail block. + // + // The index is the position, except that mrope gives every token of one image the same + // position. The cells are then ranked in sequence order and keyed on the rank instead. + GGML_ASSERT(r <= 64); + const uint64_t slots_full = r == 64 ? ~uint64_t(0) : ((uint64_t(1) << r) - 1); + + // only complete groups get an id, and a complete group owns r distinct cells, so at + // most n_kv/r ids are needed and no tensor changes shape + 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; // the cell in slot 0, which carries the block's rope position + std::vector grp_slots; + std::vector grp_bid; + std::vector bid_idx; // first index of the block + std::vector bid_cell; // one cell of the block, carries its sequence set + std::vector bid_slot0; + + // only used when the positions repeat + 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 +415,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(); + + // cells that all carry one sequence cannot mix, so skip the set test there + 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; + // two cells of one group wanted the same slot, so the index does not name them apart + bool dup = false; + + bool ranked = false; + + auto group_cells = [&]() { + // an incomplete block cannot be pooled; the bias below forces those tail cells in + // -1 means no usable block + 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 a whole image, so the position stops naming a + // cell. Rank the cells in sequence order and key on the rank: that is the token + // index the reference groups on. A rank only orders cells inside one sequence, so + // this needs a cache holding one; otherwise the short groups are dropped as above. + 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 order the mrope causal mask uses: pos first, then the 2d position + 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; + } + + // a rank is always below the cell count, so no bucket can fall outside + 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 + // ids in bucket order, so one sequence starting at position zero keeps its old numbering + 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); + + // the pooled key is roped with the block's first cell, as the reference does. Text + // repeats one position in every mrope section, and that is the bucket start already + 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 past the last real block, where the per-block bias is -inf. + // that id exists whenever such a cell does, because then the real blocks do not fill + // n_blocks. per-cell mode carries the -inf itself and only needs the gather in range + 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]; + + // the query in the index the blocks are keyed on + int64_t q = ubatch->pos[i]; + + if (ranked) { + // rank of the query: how many cells are at or before it in the order above + 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 +647,14 @@ 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) { + // a block outside this sequence, or past the last real one, is not selectable + 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 +665,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 e3472646d0f..e205d4afe48 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 abf6a0502fb..f28e8cc5609 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -6,6 +6,31 @@ #include #include +// Malformed metadata has to surface as a catchable error rather than an assert: +// GGML_ASSERT aborts the process, so an embedder that hands a user-supplied GGUF +// to llama_model_load_from_file has no way to report the bad file and carry on. +// These only reject values the arch genuinely cannot represent, so a file that +// loads today still loads. +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() into a fixed std::array already rejects an array that is too long, +// but a short one is copied as-is and leaves the tail of the destination alone. +// hparams is value-initialized, so that tail reads as zero and the n-gram hash +// quietly drops those positions instead of failing. Require the length the other +// hyperparameters imply. +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 +43,32 @@ 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); + // each of these sizes a GDN tensor dimension; zero makes the linear-attention layers degenerate + 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); + // every reference rejects a count of 1 outright: there is nothing to mix with a single + // stream. transformers configuration_qwen4_exp.py:196, the vLLM port's config.py:49 and + // the SGLang port's 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 +80,12 @@ 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) { + // the reference permits several PLE modules, but hparams holds one set of hash + // constants and the converter emits one, so more than one cannot be represented here + 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 +100,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 +112,12 @@ 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)); } + // the hash reads one multiplier per n-gram position and one range per head; + // a shorter array would leave the tail unwritten and hash to the wrong rows + 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 +141,8 @@ 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); + // the key is optional, but if present it divides the layer index below + 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 +605,14 @@ 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); + // rms_norm maps ne1 to gridDim.x but ne2 to gridDim.y, which stops at 65535, and a + // 262144 cell cache at ratio 4 has 65536 blocks. So count blocks along ne1 here. + 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. + // it flattens every row above ne0 into gridDim.x, so ne2 is safe here. 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 70a88dbb94859d7c9aa20ac5cb95a6d5a8828430 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 29 Aug 2026 08:43:45 +0000 Subject: [PATCH 2/2] qwen4exp: trim comments Cut the added comment lines from 64 to 19. Keeps the invariants that the code does not state, the reference citations, and the two traps, and drops restatements, spacers and measured numbers that the pull request description already carries. Also corrects two comments in llama-memory-hybrid-idx.h that became wrong when the update context started building ctx_idx. --- src/llama-kv-cells.h | 4 +-- src/llama-memory-hybrid-idx.cpp | 51 ++++++++------------------------- src/models/qwen4exp.cpp | 28 ++++-------------- 3 files changed, 19 insertions(+), 64 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 9af6c7d21c4..f13ee545ac3 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 whole set of sequences occupying the cell. Two cells that carry the - // same set are visible to exactly the same sequences, which is what lets a - // caller group cells without having to test one sequence at a time. + // 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()); diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index b272d86a2df..04f6ad6d455 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -50,8 +50,8 @@ 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 + // 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); @@ -300,8 +300,7 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( bool optimize) : llama_memory_hybrid_context(mem, lctx, optimize), mem(mem), - // llama_kv_cache::update() is what applies a pending cross-stream seq_cp to the key buffer - // without this the copied sequence keeps the destination stream's old indexer keys + // 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)) {} @@ -374,34 +373,23 @@ 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; - // a block is r cells one sequence holds at consecutive indices, so it is keyed on - // (sequence set, index bucket). A unified cache shares one cells array between all - // sequences and every sequence counts positions from zero, so the bucket alone pools - // different sequences into one block. Cells with the same sequence set are visible to - // the same sequences, so that set is the coarsest key that never mixes two of them. - // A bucket whose cells disagree on the set splits into short groups, and a short group - // is dropped like the partial tail block. - // - // The index is the position, except that mrope gives every token of one image the same - // position. The cells are then ranked in sequence order and keyed on the rank instead. + // 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); - // only complete groups get an id, and a complete group owns r distinct cells, so at - // most n_kv/r ids are needed and no tensor changes shape 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; // the cell in slot 0, which carries the block's rope position + std::vector grp_slot0; std::vector grp_slots; std::vector grp_bid; - std::vector bid_idx; // first index of the block - std::vector bid_cell; // one cell of the block, carries its sequence set + std::vector bid_idx; + std::vector bid_cell; std::vector bid_slot0; - // only used when the positions repeat std::vector order; std::vector rank; @@ -421,7 +409,6 @@ void llama_memory_hybrid_idx_context::set_input_qsa( bid_cell .clear(); bid_slot0.clear(); - // cells that all carry one sequence cannot mix, so skip the set test there int n_seq_present = 0; for (int sq = 0; sq < LLAMA_MAX_SEQ && n_seq_present < 2; ++sq) { @@ -436,14 +423,12 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // every cache path keeps the position below the cell window, so this stays false bool oor = false; - // two cells of one group wanted the same slot, so the index does not name them apart bool dup = false; bool ranked = false; auto group_cells = [&]() { - // an incomplete block cannot be pooled; the bias below forces those tail cells in - // -1 means no usable block + // -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); @@ -506,10 +491,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( group_cells(); - // mrope repeats one position across a whole image, so the position stops naming a - // cell. Rank the cells in sequence order and key on the rank: that is the token - // index the reference groups on. A rank only orders cells inside one sequence, so - // this needs a cache holding one; otherwise the short groups are dropped as above. + // 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); @@ -520,7 +502,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } - // same order the mrope causal mask uses: pos first, then the 2d position + // 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); @@ -540,7 +522,6 @@ void llama_memory_hybrid_idx_context::set_input_qsa( rank[order[k]] = (int32_t) k; } - // a rank is always below the cell count, so no bucket can fall outside ranked = true; group_cells(); @@ -548,7 +529,6 @@ void llama_memory_hybrid_idx_context::set_input_qsa( GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); - // ids in bucket order, so one sequence starting at position zero keeps its old numbering int32_t n_bid = 0; for (int64_t pb = 0; pb < n_blocks; ++pb) { @@ -567,8 +547,6 @@ void llama_memory_hybrid_idx_context::set_input_qsa( GGML_ASSERT(n_bid <= n_blocks); - // the pooled key is roped with the block's first cell, as the reference does. Text - // repeats one position in every mrope section, and that is the bucket start already 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] }; @@ -588,9 +566,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } - // unpooled cells point past the last real block, where the per-block bias is -inf. - // that id exists whenever such a cell does, because then the real blocks do not fill - // n_blocks. per-cell mode carries the -inf itself and only needs the gather in range + // 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) { @@ -611,11 +587,9 @@ void llama_memory_hybrid_idx_context::set_input_qsa( const int64_t i = s*n_tps + ii; const llama_seq_id seq_id = ubatch->seq_id[i][0]; - // the query in the index the blocks are keyed on int64_t q = ubatch->pos[i]; if (ranked) { - // rank of the query: how many cells are at or before it in the order above 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]; @@ -647,7 +621,6 @@ 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) { - // a block outside this sequence, or past the last real one, is not selectable if (b >= n_bid || !cells.seq_has((uint32_t) bid_cell[b], seq_id)) { cur_blk_bias[b] = -INFINITY; continue; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index f28e8cc5609..f2a189ef844 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -6,22 +6,14 @@ #include #include -// Malformed metadata has to surface as a catchable error rather than an assert: -// GGML_ASSERT aborts the process, so an embedder that hands a user-supplied GGUF -// to llama_model_load_from_file has no way to report the bad file and carry on. -// These only reject values the arch genuinely cannot represent, so a file that -// loads today still loads. +// 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() into a fixed std::array already rejects an array that is too long, -// but a short one is copied as-is and leaves the tail of the destination alone. -// hparams is value-initialized, so that tail reads as zero and the n-gram hash -// quietly drops those positions instead of failing. Require the length the other -// hyperparameters imply. +// 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); @@ -43,7 +35,6 @@ 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); - // each of these sizes a GDN tensor dimension; zero makes the linear-attention layers degenerate 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); @@ -53,9 +44,8 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { // 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); - // every reference rejects a count of 1 outright: there is nothing to mix with a single - // stream. transformers configuration_qwen4_exp.py:196, the vLLM port's config.py:49 and - // the SGLang port's configs/qwen4_exp.py:38 all raise on hc_count <= 1 + // 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)); @@ -81,8 +71,7 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { std::vector ple_layers; ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers); if (n_ple != 1) { - // the reference permits several PLE modules, but hparams holds one set of hash - // constants and the converter emits one, so more than one cannot be represented here + // 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)); } @@ -112,8 +101,6 @@ 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)); } - // the hash reads one multiplier per n-gram position and one range per head; - // a shorter array would leave the tail unwritten and hash to the wrong rows 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); @@ -141,7 +128,6 @@ 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); - // the key is optional, but if present it divides the layer index below 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); @@ -605,13 +591,11 @@ 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); - // rms_norm maps ne1 to gridDim.x but ne2 to gridDim.y, which stops at 65535, and a - // 262144 cell cache at ratio 4 has 65536 blocks. So count blocks along ne1 here. + // 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. - // it flattens every row above ne0 into gridDim.x, so ne2 is safe here. pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,