Why we use q,k,k_max = (64, 64, 32) / (64, 128,128) / (64, 128, 2e16) for sm80?

[ZhangDY-6483]

In generate_kernel.py, we generate one specific arch's kernel with some fixed shapes, such as q,k,k_max = (64, 64, 32) / (64, 128,128) / (64, 128, 2e16) for sm80. I'm not sure why these three shapes are selected but not others?
(Did we emulate all the possible shape and test the performance of them and select the best combination? )

[danthe3rd]

Hi,
You might get better performance if you know ahead of time the embedding dimension and hardcode it in the kernel indeed.
However, there is a tradeoff because pre-building a lot of kernels will take a lot of space and compilation time.

[ZhangDY-6483]

Hi, thanks for your reply ~ And why we select three fixed shape as patterns during generating kernel (these three shapes are slected by xformers): xformers/csrc/attention/cuda/fmha/generate_kernels.py ?

[danthe3rd]

The sizes we select are aligned with what researchers usually use, and what is best handled by the kernel

[ZhangDY-6483]

Got it ! Thank u very much~