engine/crates/kernels/src/linear.rs

//! FP8 W8A8 linears (cuda/linear.cu). Weights: e4m3 bytes [N,K] row-major + fp32 per-row scale (checkpoint layout, no repack).
//! Decode (T <= 16): `gemv_pqd` (pre-quantised activations, bandwidth-optimal) or `gemv_fused` (prologue: optional RMSNorm + per-token quant inside the kernel).
//! Prefill: `quant_rows` + `gemm` / `gemm_glu`.  All pointers are device addresses (u64) as everywhere in dllm-kernels; `y` is bf16 [T, ldy].
use crate::backend::{Arg, Backend};
use anyhow::Result;
use std::sync::Arc;
use Arg::{F, I, P};

/// one consecutive output segment: weight [n,k] e4m3, fp32 row scales [n]
#[derive(Debug, Clone, Copy)]
pub struct Seg {
    pub w: u64,
    pub s: u64,
    pub n: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Prologue {
    /// x bf16 raw -> per-token quant
    Bf16,
    /// rmsnorm(x; nw, eps) (nw = 0 -> unweighted) -> bf16 round -> per-token quant
    Rms,
}

/// measured default K-split (warps per CTA) for the direct kernel
pub fn default_ks(tiles: usize, t: usize) -> usize {
    let t = if super::row_invariant() { 1 } else { t };
    if tiles < 160 {
        8
    } else if tiles < 400 {
        4
    } else if t <= 2 {
        4
    } else {
        2
    }
}

fn segs3(segs: &[Seg]) -> [Seg; 3] {
    assert!(!segs.is_empty() && segs.len() <= 3);
    let mut o = [Seg { w: 0, s: 0, n: 0 }; 3];
    for (i, s) in segs.iter().enumerate() {
        assert!(s.n % 16 == 0);
        o[i] = *s;
    }
    o
}

/// y[t, off..] = (xq,xs) . W^T for up to 3 consecutive segments (q,k,v), or GLU (segs = [gate, up]: y[t,i] = gelu_tanh(gate)*up, ldy = n).  xq u8 [T,K], xs f32 [T]. K % 64 == 0, T <= 16.
#[allow(clippy::too_many_arguments)]
pub fn gemv_pqd(be: &Arc<dyn Backend>, y: u64, ldy: usize, xq: u64, xs: u64, t: usize, k: usize, segs: &[Seg], glu: bool, ks: Option<usize>) -> Result<()> {
    assert!(t >= 1 && t <= 16 && k % 64 == 0);
    let tiles = if glu { segs[0].n / 16 } else { segs.iter().map(|s| s.n / 16).sum() };
    let ks = ks.unwrap_or_else(|| default_ks(tiles, t));
    let [a, b, c] = segs3(segs);
    let name = format!("lin_gemv_{}_ks{}_pqd{}", if t <= 8 { "t8" } else { "t16" }, ks, if glu { "_glu" } else { "" });
    be.launch("linear", &name, [tiles as u32, 1, 1], [(ks * 32) as u32, 1, 1],
        &[P(y), I(ldy as i32), P(xq), P(xs), I(t as i32), I(k as i32), P(a.w), P(a.s), I(a.n as i32), P(b.w), P(b.s), I(b.n as i32), P(c.w), P(c.s), I(c.n as i32)])
}

/// fused-prologue GEMV: x bf16 [T,K] (+ optional rmsnorm weight `nw` bf16 [K], 0 = none) quantised inside the kernel.
#[allow(clippy::too_many_arguments)]
pub fn gemv_fused(be: &Arc<dyn Backend>, y: u64, ldy: usize, x: u64, pro: Prologue, nw: u64, eps: f32, t: usize, k: usize, segs: &[Seg], glu: bool, ks: Option<usize>) -> Result<()> {
    assert!(t >= 1 && t <= 16 && k % 64 == 0);
    let tiles = if glu { segs[0].n / 16 } else { segs.iter().map(|s| s.n / 16).sum() };
    let ks = ks.unwrap_or(8);
    let [a, b, c] = segs3(segs);
    let name = format!("lin_gemv_{}_ks{}_{}{}", if t <= 8 { "t8" } else { "t16" }, ks, if pro == Prologue::Rms { "rms" } else { "bf" }, if glu { "_glu" } else { "" });
    be.launch("linear", &name, [tiles as u32, 1, 1], [(ks * 32) as u32, 1, 1],
        &[P(y), I(ldy as i32), P(x), P(0), P(0), P(nw), F(eps), I(t as i32), I(k as i32), P(a.w), P(a.s), I(a.n as i32), P(b.w), P(b.s), I(b.n as i32), P(c.w), P(c.s), I(c.n as i32)])
}

/// per-token e4m3 quantisation (optionally fused rmsnorm): q u8 [rows,K], s f32 [rows]
#[allow(clippy::too_many_arguments)]
pub fn quant_rows(be: &Arc<dyn Backend>, q: u64, s: u64, x: u64, nw: u64, eps: f32, rows: usize, k: usize, rms: bool) -> Result<()> {
    be.launch("linear", "lin_quant_rows", [rows as u32, 1, 1], [256, 1, 1], &[P(q), P(s), P(x), P(nw), F(eps), I(k as i32), I(rms as i32)])
}

/// CTA tile of the prefill GEMM: (MT, WM, WN) = warp tile (MT*16) x 32, WM x WN warps (BM = WM*MT*16, BN = WN*32). All tiles give bit-identical results.
pub type Tile = (usize, usize, usize);
/// every compiled tile variant (kernel suffix `_v{MT}{WM}{WN}`; (4,2,4) is the plain `lin_gemm` / `lin_gemm_glu`)
/// (the tile grid of round 3 measured 12 shapes; these four cover the fitted rule of `pick_tile`)
pub const TILES: [Tile; 4] = [(4, 2, 4), (4, 2, 2), (2, 2, 2), (1, 2, 2)];

fn tile_name(base: &str, t: Tile) -> String {
    if t == (4, 2, 4) { base.to_string() } else { format!("{base}_v{}{}{}", t.0, t.1, t.2) }
}

/// tile choice for an [m, n] output (glu: n = output columns): the 128x128 tile once it fills the GPU (>= ~2 waves of CTAs), smaller tiles for
/// short prefills (M ~ 32-512), where 128x128 leaves most of the SMs idle
pub fn pick_tile(m: usize, n: usize, glu: bool) -> Tile {
    let nc = if glu { n * 2 } else { n };
    let ctas = |t: Tile| ((nc + t.2 * 32 - 1) / (t.2 * 32)) * ((m + t.0 * t.1 * 16 - 1) / (t.0 * t.1 * 16));
    // fitted on the RTX 5090 grid of tests/gemm_tiles.rs (q/k/o/down/gate-up shapes, M 16..8192): mean 0.7 % above the per-shape best tile
    for (t, min_ctas) in [((4, 2, 4), 2000), ((4, 2, 2), 600), ((2, 2, 2), 200)] {
        if ctas(t) >= min_ctas {
            return t;
        }
    }
    (1, 2, 2)
}

/// prefill GEMM  y[M,N] = xq[M,K] . W[N,K]^T (K % 64 == 0)
#[allow(clippy::too_many_arguments)]
pub fn gemm(be: &Arc<dyn Backend>, y: u64, xq: u64, xs: u64, w: u64, ws: u64, m: usize, n: usize, k: usize) -> Result<()> {
    gemm_tile(be, pick_tile(m, n, false), y, xq, xs, w, ws, m, n, k)
}

#[allow(clippy::too_many_arguments)]
pub fn gemm_tile(be: &Arc<dyn Backend>, t: Tile, y: u64, xq: u64, xs: u64, w: u64, ws: u64, m: usize, n: usize, k: usize) -> Result<()> {
    assert!(k % 64 == 0);
    let (bm, bn) = (t.0 * t.1 * 16, t.2 * 32);
    be.launch("linear", &tile_name("lin_gemm", t), [((n + bn - 1) / bn) as u32, ((m + bm - 1) / bm) as u32, 1], [(t.1 * t.2 * 32) as u32, 1, 1], &[P(y), P(xq), P(xs), P(w), P(ws), I(m as i32), I(n as i32), I(k as i32)])
}

/// prefill dense-MLP gate/up with fused GELU*up: y[M,N] = bf16(gelu(bf16 g) * bf16 u)
#[allow(clippy::too_many_arguments)]
pub fn gemm_glu(be: &Arc<dyn Backend>, y: u64, xq: u64, xs: u64, wg: u64, wsg: u64, wu: u64, wsu: u64, m: usize, n: usize, k: usize) -> Result<()> {
    gemm_glu_tile(be, pick_tile(m, n, true), y, xq, xs, wg, wsg, wu, wsu, m, n, k)
}

#[allow(clippy::too_many_arguments)]
pub fn gemm_glu_tile(be: &Arc<dyn Backend>, t: Tile, y: u64, xq: u64, xs: u64, wg: u64, wsg: u64, wu: u64, wsu: u64, m: usize, n: usize, k: usize) -> Result<()> {
    assert!(k % 64 == 0);
    let (bm, bn2) = (t.0 * t.1 * 16, t.2 * 16);
    be.launch("linear", &tile_name("lin_gemm_glu", t), [((n + bn2 - 1) / bn2) as u32, ((m + bm - 1) / bm) as u32, 1], [(t.1 * t.2 * 32) as u32, 1, 1],
        &[P(y), P(xq), P(xs), P(wg), P(wsg), P(wu), P(wsu), I(m as i32), I(n as i32), I(k as i32)])
}