kv-cache : optimize restoring non-contiguous cells (#27991)
* kv cache : batch state restore scatter reads per contiguous run
When restoring state into non-contiguous destination cells (e.g. a
prompt-cache snapshot into a fragmented ring), state_read_data issued
one small copy per KV cell - ~1.4M copies of a few KiB each for a
40k+ token restore, taking 25-63 s on the CUDA backend.
The snapshot stores cell rows in cell order, so a maximal run of
consecutive destination indices maps to one contiguous block and can
be restored with a single copy. Precompute the runs once and use them
in all three scatter loops (K, V, transposed V). Byte-identical.
The on-device reader copies with a byte cursor when the read and
write chunking differs, so the batched reads are safe for it as well.
Batching makes equal tensor counts with a different split reachable
(save ranges [2,1] vs restore runs [1,2]); the next commit teaches the
reader's 1:1 path to fall back to the byte cursor in that case.
Verified in a production setup: 1,363,616 copies / 25-63 s -> 224
copies / 221-424 ms for the same restores (42,603 cells, 4 runs).
Assisted-by: Claude Code (unsloth/qwen3.8-27b)
* context : fall back to the byte cursor when read and write chunking differ
the on-device reader copies saved state back with a 1:1 copy by tensor
index whenever the write and read sides recorded the same number of
tensors, guarded by a per-tensor size assert.
equal tensor counts do not imply equal chunking: a state restore may
batch its reads per contiguous run of destination cells while the save
used per-range reads, so both sides can record two tensors that split
the same data differently, and the assert aborts in all builds.
compare the per-tensor sizes and only take the 1:1 path when the
chunking actually matches, otherwise fall through to the existing
byte-cursor copy. both sides enumerate the same logical data in the
same order, so the cursor copy is well-defined across tensor
boundaries.
Assisted-by: Claude Code (unsloth/qwen3.8-27b)
* tests : cover state restore scatter reads on host and on-device paths
decode the same prefix on two sequences, interleaving the seq 0 cells
between the seq 1 cells, so the seq 1 cells are isolated from each
other in the kv cache (three cells, two saved ranges). save the seq 1
state, free the interleaved seq 0 cells, and restore: the destination
is then non-contiguous (two runs), and the restore-side chunking has
the same tensor count as the save-side with a different split, so the
scatter path is batched per contiguous run and the on-device reader's
byte-cursor fallback is exercised.
the restored state is saved again on the host and compared byte for
byte with the first save: the blob is serialized in sequence cell
order, so the two saves are identical if and only if the scatter
restore wrote exactly the same KV content. this documents the
byte-identical guarantee of the run-batched scatter reads.
one test per io backend: the host (CPU) path and the on-device path.
Assisted-by: Claude Code (unsloth/qwen3.8-27b)
