Consolidation of tensor copies to backend to reduce API overhead

[agray3]

Prerequisites

• I am running the latest code. Mention the version if possible as well.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

Before the generation of each token, tensors are copied from the host CPU to the backend GPU. There currently exists a separate copy for each iteration of a loop:
for (int input_id = 0; input_id < split->n_inputs; input_id++)
where, e.g. for llama3 8B Q4_K_M n_inputs is 6, resulting in 6 separate copies which has significant CUDA API overhead (when running on an NVIDIA GPU) caused by the multiple memcpy and streamSync calls:

These can be consolidated into a single copy, substantially reducing the overheads:

This can be done in an isolated change which involves introducing buffers to hold the multiple tensors in consecutive memory on host and device.

Motivation

The optimisation improves inference performance. At the moment the improvement is modest but significant (e.g. around 1% on Blackwell RTX Pro 6000 for llama3 8B Q4_K_M ). Optimisation of such CPU-side API overheads will be come increasingly important as GPUs continue to get faster while CPU performance is largely stagnant.

Implementation

PR: #15750

[agray3]

Closing as per comments in #15750

[slaren]

Despite the PR being rejected, I still think that this would be a worthwhile optimization to implement after a refactor of ggml_backend_sched is done. However, it may also require changes to the ggml-backend interface, since at the moment aliasing tensor data between different types is not guaranteed to work.

It may also be worth trying first if removing the synchronization between each tensor copy (so that all the copies can be queued asynchronously) could achieve similar results, since that would be a lot easier to implement.

[agray3]

Despite the PR being rejected, I still think that this would be a worthwhile optimization to implement after a refactor of ggml_backend_sched is done. However, it may also require changes to the ggml-backend interface, since at the moment aliasing tensor data between different types is not guaranteed to work.

It may also be worth trying first if removing the synchronization between each tensor copy (so that all the copies can be queued asynchronously) could achieve similar results, since that would be a lot easier to implement.

There are currently 2 cudaStreamSynchronize calls between every tensor copy (as shown in the above profile), one from a call to ggml_backend_synchronize before the copy, and then one within ggml_backend_cuda_buffer_set_tensor in ggml-cuda.cu after the copy . At least for the CUDA backend, none of these should be required since everything is in the same stream. We only need a cudaStreamSynchronize or event-based device-side sync after all tensor copies, before the CUDA Graph executes (since that is in a different stream). Yes, removing all (or some) of these should indeed give a benefit, but I'm not sure if that can be done in a trivial way, or if there is a risk removing any of these could harm other backends (or split backends).

[wishstudio]

I think the cleanest approach is to add support for memcpy operations as a ggml op. The op will become part of the graph such that any launch overhead will be eliminated even for the case of a single memcpy.

We can even make the current partial MoE experts loading logic into a separate op, enabling us to remove the code from the cluttered sched code. After that the compute function will become a simple for loop calling each split to compute.

If this is too hard to implement, I think the middle ground is to add a backend interface for batched memcpy and let the backend handle batching logic (cudaMemcpyBatchAsync in CUDA, vkCmdCopyBuffer in vulkan). We may want to implement this interface anyway as it can help code outside graph compute, such as model loading.