candding 0.1.0

//! Additive attention terms built once per forward pass and handed to every layer.
//! Padding, causal masks and per-head biases (MPNet, ALiBi) all enter here, so attention has no per-family code.

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;

use crate::error::CanddingError;
use crate::error::Result;

/// Most negative finite value for a dtype; `-inf` would turn fully masked rows into NaN.
pub fn mask_min(dtype: DType) -> f64 {
  match dtype {
    DType::F16 => -65504.0,
    DType::BF16 => -3.3895e38,
    _ => f32::MIN as f64,
  }
}

/// `additive` broadcasts to `(batch, heads, seq, seq)` and already has that shape, as a view, where Metal SDPA runs.
/// `causal` is applied by the Metal kernel itself and folded into `additive` on every other path.
pub struct AttentionMask {
  additive: Option<Tensor>,
  causal: bool,
}

impl AttentionMask {
  pub fn none() -> Self {
    Self {
      additive: None,
      causal: false,
    }
  }

  pub fn additive(&self) -> Option<&Tensor> {
    self.additive.as_ref()
  }

  pub fn is_causal(&self) -> bool {
    self.causal
  }

  /// Bidirectional padding term, `(batch, 1, 1, seq)`, or none when no row is padded; one `lengths` entry per row.
  pub fn padding(
    attention_mask: &Tensor,
    lengths: &[usize],
    num_heads: usize,
    head_dim: usize,
    dtype: DType,
  ) -> Result<Self> {
    let (batch, seq) = attention_mask.dims2()?;
    check_lengths(lengths, batch)?;
    if lengths.iter().all(|&len| len == seq) {
      return Ok(Self::none());
    }
    let inverted = attention_mask.to_dtype(DType::F32)?.affine(-1.0, 1.0)?;
    let mask = inverted
      .affine(mask_min(dtype), 0.0)?
      .to_dtype(dtype)?
      .reshape((batch, 1, 1, seq))?;
    Ok(Self {
      additive: Some(broadcast_for_device(mask, batch, num_heads, head_dim, seq)?),
      causal: false,
    })
  }

  /// Padding composed with a per-head F32 `bias`, `(1, heads, seq, seq)`; never `none()`, the bias reaches every layer.
  /// A selection, not an addition: masked keys are exactly `mask_min(dtype)`, since bias plus floor nears `-inf`.
  pub fn padding_with_bias(
    attention_mask: &Tensor,
    lengths: &[usize],
    bias: &Tensor,
    num_heads: usize,
    head_dim: usize,
    dtype: DType,
  ) -> Result<Self> {
    let (batch, seq) = attention_mask.dims2()?;
    check_lengths(lengths, batch)?;
    let combined = if lengths.iter().all(|&len| len == seq) {
      bias.to_dtype(dtype)?
    } else {
      let keep = attention_mask
        .to_dtype(DType::F32)?
        .reshape((batch, 1, 1, seq))?;
      let floor = keep.affine(-1.0, 1.0)?.affine(mask_min(dtype), 0.0)?;
      bias
        .broadcast_mul(&keep)?
        .broadcast_add(&floor)?
        .to_dtype(dtype)?
    };
    Ok(Self {
      additive: Some(broadcast_for_device(
        combined, batch, num_heads, head_dim, seq,
      )?),
      causal: false,
    })
  }

  /// Causal attention: no tensor where the Metal kernel runs, else a `(1, 1, seq, seq)` upper-triangular floor.
  /// No padding term: with right padding a real token never attends to a padding position.
  pub fn causal(seq: usize, head_dim: usize, dtype: DType, device: &Device) -> Result<Self> {
    if crate::layers::attention::metal_sdpa_applies(device, head_dim, dtype) {
      return Ok(Self {
        additive: None,
        causal: true,
      });
    }
    let floor = mask_min(dtype) as f32;
    let values: Vec<f32> = (0..seq)
      .flat_map(|i| (0..seq).map(move |j| if j > i { floor } else { 0.0 }))
      .collect();
    let mask = Tensor::from_vec(values, (1, 1, seq, seq), device)?.to_dtype(dtype)?;
    Ok(Self {
      additive: Some(mask),
      causal: false,
    })
  }

  /// Bidirectional sliding window with padding: `i` sees `j` when `|i - j| < window` and `j` is real (Gemma3's rule).
  /// Every position keeps its own key, or a far padding row would have none and Metal's softmax would give NaN.
  pub fn sliding(
    attention_mask: &Tensor,
    lengths: &[usize],
    window: usize,
    num_heads: usize,
    head_dim: usize,
    dtype: DType,
  ) -> Result<Self> {
    let (batch, seq) = attention_mask.dims2()?;
    check_lengths(lengths, batch)?;
    if seq <= window && lengths.iter().all(|&len| len == seq) {
      return Ok(Self::none());
    }
    let device = attention_mask.device();
    let positions = Tensor::arange(0u32, seq as u32, device)?.to_dtype(DType::F32)?;
    let distance = positions
      .reshape((seq, 1))?
      .broadcast_sub(&positions.reshape((1, seq))?)?
      .abs()?;
    let band = distance.lt(window as f32)?.reshape((1, 1, seq, seq))?;
    let keys = attention_mask
      .to_dtype(DType::U8)?
      .reshape((batch, 1, 1, seq))?;
    let own_key = Tensor::eye(seq, DType::U8, device)?.reshape((1, 1, seq, seq))?;
    let visible = band.broadcast_mul(&keys)?.broadcast_maximum(&own_key)?;
    let mask = visible
      .to_dtype(DType::F32)?
      .affine(-1.0, 1.0)?
      .affine(mask_min(dtype), 0.0)?
      .to_dtype(dtype)?;
    Ok(Self {
      additive: Some(broadcast_for_device(mask, batch, num_heads, head_dim, seq)?),
      causal: false,
    })
  }
}

/// A short `lengths` makes `all()` vacuously true and unmasks padding, so this is an error, not a `debug_assert`.
pub(crate) fn check_lengths(lengths: &[usize], batch: usize) -> Result<()> {
  if lengths.len() != batch {
    return Err(CanddingError::LengthsMismatch {
      lengths: lengths.len(),
      rows: batch,
    });
  }
  Ok(())
}

/// Metal SDPA wants the full `(batch, heads, seq, seq)` mask but honours strides, so a zero-stride broadcast view
/// avoids `heads * seq` copies; the matmul path wants the narrowest shape.
fn broadcast_for_device(
  mask: Tensor,
  batch: usize,
  num_heads: usize,
  head_dim: usize,
  seq: usize,
) -> Result<Tensor> {
  if crate::layers::attention::metal_sdpa_applies(mask.device(), head_dim, mask.dtype()) {
    return Ok(mask.broadcast_as((batch, num_heads, seq, seq))?);
  }
  Ok(mask)
}

#[cfg(test)]
mod tests;