candding 0.1.0

//! Masked scaled dot-product attention shared by every family, with grouped-query support.

use candle_core::DType;
use candle_core::Device;
use candle_core::Result;
use candle_core::Tensor;
use candle_nn::ops::softmax_last_dim;

use crate::layers::mask::AttentionMask;

/// Head sizes candle's Metal SDPA kernel accepts, minus 512, whose F32 tile overflows threadgroup memory.
const SDPA_HEAD_DIMS: &[usize] = &[32, 64, 72, 80, 96, 128, 256];

/// At head size 256 an F32 tile fills the 32 KB threadgroup memory and Metal fails to build the pipeline.
const SDPA_F32_MAX_HEAD_DIM: usize = 128;

/// The device-free half of `metal_sdpa_applies`, so tests can assert it without Metal.
fn sdpa_accepts(head_dim: usize, dtype: DType) -> bool {
  SDPA_HEAD_DIMS.contains(&head_dim) && (dtype != DType::F32 || head_dim <= SDPA_F32_MAX_HEAD_DIM)
}

pub(crate) fn metal_sdpa_applies(device: &Device, head_dim: usize, dtype: DType) -> bool {
  device.is_metal() && sdpa_accepts(head_dim, dtype)
}

/// `q` is `(batch, heads, seq, head_dim)`, `k` and `v` `(batch, kv_heads, seq, head_dim)`; `kv_heads` divides `heads`.
pub fn scaled_dot_product(
  q: &Tensor,
  k: &Tensor,
  v: &Tensor,
  mask: &AttentionMask,
  scale: f64,
) -> Result<Tensor> {
  let head_dim = q.dim(3)?;
  if metal_sdpa_applies(q.device(), head_dim, q.dtype()) {
    return candle_nn::ops::sdpa(
      &q.contiguous()?,
      &k.contiguous()?,
      &v.contiguous()?,
      mask.additive(),
      mask.is_causal(),
      scale as f32,
      1.0,
    );
  }
  matmul_attention(q, k, v, mask, scale).map(|(context, _weights)| context)
}

/// The one matmul-path formula behind both public entry points, so the two cannot drift apart.
fn matmul_attention(
  q: &Tensor,
  k: &Tensor,
  v: &Tensor,
  mask: &AttentionMask,
  scale: f64,
) -> Result<(Tensor, Tensor)> {
  debug_assert!(
    mask.additive().is_some() || !mask.is_causal(),
    "a causal mask without a tensor reached the matmul path"
  );
  let (heads, kv_heads) = (q.dim(1)?, k.dim(1)?);
  if heads % kv_heads != 0 {
    candle_core::bail!("{heads} query heads are not a multiple of {kv_heads} kv heads");
  }
  let groups = heads / kv_heads;
  let k = repeat_kv(k, groups)?;
  let v = repeat_kv(v, groups)?;
  let scores = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * scale)?;
  let scores = match mask.additive() {
    Some(additive) => scores.broadcast_add(additive)?,
    None => scores,
  };
  let weights = softmax_last_dim(&scores)?;
  let context = weights.matmul(&v.contiguous()?)?;
  Ok((context, weights))
}

/// `scaled_dot_product` plus its softmax weights; always the matmul path, as the Metal kernel does not expose them.
pub fn scaled_dot_product_with_weights(
  q: &Tensor,
  k: &Tensor,
  v: &Tensor,
  mask: &AttentionMask,
  scale: f64,
) -> Result<(Tensor, Tensor)> {
  matmul_attention(q, k, v, mask, scale)
}

/// Tiles the kv heads to the query heads for the matmul path, kv head `h` serving query heads `h·groups..(h+1)·groups`.
pub(crate) fn repeat_kv(xs: &Tensor, groups: usize) -> Result<Tensor> {
  if groups == 1 {
    return Ok(xs.clone());
  }
  let (batch, kv_heads, seq, head_dim) = xs.dims4()?;
  xs.unsqueeze(2)?
    .expand((batch, kv_heads, groups, seq, head_dim))?
    .reshape((batch, kv_heads * groups, seq, head_dim))
}

/// Splits a fused `(3 * hidden, hidden)` projection by rows into query, key and value, `three` outermost;
/// reading it per head as `(heads, three, head_dim)` would mix the three with no visible sign.
pub fn split_fused_qkv(weight: &Tensor, hidden: usize) -> Result<(Tensor, Tensor, Tensor)> {
  let q = weight.narrow(0, 0, hidden)?.contiguous()?;
  let k = weight.narrow(0, hidden, hidden)?.contiguous()?;
  let v = weight.narrow(0, 2 * hidden, hidden)?.contiguous()?;
  Ok((q, k, v))
}

#[cfg(test)]
mod tests;