kv-cache : use ggml_set_rows

[ggerganov]

cont #14274

Utilize ggml_set_rows() for updating the KV cache.

• Make the graph static with respect to the KV cells head offset
• Relax the requirement for continuous KV slots of the input ubatch, making the defrag logic almost obsolete

Currently enabled only if the environment variable LLAMA_SET_ROWS is defined. If not, we fallback to the original way of updating the KV cache using a view + cpy of continuous slots. This is needed until the ggml_set_rows() implementation is finalized and supported by all backends.

Testing

# regular
LLAMA_SET_ROWS=1 ./bin/llama-cli -m ../models/llama-3.2-3b-instruct/ggml-model-q8_0.gguf \
     -p "I believe the meaning of life is" -n 32 --top-k 1 -fa

# SWA
LLAMA_SET_ROWS=1 ./bin/llama-cli -m ../models/gemma-3-4b/ggml-model-q8_0.gguf \
     -p "I believe the meaning of life is" -n 32 --top-k 1 -fa

Next PRs

• Introduce the concept of "virtual sequences"
• Extend llama_kv_cache_unified to support virtual sequences
• Extend the attention graph to support virtual sequences
• Enable efficient multi-sequence decoding
• Enable compute graph reuse

[rgerganov]

I tried this PR with the following change in the RPC backend:

diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp
index f468f796..dcbede89 100644
--- a/ggml/src/ggml-rpc/ggml-rpc.cpp
+++ b/ggml/src/ggml-rpc/ggml-rpc.cpp
@@ -761,6 +761,8 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g
     ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;
     std::vector<uint8_t> input;
     serialize_graph(cgraph, input);
+    auto graph_hash = fnv_hash(input.data(), input.size());
+    printf("RPC graph compute: hash = 0x%" PRIx64 ", size = %zu\n", graph_hash, input.size());
     rpc_msg_graph_compute_rsp response;
     auto sock = get_socket(rpc_ctx->endpoint);
     bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size(), &response, sizeof(response));

The compute graph doesn't change and produces the same hash with gpt2, tinyllama and mistral-7b models. However, the hash does change with gemma3 models. The serialized graph includes tensor addresses, so it's possible that we rebuild same tensors on different addresses resulting in different graph hash.

But in any way this looks like a great progress!

[ggerganov]

However, the hash does change with gemma3 models.

Yes, this is expected. I've applied the change only for the unified cache. For the unified+iswa, it is still using the ggml_cpy:

https://github.com/ggml-org/llama.cpp/pull/14285/files#diff-9be9eea14f4aefce7375482c05968900192634e88e92ac263cedb955a64ad7feR1290-R1291

[ggerganov]

Should work with Gemma now as well.

[ggerganov]

The non-FA path is now also supported, though I am not 100% sure this is the best way to do it.

[ggerganov]

The non-FA path is now also supported, though I am not 100% sure this is the best way to do it.

I don't observe any performance regression with CPU-only build, so the implementation should be good enough I think.

[rgerganov]

@ggerganov the following test segfaults on my machine:

$ bin/test-backend-ops -b CPU -p "type=iq2_xxs,n=256,m=5,r=4,b0=1,bs=1,v=0"
Testing 1 devices

Backend 1/1: CPU
  Device description: 11th Gen Intel(R) Core(TM) i7-11850H @ 2.50GHz
  Device memory: 31835 MB (31835 MB free)

  SET_ROWS(type=iq2_xxs,n=256,m=5,r=4,b0=1,bs=1,v=0): [1]    786752 segmentation fault (core dumped)  bin/test-backend-ops -b CPU -p "type=iq2_xxs,n=256,m=5,r=4,b0=1,bs=1,v=0"

[ggerganov]

Apply this patch:

diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp
index 735ef3f01..cc9b922fa 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.cpp
+++ b/ggml/src/ggml-cpu/ggml-cpu.cpp
@@ -416,6 +416,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
 
     switch (op->op) {
         case GGML_OP_CPY:
+        case GGML_OP_SET_ROWS:
             return
                 op->type != GGML_TYPE_IQ3_XXS &&
                 op->type != GGML_TYPE_IQ3_S   &&

[ggerganov]

@slaren This is ready to for a detailed review. I've prototyped 2 important use cases that can be enabled by adopting this change:

• llama : reuse compute graphs #14482 for reusing a graph from previous ubatch
• llama : add high-throughput mode #14363 for decoupling the KV cache buffers of separate sequences in order to improve multi-sequence performance

These still need some more work, which if we accept the current PR, I will do in next PRs.

For now, this PR updates the KV cache logic to start using ggml_set_rows() and splits the kv_idxs into 2 tensors: k_idxs and v_idxs. The latter is currently not relevant, but it will be needed for #14363.

Without setting the LLAMA_SET_ROWS environment variable, this branch should behave identical to master.

[slaren]

Replacing this with reserve and using push_back would remove the possibility of an OOB write.

[slaren]

Might be better to use at here for safety.