Could you elaborate more on this?

-7 as it will be padded to 8 anyway

A quick question: How did you choose 512 for the buffer size?

I think I got (64 + 256) * sizeof(float32) for other __shared__ variables by reading the CUDA kernel, so I just rounded it up to 512 * sizeof(float32) to be safe. But it may be too conservative.

@WoosukKwon would appreciate if you could provide a more accurate measurement :)

@Yard1 Do you mean

__shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];

and

__shared__ float red_smem[2 * NUM_WARPS];

?

yes - they are included in the shared memory usage - we should set the buffer to upper bound of those

If my calculation is correct, the size of q_vecs is head_size * sizeof(scalar_t) <= 256 * 4 = 1024. The size of red_smem is obviously 64 * 4 = 256. In total, it's 1280 bytes (=320 float elements). So 512 is actually a bit conservative upper bound. However, I think this is acceptable.

Isn't this always 4?

It should technically be, but that way it ensures it's always true irrespective of the platform/implementation and is also self documenting

To my knowledge, the size of float is defined as an IEEE standard and is independent from the underlying machine architecture (unlike integer types). That being said, I like that this is self-documenting. Let's keep it!