Method causal_mask

cake-core/src/backends/mod.rs:533–553

Create a causal attention mask. Returns a U8 tensor of shape `(seq_len, kv_len)` where 1 = masked (future position), 0 = attend. Callers use `masked_fill` or `where_cond` to apply the mask.

(
        &self,
        seq_len: usize,
        kv_len: usize,
        device: &Device,
    )

    /// where 1 = masked (future position), 0 = attend.
    /// Callers use `masked_fill` or `where_cond` to apply the mask.
    fn causal_mask(
        &self,
        seq_len: usize,
        kv_len: usize,
        device: &Device,
    ) -> Result<Tensor> {
        if seq_len == 1 {
            return Tensor::zeros((1, kv_len), DType::U8, device);
        }
        // Build full (seq_len, kv_len) mask in one allocation — no Tensor::cat.
        let prefix = kv_len.saturating_sub(seq_len);
        let mut mask = vec![0u8; seq_len * kv_len];
        for i in 0..seq_len {
            let row_start = i * kv_len + prefix + i + 1;
            let row_end = (i + 1) * kv_len;
            if row_start < row_end {
                mask[row_start..row_end].fill(1);
            }
        }
        Tensor::from_vec(mask, (seq_len, kv_len), device)
    }

    /// Top-K selection: returns `(values, indices)` for the largest K elements
    /// along the last dimension. Uses partial sort O(N + K log K) for large N.