llama : reuse compute graphs

[ggerganov]

target #14285

Reuse computation graphs from the previous ubatch when possible. Works with any batch size and any model.

Note

The functionality currently requires the LLAMA_SET_ROWS from #14285

• CPU
• Metal
• CUDA
• Vulkan
• OpenCL
• SYCL

This functionality requires the ggml_set_rows() operator to be supported (see #14285). In order to be able to reuse a compute graph, its topology (shapes, strides, parameters, etc.) has to be entirely defined by the set of input tensors (e.g. inp_embd, inp_pos, inp_attn, etc.).

This PR adds logic to update a previous llm_graph_result by verifying that the new llm_graph_params would result in the same tensor shapes. For this to work, we should no longer preemptively reset the scheduler after processing a batch so that all buffers from the previous graph remain allocated and ready for reuse, in case the new ubatch is compatible. See the new llm_graph_result::update() method:

llama.cpp/src/llama-graph.h

Lines 506 to 525 in fc4fdf6

	// try to update the existing graph result using the new graph parameters
	// this can only be done if we determine that the resulting graph using the new graph parameters
	// would be identical to the existing graph. in that case, we simply have to update the memory
	// contexts of the input tensors of the graph and we can reuse it for another computation
	// return true if the graph was updated and can be reused
	bool update(const llm_graph_params & params) override {
	if (!this->params.is_same(params)) {
	return false;
	}
	bool res = true;
	for (auto & input : inputs) {
	res &= input->update(params);
	}
	return res;
	}

The other change that is needed is to introduce a way to swap the llama_memory_context of all graph inputs, so that the new call to llm_graph_result_i::set_inputs() uses the correct context from the current ubatch. This is performed by calling the llm_graph_input_i::update() method of all input tensors.

To enable this feature, define the LLAMA_SET_ROWS environment variable and add the new --graph-reuse CLI argument to the llama.cpp tools.

API Changes

• Add bool llama_context_params::graph_reuse. Default is false

Tests

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 -gr

LLAMA_SET_ROWS=1 ./bin/llama-parallel -m ../models/qwen2.5-3b-coder/ggml-model-q8_0.gguf -np 8 -ns 128 -s 1 -c 4096 -fa -n 128 -gr

Benchmark on M2 Ultra:

LLAMA_ARG_GRAPH_REUSE=1 LLAMA_SET_ROWS=1 ./scripts/compare-commits.sh gg/kv-cache-use-set-rows gg/llama-reuse-graphs -m models/qwen2.5-3b-coder/ggml-model-q8_0.gguf -m models/qwen2.5-3b-coder/ggml-model-q4_0.gguf -m models/qwen2.5-1.5b-coder/ggml-model-q4_0.gguf -m models/qwen2.5-1.5b-coder/ggml-model-q8_0.gguf -m models/gemma-3-4b/ggml-model-q4_0.gguf -m models/llama-3.2-1b-instruct/ggml-model-q8_0.gguf -fa 0,1 -t 1 -r 10 -n 1,32 -p 0

Model	Test	t/s master	t/s gg/llama-reuse-graphs	Speedup
gemma3 4B Q4_0	tg32	113.56	116.19	1.02
llama 1B Q8_0	tg32	275.25	280.95	1.02
qwen2 1.5B Q4_0	tg32	206.63	215.55	1.04
qwen2 1.5B Q8_0	tg32	174.63	180.79	1.04
qwen2 3B Q4_0	tg32	141.91	147.95	1.04
qwen2 3B Q8_0	tg32	111.40	114.48	1.03

TODO

• Clean-up and improve new interfaces and members
• Avoid graph input dynamic casts in is_same methods?
• Allow to reuse more models
• Manual user option to force disable of graph reuse?

Next PRs

• Remove llama_graph_result_i interface - does not seem to have any purpose
• Be able to compare the unique sequence ids of 2 ubatches
• Avoid passing ggml_cgraph * gf everywhere. Simply move it to llm_graph_context
• Try to reuse Metal graphs via MTLIndirectCommandBuffer
• Make CUDA reuse CUDA graphs using this new mechanism

[ggerganov]

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.

[compilade]

If update() is implemented for the recurrent cache, I think it could work even without adapting it to ggml_set_rows, because the head offset tends to be the same for similar consecutive ubatches in find_slot.

That might not work as well once multiple recurrent state cells per sequence are implemented (because they won't get re-used as much), but at that point it should be possible to use ggml_set_rows.

[ggerganov]

Yes, it has to work. As long as the check for the head is correctly added to the update() of the respective inputs it should be good.

[ggerganov]

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.

[am17an]

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

Model	FlashAttention	Test	t/s 99af79e	t/s cuda_set_rows	Speedup
gemma3 4B Q8_0	No	tg1	107.63	109.06	1.01
gemma3 4B Q8_0	No	tg32	110.20	110.77	1.01
gemma3 4B Q8_0	Yes	tg1	104.97	106.19	1.01
gemma3 4B Q8_0	Yes	tg32	106.32	106.32	1.00
llama 13B Q5_K_M	No	tg1	71.82	54.37	0.76
llama 13B Q5_K_M	No	tg32	71.49	68.12	0.95
llama 13B Q5_K_M	Yes	tg1	73.12	64.46	0.88
llama 13B Q5_K_M	Yes	tg32	72.19	66.93	0.93
llama 7B Q5_K_M	No	tg1	122.50	124.71	1.02
llama 7B Q5_K_M	No	tg32	119.88	121.22	1.01
llama 7B Q5_K_M	Yes	tg1	126.75	127.27	1.00
llama 7B Q5_K_M	Yes	tg32	125.98	126.06	1.00
qwen2vl 3B Q4_K_M	No	tg1	87.15	185.50	2.13
qwen2vl 3B Q4_K_M	No	tg32	191.90	192.95	1.01
qwen2vl 3B Q4_K_M	Yes	tg1	186.79	187.19	1.00
qwen2vl 3B Q4_K_M	Yes	tg32	187.53	187.20	1.00

To be sure, I ran it again

Model	FlashAttention	Test	t/s 99af79e	t/s cuda_set_rows	Speedup
gemma3 4B Q8_0	No	tg1	104.31	105.15	1.01
gemma3 4B Q8_0	No	tg32	104.44	110.46	1.06
gemma3 4B Q8_0	Yes	tg1	103.53	103.42	1.00
gemma3 4B Q8_0	Yes	tg32	104.16	95.79	0.92
llama 13B Q5_K_M	No	tg1	68.89	53.73	0.78
llama 13B Q5_K_M	No	tg32	66.89	70.87	1.06
llama 13B Q5_K_M	Yes	tg1	68.92	72.69	1.05
llama 13B Q5_K_M	Yes	tg32	71.02	69.24	0.98
llama 7B Q5_K_M	No	tg1	108.22	116.91	1.08
llama 7B Q5_K_M	No	tg32	113.81	117.13	1.03
llama 7B Q5_K_M	Yes	tg1	120.66	121.00	1.00
llama 7B Q5_K_M	Yes	tg32	124.38	121.43	0.98
qwen2vl 3B Q4_K_M	No	tg1	52.09	126.43	2.43
qwen2vl 3B Q4_K_M	No	tg32	188.24	191.41	1.02
qwen2vl 3B Q4_K_M	Yes	tg1	175.98	190.76	1.08
qwen2vl 3B Q4_K_M	Yes	tg32	180.16	188.93	1.05

[ggerganov]

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.

[am17an]

@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

Model	FlashAttention	Test	t/s 99af79e	t/s faa60717	Speedup
gemma3 4B Q8_0	No	tg1	111.89	110.63	0.99
gemma3 4B Q8_0	No	tg64	109.37	104.83	0.96
gemma3 4B Q8_0	No	tg128	105.82	112.28	1.06
gemma3 4B Q8_0	Yes	tg1	108.19	106.61	0.99
gemma3 4B Q8_0	Yes	tg64	108.57	101.07	0.93
gemma3 4B Q8_0	Yes	tg128	103.60	110.38	1.07
llama 13B Q5_K_M	No	tg1	57.01	58.88	1.03
llama 13B Q5_K_M	No	tg64	70.34	73.04	1.04
llama 13B Q5_K_M	No	tg128	69.97	68.06	0.97
llama 13B Q5_K_M	Yes	tg1	77.40	72.23	0.93
llama 13B Q5_K_M	Yes	tg64	71.20	73.46	1.03
llama 13B Q5_K_M	Yes	tg128	71.88	71.90	1.00
llama 7B Q5_K_M	No	tg1	119.34	129.08	1.08
llama 7B Q5_K_M	No	tg64	117.72	128.20	1.09
llama 7B Q5_K_M	No	tg128	115.36	120.00	1.04
llama 7B Q5_K_M	Yes	tg1	128.29	126.52	0.99
llama 7B Q5_K_M	Yes	tg64	125.85	131.65	1.05
llama 7B Q5_K_M	Yes	tg128	116.00	121.91	1.05
qwen2vl 3B Q4_K_M	No	tg1	141.47	136.11	0.96
qwen2vl 3B Q4_K_M	No	tg64	178.99	195.90	1.09
qwen2vl 3B Q4_K_M	No	tg128	175.23	193.03	1.10
qwen2vl 3B Q4_K_M	Yes	tg1	172.34	187.55	1.09
qwen2vl 3B Q4_K_M	Yes	tg64	185.92	175.19	0.94
qwen2vl 3B Q4_K_M	Yes	tg128	187.77	192.94	1.03

[slaren]

Are there any downsides of enabling this option, besides the removal of the preemptive scheduler reset?

[slaren]

The name of this function doesn't tell me much. I would expect it to be an operator==, but it clearly isn't. Why is it necessary for hparams and cparams? Shouldn't the graph result be associated to a llama_context?

[slaren]

need_update, can_reuse, ..