Method attention

cake-core/src/backends/rocm/mod.rs:380–396

(&self, q: &Tensor, k: &Tensor, v: &Tensor, scale: f32, causal: bool)

    }

    fn attention(&self, q: &Tensor, k: &Tensor, v: &Tensor, scale: f32, causal: bool) -> Result<Tensor> {
        let orig = q.dtype();
        let (qf, kf, vf);
        let (q, k, v) = if orig == DType::F32 { (q, k, v) }
            else { qf = q.to_dtype(DType::F32)?; kf = k.to_dtype(DType::F32)?; vf = v.to_dtype(DType::F32)?; (&qf, &kf, &vf) };
        let kt = k.t()?;
        let attn = (self.tensor_matmul(q, &kt)? * scale as f64)?;
        let attn = if causal {
            let (sl, kl) = (q.dim(2)?, k.dim(2)?);
            let mut mask = vec![0u8; sl*kl];
            for i in 0..sl { for j in 0..=kl.saturating_sub(sl)+i { if j<kl { mask[i*kl+j]=1; } } }
            let m = Tensor::from_vec(mask, (1,1,sl,kl), q.device())?;
            let ni = Tensor::full(f32::NEG_INFINITY, attn.shape(), q.device())?;
            m.broadcast_as(attn.shape())?.where_cond(&attn, &ni)?
        } else { attn };
        self.tensor_matmul(&candle_nn::ops::softmax_last_dim(&attn)?, v)?.to_dtype(orig)
    }

    fn silu_mul(&self, g: &Tensor, u: &Tensor) -> Result<Tensor> { (candle_nn::ops::silu(&g.contiguous()?)? * u.contiguous()?)?.contiguous() }
    fn stable_softplus(&self, x: &Tensor) -> Result<Tensor> { let t=Tensor::full(88f32,x.shape(),x.device())?.to_dtype(x.dtype())?; x.maximum(&(x.minimum(&t)?.exp()?+1.0)?.log()?) }