Comment by ryao | Hacker News Reader

Comment by ryao | Hacker News Reader

Parent Full thread

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.

Interesting! CUDA graphs indeed look very promising (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!

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.

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-in...

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.

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-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

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

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).

> 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.

Oh 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/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.

I plan to write my own GEMM implementations so I can do kernel fusion. I am using GEMM implementations from Nvidia and Intel until I have my own performant substitutes. So far, my initial attempts at writing my own GEMM implementations have not been performance competitive. Intel and Nvidia use tiling algorithms. I know (unintentionally) from profiling that Intel is doing a memcpy trick that is presumably similar to the one described here:

https://youtu.be/vPhG9KSxguM

It is also presumably similar to the packing technique mentioned here:

https://salykova.github.io/matmul-cpu

As for the annotations, some might be outdated. Much like the GPU version, I had published the CPU version at the request of another OSS developer and before I had reached the point where I would make a cleanup pass to address any stale/obsolete/missing comments. I made some effort to have good comments for my own benefit during the current stage of development, but I was not disciplined enough to make a cleanup pass unnecessary. That is partly why I claimed it is messy. I am happy to hear that they look good even before I have done a cleanup pass.

As for xb being in L3, that is part of it. I suspect that Intel has multiple threads doing “tiles” from both GEMM (GEMV in GEMM) operations in parallel and the reads and prefetches on xb are being coalesced by L3 cache. I also had the idea of adapting my handwritten matmul to operate on 2 outputs at the same time, and even tried it, but it did not work. At the time, I had thought this was sue to register spilling occurred because I ran out of ymm registers on Zen 3. Presumably, this experiment could be attempted on Zen 4 using the additional ymm registers that AVX-512 gives to AVX2.

However, two further thoughts occurred to me while writing this response. The first is that this experiment could be done on Zen 3 by doing do 6 rows at a time per matrix instead of 8. The second is that this will be throwing a larger number of parallel linear accesses at the core’s memory prefetcher than what I suspect Intel does, so this might fail to work even if the underlying idea about how Intel is getting better performance is accurate. Honestly, I had been surprised to see Zen 3 handle 9 linear accesses per loop iteration, as I had felt that doing 9 (8 rows + xb) had already been asking too much, but to my happy surprise, it was not.

Right now, some GPU-based experiments are at the top of my mental to do list. I will likely circle back to the CPU side to try to understand how to get this performance from my own GEMM/GEMV kernels at some point after doing them.

+1, rounding f32 to bf16 is helpful. For the other direction, the approach we take in Highway/gemma.cpp is to load a full vector of bf16, then shift/AND to isolate the odd/even elements and convert to float. These can execute two per cycle, whereas promoting 16->32 bit is often just one per cycle (though a different port than FMA).

Reply on news.ycombinator.com