Function burn_attention 

pub fn burn_attention(
    q: Tensor<4>,
    k: Tensor<4>,
    v: Tensor<4>,
    scale: f32,
    mask: Option<Tensor<4>>,
) -> Tensor<4>

Scaled dot-product attention with additive mask. q/k/v: [batch, heads, seq_q/kv, head_dim].