kv cache : optimize restoring non-contiguous cells

[itsnotoger]

Optimize KV cache for non-contiguous cells (agentic use)

Overview

I have been running qwen3.5+ models via Claude Code for many weeks, bottlenecked by some phase after tool results are visible in the harness, but before the PP progresses. Any tool-calling heavy conversation was dominated by this phase.
Config: --kv-unified, --cache-idle-slots (default), --parallel 2.

Performance Analysis

GPU utilization before (left half is representative of multi-turn tool usages). GPU utilization after, with the fix applied.

Fix

Commit 4529c66 introduced the slow path, to enable non-contiguous cell support, but throughput remains very slow.
To improve this, we identify contiguous runs in the non-contiguous data, and upload each run as a single read operation.

Tests

The second commit adds test cases for non-contiguous cells (AI-authored). Case 6 tests scattered cells restoring from host, case 7 on device.

Blast radius

Prior to this change, the save/restore API could expect equal read/write counts in llama-context, which no longer holds. A test case for this is also added.

Additional information

From experiments with additional logging:

596.37.218.629 I srv  oaicompat_ch:   history: 32 messages, 14 tool_calls, 14 tool results (45281 bytes total)
596.37.260.175 I srv  oaicompat_ch: chat template apply: 0.041s (tools=0, grammar_bytes=0, parser_bytes=765)
596.37.326.577 I srv  oaicompat_ch: tools: 25 definitions, 100862 bytes (raw json)
596.37.327.139 I srv  oaicompat_ch:   history: 32 messages, 14 tool_calls, 14 tool results (45281 bytes total)
596.37.390.540 I srv  oaicompat_ch: chat template apply: 0.063s (tools=25, grammar_bytes=64204, parser_bytes=130806)
596.37.466.421 I srv   operator (): chat_parser load (PEG compile): 0.004s (130806 bytes)
596.37.467.544 I slot get_availabl: id  1 | task -1 | selected slot by LRU, t_last = 35782129779
597.39.846.338 I srv          load: load: lcp scan 0.05 ms, set_data_ext 62325.56 ms, clear 53.84 ms (size = 2778.599 MiB)
597.40.154.192 I srv  get_availabl: prompt cache update: save 0.00 ms, load 62687.31 ms, total 62687.31 ms
597.40.154.236 I common_sampler_init: chain_init 0.002 ms (grammar_bytes=64204, triggers=1)
597.40.154.239 I common_sampler_init: pre_grammar (trigger loop) 0.006 ms
597.40.158.350 I slot launch_slot_: id  1 | task 176954 | processing task, is_child = 0
597.40.280.832 I srv         alloc: alloc: obsolete+evict 0.00 ms, resize 119.39 ms (tgt = 1442.834 MiB, drft = 152.797 MiB)
597.42.012.616 I srv   prompt_save: prompt_save: alloc 238.63 ms, get_data_ext 1612.57 ms (tgt = 1442.834 MiB, drft = 152.797 MiB)

Notably: set_data_ext takes 62 seconds. The restore-time correlates strongly with cell-count.
In a production setting (Windows, 200k context), transfers usually take from 20 to 60 seconds.

From there on, the bottleneck was identified to the documented "slow path". When this path is hit, potentially millions of io.read_tensor() calls are being performed individually, resulting in effective memory transfer speeds around ~50 MB/s.

Fixes #20854

Discussion

I believe my use case is the common one, yet I don't see this bottleneck talked about much. A harness usually comes with: 1. a main conversation thread, and 2. an auto-classifier for tool call safety evaluation. This is essentially the most basic use case, even if you do not use subagents. Although you can configure a harness with longer timeouts and a single slot, I believe allowing parallel requests makes sense here.
Moreover, unified KV cache also makes sense, since many users will run models just at the edge of what their VRAM allows.
Due to aforementioned at least two conversation threads perpetually occurring, the allocator produces non-contiguous cells in the KV cache. My understanding is that (large, ~2 GB) RAM cached checkpoints need host/device transfers, which is where this bottleneck lies. I want to make the case that this is not a minor optimization for specific use cases, but an important general one.
The commonly user-cited metrics PP and TG do not capture this bottleneck, which I think shadows the importance of this optimization.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes. AI was used to track down the bottleneck with debug instrumentation, and to author the fix. AI also reviewed the fix branch extensively.
  • I tested performance of each build manually.
  • I ensured that performance of my local build matches the prebuilt release prior to and after optimizations.
  • I reviewed the diff. It is a straightforward identification of contiguous runs, and upload thereof in three places, replacing the previous slow path logic (excluding the on-device case).
  • But I am not deeply familiar with C++ or this codebase.
• Since I am on Windows I cannot easily run full local CI, but the new test cases were extensively ran, with both CPU and CUDA, under differnt FA and -ngl settings.
  • Perplexity/llama-bench are not expected to be affected by this change, since they don't seem to be touching restore API

[ggerganov]

I think you need to gate the new tests only for COMMON_CONTEXT_SEQ_RM_TYPE_PART contexts. For example the DSv4 memory does not support partial sequence removal so it fails the tests:

https://github.com/ggml-org/llama.cpp/actions/runs/33293259535/job/99212684620?pr=27991#step:8:12896

[ggerganov]

That also works and it's actually better than what I suggested.

[ggerganov]

Hm, the webgpu error would need an extra look: https://github.com/ggml-org/llama.cpp/actions/runs/33296367390/job/99216781100?pr=27991#step:4:3861

[ggerganov]

cc @ggml-org/ggml-webgpu or @fairydreaming in case you have a suggestion

[fairydreaming]

cc @ggml-org/ggml-webgpu or @fairydreaming in case you have a suggestion

@ggerganov From what I see WebGPU CopyBufferToBuffer requires sizes and offsets to be multiplies of 4, while (I suppose) ggml_backend_tensor_get() that uses it internally has no such restrictions. In the implementation size is rounded to multiplies of 4, but offset is not - that's why it crashes.

I tested this with this fix that rounds global offset to a multiple of 4 and adds local offset that is used during the final memcpy:

diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index b953118a7a3..75a21a33610 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -3713,11 +3713,18 @@ static void ggml_backend_webgpu_buffer_get_tensor(ggml_backend_buffer_t buffer,
 
     size_t total_offset = ggml_webgpu_tensor_offset(tensor) + offset;
 
-    size_t final_size = size;
-    if (size % 4 != 0) {
+    size_t local_offset = total_offset % 4;
+    if (local_offset != 0) {
+        // If offset is not a multiple of 4, we need to round it down to the previous
+        // multiple of 4
+        total_offset = total_offset - local_offset;
+    }
+
+    size_t final_size = size + local_offset;
+    if (final_size % 4 != 0) {
         // If size is not a multiple of 4, we need to round it up to the next
         // multiple of 4
-        final_size = size + (4 - (size % 4));
+        final_size = final_size + (4 - (final_size % 4));
     }
 
     std::lock_guard<std::recursive_mutex> lock(buf_ctx->global_ctx->mutex);
@@ -3748,7 +3755,7 @@ static void ggml_backend_webgpu_buffer_get_tensor(ggml_backend_buffer_t buffer,
     const void * mapped_range = buf_ctx->global_ctx->get_tensor_staging_buf.GetConstMappedRange(0, final_size);
 
     // Copy the data from the mapped range to the output buffer
-    std::memcpy(data, mapped_range, size);
+    std::memcpy(data, (const void *) ((const char *) mapped_range + local_offset), size);
     buf_ctx->global_ctx->get_tensor_staging_buf.Unmap();
     WEBGPU_CPU_PROFILE_TOTAL_END(get_tensor, buf_ctx->global_ctx);
 }

$ ctest -R test-save-load-state
Test project /home/phm/projects/llama.cpp/build-webgpu
    Start 28: test-generate-models
1/2 Test #28: test-generate-models .............   Passed   34.36 sec
    Start 32: test-save-load-state
2/2 Test #32: test-save-load-state .............   Passed   74.64 sec

100% tests passed, 0 tests failed out of 2

Edit: fixed a bug in the fix.

[fairydreaming]

Added #28045 with the WebGPU fix, let's see how it goes.

[fairydreaming]

@itsnotoger merge the current master into your branch, WebGPU fix is already there

[gianlorenzomungiovino]

Here qwen3.6 35b a3b mtp. I noticed a little bit gain in eval rate (just few t/s) and prefill (about 50t/s) velocities. Thank you