This should be ready for review. Currently, there is some small gain for Metal where ggml_set_rows() is available. We basically save the time for creating a new ggml_cgraph for each ubatch.

It would be interesting to try to reuse the Metal command buffers to speed this up even further on the backend side. Currently, we use MTLCommandBuffer to encode the compute commands and these objects do not allow the commands to be reused. However, according to the Apple documentation, the MTLIndirectCommandBuffer can be reused multiple times, so it seems to be what we need. It's still not clear to me how to encode compute commands to it (the docs only show examples of encoding rendering commands), but it might be possible. Any hints would be highly appreciated.

It would be interesting to try to reuse the Metal command buffers to speed this up even further on the backend side. Currently, we use MTLCommandBuffer to encode the compute commands and these objects do not allow the commands to be reused. However, according to the Apple documentation, the MTLIndirectCommandBuffer can be reused multiple times, so it seems to be what we need. It's still not clear to me how to encode compute commands to it (the docs only show examples of encoding rendering commands), but it might be possible. Any hints would be highly appreciated.

An update on this - it seems that MTLIndirectCommandBuffer won't cut it. However, I think an alternative simple approach of preparing 2 MTLCommandBuffer for each compute would work. The second command buffer will be prepared asynchronously after we commit the first one and on the next iteration, if the application informed us that we can reuse the graph, we will directly commit the prepared command buffer from the previous iteration and asynchronously prepare a new one for the next iteration. I will attempt an implementation of this in a follow-up PR and try to utilize ggml_backend_graph_plan_t in the process.

Edit: prototyped this in #14570. Does not seem worth pursuing as the gains are microscopic.

I tested this on CUDA on various models I had lying around and I see there a perf regression on a larger model, not sure if I'm doing something wrong. I cherry-picked this PR on top of #14551 and compared with commit1 = 14551, and commit2 = 14551 + this PR. Also interesting is 2x(?) speed up on qwen2lvl 3B

To be sure, I ran it again

The tg1 tests are irrelevant - they are just for warming up the next test. So the x2 is just fluctuation.

The graph_reuse parameter in llama-bench is disabled on this branch. I pushed a change to toggle it with -gr <0|1> CLI arg, but it won't work with the compare-commits.sh script.

But before benchmarking, I think you should make sure that the generated results when graph reuse is enabled are coherent (see the llama-cli and llama-parallel commands). Make sure that they report successful graph reusing at the end of the execution:

0.01.045.581 I llama_perf_context_print: prompt eval time =      36.08 ms /     8 tokens (    4.51 ms per token,   221.75 tokens per second)
0.01.045.582 I llama_perf_context_print:        eval time =     399.58 ms /    31 runs   (   12.89 ms per token,    77.58 tokens per second)
0.01.045.582 I llama_perf_context_print:       total time =     441.40 ms /    39 tokens
0.01.045.584 I llama_perf_context_print:    graphs reused =         30        <--- pay attention to this number
0.01.045.992 I ggml_metal_free: deallocating

And also from the results that you posted it seems that there is a lot of variability on your system for this tg32 test. Try to narrow down what is causing it and make sure the results are stable before benchmarking the graph reuse. Or use larger tg128 test and more repetitions.

@ggerganov I re-ran with r=100, and tg 64 and 128, I see quite a bit of variability at tg128, but tg64 is pretty tight (<1% variability). Also confirming the llama_perf_context_print

I am not sure what would happen if e.g. in a call to kv_self_update the scheduler is reset. Is this detected to prevent reusing the graph?

In acaf4b7 I made a relatively significant change to how the llama_ubatch data buffers are managed. The old implementation was keeping the ubatch data buffers in the llama_batch_allocr with the idea of avoiding memory allocations for each new llama_ubatch. Hence the llama_ubatch was just pointing to data in llama_batch_allocr. However, this logic was making it complicated to check if graph reusing is possible, since the old llama_ubatch could be pointing to buffers that are now being used for the new llama_ubatch. So I changed this to no longer keep the llama_ubatches data inside the balloc, but instead to carry it along and reference count it with std::shared_ptr.

@slaren Might be worth to take another look just in case. Even if not perfect, the results should be correct and graph reuse with split KV cache should now be functional. I will be improving the implementation in a follow-up PR.