diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4b52ba5ea3..31832aa144 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3000,9 +3000,8 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } -// Restrict a KQ mask to the positions named by top_k. Lifted verbatim out of the -// DSA build_attn below so the sparse-attention architectures that do not use MLA -// can share it; the node sequence is unchanged. +// Restrict a KQ mask to the positions named by top_k. Lifted verbatim from the DSA +// build_attn below so non-MLA sparse architectures can share it; node sequence unchanged. ggml_tensor * llm_graph_context::build_attn_mask_top_k( ggml_tensor * kq_mask, ggml_tensor * top_k) const { diff --git a/src/llama-graph.h b/src/llama-graph.h index 24b6100a5a..deb0bd444d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1206,8 +1206,8 @@ struct llm_graph_context { float kq_scale, int il) const; - // as above, but attending only to the positions named by top_k. Used by - // architectures whose sparse attention is plain GQA rather than MLA. + // as above, but attending only to positions named by top_k; for sparse attention + // architectures that are plain GQA rather than MLA ggml_tensor * build_attn( llm_graph_input_attn_kv * inp, ggml_tensor * wo, diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 2dd77ed3e3..7a0c4bd721 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1811,19 +1811,17 @@ void llama_kv_cache::set_input_qsa( int32_t * dst_blk_pos = (int32_t *) blk_pos->data; float * dst_bias = (float *) bias->data; - // block b is positions [b*ratio, (b+1)*ratio), so its first token sits at - // b*ratio. The three mrope sections all carry that position: exact for text, - // an approximation for the interleaved t/h/w positions of image tokens. + // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio. All three + // mrope sections carry that position: exact for text, approximate for image tokens. for (int64_t s = 0; s < 4; ++s) { for (int64_t b = 0; b < n_blocks; ++b) { dst_blk_pos[s*n_blocks + b] = (int32_t) (b*r); } } - // a block that is not completely populated cannot be pooled. Those cells are - // exactly the tail of the sequence, which the bias below forces in whatever - // score they carry, so they are pointed at block 0 only to keep the gather - // in range. -1 marks a cell with no usable block at all. + // an incompletely populated block cannot be pooled. Those cells are the tail, which + // the bias below forces in anyway, so they point at block 0 just to keep the gather + // in range. -1 marks a cell with no usable block. std::vector blk_of(n_kv, -1); std::vector filled(n_blocks, 0); @@ -1857,9 +1855,8 @@ void llama_kv_cache::set_input_qsa( const llama_seq_id seq_id = ubatch->seq_id[i][0]; const llama_pos q = ubatch->pos[i]; - // everything from here on is inside an incomplete block and is always - // attended to, which is what makes the whole selection land on block - // boundaries the way the reference implementation does + // from here on we are inside an incomplete block, always attended to, which is what + // lands the selection on block boundaries as the reference does const llama_pos tail_start = (q + 1)/r*r; for (int64_t j = 0; j < n_kv; ++j) { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 8511dc4186..37616e7b35 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -216,13 +216,12 @@ public: void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; - // block-compressed sparse attention (qwen4exp QSA), computed over this - // cache's cells. Blocks are cuts of the *position* line, not of the cell - // array, so nothing here assumes the cache is laid out contiguously: - // cell_blk I32 [n_kv] block each cell belongs to - // blk_cells I32 [ratio*n_blocks] the cells making up each block - // blk_pos I32 [4*n_blocks] mrope position rows of each block's first token - // bias F32 [n_kv, n_tokens] -inf where invisible, large where always visible + // block-compressed sparse attention (qwen4exp QSA) over this cache's cells. Blocks + // cut the *position* line, not the cell array, so nothing assumes contiguous layout: + // cell_blk I32 [n_kv] block each cell belongs to + // blk_cells I32 [ratio*n_blocks] cells making up each block + // blk_pos I32 [4*n_blocks] mrope position rows of each block's first token + // bias F32 [n_kv, n_tokens] -inf where invisible, large where always visible 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) const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 55054bb9d2..1e48c30ec4 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -66,8 +66,7 @@ llama_memory_hybrid::llama_memory_hybrid( : filter_recr )), mem_idx(filter_idx == nullptr ? nullptr : [&] { - // MQA with a single key head of indexer_head_size, the same shaping - // llama_kv_cache_dsa applies to its lightning-indexer cache + // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own 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; @@ -124,20 +123,21 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba } // prepare the attention cache - llama_kv_cache::slot_info_vec_t heads_idx; - if (mem_idx) { - heads_idx = mem_idx->prepare(ubatches); - if (heads_idx.empty()) { - break; - } - } - auto heads_attn = mem_attn->prepare(ubatches); if (heads_attn.empty()) { LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } + // The indexer cache is a side buffer addressed by the attention cache's cells, so it + // takes that slot layout rather than finding its own. Allocating separately let the + // two drift once context was rewritten between turns, pointing QSA top-k at the + // wrong cells. + llama_kv_cache::slot_info_vec_t heads_idx; + if (mem_idx) { + heads_idx = heads_attn; + } + return std::make_unique( this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches)); } while(false); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 2c03ee7ca5..5a167fef18 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -92,8 +92,7 @@ public: private: const llama_hparams & hparams; - // geometry for the indexer cache: MQA with a single key head of - // indexer_head_size, mirroring how llama_kv_cache_dsa builds its own + // indexer cache geometry: MQA, one key head of indexer_head_size, as in llama_kv_cache_dsa llama_hparams hparams_idx; const std::unique_ptr mem_attn; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index dad7c0b81b..1b058d82bc 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2434,8 +2434,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; - // left null for every architecture but the sparse-attention - // ones, which is what keeps the indexer cache from existing + // null for every architecture but the sparse-attention ones, which is what keeps + // the indexer cache from existing llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; diff --git a/src/models/models.h b/src/models/models.h index 6f024896d2..b952790e8c 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2314,8 +2314,7 @@ struct llama_model_qwen4exp : public llama_model_base { int * sections, int il); - // QSA: the token indices this layer's queries may attend to, or nullptr - // to attend densely + // QSA: token indices this layer's queries may attend to, or nullptr for dense ggml_tensor * build_qsa_top_k( const llama_kv_cache_context * mctx_idx, ggml_tensor * cur, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 6a5247e7c0..49028a674c 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -266,8 +266,9 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa auto * inp = build_inp_mem_hybrid(); - // present only when the GGUF carries indexer tensors, so a model without - // them still builds a dense graph + // present only when the GGUF carries indexer tensors, so a model without them still + // builds a dense graph. The indexer cache takes the attention cache's slot layout, + // so the two agree cell for cell by construction. const llama_kv_cache_context * mctx_idx = inp->mctx->get_idx(); if (mctx_idx) { GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() && @@ -372,14 +373,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( return ggml_mul(ctx0, normalized, gated); } -// QSA attends to a budget of whole blocks of `compress_ratio` tokens, chosen by -// scoring one mean-pooled indexer key per block, plus the tail of tokens that do -// not yet make up a complete block, which is always visible. Below -// indexer_top_k + compress_ratio - 1 cached tokens every block fits in the -// budget and the result is exactly dense attention. -// -// Everything that depends on the cache layout - which cells make up a block, -// which blocks a query may see - is computed host-side in set_input, so the +// QSA attends to a budget of whole blocks of `compress_ratio` tokens, scored by one +// mean-pooled indexer key each, plus the always-visible incomplete tail. Below +// indexer_top_k + compress_ratio - 1 cached tokens this is exactly dense attention. +// Everything depending on cache layout is computed host-side in set_input; the // graph only gathers, pools and scores. class llm_graph_input_qsa : public llm_graph_input_i { public: @@ -432,8 +429,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( llm_graph_input_qsa * inp = qsa.get(); res->add_input(std::move(qsa)); - // the cached indexer keys are raw: pooling happens before the norm and the - // rotation, so neither may be applied on the way in + // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens); cb(k_raw, "indexer_k_raw", il); @@ -447,8 +443,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); members = ggml_reshape_3d(ctx0, members, idx_dim, r, n_blocks); - // mean over the block's members. compress_ratio is small, so summing the - // slices costs less than transposing to reach ggml_sum_rows + // mean over the block's members; compress_ratio is small, so summing slices beats + // transposing to reach ggml_sum_rows ggml_tensor * pooled = nullptr; for (int64_t i = 0; i < r; ++i) { ggml_tensor * slice = ggml_cont(ctx0, @@ -474,10 +470,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ext_factor, attn_factor, beta_fast, beta_slow); cb(q, "indexer_q", il); - // Each head's dot product is rectified before the heads are summed, as in - // DeepSeek's lightning indexer, but with no learned per-head weight. The - // reference then divides by a constant, which cannot reorder anything, so - // that is left out. + // Each head's dot product is rectified before summing, as in DeepSeek's lightning + // indexer but with no learned per-head weight. The reference's constant divisor + // cannot reorder anything, so it is left out. ggml_tensor * score = ggml_mul_mat(ctx0, pooled, ggml_reshape_2d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tokens)); score = ggml_reshape_3d(ctx0, score, n_blocks, n_idx_h, n_tokens); @@ -487,19 +482,18 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( score = ggml_reshape_2d(ctx0, score, n_blocks, n_tokens); cb(score, "indexer_score", il); - // Give every token of a block its block's score, rather than expanding the - // selected block indices: that would need an integer multiply-add, which - // ggml has no op for. Because the budget is a whole number of blocks and the - // members of a block tie exactly, the cut still lands on a block boundary. - // get_rows gathers rows, so the scores are transposed for the gather. + // Give every token of a block its block's score rather than expanding the block + // indices, which would need an integer multiply-add ggml has no op for. The + // budget is a whole number of blocks and members tie, so the cut still lands on + // a block boundary. get_rows gathers rows, so scores are transposed first. ggml_tensor * expanded = ggml_get_rows(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, score)), inp->cell_blk); expanded = ggml_cont(ctx0, ggml_transpose(ctx0, expanded)); expanded = ggml_add(ctx0, expanded, inp->bias); cb(expanded, "indexer_score_tokens", il); - // the reference returns indexer_top_k + compress_ratio - 1 tokens: a whole - // budget of blocks plus the incomplete tail + // the reference returns indexer_top_k + compress_ratio - 1: a whole budget of + // blocks plus the incomplete tail const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); @@ -518,9 +512,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); - // The indexer reads the same block input as q/k/v. Without an indexer cache - // this falls back to dense attention, which is what the model computes - // anyway for anything shorter than the budget. + // The indexer reads the same block input as q/k/v; with no indexer cache this + // falls back to dense, which is what the model computes below the budget anyway. ggml_tensor * top_k = mctx_idx ? build_qsa_top_k(mctx_idx, cur, inp_pos, sections, il) : nullptr; // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention