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.