pub fn attn_flash_threads(hd: usize, vd: usize) -> usize
The threads of a flash attention block: 256 over 64 queries (a thread 4 queries by 4 keys), 128 over 32 — where the 3060 measured each fastest.