3 ms·
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.
by 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).
- ryao 2y ago> we would only need to worry about this when converting from BF16 to F32 This should be: > we would only need to worry about this when converting from F32 to BF16 My apologies for the typo.
- deleted 2y ago[deleted]
- reasonableklout 2y agoOh right, I forgot that BF16 allocates the same number of sign and exponent bits as F32, so we can convert from BF16 to F32 by just padding zeros to the fraction bits. I'd guess conversion in the opposite direction would require rounding rather than truncation if we want to preserve as much intelligence as possible, so something like (https://hhhhhojeihsu.github.io/tensorflow_1.8_woboq/tensorflow_1.8_xla/tensorflow/tensorflow/core/lib/bfloat16/bfloat16.h.html#_ZN10tensorflow8bfloat1617round_to_bfloat16Ef https://hhhhhojeihsu.github.io/tensorflow_1.8_woboq/tensorfl...). In any case, a constraint of my exercise was to compare a hand-written inference implementation to existing inference engines (calm/llama.cpp) which both use FP16 for weight and KV cache quantization by default rather than BF16, so a BF16 backend will be an interesting addition but not so relevant to the blog post. Yes, we can totally use cuBLAS, and it's likely what I'd go with in a production implementation for any kernels which I wasn't planning on fusing, but as mentioned in the post, defeats the point of the exercise and usually less fun than doing things from scratch. That said, I do plan on exploring for myself at some point what kind of gain I could get from switching existing hand-rolled ops to cuBLAS/cuDNN (and then trying to match that). Thanks for linking your prefill code with GEMMs! Actually, your code looks pretty readable and well-annotated. It'll be cool digging in as your CPU backend seems better developed than mine. Cool observation about how using batched-GEMM instead of GEMV speeds things up in the single-batch forward() function. I am curious about your intuition behind the reduction in memory bandwidth ("if multiple cores read the same locations in xb..."). Is this about how xb might be kept in the L3 cache which is shared between cores in Zen 3? I wonder if you could validate this idea by adapting your handwritten matmul() function to operate on 2 matrices and outputs at the same time.
- ryao 2y ago