engine/crates/kernels/src/ops.rs

//! `Gpu`: typed launchers of the small bf16 kernels (embedding gather, RMSNorm family, the draft's q-only RoPE prep, MoE prefill routing,
//! full-vocab lm_head scan + argmax, option-letter readout). The decoder's heavy kernels have their own modules (linear, moe, attn, head, norm).
//! All pointers are u64 device addresses; row-major bf16 unless stated.
use crate::backend::{Arg, Backend};
use anyhow::Result;
use std::sync::Arc;

use Arg::{F, I, P};

#[derive(Clone)]
pub struct Gpu {
    pub be: Arc<dyn Backend>,
}

fn ceil(a: usize, b: usize) -> u32 {
    ((a + b - 1) / b) as u32
}

impl Gpu {
    pub fn new(be: Arc<dyn Backend>) -> Self {
        Gpu { be }
    }
    fn l(&self, file: &str, sym: &str, grid: [u32; 3], block: u32, args: &[Arg]) -> Result<()> {
        self.be.launch(file, sym, grid, [block, 1, 1], args)
    }

    /// embedding lookup from the fp8 per-row tied table (dequantised row rounded to bf16, times the embedding scale, rounded again)
    pub fn embed_gather_fp8(&self, out: u64, table: u64, row_scale: u64, ids: u64, rows: usize, h: usize, scale: f32) -> Result<()> {
        self.l("embed", "embed_gather_fp8", [rows as u32, 1, 1], 256, &[P(out), P(table), P(row_scale), P(ids), I(h as i32), F(scale)])
    }
    pub fn gather_rows(&self, out: u64, src: u64, idx: u64, rows: usize, h: usize) -> Result<()> {
        self.l("embed", "gather_rows_bf16", [rows as u32, 1, 1], 256, &[P(out), P(src), P(idx), I(h as i32)])
    }
    /// w = 0 -> unscaled
    pub fn rmsnorm(&self, out: u64, x: u64, w: u64, rows: usize, dim: usize, eps: f32) -> Result<()> {
        self.l("rmsnorm", "rmsnorm_bf16", [rows as u32, 1, 1], dim.min(1024) as u32, &[P(out), P(x), P(w), I(dim as i32), F(eps)])
    }
    /// resid_out = resid + rmsnorm(x; w_post) ; normed = rmsnorm(resid_out; w_pre)  (w_pre = 0 skips the second output)
    #[allow(clippy::too_many_arguments)]
    pub fn sandwich(&self, resid_out: u64, normed: u64, resid: u64, x: u64, w_post: u64, w_pre: u64, rows: usize, dim: usize, eps: f32) -> Result<()> {
        self.l("rmsnorm", "post_norm_add_pre_norm_bf16", [rows as u32, 1, 1], dim.min(1024) as u32, &[P(resid_out), P(normed), P(resid), P(x), P(w_post), P(w_pre), I(dim as i32), F(eps)])
    }
    /// out = (mid + rmsnorm(x; w)) * layer_scalar
    #[allow(clippy::too_many_arguments)]
    pub fn layer_out(&self, out: u64, mid: u64, x: u64, w: u64, layer_scalar: u64, rows: usize, dim: usize, eps: f32) -> Result<()> {
        self.l("rmsnorm", "layer_out_bf16", [rows as u32, 1, 1], dim.min(1024) as u32, &[P(out), P(mid), P(x), P(w), P(layer_scalar), I(dim as i32), F(eps)])
    }

    /// fused q/k/v norm + RoPE + paged KV write (`nkv` = 0: q only, as used by the draft head). `v_raw` may equal `k_raw` (k_eq_v). cos_sin: [maxpos][hd/2][2] f32
    #[allow(clippy::too_many_arguments)]
    pub fn qkv_prep(&self, q_out: u64, k_pool: u64, v_pool: u64, q_raw: u64, k_raw: u64, v_raw: u64, q_norm: u64, k_norm: u64, cos_sin: u64, pos: u64, slot: u64, rows: usize, nh: usize, nkv: usize, hd: usize, eps: f32, pmap: u64) -> Result<()> {
        self.l("rope_kv_write", "qkv_prep_bf16", [rows as u32, (nh + nkv) as u32, 1], 256,
            &[P(q_out), P(k_pool), P(v_pool), P(q_raw), P(k_raw), P(v_raw), P(q_norm), P(k_norm), P(cos_sin), P(pos), P(slot), I(nh as i32), I(nkv as i32), I(hd as i32), F(eps), P(pmap)])
    }

    /// MoE prefill routing: expert counts -> offsets -> expert-sorted slots (pair_of_slot / slot_of_pair) from topk_ids [rows, k]
    #[allow(clippy::too_many_arguments)]
    pub fn moe_route(&self, counts: u64, offsets: u64, cursor: u64, pair_of_slot: u64, slot_of_pair: u64, topk_ids: u64, rows: usize, k: usize, e: usize) -> Result<()> {
        let n = rows * k;
        self.be.memset0(counts, e * 4)?;
        self.l("moe_route", "moe_count", [ceil(n, 256), 1, 1], 256, &[P(counts), P(topk_ids), I(n as i32)])?;
        self.l("moe_route", "moe_offsets", [1, 1, 1], 32, &[P(offsets), P(cursor), P(counts), I(e as i32)])?;
        self.l("moe_route", "moe_scatter", [ceil(n, 256), 1, 1], 256, &[P(pair_of_slot), P(slot_of_pair), P(cursor), P(offsets), P(topk_ids), I(n as i32)])
    }

    /// out[m, v] fp32 = softcap(h[m] . E[v]) over the fp8 tied table (e4m3 [V,H] + per-row scale); m <= 16
    #[allow(clippy::too_many_arguments)]
    pub fn lm_head_logits_fp8(&self, out: u64, h: u64, e8: u64, row_scale: u64, m: usize, vocab: usize, hdim: usize, softcap: f32) -> Result<()> {
        assert!(m <= 16 && hdim % 8 == 0);
        self.l("lm_head", "lm_head_logits_fp8", [ceil(vocab, 4), 1, 1], 128, &[P(out), P(h), P(e8), P(row_scale), I(m as i32), I(vocab as i32), I(hdim as i32), F(softcap)])
    }
    /// per-row argmax of fp32 logits [m, vocab], never returning ex0 / ex1 (-1 = none)
    pub fn argmax_rows(&self, idx: u64, val: u64, logits: u64, m: usize, vocab: usize, ex0: i32, ex1: i32) -> Result<()> {
        self.l("lm_head", "argmax_rows", [m as u32, 1, 1], 1024, &[P(idx), P(val), P(logits), I(vocab as i32), I(ex0), I(ex1)])
    }
    /// out[row, j] fp32 = softcap(hidden[row] . E[label_ids[j]]) over the fp8 tied table
    #[allow(clippy::too_many_arguments)]
    pub fn letter_logits_fp8(&self, out: u64, hidden: u64, e8: u64, row_scale: u64, label_ids: u64, rows: usize, n_labels: usize, hdim: usize, softcap: f32) -> Result<()> {
        self.l("letter_readout", "letter_logits_fp8", [rows as u32, 1, 1], 256, &[P(out), P(hidden), P(e8), P(row_scale), P(label_ids), I(n_labels as i32), I(hdim as i32), F(softcap)])
    }
}