Struct Sdpa 

pub struct Sdpa;

Implementations

impl Sdpa

pub fn run_attention( &self, q: &Tensor, k: &Tensor, v: &Tensor, mask: &AttentionMask, flash_params: Option<&FlashParams>, sdpa_params: &SdpaParams, ) -> Result<Tensor>

Computes softmax(QK^T*sqrt(d_k))V

Inputs:

• q: (b_sz, n_attn_heads, q_len, head_dim)
• k: (b_sz, n_kv_heads, q_len, head_dim)
• v: (b_sz, n_kv_heads, q_len, head_dim)

Dispatch attention based on the AttentionMask variant:

• AttentionMask::CausalFlash: flash attention with is_causal = true
• AttentionMask::None: flash if available (decode), else eager without mask
• AttentionMask::Custom: eager attention with the explicit mask tensor

pub fn run_attention_noflash( &self, q: &Tensor, k: &Tensor, v: &Tensor, mask: Option<&Tensor>, sdpa_params: &SdpaParams, causal: bool, ) -> Result<Tensor>

Same as run_attention, but skips the flash-attention dispatch.

causal tells the Metal SDPA-full kernel to enable its upper-triangle skip (do_causal=true). Pass true only when the caller’s mask is causal-or-stricter. Pass false` for bidirectional masks (e.g. vision attention).