Function tmp_buffer_bytes 

pub fn tmp_buffer_bytes(num_heads: u32, head_dim: u32, q_seq_len: u32) -> usize

Compute the temp buffer size needed for tree-attention. Identical to flash_attn_vec’s tmp buffer (same output layout) — kept as a separate function for clarity at call sites.

ADR-037 Phase E4b.6 codex /cfa (2026-05-22): uses saturating_mul so adversarial inputs above usize::MAX clamp to usize::MAX instead of wrapping. Callers see a clearly-too-large requirement which fails any subsequent alloc/comparison rather than silently passing with a small wrapped value.