Skip to main content

tmp_buffer_bytes

mlx_native::ops::flash_attn_vec

Function tmp_buffer_bytes 

Source 
pub fn tmp_buffer_bytes(num_heads: u32, head_dim: u32) -> usize
Expand description

Compute the size in bytes of the temporary buffer needed for flash_attn_vec.

The temp buffer stores partial results from NWG workgroups:

  • nrows * head_dim * NWG floats for the partial output vectors
  • nrows * 2 * NWG floats for the S and M values

ADR-034 task #89 (2026-05-21) — nrows = num_heads * q_seq_len. At the default q_seq_len = 1 this returns the same value as pre-task-#89 (num_heads * NWG * (head_dim + 2) * 4). For q_seq_len > 1 the buffer scales linearly with qL.