oa::ml::matrix::scaled_dot_product_attention_causal

Apply multi-head causal scaled dot-product attention to packed [B*S, D] values. Errors Returns an error unless Q, K, and V are equal nonempty same-engine F32 matrices, their row count is divisible by sequence_length, their width is divisible by num_heads, or runtime recording fails.

FnRust public API

Function Documentation

pub fn scaled_dot_product_attention_causal( query: &Matrix, key: &Matrix, value: &Matrix, sequence_length: usize, num_heads: usize ) -> Result<Matrix>
Apply multi-head causal scaled dot-product attention to packed [B*S, D] values. Errors Returns an error unless Q, K, and V are equal nonempty same-engine F32 matrices, their row count is divisible by sequence_length, their width is divisible by num_heads, or runtime recording fails.