get_prev_tokens() rebuilt a (seq, pos) -> token hash map on every
ubatch by walking all used cells, while llama_kv_cells already keeps
an ordered index of the positions of each sequence in seq_pos, updated
on every cell mutation to serve seq_pos_min() and seq_pos_max().
The index now stores (pos, cell) pairs in a std::set instead of a
position -> count map, so a repeated position (cache reuse via rm + add,
vision inputs with shared positions) yields distinct entries and the
removal of a cell erases its own pair. The new seq_pos_tok_le() returns
the token of the cell at the largest position <= p in logarithmic time,
which is exactly what the old window lookup and its M-RoPE gap fallback
computed together.
get_prev_tokens() shrinks to a direct lookup per (token, offset) and
for_each_token_in() goes away with its only caller. The kv-cache keeps
no n-gram logic of its own.
Measured on Qwen3.8-Flash-Next UD-Q4_K_XL at 71k context, alternating
two binaries with the first run discarded: tg 69.3 -> 72.7 t/s (+4.9%),
pp unchanged at ~2720 t/s, greedy output identical, needle retrieved.
* qwen4exp: follow up fixes
* -kvu NaN collapse fix
Assisted-by: Claude
* indexer cache ext.x/ext.y restore fix
Assisted-by: Claude
* kv-cells: rename seq_set to seq_get_all
seq_get is already taken by the single-id getter, so the suggested name
cannot be overloaded on return type alone.
Assisted-by: Claude
* memory-hybrid-idx: implement set_input_qsa on the memory class
The context held the whole implementation, where the pattern elsewhere is a
thin context forwarding to the memory class, as llama_kv_cache_context does
for set_input_kq_mask. The body reads no context state, so it moves unchanged
and the context keeps a forwarder.
Also shortens the seq_get_all comment as suggested.
* tests: check that a sequence state survives a save/restore round-trip
Saves seq 0, erases it, restores the blob and saves again, requiring the two
blobs to match. Compares blobs rather than generated text, which cannot see a
field dropped on the way back in.
Note this passes on master for qwen4exp, so it does not demonstrate the
ext.x/ext.y drop this PR fixes; reaching that needs 2D mrope content.
* tests: give the synthetic qwen4exp a PLE so the state test bites
has_cell_ext() is n_pos_per_embd() > 1 || ple_n_heads > 0, and the indexer
cache sets rope_type = NONE, so without a PLE it serializes no cell ext at
all and the round-trip test cannot see a dropped ext.x/ext.y. With one,
removing the ext_set restore in state_read_meta fails the test: 198 of
335692 bytes differ, first at offset 282092.
Loading such a model needed two fixes:
- the row count of per_layer_token_embd came from require_weight(), which a
model synthesised from metadata alone has no file to answer. Derive it
from the head ranges and prefer the file's padded count where there is one.
- the PLE conv history is a row of the recurrent cache, so a PLE on a full
attention layer dereferenced a null p_l. Reject it at load time instead.
The meta mirror is skipped for qwen4exp. It returned NaN logits before this
fixture carried a PLE, which the nmse check passes since a NaN comparison is
false, and aborts with one. -sm tensor on real devices works.
Assisted-by: Claude
* llama: disable -sm tensor for qwen4exp
test-llama-archs skipped the tensor split for this arch from inside the
test, so the arch still advertised support it does not have. Declare it in
llm_arch_supports_sm_tensor instead and drop the test-side exception; the
existing llm_arch_supports_sm_tensor branch then does the skipping.
Assisted-by: Claude
for_each_token_in tested all LLAMA_MAX_SEQ sequences for every used cell,
while a cell almost always belongs to one. The scan now stops once the
cell's own sequences have been seen. Same visit order, same callback
arguments, so behaviour is unchanged.
get_prev_tokens is the only caller, so this affects the n-gram path.
RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs:
55k context generation 56.3 -> 74.3 t/s
132k context generation 33.6 -> 50.9 t/s
Prompt processing is unchanged, the scan is amortised over the ubatch
there. The gain follows the number of used cells, so it grows with
context and is invisible on short prompts.
* kv-cache : use ggml_set_rows
ggml-ci
* graph : separate k and v indices
ggml-ci
* cont : remove redundant ifs
ggml-ci
* kv-cache : improve find_slot impl
* kv-cache : bounds-check when accessing slot_info indices
* kv-cache : add comments
ggml-ci
* ggml : add TODOs for adding GGML_OP_SET_ROWS support in the backends
ggml-ci
* kv-cache : simplify the "struct llama_kv_cache" interface
ggml-ci
* kv-cache : revert the (n_swa + n_ubatch) change (for next PR)
ggml-ci
* kv-cache : some comments
ggml-ci
* context : fix graph reserve for multiple sequences
ggml-ci
* kv-cache : fix typo [no ci]
* kv-cache : fix find_slot() logic for free slots
ggml-ci
* llama : add TODO for deprecating the defrag API in the future
* kv-cache : improve find_slot() using min/max seq pos info
ggml-ci
* llama : handle aborts and compute errors
ggml-ci
* memory : extract state into llama_memory_state
ggml-ci
* kv-cache : add comments
ggml-ci
* server : update batching logic to reset n_batch on successful decode
* server : upon full re-processing, remove the sequence from the cache
* kv-cache : add TODO for doing split_equal when split_simple fails
ggml-ci