Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 17:24:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama: hold the qwen4exp indexer cache in a new llama_memory_hybrid_idx

Browse Source 
The indexer key cache was added by extending llama_memory_hybrid with an
optional third cache, and the host-side cell/block mapping that drives QSA was
added as set_input_qsa on llama_kv_cache. Both are shared classes that every
hybrid and every attention model goes through.

Move both into a new memory type, llama_memory_hybrid_idx, following
llama_kv_cache_msa: the indexer cache and the pos<->cell translation live with
the sparse-attention memory rather than in the classes that serve every other
architecture. llama-kv-cache.{h,cpp} and llama-memory-hybrid.{h,cpp} are
restored to their unmodified state.

init_batch is repeated from llama_memory_hybrid because the indexer cache has to
be handed the attention cache's slot infos, and those are not reachable through
the context the base returns. Allocating them separately lets the two caches
drift, which is what pointed QSA's top-k at the wrong cells before.

The context derives from llama_memory_hybrid_context so build_inp_mem_hybrid
keeps working unchanged, and get_n_stream is computed from the slot infos
exactly as llama_kv_cache_context did.

Behaviour is unchanged: logits over an 8192-token sequence are bit-identical to
the previous implementation, sparse and dense alike.
This commit is contained in:
Daniel Han
2026-08-26 14:43:13 +00:00
parent 7fee670a95
commit 035e22731a
10 changed files with 597 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CMakeLists.txt
+1
View File
@@ -31,6 +31,7 @@ add_library(llama
llama-memory.cpp
llama-memory-hybrid.cpp
llama-memory-hybrid-iswa.cpp
llama-memory-hybrid-idx.cpp
llama-memory-recurrent.cpp
llama-mmap.cpp
llama-model-loader.cpp
src/llama-kv-cache.cpp
-118
View File
@@ -1787,109 +1787,6 @@ void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch
}
}
void llama_kv_cache::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 {
GGML_ASSERT(ratio > 0);
GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
const int64_t n_kv = cell_blk->ne[0];
const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch
const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
const int64_t n_tokens = ubatch->n_tokens;
const int64_t r = ratio;
GGML_ASSERT(n_tokens % n_ns == 0);
const int64_t n_tps = n_tokens/n_ns; // tokens per stream
int32_t * dst_cell_blk = (int32_t *) cell_blk->data;
int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
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 three
// mrope sections carry it: exact for text, approximate for images. Positions repeat per stream.
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.
// n_ns == 1 is the single-stream behaviour this replaced.
std::vector<int32_t> blk_of(n_kv);
std::vector<int32_t> filled(n_blocks);
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 = v_cells[seq_to_stream[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);
// an incomplete block cannot be pooled: those tail cells are forced in by the bias
// below, so block 0 only keeps the gather in range. -1 = no usable block.
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);
for (int64_t j = 0; j < n_kv; ++j) {
if (cells.is_empty(j)) {
continue;
}
const llama_pos p = cells.pos_get(j);
const int64_t b = p/r;
if (b >= n_blocks) {
continue;
}
blk_of[j] = (int32_t) b;
cur_blk_cells[b*r + (p%r)] = (int32_t) j;
filled[b]++;
}
for (int64_t j = 0; j < n_kv; ++j) {
if (blk_of[j] >= 0 && filled[blk_of[j]] < r) {
blk_of[j] = -1;
}
cur_cell_blk[j] = blk_of[j] < 0 ? 0 : 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 rest is an incomplete block, always attended to, which is what lands the
// selection on block boundaries like the reference
const llama_pos tail_start = (q + 1)/r*r;
float * cur_bias = dst_bias + i*n_kv;
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);
}
cur_bias[j] = v;
}
}
}
}
void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const {
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
@@ -2688,11 +2585,6 @@ uint32_t llama_kv_cache_context::get_n_kv() const {
return n_kv;
}
uint32_t llama_kv_cache_context::get_n_stream() const {
// streams in the current slot info, matching get_k/get_v's `ns`. 1 if unified.
return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1;
}
ggml_type llama_kv_cache_context::type_k() const {
return kv->type_k();
}
@@ -2753,16 +2645,6 @@ void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama
kv->set_input_pos_bucket(dst, ubatch);
}
void llama_kv_cache_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) const {
kv->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio);
}
void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const {
kv->set_input_k_rot(dst);
}
src/llama-kv-cache.h
-15
View File
@@ -216,15 +216,6 @@ 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) 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;
void set_input_k_rot(ggml_tensor * dst) const;
void set_input_v_rot(ggml_tensor * dst) const;
@@ -375,9 +366,6 @@ public:
uint32_t get_n_kv() const;
// streams in the current slot info; 1 for a unified cache
uint32_t get_n_stream() const;
ggml_type type_k() const;
ggml_type type_v() const;
@@ -410,9 +398,6 @@ 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;
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;
void set_input_k_rot(ggml_tensor * dst) const;
void set_input_v_rot(ggml_tensor * dst) const;
src/llama-memory-hybrid-idx.cpp
+380
View File
@@ -0,0 +1,380 @@
#include "llama-memory-hybrid-idx.h"
#include "llama-impl.h"
#include "llama-batch.h"
#include "llama-model.h"
#include <algorithm>
#include <cassert>
#include <cmath>
//
// llama_memory_hybrid_idx
//
llama_memory_hybrid_idx::llama_memory_hybrid_idx(
const llama_model & model,
/* attn */
ggml_type type_k,
ggml_type type_v,
bool v_trans,
uint32_t kv_size,
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
/* recurrent */
ggml_type type_r,
ggml_type type_s,
uint32_t rs_size,
/* common */
uint32_t n_seq_max,
uint32_t n_rs_seq,
bool offload,
bool unified,
/* layer filters */
const layer_filter_cb & filter_attn,
const layer_filter_cb & filter_recr,
const layer_filter_cb & filter_idx) :
llama_memory_hybrid(
model,
type_k, type_v, v_trans, kv_size, n_pad, n_swa, swa_type,
type_r, type_s, rs_size,
n_seq_max, n_rs_seq, offload, unified,
filter_attn, filter_recr),
hparams_idx(model.hparams),
mem_idx(filter_idx == nullptr ? nullptr : [&] {
// 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;
LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
return new llama_kv_cache(
model, hparams_idx, type_k, type_v, v_trans, offload, unified,
kv_size, n_seq_max, n_pad, n_swa, swa_type,
nullptr, filter_idx, nullptr, nullptr);
}()) {}
llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
// note: this repeats llama_memory_hybrid::init_batch because the indexer cache has to be
// handed the attention cache's slot infos, and those are not reachable through the
// llama_memory_hybrid_context that the base implementation returns
do {
balloc.split_reset();
// follow the recurrent pattern for creating the ubatch splits
std::vector<llama_ubatch> ubatches;
while (true) {
llama_ubatch ubatch;
if (embd_all) {
// if all tokens are output, split by sequence
ubatch = balloc.split_seq(n_ubatch);
} else {
// Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
const bool unified = (get_mem_attn()->get_n_stream() == 1);
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
// so that the rollback snapshots remain valid
const uint32_t n_rs_seq = get_mem_recr()->n_rs_seq;
ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
}
if (ubatch.n_tokens == 0) {
break;
}
ubatches.push_back(std::move(ubatch)); // NOLINT
}
if (balloc.get_n_used() < balloc.get_n_tokens()) {
// failed to find a suitable split
break;
}
// prepare the recurrent batches first
if (!get_mem_recr()->prepare(ubatches)) {
// TODO: will the recurrent cache be in an undefined context at this point?
LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__);
return std::make_unique<llama_memory_hybrid_idx_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}
// prepare the attention cache
auto heads_attn = get_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_hybrid_idx_context>(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<llama_memory_hybrid_idx_context>(
this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches));
} while(false);
return std::make_unique<llama_memory_hybrid_idx_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}
llama_memory_context_ptr llama_memory_hybrid_idx::init_full() {
return std::make_unique<llama_memory_hybrid_idx_context>(this);
}
llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lctx, bool optimize) {
return std::make_unique<llama_memory_hybrid_idx_context>(this, lctx, optimize);
}
void llama_memory_hybrid_idx::clear(bool data) {
llama_memory_hybrid::clear(data);
if (mem_idx) {
mem_idx->clear(data);
}
}
bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
// same order as llama_memory_hybrid::seq_rm: try the recurrent cache first since it is the
// one that may refuse, and if it does the caches are left untouched
if (!get_mem_recr()->seq_rm(seq_id, p0, p1)) {
return false;
}
if (mem_idx) {
mem_idx->seq_rm(seq_id, p0, p1);
}
return get_mem_attn()->seq_rm(seq_id, p0, p1);
}
void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1);
if (mem_idx) {
mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);
}
}
void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {
llama_memory_hybrid::seq_keep(seq_id);
if (mem_idx) {
mem_idx->seq_keep(seq_id);
}
}
void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
llama_memory_hybrid::seq_add(seq_id, p0, p1, shift);
if (mem_idx) {
mem_idx->seq_add(seq_id, p0, p1, shift);
}
}
void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
llama_memory_hybrid::seq_div(seq_id, p0, p1, d);
if (mem_idx) {
mem_idx->seq_div(seq_id, p0, p1, d);
}
}
std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid_idx::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> mb = llama_memory_hybrid::memory_breakdown();
if (mem_idx) {
for (const auto & buft_size : mem_idx->memory_breakdown()) {
mb[buft_size.first] += buft_size.second;
}
}
return mb;
}
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<uint32_t> llama_memory_hybrid_idx_ns(const llama_kv_cache::slot_info_vec_t & sinfos) {
std::vector<uint32_t> 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) {}
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) {}
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<llama_ubatch> 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<const llama_kv_cache_context *>(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) const {
GGML_ASSERT(ratio > 0);
GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
const int64_t n_kv = cell_blk->ne[0];
const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch
const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
const int64_t n_tokens = ubatch->n_tokens;
const int64_t r = ratio;
GGML_ASSERT(n_tokens % n_ns == 0);
const int64_t n_tps = n_tokens/n_ns; // tokens per stream
int32_t * dst_cell_blk = (int32_t *) cell_blk->data;
int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
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 three
// mrope sections carry it: exact for text, approximate for images. Positions repeat per stream.
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.
// n_ns == 1 is the single-stream behaviour this replaced.
std::vector<int32_t> blk_of(n_kv);
std::vector<int32_t> filled(n_blocks);
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);
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: those tail cells are forced in by the bias
// below, so block 0 only keeps the gather in range. -1 = no usable block.
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);
for (int64_t j = 0; j < n_kv; ++j) {
if (cells.is_empty(j)) {
continue;
}
const llama_pos p = cells.pos_get(j);
const int64_t b = p/r;
if (b >= n_blocks) {
continue;
}
blk_of[j] = (int32_t) b;
cur_blk_cells[b*r + (p%r)] = (int32_t) j;
filled[b]++;
}
for (int64_t j = 0; j < n_kv; ++j) {
if (blk_of[j] >= 0 && filled[blk_of[j]] < r) {
blk_of[j] = -1;
}
cur_cell_blk[j] = blk_of[j] < 0 ? 0 : 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 rest is an incomplete block, always attended to, which is what lands the
// selection on block boundaries like the reference
const llama_pos tail_start = (q + 1)/r*r;
float * cur_bias = dst_bias + i*n_kv;
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);
}
cur_bias[j] = v;
}
}
}
}
src/llama-memory-hybrid-idx.h
+157
View File
@@ -0,0 +1,157 @@
#pragma once
#include "llama-memory-hybrid.h"
#include <memory>
#include <vector>
//
// llama_memory_hybrid_idx
//
// llama_memory_hybrid plus a third cache holding one indexer key per token, for hybrid
// architectures whose attention layers are block-sparse (qwen4exp QSA).
//
// this is a separate llama_memory type rather than an option on llama_memory_hybrid so that
// nothing in the hybrid path used by the other architectures changes. it duplicates
// llama_memory_hybrid::init_batch because the indexer cache must be given the attention
// cache's slot layout instead of finding its own, and that layout is not observable from
// outside the returned context.
//
// the indexer cache is a side buffer addressed by the attention cache's cells: same size,
// same padding, same stream count, same slots, so cell j means the same token in both.
// everything that depends on that layout is computed host-side in set_input_qsa; the model
// graph only gathers, pools and scores.
class llama_memory_hybrid_idx : public llama_memory_hybrid {
public:
llama_memory_hybrid_idx(
const llama_model & model,
/* attn */
ggml_type type_k,
ggml_type type_v,
bool v_trans,
uint32_t kv_size,
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
/* recurrent */
ggml_type type_r,
ggml_type type_s,
uint32_t rs_size,
/* common */
uint32_t n_seq_max,
uint32_t n_rs_seq,
bool offload,
bool unified,
/* layer filters */
const layer_filter_cb & filter_attn,
const layer_filter_cb & filter_recr,
/* the indexer cache exists only if this is given */
const layer_filter_cb & filter_idx);
~llama_memory_hybrid_idx() = default;
//
// llama_memory_i
//
llama_memory_context_ptr init_batch(
llama_batch_allocr & balloc,
uint32_t n_ubatch,
bool embd_all) override;
llama_memory_context_ptr init_full() override;
llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;
void clear(bool data) override;
bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override;
void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;
void seq_keep(llama_seq_id seq_id) override;
void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override;
void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override;
std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;
//
// llama_memory_hybrid_idx specific API
//
llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer
private:
// the indexer cache stores only one key head per layer, so it needs its own hparams
// instance: llama_kv_cache keeps a reference to whatever it is given
llama_hparams hparams_idx;
const std::unique_ptr<llama_kv_cache> mem_idx;
};
class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
public:
using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
// used for errors
explicit llama_memory_hybrid_idx_context(llama_memory_status status);
// used to create a full-cache context
explicit llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem);
// used to create an update context
llama_memory_hybrid_idx_context(
llama_memory_hybrid_idx * mem,
llama_context * lctx,
bool optimize);
// used to create a batch processing context from a batch
llama_memory_hybrid_idx_context(
llama_memory_hybrid_idx * mem,
slot_info_vec_t sinfos_attn,
slot_info_vec_t sinfos_idx,
std::vector<llama_ubatch> ubatches);
~llama_memory_hybrid_idx_context() = default;
//
// llama_memory_context_i
//
bool next() override;
bool apply() override;
//
// llama_memory_hybrid_idx_context specific API
//
// nullptr when the model carries no indexer, and for the full and update contexts,
// which do not drive the sparse-attention graph
const llama_kv_cache_context * get_idx() const;
// streams in the current slot info, matching get_k/get_v's `ns`. 1 if unified.
uint32_t get_n_stream() const;
// block-compressed sparse attention (qwen4exp QSA) over the indexer cache's cells.
// blocks cut the *position* line, not the cell array, so nothing 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
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;
private:
const llama_memory_hybrid_idx * mem = nullptr;
// streams per ubatch, taken from the slot infos before they are handed to ctx_idx.
// declared first so that it is initialised while sinfos_idx is still intact
const std::vector<uint32_t> ns_ubatch;
// null unless the model has an indexer and this is a batch context
const llama_memory_context_ptr ctx_idx;
// mirrors the base class's ubatch cursor, which is private there
size_t i_cur = 0;
};
src/llama-memory-hybrid.cpp
+3 -45
View File
@@ -29,10 +29,8 @@ llama_memory_hybrid::llama_memory_hybrid(
bool unified,
/* layer filters */
const layer_filter_cb & filter_attn,
const layer_filter_cb & filter_recr,
const layer_filter_cb & filter_idx) :
const layer_filter_cb & filter_recr) :
hparams(model.hparams),
hparams_idx(model.hparams),
mem_attn(new llama_kv_cache(
model,
model.hparams,
@@ -64,19 +62,7 @@ llama_memory_hybrid::llama_memory_hybrid(
filter_recr == nullptr ?
[&](int32_t il) { return hparams.is_recr(il); }
: filter_recr
)),
mem_idx(filter_idx == nullptr ? nullptr : [&] {
// 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;
LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
return new llama_kv_cache(
model, hparams_idx, type_k, type_v, v_trans, offload, unified,
kv_size, n_seq_max, n_pad, n_swa, swa_type,
nullptr, filter_idx, nullptr, nullptr);
}()) {}
)) {}
llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
do {
@@ -129,17 +115,8 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
return std::make_unique<llama_memory_hybrid_context>(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<llama_memory_hybrid_context>(
this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches));
this, std::move(heads_attn), std::move(ubatches));
} while(false);
return std::make_unique<llama_memory_hybrid_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
@@ -160,7 +137,6 @@ bool llama_memory_hybrid::get_can_shift() const {
void llama_memory_hybrid::clear(bool data) {
mem_attn->clear(data);
if (mem_idx) mem_idx->clear(data);
mem_recr->clear(data);
}
@@ -170,31 +146,26 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
if (!mem_recr->seq_rm(seq_id, p0, p1)) {
return false;
}
if (mem_idx) mem_idx->seq_rm(seq_id, p0, p1);
return mem_attn->seq_rm(seq_id, p0, p1);
}
void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1);
if (mem_idx) mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);
mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1);
}
void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) {
mem_attn->seq_keep(seq_id);
if (mem_idx) mem_idx->seq_keep(seq_id);
mem_recr->seq_keep(seq_id);
}
void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
mem_attn->seq_add(seq_id, p0, p1, shift);
if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift);
mem_recr->seq_add(seq_id, p0, p1, shift);
}
void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
mem_attn->seq_div(seq_id, p0, p1, d);
if (mem_idx) mem_idx->seq_div(seq_id, p0, p1, d);
mem_recr->seq_div(seq_id, p0, p1, d);
}
@@ -234,10 +205,6 @@ llama_kv_cache * llama_memory_hybrid::get_mem_attn() const {
return mem_attn.get();
}
llama_kv_cache * llama_memory_hybrid::get_mem_idx() const {
return mem_idx.get();
}
llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const {
return mem_recr.get();
}
@@ -262,14 +229,11 @@ llama_memory_hybrid_context::llama_memory_hybrid_context(
llama_memory_hybrid_context::llama_memory_hybrid_context(
llama_memory_hybrid * mem,
slot_info_vec_t sinfos_attn,
slot_info_vec_t sinfos_idx,
std::vector<llama_ubatch> ubatches) :
ubatches(std::move(ubatches)),
// note: here we copy the ubatches. not sure if this is ideal
ctx_attn(new llama_kv_cache_context(mem->get_mem_attn(), std::move(sinfos_attn), this->ubatches)),
ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)),
ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), this->ubatches)),
status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {
}
@@ -278,7 +242,6 @@ bool llama_memory_hybrid_context::next() {
ctx_attn->next();
ctx_recr->next();
if (ctx_idx) ctx_idx->next();
if (++i_next >= ubatches.size()) {
return false;
@@ -294,7 +257,6 @@ bool llama_memory_hybrid_context::apply() {
res = res & ctx_attn->apply();
res = res & ctx_recr->apply();
if (ctx_idx) res = res & ctx_idx->apply();
return res;
}
@@ -315,7 +277,3 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const {
const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const {
return static_cast<const llama_memory_recurrent_context *>(ctx_recr.get());
}
const llama_kv_cache_context * llama_memory_hybrid_context::get_idx() const {
return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
}
src/llama-memory-hybrid.h
+1 -14
View File
@@ -39,12 +39,7 @@ public:
bool unified,
/* layer filters */
const layer_filter_cb & filter_attn = nullptr,
const layer_filter_cb & filter_recr = nullptr,
/* optional per-token indexer key cache, for hybrid
models whose attention layers are sparse. absent
unless filter_idx is given, so every existing
architecture is unaffected. */
const layer_filter_cb & filter_idx = nullptr);
const layer_filter_cb & filter_recr = nullptr);
~llama_memory_hybrid() = default;
@@ -87,17 +82,12 @@ public:
llama_kv_cache * get_mem_attn() const;
llama_memory_recurrent * get_mem_recr() const;
llama_kv_cache * get_mem_idx() const; // nullptr when the model has no indexer
private:
const llama_hparams & hparams;
// 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<llama_kv_cache> mem_attn;
const std::unique_ptr<llama_memory_recurrent> mem_recr;
const std::unique_ptr<llama_kv_cache> mem_idx;
};
class llama_memory_hybrid_context : public llama_memory_context_i {
@@ -120,7 +110,6 @@ public:
llama_memory_hybrid_context(
llama_memory_hybrid * mem,
slot_info_vec_t sinfos_attn,
slot_info_vec_t sinfos_idx,
std::vector<llama_ubatch> ubatches);
~llama_memory_hybrid_context() = default;
@@ -137,7 +126,6 @@ public:
const llama_kv_cache_context * get_attn() const;
const llama_memory_recurrent_context * get_recr() const;
const llama_kv_cache_context * get_idx() const; // nullptr without an indexer
private:
// the index of the next ubatch to process
@@ -147,7 +135,6 @@ private:
const llama_memory_context_ptr ctx_attn;
const llama_memory_context_ptr ctx_recr;
const llama_memory_context_ptr ctx_idx; // null unless the model has an indexer
const llama_memory_status status;
};
src/llama-model.cpp
+27 -4
View File
@@ -16,6 +16,7 @@
#include "llama-kv-cache-dsv4.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-iswa.h"
#include "llama-memory-hybrid-idx.h"
#include "llama-memory-recurrent.h"
#include "llama.h"
@@ -2434,9 +2435,10 @@ 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;
// null for every architecture but the sparse-attention ones, which is what keeps
// the indexer cache from existing
// llama_memory_hybrid_idx is used only by the sparse-attention architectures;
// filter_idx null within it means the GGUF carries no indexer tensors
llama_memory_hybrid::layer_filter_cb filter_idx = nullptr;
const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP);
if (arch == LLM_ARCH_FALCON_H1) {
filter_attn = [&](uint32_t) { return true; };
filter_recr = [&](uint32_t) { return true; };
@@ -2483,8 +2485,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
/* unified */ cparams.kv_unified,
/* filter_attn */ std::move(filter_attn),
/* filter_recr */ std::move(filter_recr));
} else {
res = new llama_memory_hybrid(
} else if (needs_mem_idx) {
// sparse attention over a per-token indexer cache: a separate memory
// type, so the plain hybrid path is untouched
res = new llama_memory_hybrid_idx(
/* model */ *this,
/* attn_type_k */ params.type_k,
/* attn_type_v */ params.type_v,
@@ -2503,6 +2507,25 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
/* filter_attn */ std::move(filter_attn),
/* filter_recr */ std::move(filter_recr),
/* filter_idx */ std::move(filter_idx));
} else {
res = new llama_memory_hybrid(
/* model */ *this,
/* attn_type_k */ params.type_k,
/* attn_type_v */ params.type_v,
/* attn_v_trans */ !cparams.flash_attn,
/* attn_kv_size */ cparams.n_ctx_seq,
/* attn_n_pad */ 1,
/* attn_n_swa */ hparams.n_swa,
/* attn_swa_type */ hparams.swa_type,
/* recurrent_type_k */ GGML_TYPE_F32,
/* recurrent_type_v */ GGML_TYPE_F32,
/* recurrent_kv_size */ std::max((uint32_t) 1, cparams.n_seq_max),
/* n_seq_max */ cparams.n_seq_max,
/* n_rs_seq */ cparams.n_rs_seq,
/* offload */ cparams.offload_kqv,
/* unified */ cparams.kv_unified,
/* filter_attn */ std::move(filter_attn),
/* filter_recr */ std::move(filter_recr));
}
} else {
llama_kv_cache::layer_filter_cb filter = nullptr;
src/models/models.h
+5 -3
View File
@@ -8,6 +8,8 @@
#include <cmath>
#include <map>
class llama_memory_hybrid_idx_context;
//
// base classes
//
@@ -2307,8 +2309,8 @@ struct llama_model_qwen4exp : public llama_model_base {
int il);
ggml_tensor * build_layer_attn(
llm_graph_input_attn_kv * inp_attn,
const llama_kv_cache_context * mctx_idx,
llm_graph_input_attn_kv * inp_attn,
const llama_memory_hybrid_idx_context * mctx_hyb,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int * sections,
@@ -2328,7 +2330,7 @@ struct llama_model_qwen4exp : public llama_model_base {
// 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,
const llama_memory_hybrid_idx_context * mctx_hyb,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int * sections,
src/models/qwen4exp.cpp
+23 -17
View File
@@ -1,5 +1,5 @@
#include "models.h"
#include "llama-memory-hybrid.h"
#include "llama-memory-hybrid-idx.h"
#include "llama-memory-recurrent.h"
void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
@@ -275,7 +275,11 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
// 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();
// qwen4exp always builds llama_memory_hybrid_idx, so this downcast is total; the
// indexer cache inside it is absent when the GGUF carries no indexer tensors
const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(inp->mctx);
const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
if (mctx_idx) {
GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() &&
"the indexer cache must track the attention cache cell for cell");
@@ -310,7 +314,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
if (hparams.is_recr(il)) {
cur = build_layer_attn_linear(inp->get_recr(), cur, il);
} else {
cur = build_layer_attn(inp->get_attn(), mctx_idx, cur, inp_pos, sections, il);
cur = build_layer_attn(inp->get_attn(), mctx_hyb, cur, inp_pos, sections, il);
}
res_hc = build_hc_combine(res_hc, cur, inject, il);
@@ -386,12 +390,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated(
// graph only gathers, pools and scores.
class llm_graph_input_qsa : public llm_graph_input_i {
public:
llm_graph_input_qsa(const llama_kv_cache_context * mctx, uint32_t ratio) :
llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio) :
mctx(mctx), ratio(ratio) {}
virtual ~llm_graph_input_qsa() = default;
void set_input(const llama_ubatch * ubatch) override {
mctx->set_input_k_idxs(k_idxs, ubatch);
mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);
mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio);
}
@@ -402,16 +406,18 @@ public:
ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream]
ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens/n_stream, n_stream]
const llama_kv_cache_context * mctx;
const llama_memory_hybrid_idx_context * mctx;
const uint32_t ratio;
};
ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
const llama_kv_cache_context * mctx_idx,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int * sections,
int il) {
const llama_memory_hybrid_idx_context * mctx_hyb,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int * sections,
int il) {
const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
const int64_t idx_dim = hparams.indexer_head_size;
const int64_t n_idx_h = hparams.indexer_n_head;
const int64_t r = hparams.dsv4_compress_ratios[il];
@@ -421,12 +427,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
const int64_t n_blocks = (n_kv + r - 1)/r;
// n_tps: tokens divide evenly across streams, as build_attn_mask_top_k and the KQ mask assume.
const int64_t n_stream = mctx_idx->get_n_stream();
// n_tps: tokens divide evenly across streams, as build_attn_qsa and the KQ mask assume.
const int64_t n_stream = mctx_hyb->get_n_stream();
GGML_ASSERT(n_tokens % n_stream == 0);
const int64_t n_tps = n_tokens/n_stream;
auto qsa = std::make_unique<llm_graph_input_qsa>(mctx_idx, (uint32_t) r);
auto qsa = std::make_unique<llm_graph_input_qsa>(mctx_hyb, (uint32_t) r);
qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch);
qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream);
@@ -595,7 +601,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
llm_graph_input_attn_kv * inp,
const llama_kv_cache_context * mctx_idx,
const llama_memory_hybrid_idx_context * mctx_hyb,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int * sections,
@@ -604,9 +610,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
// indexer reads the same block input as q/k/v; no cache or no ratio means dense
const bool qsa = mctx_idx != nullptr && hparams.dsv4_compress_ratios[il] > 0;
const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0;
ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_idx, cur, inp_pos, sections, il) : nullptr;
ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, sections, il) : nullptr;
// Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention