Graphs in llama.cpp

[abhinav9629]

Hi
I'm trying to debug and understand the codebase of llama.cpp
I have a couple of questions regarding graph creation.

• In llama_new_context_with_model whats the need of two graphs pp and tg graph?
• What exactly is the worst case scenario for which we reserve a worst case graph?
• Why we need to reserve a worst case graph again in case kv cache defragments or K shift is there?

[danbev]

I'll try to answer and hopefully others can correct me if I'm wrong about anything here.

The two computation graphs are for the prefill prompt (pp), which is the initial prompt and usually consists of multiple tokens, and the token generation (tg) computation graph (single token).

The worst case is for the case where the prefill prompt is completely filled with the maximum number of tokens.
The token generation reservation I believe is for getting the number of graph splits and the number of nodes. And after this the pp graph is reserved again so that the reservation does not have to be done at inference time, as the tensor allocation sizes will be that of the tg graph nodes without this "re-reservation".

Regarding the kv cache I think this because if a k-shift is needed or a defragmentation, then the scheduler is reset and a new graph is built (llama_build_graph_k_shift or llama_build_graph_defrag), the scheduler is allocated with the new graph which is then computed. After this the computation graph for pp is again built and reserved to reset the scheduler to the state it was before the k-shift/defrag.