//! Greedy lm_head (cuda/head.cu): full-vocab scan of the fp8 tied table with a fused argmax (bf16 activations for <= 8 rows, per-token fp8 activations up to 16 rows).
//! The tied table is ONE fp8 per-row table: e4m3 bytes [V,H] + fp32 row_scale [V].
use crate::backend::{Arg, Backend, DBuf};
use anyhow::Result;
use std::sync::Arc;
use Arg::{F, I, P};

pub const NBLK_HEAD: usize = 170;

/// partial-result scratch of the lm_head scan (CTA partial max/argmax)
pub struct HeadScratch {
    pub nblk: usize,
    pv: DBuf,
    pi: DBuf,
}
impl HeadScratch {
    pub fn new(be: &Arc<dyn Backend>) -> Result<Self> {
        Ok(HeadScratch { nblk: NBLK_HEAD, pv: DBuf::zeros(be, NBLK_HEAD * 16 * 4)?, pi: DBuf::zeros(be, NBLK_HEAD * 16 * 4)? })
    }
}

/// idx i32 [T], val f32 [T] (softcapped max logit) = argmax over the vocab of h[t].E[v] (exclusions ex0/ex1, -1 = none). `w8a8 = false`: bf16 activations (T <= 8, default); true: hidden quantised per token (T <= 16).
#[allow(clippy::too_many_arguments)]
pub fn lm_head_argmax(be: &Arc<dyn Backend>, sc: &HeadScratch, idx: u64, val: u64, h: u64, table: u64, row_scale: u64, t: usize, v: usize, hdim: usize, ex0: i32, ex1: i32, softcap: f32, w8a8: bool) -> Result<()> {
    assert!(hdim % 64 == 0 && hdim <= 2816 && v % 16 == 0);
    let (name, stride) = if !w8a8 {
        assert!(t <= 8);
        ("head_lm_w8a16_t8", 8)
    } else if t <= 8 {
        ("head_lm_w8a8_t8", 16)
    } else {
        assert!(t <= 16);
        ("head_lm_w8a8_t16", 16)
    };
    be.launch("head", name, [sc.nblk as u32, 1, 1], [256, 1, 1], &[P(sc.pv.ptr), P(sc.pi.ptr), P(h), P(table), P(row_scale), I(t as i32), I(v as i32), I(hdim as i32), I(ex0), I(ex1)])?;
    be.launch("head", "head_argmax_final", [t as u32, 1, 1], [256, 1, 1], &[P(idx), P(val), P(sc.pv.ptr), P(sc.pi.ptr), I(sc.nblk as i32), I(stride), F(softcap)])
}