Interface ComputeBackend

cake-core/src/backends/mod.rs:40–685

Abstraction over compute backends (CPU, CUDA, Metal, Vulkan). Each method has a CPU-based default via [`CpuBackend`]. GPU backends override specific methods for acceleration while falling back to CPU for unimplemented ops.

/// Each method has a CPU-based default via [`CpuBackend`]. GPU backends override
/// specific methods for acceleration while falling back to CPU for unimplemented ops.
pub trait ComputeBackend: Send + Sync + std::fmt::Debug {
    /// Human-readable backend name for logging.
    fn name(&self) -> &str;

    /// The candle device this backend operates on.
    fn device(&self) -> &Device;

    // ── Attention ────────────────────────────────────────────────────

    /// Scaled dot-product attention.
    ///
    /// Backends may use flash-attn (CUDA), fused SDPA (Metal), wgpu matmul (Vulkan),
    /// or manual matmul + softmax (CPU).
    ///
    /// Input layout: `(batch, heads, seq_len, head_dim)`.
    fn attention(
        &self,
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        scale: f32,
        causal: bool,
    ) -> Result<Tensor>;

    /// Fused scaled dot-product attention (Metal SDPA, Flash-Attn style).
    /// Returns `softmax(Q @ K^T * scale + mask) @ V`.
    /// Default: delegates to `candle_nn::ops::sdpa`.
    fn sdpa(
        &self,
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        mask: Option<&Tensor>,
        causal: bool,
        scale: f32,
    ) -> Result<Tensor> {
        candle_nn::ops::sdpa(q, k, v, mask, causal, scale, 1.0)
    }

    // ── Fused activations ────────────────────────────────────────────

    /// `silu(gate) * up` — MLP activation gate.
    fn silu_mul(&self, gate: &Tensor, up: &Tensor) -> Result<Tensor>;

    /// `ln(1 + exp(clamp(x, -inf, 88)))` with `max(x, result)` — GDN gate.
    fn stable_softplus(&self, x: &Tensor) -> Result<Tensor>;

    // ── Fused normalization ──────────────────────────────────────────

    /// `rms_norm(x, weight, eps) * silu(z)` — GDN output gating.
    fn rms_norm_gated(
        &self,
        x: &Tensor,
        z: &Tensor,
        weight: &Tensor,
        eps: f32,
    ) -> Result<Tensor>;