bench : cache the llama_context state at computed depth

[ggerganov]

See #16944 (comment)

Sample commands:

make -j && ./bin/llama-bench -m ../models/gpt-oss-20b/ggml-model-mxfp4.gguf -t 1 -fa 1 -b 16384 -ub 2048 -d 0,1024,2048,4096,8192,16384,32768 -n 32 -p 2048

make -j && ./bin/llama-bench -m ../models/qwen2.5-3b-coder/ggml-model-q4_k.gguf -fa 1 -d 1024 -p 512 -ctk f16,q8_0

[JohannesGaessler]

I should clarify, for a MoE model this is not going to work correctly. Because the expert selection depends on the numerical input values leaving the memory of the KV cache uninitialized is going to bias the results.

My understanding is that recently functionality was added to the server which swaps on-device KV cache with RAM. If that could be repurposed for an implementation like this it would I think already be fast enough:

1. Do depth run with default batch size of 512.
2. Save the state of the KV cache.
3. For each benchmark run, load the KV cache state first instead of re-calculating it.

[ggerganov]

@JohannesGaessler Good idea - pushed a version that I think does what you described.

[JohannesGaessler]

Thank you, this version seems to be working correctly. One caveat is that the order in which you do the batch sizes matters. Preferably you would run the large batch sizes first as they are going to process the first, uncached KV context much faster. However, the syntax -ub "1-512*2" will do the batch sizes in a suboptimal order (for my purposes this doesn't matter because I'm going to script it anyways). One solution would be to do the depth run always with a constant batch size but previously when we tried that that was causing issues (and I'm not sure it would be worth the opportunity cost to fix).

[ggerganov]

One solution would be to do the depth run always with a constant batch size but previously when we tried that that was causing issues (and I'm not sure it would be worth the opportunity cost to fix).

Can you point me to this previous attempt?

[JohannesGaessler]

#13096

[ggerganov]

One solution would be to do the depth run always with a constant batch size but previously when we tried that that was causing issues (and I'm not sure it would be worth the opportunity cost to fix).

Hm yeah, not sure what's the best way to do that.

[slaren]

Wouldn't this need also to check if the model is the same, using the same KV type, or other parameters that may make the cache incompatible with the current test?

[ggerganov]

llama_state_seq_set_data should (in theory) return error (i.e. 0) when the state is incompatible with the current llama_context. I did try a few cases (different models, different KV cache types) and it seems to work as expected.

But it is a bit risky if somehow it's internal checks fail to detect an incompatibility, which can lead to invalid benches. So not sure - we can simply the logic to just reuse the state for the repetitions of the same test?

[ggerganov]

Did some extra tests with SSMs and it works as expected. Let's keep an eye out just in case - using --progress is useful for debugging this feature.