3 ms·

You can probably get this to run slightly faster using CUDA graphs, which should reduce the amount of communication between the GPU and CPU. I notice that llama

by ryao 2y ago

You can probably get this to run slightly faster using CUDA graphs, which should reduce the amount of communication between the GPU and CPU. I notice that llama.cpp uses them, but glancing at your blog post, it appears that you do not. You might want to consider switching to bfloat16. The model weights are published as bfloat16. LLMs are more sensitive to the exponent bits than the mantissa bits, so using fp16 over bfloat16 should slightly hurt the model’s intelligence. Using bfloat16 would require you to use the tensor cores instead of the CUDA cores. Interestingly, llama.cpp performs very poorly when asked to use bfloat16, so performance comparisons of a bfloat16 would be best done against llama.cpp using fp16, since there is no reason for performance of fp16 and bfloat16 to be different, other than a lack of optimization. Using the tensor cores should not make your forward pass any faster despite the additional computational power since the GEMV operations that your forward pass uses should be memory bandwidth bound, although if you ever decide to try doing multiple queries in parallel, you will see a performance boost from being able to use that additional compute via GEMM operations. If you write a prompt processing version of the forward pass, you can use GEMM operations to process all prompt tokens simultaneously even without doing multiple queries in parallel. This would give you fast prompt processing speeds and open another front on which you can try to outdo llama.cpp.

reasonableklout 2y ago
Interesting! CUDA graphs indeed look very promising (https://developer.nvidia.com/blog/cuda-graphs/ https://developer.nvidia.com/blog/cuda-graphs/). I'd expected that kernel dispatch is not a problem because it is done asynchronously by the host and profiles suggested that the device thread is rarely idle, but from the NVIDIA blog I see the effective time per kernel goes from 3.8usec to 3.4usec. I guess there is still dispatch overhead that occurs on the device? Re: bfloat16. Yes, this is a good suggestion and I'm aware that bfloat16 is better suited for ML applications for the reason you mentioned (the range being more important than the precision) while offering the same footprint as float16. TBH, the main reason I haven't done this is because of lack of native support on my CPU (EPYC 7702P). This makes it a bit more annoying to maintain the format parity for the CPU and GPU backends, which is nice because it allows for easy testing, although if I wanted to preserve that I could just write non-performant test-only CPU code for bfloat16. Prompt pre-filling optimizations using GEMMs are great and I plan to explore this when I have free time. I talk about this in the "What's next" section briefly. Optimizing GEMM is a rabbit hole, as you point out in some of the other blog posts you shared. Will be a good next challenge!
- saagarjha 2y ago
  FWIW CUDA graphs are only really useful when you have a lot of kernels you want to launch. Otherwise they seem to perform similarly to just launching all the kernels in a loop on the host.
  - deleted 2y ago
    [deleted]
  - ryao 2y ago
    llama.cpp saw a 10% performance improvement from using CUDA graphs, so inference does benefit from it: https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-inference-with-cuda-graphs/ https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-in...
    - saagarjha 2y ago
      Yeah I mean I am sure there are workloads that it helps but there are also a lot where you'd think it would help and the driver will actually just fail to cooperate.
- ryao 2y ago
  I have not explored CUDA graphs yet in my own code, but llama.cpp saw a 10% performance improvement from it, so it certainly does something: https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-inference-with-cuda-graphs/ https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-in... As for why this is the case, anything I say would be pure speculation, but my current guess is that the standard API is actually synchronous between the GPU and CPU, and is only asynchronous on the CPU between userland and the kernel. This allows CUDA graphs to improve performance by turning all of the synchronous calls between the GPU and CPU into a single synchronous call. Again, this is purely a guess, but if it is right, it would explain how CUDA graphs are able to improve the performance of an asynchronous API. As for doing BF16 on CPUs that do not support BF16, you can emulate it in F32 using a widening instruction and a bitshift to load BF16 into F32 SIMD registers. Load the BF16 values into a 128-bit SIMD register. Then use _mm256_cvtepu16_epi32() to widen the 16-bit values into 32-bit values in a 256-bit register and do a left bitshift by 16 on them. This trick disregards subnormal numbers, but BF16 does not support subnormal numbers, so we would only need to worry about this when converting from BF16 to F32. F32 does support subnormal numbers, but on amd64 processors, I believe you can set bit 15 in the MXCSR register to enable flush to zero to zero subnormal results. Then a shuffle and permute could be used to convert F32 numbers to BF16 numbers like this: https://godbolt.org/z/dWPaW61Wf https://godbolt.org/z/dWPaW61Wf Unfortunately, I do not have an example of using shuffle and permutate to go in the opposite direction, so that will be an exercise for the reader. I will likely implement this in the CPU version of my llama3.c fork to make a BF16 version after I figure out how to write performant GEMM kernels, since I would want to have the fast prompt processing with BF16. As for GEMM, you can skip having to optimize it by using cuBLAS. Then implementing your own GEMM kernels is only needed if you want to do kernel fusion to for example, compute the max values for rmsnorm as part of the previous GEMM operation, which would reduce memory bandwidth usage slightly. You could look at how I generalized the forward() function into the precompute_input_logits() function to calculate all prompt processing via GEMM rather than GEMV: https://github.com/ryao/llama3.c/blob/master/run.c https://github.com/ryao/llama3.c/blob/master/run.c It was fairly straightforward, although the code is somewhat messy (as it is from a WIP personal research project). My precompute_input_logits() function also needs to be changed to stop assuming the starting position is 0, as hypothetically speaking, if this were made into a library and a chat frontend were put in front, end users could modify the chat history in the center, and rather than recompute everything from scratch, we only need to precompute tokens from the first changed token. However, that is a future improvement to explore. You might wonder why I am using the batched GEMM function in forward() when everything is on the CPU in that version. All I should need is GEMV there and my own handwritten AVX2 matmul() implementing GEMV is actually faster than the Intel MKL GEMV. However, shoehorning the GEMV into the Intel MKL batched GEMM makes is even faster. I believe it is reducing memory bandwidth requirements in the forward() pass by having matrix multiplication operations that have locality run in parallel. For computing K and V, the xb vector is shared, so if multiple cores read the same locations in xb, then only 1 read needs to go to system memory. Similarly, in the attention calculations, the matrices rows are interleaved, so the CPU is able to read rows for multiple cores simultaneously, and what would have been 32 separate passes over the weight matrices become 2 passes on my Ryzen 7 5800X. At least, this is my best guess why the batched API is faster than calling a GEMV routine sequentially 32 times there. Note that using more cores on my Ryzen 7 5800X does not increase memory bandwidth, since a single Zen 3 core in the Ryzen 7 5800X is able to use the CPU's full memory bandwidth (unlike cores in most other processors).

by ryao 2y ago

reasonableklout 2y ago

Interesting! CUDA graphs indeed look very promising (https://developer.nvidia.com/blog/cuda-graphs/ https://developer.nvidia.com/blog/cuda-graphs/). I'd expected that kernel dispatch is not a problem because it is done asynchronously by the host and profiles suggested that the device thread is rarely idle, but from the NVIDIA blog I see the effective time per kernel goes from 3.8usec to 3.4usec. I guess there is still dispatch overhead that occurs on the device? Re: bfloat16. Yes, this is a good suggestion and I'm aware that bfloat16 is better suited for ML applications for the reason you mentioned (the range being more important than the precision) while offering the same footprint as float16. TBH, the main reason I haven't done this is because of lack of native support on my CPU (EPYC 7702P). This makes it a bit more annoying to maintain the format parity for the CPU and GPU backends, which is nice because it allows for easy testing, although if I wanted to preserve that I could just write non-performant test-only CPU code for bfloat16. Prompt pre-filling optimizations using GEMMs are great and I plan to explore this when I have free time. I talk about this in the "What's next" section briefly. Optimizing GEMM is a rabbit hole, as you point out in some of the other blog posts you shared. Will be a good next challenge!

saagarjha 2y ago
FWIW CUDA graphs are only really useful when you have a lot of kernels you want to launch. Otherwise they seem to perform similarly to just launching all the kernels in a loop on the host.
- deleted 2y ago
  [deleted]
- ryao 2y ago
  llama.cpp saw a 10% performance improvement from using CUDA graphs, so inference does benefit from it: https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-inference-with-cuda-graphs/ https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-in...
  - saagarjha 2y ago
    Yeah I mean I am sure there are workloads that it helps but there are also a lot where you'd think it would help and the driver will actually just fail to cooperate.
ryao 2y ago
I have not explored CUDA graphs yet in my own code, but llama.cpp saw a 10% performance improvement from it, so it certainly does something: https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-inference-with-cuda-graphs/ https://developer.nvidia.com/blog/optimizing-llama-cpp-ai-in... As for why this is the case, anything I say would be pure speculation, but my current guess is that the standard API is actually synchronous between the GPU and CPU, and is only asynchronous on the CPU between userland and the kernel. This allows CUDA graphs to improve performance by turning all of the synchronous calls between the GPU and CPU into a single synchronous call. Again, this is purely a guess, but if it is right, it would explain how CUDA graphs are able to improve the performance of an asynchronous API. As for doing BF16 on CPUs that do not support BF16, you can emulate it in F32 using a widening instruction and a bitshift to load BF16 into F32 SIMD registers. Load the BF16 values into a 128-bit SIMD register. Then use _mm256_cvtepu16_epi32() to widen the 16-bit values into 32-bit values in a 256-bit register and do a left bitshift by 16 on them. This trick disregards subnormal numbers, but BF16 does not support subnormal numbers, so we would only need to worry about this when converting from BF16 to F32. F32 does support subnormal numbers, but on amd64 processors, I believe you can set bit 15 in the MXCSR register to enable flush to zero to zero subnormal results. Then a shuffle and permute could be used to convert F32 numbers to BF16 numbers like this: https://godbolt.org/z/dWPaW61Wf https://godbolt.org/z/dWPaW61Wf Unfortunately, I do not have an example of using shuffle and permutate to go in the opposite direction, so that will be an exercise for the reader. I will likely implement this in the CPU version of my llama3.c fork to make a BF16 version after I figure out how to write performant GEMM kernels, since I would want to have the fast prompt processing with BF16. As for GEMM, you can skip having to optimize it by using cuBLAS. Then implementing your own GEMM kernels is only needed if you want to do kernel fusion to for example, compute the max values for rmsnorm as part of the previous GEMM operation, which would reduce memory bandwidth usage slightly. You could look at how I generalized the forward() function into the precompute_input_logits() function to calculate all prompt processing via GEMM rather than GEMV: https://github.com/ryao/llama3.c/blob/master/run.c https://github.com/ryao/llama3.c/blob/master/run.c It was fairly straightforward, although the code is somewhat messy (as it is from a WIP personal research project). My precompute_input_logits() function also needs to be changed to stop assuming the starting position is 0, as hypothetically speaking, if this were made into a library and a chat frontend were put in front, end users could modify the chat history in the center, and rather than recompute everything from scratch, we only need to precompute tokens from the first changed token. However, that is a future improvement to explore. You might wonder why I am using the batched GEMM function in forward() when everything is on the CPU in that version. All I should need is GEMV there and my own handwritten AVX2 matmul() implementing GEMV is actually faster than the Intel MKL GEMV. However, shoehorning the GEMV into the Intel MKL batched GEMM makes is even faster. I believe it is reducing memory bandwidth requirements in the forward() pass by having matrix multiplication operations that have locality run in parallel. For computing K and V, the xb vector is shared, so if multiple cores read the same locations in xb, then only 1 read needs to go to system memory. Similarly, in the attention calculations, the matrices rows are interleaved, so the CPU is able to read rows for multiple cores simultaneously, and what would have been 32 separate passes over the weight matrices become 2 passes on my Ryzen 7 5800X. At least, this is my best guess why the batched API is faster than calling a GEMV routine sequentially 32 times there. Note that using more cores on my Ryzen 7 5800X does not increase memory bandwidth, since a single Zen 3 core in the Ryzen 7 5800X is able to use the CPU's full memory bandwidth (unlike cores in most other processors).