diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 5167c037db3..f13ee545ac3 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()); @@ -510,8 +519,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..04f6ad6d455 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 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..f2a189ef844 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);