zsfm_burn/
attn.rs

1use burn::tensor::{activation, Tensor};
2
3/// Scaled dot-product attention with additive mask.
4/// `q/k/v`: [batch, heads, seq_q/kv, head_dim].
5pub fn burn_attention(
6    q: Tensor<4>,
7    k: Tensor<4>,
8    v: Tensor<4>,
9    scale: f32,
10    mask: Option<Tensor<4>>,
11) -> Tensor<4> {
12    let scores = q.matmul(k.permute([0, 1, 3, 2])).mul_scalar(scale);
13    let scores = match mask {
14        Some(m) => scores + m,
15        None => scores,
16    };
17    let probs = activation::softmax(scores, 3);
18    probs.matmul(v)
19}
20
21#[cfg(test)]
22mod tests {
23    use super::*;
24    use burn::tensor::{Device, TensorData};
25
26    #[test]
27    fn uniform_attention_averages_values() {
28        let device = Device::flex();
29        let q = Tensor::<4>::zeros([1, 1, 2, 1], &device);
30        let k = Tensor::<4>::zeros([1, 1, 2, 1], &device);
31        let v = Tensor::<4>::from_data(TensorData::from([[[[2.0f32], [4.0]]]]), &device);
32        let out = burn_attention(q, k, v, 1.0, None);
33        let vals: Vec<f32> = out.to_data().try_to_vec().unwrap();
34        assert!((vals[0] - 3.0).abs() < 1e-5);
35        assert!((vals[1] - 3.0).abs() < 1e-5);
36    }
37}
zsfm_burn/attn.rs

zsfm_burn/
attn.rs
