llama: add token ID tracking to KV cell by ngxson · Pull Request #27762 · ggml-org/llama.cpp

ngxson · 2026-08-26T19:27:19Z

Overview

Extract a generic change from github.com//pull/19167 and #27742

Save token id to KV cell, so that n-gram hash can read it later

Migration script https://gist.github.com/ngxson/e8ef16ebb62d3451f0c4d1fa47f82de5

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: yes

ggerganov · 2026-08-26T19:51:46Z

+void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
+    const uint32_t n_tokens = ubatch.n_tokens;
+
+    res.clear();
+    res.resize(n_tokens*n, LLAMA_TOKEN_NULL);
+
+    if (n == 0) {
+        return;
+    }
+
+    // note: apply_ubatch() has already stored the current ubatch
+    //       the window below thus covers tokens of this very ubatch as well, which is what we want
+    llama_pos p_min = std::numeric_limits<llama_pos>::max();
+    llama_pos p_max = std::numeric_limits<llama_pos>::min();
+
+    std::bitset<LLAMA_MAX_SEQ> seqs;
+
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        p_min = std::min(p_min, ubatch.pos[i]);
+        p_max = std::max(p_max, ubatch.pos[i]);
+    }
+
+    for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+        seqs.set(ubatch.seq_id_unq[s]);
+    }
+
+    // (seq_id, pos) -> token, for every cell that could be a predecessor of a ubatch token
+    std::unordered_map<uint64_t, llama_token> hist;
+
+    const auto key = [](llama_seq_id seq_id, llama_pos pos) {
+        return ((uint64_t) seq_id << 32) | (uint32_t) pos;
+    };
+
+    for (uint32_t s = 0; s < n_stream; ++s) {
+        v_cells[s].for_each_token_in(seqs, p_min - (llama_pos) n, p_max,
+            [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
+                hist[key(seq_id, pos)] = tok;
+            });
+    }
+
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        // TODO: a token that belongs to more than one sequence has an ambiguous history.
+        //       the n-gram architectures have to reject such batches
+        const llama_seq_id seq_id = ubatch.seq_id[i][0];
+
+        for (uint32_t j = 0; j < n; ++j) {
+            const llama_pos p = ubatch.pos[i] - (llama_pos) (n - j);
+            if (p < 0) {
+                continue;
+            }
+
+            const auto it = hist.find(key(seq_id, p));
+            if (it != hist.end()) {
+                res[i*n + j] = it->second;
+            }
+        }
+    }
+}


A bit difficult to digest, but I guess should be fine. Maybe at some point we can move the logic to llama_kv_cells and write unit+perf tests.

yeah the logic here is not very idea: because kv cells can be out-of-order, we need to firstly write them to std::unordered_map, then for each token in a batch we find the list of N predecessor of that token

I think probably at some point it's better to track an ordered list of cells (in case of unified kv), can be a refactoring in the future

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

spb5030

Why the hell did you ruin backwards compatibility with version 2?
Now my KV cache files simply won't load.
Do you even think about what you're doing?

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens (cherry picked from commit 925e117)

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

ngxson · 2026-08-27T17:34:58Z

there is no need to be aggressive towards maintainers. you can just politely ask an AI agent to migrate it

here, my agent wrote this migration script in 5 minutes: https://gist.github.com/ngxson/e8ef16ebb62d3451f0c4d1fa47f82de5

…742) We merged PR ggml-org#27742 at its llama_memory_hybrid_idx commit. Upstream merged it as 6c84c7d after ~30 more review commits. This brings those in, plus the two upstream-master PRs they depend on (ggml-org#27762 KV token tracking, ggml-org#27794 TENSOR_READ_LAZY) and ggml-org#27795, which replaced our own quantizer buffer fix. Cost and correctness, in rough order of what it is worth here: - QSA bias is now per block, not per cell. It was f32[n_kv, n_tokens], which is about 1 MB per token at -c 262144. The per-cell half is the attention mask, which the graph already has, so only the per-block half is uploaded. - Layers sharing a compress ratio now share one QSA input set. 48 layers each allocated their own cell_blk / blk_cells / blk_pos / bias. - can_reuse() on the QSA and PLE graph inputs, so decode reuses the graph. - The full memory context now carries its indexer cache. graph_reserve() walks a full context and qwen4exp builds sparse attention only when one is present, so the reserved worst case was the smaller dense graph and ggml-alloc grew the compute buffer past the size reported at load. - Output rows are trimmed at the last layer instead of after the final mixer. Our MTP sidecar reads the hyper-connection streams by raw token position, so that trim is skipped when unmasked nextn embeddings are requested and applied after the mixer instead, the same way deepseek4.cpp does it. - QSA supports a rotated/quantized KV cache: q/k/v are hadamard-rotated in build_attn_qsa rather than asserting self_k_rot is null. Inert while LLAMA_ATTN_ROT_DISABLE defaults on for TurboQuant, but the assert is gone. - The PLE n-gram history moves out of a map on llama_model and is recovered from the KV cells (llama_kv_cell_ext::tok, get_prev_tokens). A llama_model is shared by every context and the map was keyed only by seq_id, so two contexts on one model overwrote each other's window; it was also in no state blob, so a restored context hashed against EOS padding. - The PLE conv history gets its own mirrored recurrent row (cache_ple_r_l), and n_embd_r() returns to n_conv. The row exists only on PLE layers, so the recurrent footprint drops rather than grows. - The indexer cache gets its own tensor names (cache_idx_k_l%d); it was matching the attention split pattern. - Indexer state save/restore, whole-context restore clears once, mirrored slot layout is checked rather than assumed. - PLE hparams shrink: is_ple_impl becomes a bitset (2048 bytes to 64), head offsets and vocab sizes narrow to uint32. llama_hparams is held by value in llm_graph_params, which is a stack local on every worker thread. - Converter streams the 128-shard PLE table through LazyChunkedTensor; the quantizer dequantizes in row bands. Drops our -sm tensor deny for qwen4exp (1cecef9). That existed because the PR shipped no split rules for the hyper-connection, PLE and indexer tensors. Upstream wrote them, so the deny is obsolete and was blocking a capability we can use. test-llama-archs -a qwen4exp: R9700 OK 9.58e-14 roundtrip OK RX 6900 XT OK 9.95e-14 roundtrip OK CPU OK 0.00e+00 roundtrip OK Meta OK 9.81e-14 (was SKIP under the deny) No regressions across the families the kv-cache and recurrent changes touch: qwen3next qwen35 qwen35moe minimax-01 falcon-h1 mamba mamba2 nemotron_h granitehybrid lfm2 plamo2 deepseek32 glm-dsa dots3note all pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NjfM4FxGPEhoApv6yvWAzD

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens (cherry picked from commit fa53740)

Upstream 925e117 (ggml-org#27762) bumped LLAMA_STATE_SEQ_VERSION from 2 to 3, and llama_state_seq_load_file checks that word first, so the frozen v2 fixture was refused outright and every assertion built on it was meaningless. Recaptured with the documented recipe: tinyllama2 geometry (-c 512 -b 512 -np 1, temp 0, --slot-save-min-tokens 0, default 256-token block), GOLDEN_PROMPT, and a graceful stop so the shutdown flush publishes the unit. Emitted filenames are unchanged (auto-de633b3190d4d950-5efede8f57f0c198-377), so the identity hash, the block-chain hash and the token count all survived the upstream absorption. The .meta sidecar is BYTE-IDENTICAL to the previous capture: our sidecar format did not change, only the engine's .bin version word did. Captured on the GPU path, which is what the harness uses when N_GPU_LAYERS is unset. This matters and was not previously written down: the saved KV bytes are BACKEND-DEPENDENT. The same unit captured with -ngl 0 differs from the GPU capture in ~43k bytes, so a fixture captured on one backend can never satisfy test_text_only_meta_byte_identical on the other. Verified the absorption itself is byte-clean: capturing with the previous release (2026-08-23-ef3ec4a02-aot) and with this one produces .bin files that differ ONLY in the version word, on both the CPU and the GPU path. Suite goes from 26 failed / 27 passed to 22 failed / 31 passed; four tests fixed, none broken (test_text_only_meta_byte_identical, test_v1_meta_still_indexed, test_unknown_meta_version_skipped, test_missing_meta_is_transient_not_rejected).

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens (cherry picked from commit 925e117)

ngxson added 3 commits August 26, 2026 21:16

kv: track token id

980256b

rm get_prev_tokens, move it to the main pr

c0ea2b3

nits

d07a0ae

ngxson requested a review from ggerganov as a code owner August 26, 2026 19:27

ngxson changed the title ~~Xsn/kv track token~~ llama: add token ID tracking to KV cell Aug 26, 2026

ngxson mentioned this pull request Aug 26, 2026

model: add Qwen3.8-Flash-Next (qwen4exp) #27742

Merged

ggerganov self-assigned this Aug 26, 2026

ggerganov approved these changes Aug 26, 2026

View reviewed changes

add get_prev_tokens

2bd7078

ggerganov reviewed Aug 26, 2026

View reviewed changes

ngxson merged commit 925e117 into master Aug 26, 2026
22 of 26 checks passed

therealkenc pushed a commit to therealkenc/llama.cpp that referenced this pull request Aug 26, 2026

llama: add token ID tracking to KV cell (ggml-org#27762)

7baae0a

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

anthony-chaudhary mentioned this pull request Aug 26, 2026

feat(kvcache): attach token lineage to native KV cells for exact reuse witnesses anthony-chaudhary/fak#9294

Closed

spb5030 reviewed Aug 27, 2026

View reviewed changes

ppenatra pushed a commit to ppenatra/llama.cpp that referenced this pull request Aug 27, 2026

llama: add token ID tracking to KV cell (ggml-org#27762)

840d5d6

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

amangupta-tether pushed a commit to tetherto/qvac-fabric-llm.cpp that referenced this pull request Aug 28, 2026

llama: add token ID tracking to KV cell (ggml-org#27762)

fa53740

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

amangupta-tether mentioned this pull request Aug 28, 2026

Port Qwen4-next PR tetherto/qvac-fabric-llm.cpp#218

Merged

3 tasks

horvay pushed a commit to horvay/atomic-llama-cpp-turboquant that referenced this pull request Aug 30, 2026

llama: add token ID tracking to KV cell (ggml-org#27762)

4dcea2b

* kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens

ggerganov mentioned this pull request Sep 4, 2026

llama.cpp : bump version to 0.4.0 #28386

Merged

Navigation Menu

Uh oh!

llama: add token ID tracking to KV cell - #27762

llama: add token ID tracking to KV cell#27762
ngxson merged 4 commits into
masterfrom
xsn/kv_track_token_id

ngxson commented Aug 26, 2026 •

edited

Loading

Uh oh!

ggerganov Aug 26, 2026 •

edited

Loading

Uh oh!

ngxson Aug 26, 2026

Uh oh!

Uh oh!

spb5030 left a comment

Uh oh!

ngxson commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

ngxson commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Requirements

Uh oh!

ggerganov Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngxson Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

spb5030 left a comment

Choose a reason for hiding this comment

Uh oh!

ngxson commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ngxson commented Aug 26, 2026 •

edited

Loading

ggerganov Aug 26, 2026 •

edited

Loading