Add option to build and use some flashinfer cuda kernels

[WilliamTambellini]

@slaren @ggerganov
would you consider a PR adding an option to build and take advantage of these header only cuda kernels in ggml-cuda:
https://docs.flashinfer.ai/installation.html#c-api
https://github.com/flashinfer-ai/flashinfer
?

[ggerganov]

I am not very familiar with the FlashInfer implementation, so cannot say if it's a good idea. If it can be incorporated in some reasonable way (i.e. no extra (or few) dependencies, easy to deploy on different OS / hardware, compatibility with existing ggml infrastructure), then it's something to consider.

In general, it will mostly depend on how feasible this is for local inference. Correct me if I'm wrong, but the use cases of FlashInfer so far mainly for server/cloud deployment.

[WilliamTambellini]

If it can be incorporated in some reasonable way (i.e. no extra (or few) dependencies, easy to deploy on different OS / hardware, compatibility with existing ggml infrastructure), then it's something to consider.

• deps: apparently no deps if AOT (just some new cuh headers). If JIT, then likely nvrtc: https://docs.nvidia.com/cuda/nvrtc/
• easy to deploy: well a first impl could be to just use some of the their cuda kernels implementation so no change of deployment I guess

the use cases of FlashInfer so far mainly for server/cloud deployment.

Not necessarily: seems to benefit any nvidia gpu with a minimum of CUDA and Tensor Cores. Now, it s likely a cuda12+ tech.

Let us profile a bit our LLMs and identify which kernel would we favor for our usecase.
CU, tks.

[JohannesGaessler]

There are in essence only two operations that are performance critical and where transplanting CUDA code could potentially make sense:

1. Matrix multiplications: For floating point data NVIDIA already provides a highly optimized BLAS library in cuBLAS, for our custom quantized data formats FlashInfer they most likely does not have support. So for matrix multiplications there would likely be no benefit.
2. Fused attention kernels like FlashAttention: For long contexts and many parallel users in particular the attention takes up a large percentage of the runtime. The ggml CUDA code uses custom FlashAttention kernels that I wrote, it could very well be that another kernel would be faster. However, the attention performance is tightly linked to the KV cache data layout and it may very well be that for the code to work at all it would be necessary to make changes outside the CUDA backend. Also I think a very relevant question for transplanted code would be who would take over maintenance of that code. You cannot expect transplanted code to work or be faster for all combinations of tensor shapes and hardware so someone will need to put in the work to continually figure out the optimal code paths as the project develops further.

[JohannesGaessler]

I forgot: my understanding is that right now multiple sequences are handled in llama.cpp by adding the tokens to the same KV cache and then applying masking. I think that unless there is significant overlap between sequences it would be better to instead have multiple separate KV caches so that each sequence only needs to iterate over its own tokens. Is that still the case? My intuition would be that that would make a larger impact for throughput with many users than changing the CUDA code.

[ggerganov]

@JohannesGaessler Yes, this is my intuition as well and supporting separate KV caches per sequence is something that we should eventually implement. My expectation/hope is that when we do that, the performance of llama.cpp for multi-user / large context cases would become up to par with other dedicated engines, so it is definitely what we should prioritize for better CUDA performance.

My point was that we should not straight-up close the door for supporting other kernel implementations before getting familiar with the details, even if it seems very unlikely to do so.