server : support unified cache across slots

[ggerganov]

ref #4130 (reply in thread)

Current logic in this PR (subject to change):

• When using unified KV cache with -kvu, share the entire context -c N among all parallel slots of the server -np N
• When we run out of space, try to free some by purging old sequences from idle slots, one by one, in no particular order
• If we still run out of space, terminate all active slots at once
• The -np N argument is still utilized to control the max number of parallel jobs, but it is no longer used to change the per-slot context
• By default, start the server using 4 slots and unified KV cache
• llama_context now caps the n_ctx_seq to a maximum of hparams.n_ctx_train

Example:

llama-server -m model.gguf -c 8192 --jinja

TODO:

• When we run out of space, terminate the active slots one-by-one and keep trying
• Think about instead of purging, to move the slot into host-memory cache. Not sure that this is really needed thanks to the existing logic from server : host-memory prompt caching #16391
• Add tests

Future improvements:

• When run out of space, terminate slots one by one instead of all together
• Update logic for starting a new task to check that it has some extra room for generation (not very sure if needed, current logic will simply purge one of the other slots, so it should be good as it is)

[slaren]

Should this value be capped when using unified cache to avoid exceeding the model context length? I think it could be set to min(n_ctx_train, n_ctx), or add a parameter to allow the user to change it.

[ggerganov]

I guess we can cap it to n_ctx_train. The only use case for n_ctx > n_ctx_train that comes to mind is self-extend, but lately this technique seems less relevant.

We can also cap it for the non-unified case?

Suggested change

	return cparams.kv_unified ? cparams.n_ctx : cparams.n_ctx / cparams.n_seq_max;
	return stdd:min(n_ctx_train, cparams.kv_unified ? cparams.n_ctx : cparams.n_ctx / cparams.n_seq_max);

[slaren]

We can also cap it for the non-unified case?

What would happen to the leftover slots? I may be misunderstanding the way split cache works, but my assumption would be that these slots would never be used, and it would be wasted memory. So if that's capped, it should be done at context creation.

[ggerganov]

Right, we should do the capping at context creation in the llama_context constructor. Currently we have some additional logic for this in llama-model:

llama.cpp/src/llama-model.cpp

Lines 19708 to 19724 in 7863fcc

	const auto padding = llama_kv_cache::get_padding(cparams);
	uint32_t n_ctx_per_stream = cparams.n_ctx;
	if (!cparams.kv_unified) {
	n_ctx_per_stream = (cparams.n_ctx + cparams.n_seq_max - 1)/cparams.n_seq_max;
	n_ctx_per_stream = GGML_PAD(n_ctx_per_stream, padding);
	cparams.n_ctx = n_ctx_per_stream*cparams.n_seq_max;
	} else {
	n_ctx_per_stream = GGML_PAD(n_ctx_per_stream, padding);
	cparams.n_ctx = n_ctx_per_stream;
	}
	LLAMA_LOG_DEBUG("%s: n_ctx = %u (padded)\n", __func__, cparams.n_ctx);

Since we no longer need the padding logic (as of #16148 and related) we should simplify this.

I'll push a separate PR for this and then will come back to polishing this one.

[ggerganov]

This is now rebased on top of the changes in #16812. The result is that we determine the KV cache size during context creation and there should be no leftover KV cells.

Note that since we now cap the context size to the training context size, the user code is recommended to query llama_n_ctx and llama_n_ctx_seq after creating the llama_context in order to obtain the actual context size. I'll add comments in llama.h to reflect this.

Will try to clean-up this PR next and will open it for review when ready.

[ggerganov]

This branch should not be reached due to the capping above on line 117. But keeping it in case the capping logic gets changed in the future.

[ggerganov]

Ready for review. I've marked some TODOs for follow-up PRs since I think the current implementation is quite basic and at the same time gets us 90% on the way to the ideal logic. Will improve the rest of the cases from master.