Add option to build ggml-cuda as JIT using nvrtc

[WilliamTambellini]

Add option to build ggml-cuda as JIT using nvrtc.

As today ggml-cuda is AOT (Ahead Of Time):
all cuda kernels are compiled ahead of time using the local nvcc for a limited range of nvidia archs. This makes these embedded kernels only runnable on these archs. This also potentially makes ggml-cuda.so huge.

Another way to run device/npu/gpu kernels is to do JIT (Just in Time):
ggml-cuda to embed the source code of (some) kernels in the lib, link with nvrtc to compile the needed kernels/PTX, and only kernles which are needed, at runtime, just before the very first kernel execution.
Advantages: speeds up building ggml-cuda, limits the size of the lib, targets more archs dynamically, allows to parameter the kernels for the local hardware, and get better perf.
Drawbacks: the very first kernel exec is slower because needs a runtime compilation.
Refs:
https://docs.nvidia.com/cuda/nvrtc/
https://github.com/pytorch/pytorch/blob/b0a5d55c584792a504ec18600180e3d1200dfea6/torch/csrc/jit/tensorexpr/cuda_codegen.cpp#L1262
https://github.com/arrayfire/arrayfire/blob/360fefb3551a7c9f91250b0ec894aad76ec6a022/src/backend/cuda/compile_module.cpp#L153

@ggerganov
would you consider some PRs to add that option (no change of default behavior, still AOT/nvcc)?

[JohannesGaessler]

This would I think be a useful feature, particularly for builds with GGML_NATIVE=OFF. As of right now the code is compiled for compute capabilities 5.0, 6.1, 7.0, 7.5 because those are the minimum CCs for certain features. The default CMake behavior is I think to compile PTX code and PTXAS code for those CCs. However, in practice most people are using GPUs with CCs >= 8.0 so JIT compilation of PTX to PTXAS is already necessary in most cases. At the same time the AOT compilation of PTXAS is very slow. So I think it would make sense to use JIT compilation for all CCs by default if GGML_NATIVE=OFF.

[JohannesGaessler]

As of right now the code is compiled for compute capabilities 5.0, 6.1, 7.0, 7.5

Actually, I forgot that as of recently it's also being compiled for CC 8.0 but that doesn't really change anything about the conclusion.

[WilliamTambellini]

Indeed @JohannesGaessler
Confirming that today optimizing FA for a given GPU/CC would require a separate build of gggml-cuda: ref:
https://tridao.me/publications/flash2/flash2.pdf#subsection.3.3
eg many vars in
https://github.com/ggml-org/ggml/blob/master/src/ggml-cuda/fattn-tile-f16.cu
could be runtime finetuned to provide best speed:
#define FATTN_KQ_STRIDE_TILE_F16 64
constexpr int cols_per_block = 32;
...
RFC @slaren please

[JohannesGaessler]

I am assuming that what you're interested in is higher throughput for batched inference on modern hardware. The code you linked is not used for that, look at fattn-mma-f16.cuh instead. The code in that file uses a parallelization scheme that does not require any inter-warp data transfers.

Feel free to make a PR for better kernel parameters if you find any.