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

qwen4exp: support a non-unified KV cache in QSA

Browse Source 
set_input_qsa asserted n_stream == 1, so llama-server could not serve this
model with more than one slot unless -kvu was passed. With a non-unified
cache each sequence owns its own cells, and a cell index means a different
token in each stream, so a single shared mapping is wrong.

- cell_blk, blk_cells and bias gain a stream dimension. At n_stream == 1
  these collapse to the shapes they had, so the unified path is unchanged.
- Scoring is now batched over streams. ggml_mul_mat matches ne[2] on both
  operands, so stream s's queries only ever meet stream s's blocks; without
  this sequences would score against each other's context.
- set_input_qsa loops per stream and resolves cells through
  v_cells[seq_to_stream[seq_id]], following set_input_kq_mask_impl, instead
  of hardcoding v_cells[0].
- llama_kv_cache_context::get_n_stream() is added, mirroring the ns that
  get_k and get_v already derive from the slot info.

build_attn_mask_top_k needed no change: it already expects
[n_top_k, n_batch, 1, n_stream], so the top-k result is reshaped to meet it.

set_input_qsa has exactly one caller, so the blast radius is qwen4exp only.

Validation, UD-Q4_K_XL on one B200:

- unified cache unchanged within noise: 1802.9/68.85 -> 1807.2/69.11 t/s at
  batch 1, 2262.5/192.43 -> 2270.1/193.75 at batch 4.
- non-unified now runs at npl 1, 4, 16 where it previously aborted, and is
  22% faster than the -kvu workaround at batch 16 (1205 vs 984 t/s total),
  since per-stream cells avoid the cross-sequence masking a unified cache
  pays for.
- no cross-stream contamination: four concurrent sequences each carrying a
  distinct secret all recall their own and no other, on both cache modes.
- test-llama-archs green on qwen4exp, deepseek2, gemma3n, qwen3next, llama.

Note on testing: comparing concurrent output against solo output exactly is
not a valid check. It failed 0/4 with no bug present, and the unified-cache
control failed the same way, because batch composition changes the
floating-point reduction order and near-tied tokens flip. The contamination
test above is what the exit code gates on.
This commit is contained in:
danielhanchen
2026-08-26 04:09:11 +00:00
committed by Daniel Han
parent eb95d125ba
commit bea3b12dae
3 changed files with 113 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-kv-cache.cpp
+73 -50
View File
@@ -1794,80 +1794,98 @@ void llama_kv_cache::set_input_qsa(
ggml_tensor * bias,
const llama_ubatch * ubatch,
uint32_t ratio) const {
GGML_ASSERT(n_stream == 1 && "TODO: support multiple streams");
GGML_ASSERT(ratio > 0);
const auto & cells = v_cells[0];
GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
const int64_t n_kv = cell_blk->ne[0];
const int64_t n_blocks = blk_pos->ne[0]/4;
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 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);
// 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);
}
}
}
// 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<int32_t> blk_of(n_kv, -1);
std::vector<int32_t> filled(n_blocks, 0);
// 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);
std::fill(dst_blk_cells, dst_blk_cells + r*n_blocks, 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
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]];
for (int64_t j = 0; j < n_kv; ++j) {
if (cells.is_empty(j)) {
continue;
}
int32_t * cur_cell_blk = dst_cell_blk + s*n_kv;
int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
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;
dst_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;
}
dst_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j];
}
for (int64_t i = 0; i < n_tokens; ++i) {
const llama_seq_id seq_id = ubatch->seq_id[i][0];
const llama_pos q = ubatch->pos[i];
// 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;
// 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) {
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)) {
continue;
}
dst_bias[i*n_kv + j] = v;
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;
}
}
}
}
@@ -2670,6 +2688,11 @@ 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();
}
src/llama-kv-cache.h
+3
View File
@@ -375,6 +375,9 @@ 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;
src/models/qwen4exp.cpp
+37 -21
View File
@@ -9,6 +9,7 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);
ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);
ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);
@@ -22,6 +23,7 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
GGML_ASSERT(hparams.hc_low_rank > 0 && "qwen4exp needs a hyper-connection 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);
@@ -136,6 +138,7 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
const int64_t idx_dim = hparams.indexer_head_size;
layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0);
layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0);
@@ -392,11 +395,12 @@ public:
mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio);
}
// Per-stream: a cell index means a different token in each stream. n_stream 1 = the old shapes.
ggml_tensor * k_idxs = nullptr; // I32 [n_tokens]
ggml_tensor * cell_blk = nullptr; // I32 [n_kv]
ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks]
ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks]
ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens]
ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream]
ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream]
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 uint32_t ratio;
@@ -417,13 +421,18 @@ 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();
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);
qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch);
qsa->cell_blk = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_kv);
qsa->blk_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r*n_blocks);
qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks);
qsa->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens);
qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream);
qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream);
qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream);
qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream);
ggml_set_input(qsa->cell_blk);
ggml_set_input(qsa->blk_cells);
@@ -440,30 +449,33 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il));
// one key head, so the per-cell rows are contiguous and this view is dense
// one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream].
ggml_tensor * k_all = mctx_idx->get_k(ctx0, il);
k_all = ggml_view_2d(ctx0, k_all, idx_dim, n_kv, k_all->nb[2], 0);
k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0);
// gathers per stream: blk_cells row s indexes stream s's own cells
ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells);
members = ggml_reshape_3d(ctx0, members, idx_dim, r, n_blocks);
members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream);
// 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,
ggml_view_2d(ctx0, members, idx_dim, n_blocks, members->nb[2], i*members->nb[1]));
ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream,
members->nb[2], members->nb[3], i*members->nb[1]));
pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice;
}
pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r);
cb(pooled, "indexer_k_pooled", il);
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks);
// 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);
pooled = ggml_reshape_2d(ctx0, pooled, idx_dim, n_blocks);
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream);
cb(pooled, "indexer_k", il);
ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur);
@@ -475,15 +487,15 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
cb(q, "indexer_q", il);
// 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.
// no per-head weight; the constant divisor cannot reorder. mul_mat matches ne[2], so
// stream s's queries only meet stream s's blocks.
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);
ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream));
score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);
score = ggml_relu(ctx0, score);
score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
score = ggml_sum_rows(ctx0, score);
score = ggml_reshape_2d(ctx0, score, n_blocks, n_tokens);
score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream);
cb(score, "indexer_score", il);
// Give every token of a block its block's score rather than expanding the block
@@ -491,8 +503,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
// 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));
ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk);
expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3));
expanded = ggml_add(ctx0, expanded, inp->bias);
cb(expanded, "indexer_score_tokens", il);
@@ -501,6 +513,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
const int64_t width = std::min<int64_t>(n_kv, (int64_t) hparams.indexer_top_k + r - 1);
ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width));
// build_attn_mask_top_k reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask.
top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream);
cb(top_k, "indexer_top_k", il);
return top_k;
@@ -701,6 +716,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(
k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm);
// repeat to match shapes when head keys != value keys; unneeded with the fused GDN
if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) {
GGML_ASSERT(num_v_heads % num_k_heads == 0);