Function attn_flash_smem_bytes 

pub fn attn_flash_smem_bytes(hd: usize, vd: usize) -> usize

The flash attention block’s dynamic shared memory: its queries, two stages of a tile’s keys 32 values of the head at a time and of its values 16 keys at a time, and the tile’s weights — rows padded four values.