qwen4exp: keep the QSA indexer cache in lockstep with the attention cache - #135
Conversation
…ache
The indexer cache (llama_memory_hybrid_idx::mem_idx) was a separate
llama_kv_cache with its own cell metadata (positions, sequence ownership,
used-set). Although init_batch hands it the attention cache's slot infos,
the two caches' cell state can drift apart under server slot management
(seq_rm on slot release): the indexer drops cells the attention cache keeps,
so used_max_p1 -- and therefore get_n_kv() -- diverge, and the graph builder
asserts in llama_model_qwen4exp::graph::graph:
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")
This aborts deep-context (QSA budget exceeded) and multi-slot unified-KV
generation.
Fix: construct the indexer cache with the attention cache as `mem_other`, so
it shares the attention cache's cell metadata (v_cells_impl) via the existing
`other` mechanism. The two then agree cell-for-cell by construction and can
never drift. The indexer keeps its own K/V tensors (the share callback is
null), and the attn_rot/n_embd_head values inherited from the attention cache
are harmless because the indexer path never uses them (cpy_k stores raw keys).
Fixes ggml-org#27780
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 70013d030c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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, "idx_"); | ||
| get_mem_attn(), filter_idx, nullptr, nullptr, "idx_"); |
There was a problem hiding this comment.
Passing get_mem_attn() here sets llama_kv_cache::other on mem_idx. In that mode llama_kv_cache::state_write() and state_read_sinfo() immediately return, but the idx_ cache still allocates its own K/V tensors because the share callback is null; saving/loading a Qwen4Exp session therefore restores attention cell metadata without restoring the indexer keys that QSA top-k reads, producing stale or zero scores for past tokens. Please keep cell sharing from disabling indexer tensor serialization, or add a metadata-only sharing path.
Useful? React with 👍 / 👎.
| // the indexer cache is a side buffer addressed cell-for-cell by the attention | ||
| // cache: it must track the attention cache's cells exactly or QSA top-k reads | ||
| // the wrong cells. sharing the attention cache's cell metadata (v_cells_impl) | ||
| // via `other` makes the two agree by construction -- the indexer keeps its own | ||
| // K/V tensors but sees the same positions/sequences/used-set as attention, so |
There was a problem hiding this comment.
This new block is a seven-line explanation with sentences split across comment lines; the repo guidance asks comments to stay concise, usually 1-2 lines, and not to split a sentence just to wrap. In this hot memory code, condensing it to the invariant, metadata is shared with the attention cache while tensors are not, would reduce maintenance noise.
AGENTS.md reference: AGENTS.md:L73-L81
Useful? React with 👍 / 👎.
1 participant
Problem
llama-serverbuilt from the qwen4exp branch (PR ggml-org#27742)SIGABRTs when a request's context exceeds the QSA sparse-attention budget and decodes long, or when two concurrent unified-KV slots have different KV lengths:Reproduced on DGX Spark / GB10 (SM121), Q4_K_XL,
-c 262144 -ngl 48 -kvu --cache-ram -1 --parallel 2.Root cause
The QSA indexer cache (
llama_memory_hybrid_idx::mem_idx) was a separatellama_kv_cachewith its own cell metadata (v_cells_impl: positions, sequence ownership, used-set). Althoughinit_batchhands it the attention cache's slot infos, the two caches' cell state drifts apart under server slot management (seq_rmon slot release): the indexer drops cells the attention cache keeps. Becauseget_n_kv()is driven byused_max_p1(highest used cell index), the two caches'n_kvdiverge and the graph builder aborts.Fix
Construct the indexer cache with the attention cache as
mem_other, so it shares the attention cache's cell metadata (v_cells_impl) via the existingothermechanism. The two then agree cell-for-cell by construction and can never drift. The indexer keeps its own K/V tensors (thesharecallback is null), and theattn_rot/n_embd_headvalues it inherits from the attention cache are harmless because the indexer path never uses them (cpy_kstores raw keys).Verification (DGX Spark, Q4_K_XL, 2 slots)
n_kvdivergence across all decode steps; the assert no longer fires.This fix targets the qwen4exp branch in PR ggml-org#27742 (the code is not on master yet). Refs ggml-org#27780.