vision-rs 0.1.1

/*
 * Copyright 2026 Teenygrad
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */


//! Flash Attention 2 — forward and backward kernels.
//!
//! **Layout**: all tensors are stored as `[BATCH * N_HEADS, N_CTX, HEAD_DIM]`
//! row-major (contiguous).  The caller is responsible for reshaping
//! `[B, H, N, D]` PyTorch tensors to this flat 3-D layout before calling.
//!
//! **Algorithm**: each CTA processes one `(batch, head, q_row)` triple.
//! The kernel iterates over all `N_CTX_K` key/value rows with the online
//! softmax recurrence (Flash Attention paper, Dao et al. 2022/2023), so
//! the full `N_CTX_Q × N_CTX_K` attention matrix is never materialised.
//! Memory is O(N_CTX × HEAD_DIM) per CTA rather than O(N_CTX²).
//!
//! HEAD_DIM must be a power of two and is a compile-time const so the
//! inner vector loads are always fully unmasked.

use core::marker::PhantomData;
use teeny_core::dtype::Float;
use teeny_macros::kernel;
use teeny_triton::triton::{
    types::{AddOffsets, Comparison, Tensor},
    *,
};

// ── Forward ───────────────────────────────────────────────────────────────────

/// Flash Attention 2 forward pass.
///
/// Inputs  (all `[BH, N_CTX, HEAD_DIM]` row-major, where `BH = BATCH * N_HEADS`):
///   `q_ptr`, `k_ptr`, `v_ptr`
///
/// Outputs:
///   `o_ptr`  — attention output   `[BH, N_CTX_Q, HEAD_DIM]`
///   `l_ptr`  — log-sum-exp        `[BH, N_CTX_Q]`  (saved for backward)
///
/// Grid: `(N_CTX_Q, BH, 1)` — one CTA per `(batch_head, q_row)` pair.
#[kernel]
pub fn flash_attention2_forward<T: Triton, D: Float, const HEAD_DIM: i32>(
    q_ptr: T::Pointer<D>,
    k_ptr: T::Pointer<D>,
    v_ptr: T::Pointer<D>,
    o_ptr: T::Pointer<D>,
    l_ptr: T::Pointer<D>,
    n_ctx_q: i32,
    n_ctx_k: i32,
    softmax_scale: f32, // 1 / sqrt(HEAD_DIM)
    neg_inf: f32,       // f32::NEG_INFINITY — passed explicitly (no_core has no float constants)
) where
    T::I32Tensor: Tensor<i32, 1>,
    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
{
    let pid_m = T::program_id(Axis::X); // query-row index  [0, n_ctx_q)
    let pid_bh = T::program_id(Axis::Y); // (batch, head)    [0, BH)

    let kv_bh_base = pid_bh * n_ctx_k * HEAD_DIM;
    let q_row_base = pid_bh * n_ctx_q * HEAD_DIM + pid_m * HEAD_DIM;
    let o_row_base = pid_bh * n_ctx_q * HEAD_DIM + pid_m * HEAD_DIM;
    let l_row_base = pid_bh * n_ctx_q + pid_m;

    // HEAD_DIM lane offsets — no masking needed (HEAD_DIM is a power of two).
    let d = T::arange(0, HEAD_DIM);

    // Load Q[pid_bh, pid_m, :]
    let q_vec = T::load(
        q_ptr.add_offsets(d + q_row_base),
        None, None, &[], None, None, None, false,
    );

    // Online-softmax running state — all kept as [HEAD_DIM] tensors so that
    // all scf.for iter-args have the same shape (Triton requires this).
    let mut acc = T::zeros::<D>(&[HEAD_DIM]);
    let mut m_i = T::full(&[HEAD_DIM], D::from_f64(neg_inf as f64));
    let mut l_i = T::zeros::<D>(&[HEAD_DIM]);
    let scale_t = T::full(&[HEAD_DIM], D::from_f64(softmax_scale as f64));

    for k_row in 0..n_ctx_k {
        let kv_row_base = kv_bh_base + k_row * HEAD_DIM;

        let k_vec = T::load(k_ptr.add_offsets(d + kv_row_base), None, None, &[], None, None, None, false);
        let v_vec = T::load(v_ptr.add_offsets(d + kv_row_base), None, None, &[], None, None, None, false);

        // Scaled dot-product: score = sum(q · k) * scale  → [HD] (scalar replicated)
        let qk = T::sum(q_vec * k_vec, Some(0), true) * scale_t;

        // Online softmax recurrence
        let m_new = T::maximum(m_i, qk); // [HD] running max (all elements equal)
        let exp_diff = T::exp(m_i - m_new); // [HD] correction factor
        let p = T::exp(qk - m_new); // [HD] unnorm weight for this k

        l_i = exp_diff * l_i + p;
        acc = exp_diff * acc + p * v_vec;
        m_i = m_new;
    }

    // Normalise output and compute logsumexp for backward.
    let o_row = acc / l_i; // [HD] / [HD] → [HD]
    // All elements of m_i and l_i are equal (replicated scalar). Sum and divide
    // by HEAD_DIM to recover the scalar as tensor<1xf32> for the l_ptr store.
    let l_save_sum = T::sum(m_i + T::log(l_i), Some(0), false); // scalar f32
    let l_save = l_save_sum / T::full(&[1], D::from_f64(HEAD_DIM as f64)); // tensor<1xD>

    T::store(o_ptr.add_offsets(d + o_row_base), o_row, None, &[], None, None);
    T::store(l_ptr.add_offsets(T::arange(0, 1) + l_row_base), l_save, None, &[], None, None);
}

// ── Backward — dQ ─────────────────────────────────────────────────────────────

/// Flash Attention 2 backward: computes `dQ`.
///
/// Grid: `(N_CTX_Q, BH, 1)` — same grid shape as the forward pass.
#[kernel]
pub fn flash_attention2_backward_dq<T: Triton, D: Float, const HEAD_DIM: i32>(
    q_ptr: T::Pointer<D>,
    k_ptr: T::Pointer<D>,
    v_ptr: T::Pointer<D>,
    o_ptr: T::Pointer<D>,
    do_ptr: T::Pointer<D>,
    l_ptr: T::Pointer<D>,
    dq_ptr: T::Pointer<D>,
    n_ctx_q: i32,
    n_ctx_k: i32,
    softmax_scale: f32,
) where
    T::I32Tensor: Tensor<i32, 1>,
    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
{
    let pid_m = T::program_id(Axis::X);
    let pid_bh = T::program_id(Axis::Y);

    let q_row_base = pid_bh * n_ctx_q * HEAD_DIM + pid_m * HEAD_DIM;
    let kv_bh_base = pid_bh * n_ctx_k * HEAD_DIM;
    let l_row_base = pid_bh * n_ctx_q + pid_m;

    let d = T::arange(0, HEAD_DIM);
    let scale_t = T::full(&[HEAD_DIM], D::from_f64(softmax_scale as f64));

    let q_vec  = T::load(q_ptr.add_offsets(d + q_row_base),  None, None, &[], None, None, None, false);
    let o_vec  = T::load(o_ptr.add_offsets(d + q_row_base),  None, None, &[], None, None, None, false);
    let do_vec = T::load(do_ptr.add_offsets(d + q_row_base), None, None, &[], None, None, None, false);

    // D_q = rowsum(O_q * dO_q)  — scalar
    let d_q = T::sum(o_vec * do_vec, Some(0), false);

    // Load logsumexp L_q (scalar stored as 1-element vec); reduce to scalar.
    let l_q_raw = T::load(l_ptr.add_offsets(T::arange(0, 1) + l_row_base), None, None, &[], None, None, None, false);
    let l_q = T::sum(l_q_raw, Some(0), false); // scalar

    let mut dq_acc = T::zeros::<D>(&[HEAD_DIM]);

    for k_row in 0..n_ctx_k {
        let kv_row_base = kv_bh_base + k_row * HEAD_DIM;

        let k_vec = T::load(k_ptr.add_offsets(d + kv_row_base), None, None, &[], None, None, None, false);
        let v_vec = T::load(v_ptr.add_offsets(d + kv_row_base), None, None, &[], None, None, None, false);

        // Recompute attention weight: p = exp(qk * scale - L_q)
        let qk = T::sum(q_vec * k_vec, Some(0), false) * scale_t;
        let p = T::exp(qk - l_q);

        // dS = p * (dO · V_k - D_q)
        let do_dot_v = T::sum(do_vec * v_vec, Some(0), false);
        let ds = p * (do_dot_v - d_q);

        // dQ += dS * K_k * scale
        dq_acc = dq_acc + ds * k_vec * scale_t;
    }

    T::store(dq_ptr.add_offsets(d + q_row_base), dq_acc, None, &[], None, None);
}

// ── Backward — dK / dV ────────────────────────────────────────────────────────

/// Flash Attention 2 backward: computes `dK` and `dV` for one key row.
///
/// Grid: `(N_CTX_K, BH, 1)`.
#[kernel]
pub fn flash_attention2_backward_dkv<T: Triton, D: Float, const HEAD_DIM: i32>(
    q_ptr: T::Pointer<D>,
    k_ptr: T::Pointer<D>,
    v_ptr: T::Pointer<D>,
    o_ptr: T::Pointer<D>,
    do_ptr: T::Pointer<D>,
    l_ptr: T::Pointer<D>,
    dk_ptr: T::Pointer<D>,
    dv_ptr: T::Pointer<D>,
    n_ctx_q: i32,
    n_ctx_k: i32,
    softmax_scale: f32,
) where
    T::I32Tensor: Tensor<i32, 1>,
    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
{
    let pid_n = T::program_id(Axis::X); // key-row index  [0, n_ctx_k)
    let pid_bh = T::program_id(Axis::Y); // (batch, head)  [0, BH)

    let q_bh_base = pid_bh * n_ctx_q * HEAD_DIM;
    let kv_row_base = pid_bh * n_ctx_k * HEAD_DIM + pid_n * HEAD_DIM;
    let l_bh_base = pid_bh * n_ctx_q;

    let d = T::arange(0, HEAD_DIM);
    let scale_t = T::full(&[HEAD_DIM], D::from_f64(softmax_scale as f64));

    // Load K_n and V_n — fixed for this CTA.
    let k_vec = T::load(k_ptr.add_offsets(d + kv_row_base), None, None, &[], None, None, None, false);
    let v_vec = T::load(v_ptr.add_offsets(d + kv_row_base), None, None, &[], None, None, None, false);

    let mut dk_acc = T::zeros::<D>(&[HEAD_DIM]);
    let mut dv_acc = T::zeros::<D>(&[HEAD_DIM]);

    for q_row in 0..n_ctx_q {
        let q_row_base = q_bh_base + q_row * HEAD_DIM;
        let l_row_base = l_bh_base + q_row;

        let q_vec_m  = T::load(q_ptr.add_offsets(d + q_row_base),  None, None, &[], None, None, None, false);
        let o_vec_m  = T::load(o_ptr.add_offsets(d + q_row_base),  None, None, &[], None, None, None, false);
        let do_vec_m = T::load(do_ptr.add_offsets(d + q_row_base), None, None, &[], None, None, None, false);
        // Load logsumexp L_m; reduce tensor<1xf32> to scalar.
        let l_m_raw  = T::load(l_ptr.add_offsets(T::arange(0, 1) + l_row_base), None, None, &[], None, None, None, false);
        let l_m = T::sum(l_m_raw, Some(0), false); // scalar f32

        // D_m = rowsum(O_m * dO_m) — scalar
        let d_m = T::sum(o_vec_m * do_vec_m, Some(0), false);

        // Recompute p_{mn} = exp(Q_m · K_n * scale - L_m)
        let qk = T::sum(q_vec_m * k_vec, Some(0), false) * scale_t;
        let p = T::exp(qk - l_m);

        // dV += p * dO_m
        dv_acc = dv_acc + p * do_vec_m;

        // dS = p * (dO_m · V_n - D_m)
        let do_dot_v = T::sum(do_vec_m * v_vec, Some(0), false);
        let ds = p * (do_dot_v - d_m);

        // dK += dS * Q_m * scale
        dk_acc = dk_acc + ds * q_vec_m * scale_t;
    }

    T::store(dk_ptr.add_offsets(d + kv_row_base), dk_acc, None, &[], None, None);
    T::store(dv_ptr.add_offsets(d + kv_row_base), dv_acc, None, &[], None, None);
}

// ── Op wrapper ────────────────────────────────────────────────────────────────

/// Bundles the forward and backward Flash Attention 2 kernels for a given head dimension.
pub struct FlashAttention2Op<'a, D: Float + Send + Sync + 'static> {
    /// The forward-pass kernel.
    pub forward: FlashAttention2Forward<D>,
    /// The backward-pass kernel computing `dQ`.
    pub backward_dq: FlashAttention2BackwardDq<D>,
    /// The backward-pass kernel computing `dK` and `dV`.
    pub backward_dkv: FlashAttention2BackwardDkv<D>,
    _marker: PhantomData<&'a ()>,
}

impl<'a, D: Float + Send + Sync + 'static> FlashAttention2Op<'a, D> {
    /// Constructs the forward/backward kernel set for the given head dimension.
    pub fn new(head_dim: i32) -> Self {
        Self {
            forward: FlashAttention2Forward::<D>::new(head_dim),
            backward_dq: FlashAttention2BackwardDq::<D>::new(head_dim),
            backward_dkv: FlashAttention2BackwardDkv::<D>::new(head_dim),
            _marker: PhantomData,
        }
    }
}