//! Gemma-4 text decoder on the GPU: weights on the device as stored in the checkpoint, paged KV pools, ONE row-batch forward used for packed prefill and
//! for decode / speculative-verify steps (forward.rs), CUDA-graph replay of small decode batches (graphs.rs).
//! Weights: MoE experts NVFP4 W4A4 (static per-layer activation global scales; repacked into the MMA fragment layout at load), attention and dense MLP
//! FP8 W8A8 (per-row weight scale, dynamic per-token activation scale), ONE fp8 per-row tied table for the embedding, the lm_head and the letter readout,
//! router / norms / KV cache bf16.
//! KV: the full-attention layers use a pool indexed by the block-table page ids; the sliding-window layers use their OWN smaller pool, reached through a
//! device page map `smap[full page id] -> sliding page id` that the scheduler maintains (`set_sliding_map`; 0 = the never-read null page).
use anyhow::{bail, ensure, Result};
use dllm_core::config::{LayerKind, LayerSpec, TextConfig};
use dllm_core::loader::{DType, WeightStore};
use dllm_kernels::linear::Seg;
use dllm_kernels::moe::Experts;
use dllm_kernels::{DBuf, Gpu};
use half::bf16;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};

pub const PAGE: usize = 16;

#[derive(Debug, Clone)]
pub struct GpuOpts {
    /// hard row limit of one forward (workspace size)
    pub max_rows: usize,
    /// pages of the full-attention KV pool (16 tokens each)
    pub num_pages: usize,
    /// RoPE table length (max position + 1)
    pub max_pos: usize,
    /// pages of the sliding-window KV pool (+ page 0, the null page)
    pub sliding_pages: usize,
}

/// device addresses of the per-row metadata of one forward (int32 arrays): token ids, positions, KV slots, block-table offsets, last visible kv position (full / sliding layers) and the block table
#[derive(Debug, Clone, Copy)]
pub struct DevMeta {
    pub tokens: u64,
    pub pos: u64,
    pub slot: u64,
    pub bt_off: u64,
    pub kv_hi: u64,
    pub kv_hi_sl: u64,
    pub bt: u64,
}

mod forward;
mod graphs;
pub use graphs::graph_bucket;

/// FP8 W8A8 linear: e4m3 weight [n, k] + f32 per-output-row scale
pub struct Fp8Linear {
    pub n: usize,
    pub k: usize,
    w: DBuf,
    s: DBuf,
}
impl Fp8Linear {
    pub fn seg(&self) -> Seg {
        Seg { w: self.w.ptr, s: self.s.ptr, n: self.n }
    }
}

struct Moe {
    /// bf16 [E, H]
    router_proj: DBuf,
    router_scale: DBuf,
    per_expert_scale: DBuf,
    pre_ff_ln2: DBuf,
    post_ff_ln2: DBuf,
    experts: Experts,
    _bufs: Vec<DBuf>,
}

struct Layer {
    spec: LayerSpec,
    in_ln: DBuf,
    q: Fp8Linear,
    k: Fp8Linear,
    /// None: k_eq_v (values = keys)
    v: Option<Fp8Linear>,
    o: Fp8Linear,
    q_norm: DBuf,
    k_norm: DBuf,
    post_attn_ln: DBuf,
    pre_ff_ln: DBuf,
    gate: Fp8Linear,
    up: Fp8Linear,
    down: Fp8Linear,
    post_ff_ln1: DBuf,
    post_ff_ln: DBuf,
    moe: Moe,
    layer_scalar: DBuf,
    /// index into GpuModel::rope
    rope: usize,
}

/// one forward = a ragged batch of rows (tokens) with explicit KV bookkeeping
#[derive(Debug, Clone, Default)]
pub struct RowBatch {
    pub tokens: Vec<i32>,
    pub pos: Vec<i32>,
    pub slot: Vec<i32>,
    pub bt_off: Vec<i32>,
    pub block_table: Vec<i32>,
    /// last visible kv position per row on FULL layers (== pos: causal)
    pub kv_hi: Vec<i32>,
    /// ... on SLIDING layers (image blocks extend it to the end of the block: bidirectional attention inside an image)
    pub kv_hi_sliding: Vec<i32>,
    /// decode / verify step (split-KV decode attention, CUDA graphs) rather than prefill
    pub decode: bool,
    /// image soft tokens injected into the embedding stream right after the embedding gather
    pub soft: Option<SoftInject>,
}

/// soft-token injection (image features): `x[row_idx[i]] = src[src_idx[i]]` for every i, bf16 rows of the text hidden size, NO embed scale.
#[derive(Debug, Clone, Default)]
pub struct SoftInject {
    /// device pointer to bf16 `[S, H]` concatenated soft tokens
    pub src: u64,
    /// batch row of each injected token
    pub row_idx: Vec<i32>,
    /// row in `src` of each injected token
    pub src_idx: Vec<i32>,
}

impl RowBatch {
    pub fn rows(&self) -> usize {
        self.tokens.len()
    }
    /// decode/verify step: one row per (sequence, token); `seqs[i] = (pages, position of the new token)`
    pub fn decode(tokens: &[i32], seqs: &[(&[i32], usize)]) -> RowBatch {
        let mut b = RowBatch { decode: true, ..Default::default() };
        for (&t, &(pages, p)) in tokens.iter().zip(seqs) {
            let off = b.block_table.len();
            b.block_table.extend_from_slice(pages);
            b.tokens.push(t);
            b.pos.push(p as i32);
            b.slot.push(pages[p / PAGE] * PAGE as i32 + (p % PAGE) as i32);
            b.bt_off.push(off as i32);
            b.kv_hi.push(p as i32);
            b.kv_hi_sliding.push(p as i32);
        }
        b
    }
}

/// per-forward workspace (max_rows rows)
struct Ws {
    rows: usize,
    x: DBuf,
    qraw: DBuf,
    kraw: DBuf,
    vraw: DBuf,
    q: DBuf,
    ao: DBuf,
    o: DBuf,
    mid: DBuf,
    act: DBuf,
    m: DBuf,
    rn: DBuf,
    topk_ids: DBuf,
    topk_w: DBuf,
    counts: DBuf,
    offsets: DBuf,
    cursor: DBuf,
    pair_of_slot: DBuf,
    slot_of_pair: DBuf,
    moe_o: DBuf,
    hidden: DBuf,
    tokens: DBuf,
    pos: DBuf,
    slot: DBuf,
    bt_off: DBuf,
    kv_hi: DBuf,
    kv_hi_sl: DBuf,
    bt: DBuf,
    bt_cap: usize,
    /// fp8 activation scratch (q bytes [rows, max_k], scales [rows])
    aq: DBuf,
    a_s: DBuf,
    /// fp4 expert input [rows, H/2] + [rows, H/16] and expert intermediate [rows*8, I/2] + [rows*8, I/16]
    mxq: DBuf,
    mxs: DBuf,
    mhq: DBuf,
    mhs: DBuf,
    /// fused [rows <= 16, q|k|v] GEMV output of decode steps
    qkv: DBuf,
    /// fp8 dense-MLP input of `post_attn`
    fq: DBuf,
    fs: DBuf,
    /// router logits f32 [rows, 128]
    rlog: DBuf,
    /// self-resetting ticket of the decode router
    ticket: DBuf,
    /// split-KV decode partials and merge counters
    part: DBuf,
    attn_cnt: DBuf,
    /// expert-union schedule built by the decode router
    grp: DBuf,
    head: dllm_kernels::head::HeadScratch,
    head_idx: DBuf,
    head_val: DBuf,
    /// prefill attention tiles for hd256 (idx 0) / hd512 (idx 1): row0 / n per tile (device i32)
    tile_row0: [DBuf; 2],
    tile_n: [DBuf; 2],
    ntiles: [AtomicUsize; 2],
}

pub struct GpuModel {
    pub gpu: Gpu,
    pub cfg: TextConfig,
    pub opts: GpuOpts,
    /// fp8 tied table: e4m3 [V, H] + f32 row scale [V]
    embed_w: DBuf,
    embed_s: DBuf,
    final_norm: DBuf,
    layers: Vec<Layer>,
    rope: Vec<DBuf>,
    kpool: Vec<DBuf>,
    vpool: Vec<DBuf>,
    /// sliding-window pool page map (device i32 [num_pages], 0 = no sliding page)
    smap: DBuf,
    ws: Ws,
    embed_scale: f32,
    /// soft injection: device i32 [row_idx(max_rows) | src_idx(max_rows)] and the pending (src ptr, count) consumed by the next `forward_dev`
    soft_idx: DBuf,
    soft_src: AtomicU64,
    soft_n: AtomicUsize,
    /// CUDA-graph store for decode forwards (None = off), see graphs.rs
    graphs: std::sync::Mutex<Option<graphs::GraphStore>>,
    graph_busy: AtomicBool,
    force_splits: AtomicUsize,
}

fn bf16_bytes_of(st: &WeightStore, name: &str) -> Result<Vec<u8>> {
    let info = st.info(name)?;
    if info.dtype == DType::BF16 {
        return Ok(st.raw(name)?.to_vec());
    }
    let t = st.load_f32(name)?;
    Ok(t.data.iter().flat_map(|x| bf16::from_f32(*x).to_le_bytes()).collect())
}
fn f32_bytes(v: &[f32]) -> Vec<u8> {
    v.iter().flat_map(|x| x.to_le_bytes()).collect()
}

fn load_vec(gpu: &Gpu, st: &WeightStore, name: &str) -> Result<DBuf> {
    DBuf::from_bytes(&gpu.be, &bf16_bytes_of(st, name)?)
}

/// FP8 linear as stored: `<base>.weight` e4m3 [n, k] + `<base>.weight_scale` f32 (per-tensor or per-row)
fn load_fp8_linear(gpu: &Gpu, st: &WeightStore, base: &str) -> Result<Fp8Linear> {
    let wn = format!("{base}.weight");
    let info = st.info(&wn)?;
    ensure!(info.dtype == DType::F8E4M3 && info.shape.len() == 2, "{wn}: expected an e4m3 [out, in] weight (engine-native checkpoint)");
    let (n, k) = (info.shape[0], info.shape[1]);
    let sc = st.load_f32(&format!("{base}.weight_scale"))?;
    let per_row: Vec<f32> = if sc.numel() == 1 {
        vec![sc.data[0]; n]
    } else if sc.numel() == n {
        sc.data.clone()
    } else {
        bail!("{base}: fp8 scales {:?} not supported (per-tensor or per-row only)", sc.shape)
    };
    Ok(Fp8Linear { n, k, w: DBuf::from_bytes(&gpu.be, st.raw(&wn)?)?, s: DBuf::from_bytes(&gpu.be, &f32_bytes(&per_row))? })
}

/// NVFP4 W4A4 experts as stored (`experts.{gate_up_proj,down_proj}.{packed,scale,gscale}` + `experts.act_gscale`), repacked into the MMA fragment layout
fn load_experts(gpu: &Gpu, st: &WeightStore, p: &str, h: usize, im: usize, e: usize) -> Result<(Experts, Vec<DBuf>)> {
    let be = &gpu.be;
    let (gu, dn) = (format!("{p}experts.gate_up_proj"), format!("{p}experts.down_proj"));
    ensure!(st.contains(&format!("{gu}.packed")), "{p}: expected NVFP4 experts ({gu}.packed; engine-native checkpoint)");
    let up = |n: &str| DBuf::from_bytes(be, st.raw(n)?);
    let a = st.load_f32(&format!("{p}experts.act_gscale"))?.data;
    let act_g = DBuf::from_bytes(be, &f32_bytes(&[a[0], a[1]]))?;
    let (gu_g, dn_g) = (up(&format!("{gu}.gscale"))?, up(&format!("{dn}.gscale"))?);
    let (t1, t2) = (e * (2 * im / 16) * (h / 64), e * (h / 16) * (im / 64));
    let nb = [DBuf::zeros(be, t1 * 512)?, DBuf::zeros(be, t1 * 64)?, DBuf::zeros(be, t2 * 512)?, DBuf::zeros(be, t2 * 64)?];
    {
        let (gp, gs, dp, ds) = (up(&format!("{gu}.packed"))?, up(&format!("{gu}.scale"))?, up(&format!("{dn}.packed"))?, up(&format!("{dn}.scale"))?);
        gpu.repack_nvfp4(nb[0].ptr, nb[1].ptr, gp.ptr, gs.ptr, e, 2 * im, h)?;
        gpu.repack_nvfp4(nb[2].ptr, nb[3].ptr, dp.ptr, ds.ptr, e, h, im)?;
        be.sync()?;
    }
    let x = Experts { gu_data: nb[0].ptr, gu_scale: nb[1].ptr, gu_gscale: gu_g.ptr, dn_data: nb[2].ptr, dn_scale: nb[3].ptr, dn_gscale: dn_g.ptr, act_g: act_g.ptr };
    let [n0, n1, n2, n3] = nb;
    Ok((x, vec![n0, n1, n2, n3, gu_g, dn_g, act_g]))
}

impl GpuModel {
    /// load the decoder of an engine-native checkpoint (`prefix` = "model.language_model." for the multimodal checkpoint)
    pub fn load(gpu: Gpu, st: &WeightStore, cfg: &TextConfig, prefix: &str, opts: GpuOpts) -> Result<Self> {
        let be = gpu.be.clone();
        let (h, im, ne) = (cfg.hidden_size, cfg.moe_intermediate_size, cfg.num_experts);
        ensure!(cfg.enable_moe && cfg.num_experts == 128 && cfg.top_k_experts == 8 && h % 64 == 0 && im % 64 == 0, "the kernels are specialised for MoE layers with 128 experts, top-8, H and I multiples of 64");
        ensure!(opts.sliding_pages > 0, "the sliding-window KV pool needs at least one page");
        let en = format!("{prefix}embed_tokens.weight");
        ensure!(st.info(&en)?.dtype == DType::F8E4M3, "{en}: expected the fp8 tied table (engine-native checkpoint)");
        let embed_w = DBuf::from_bytes(&be, st.raw(&en)?)?;
        let embed_s = DBuf::from_bytes(&be, &f32_bytes(&st.load_f32(&format!("{prefix}embed_tokens.weight_scale"))?.data))?;
        let final_norm = load_vec(&gpu, st, &format!("{prefix}norm.weight"))?;
        let mut layers = vec![];
        let mut rope: Vec<DBuf> = vec![];
        let mut rope_keys: HashMap<(usize, u64, u64, bool), usize> = HashMap::new();
        for (i, spec) in cfg.layers.iter().enumerate() {
            let p = format!("{prefix}layers.{i}.");
            let lin = |n: &str| load_fp8_linear(&gpu, st, &format!("{p}{n}"));
            let vec = |n: &str| load_vec(&gpu, st, &format!("{p}{n}"));
            let (experts, bufs) = load_experts(&gpu, st, &p, h, im, ne)?;
            let moe = Moe {
                router_proj: vec("router.proj.weight")?,
                router_scale: vec("router.scale")?,
                per_expert_scale: vec("router.per_expert_scale")?,
                pre_ff_ln2: vec("pre_feedforward_layernorm_2.weight")?,
                post_ff_ln2: vec("post_feedforward_layernorm_2.weight")?,
                experts,
                _bufs: bufs,
            };
            let key = (spec.head_dim, spec.rope_theta.to_bits(), spec.partial_rotary.to_bits(), spec.proportional_rope);
            let ri = match rope_keys.get(&key) {
                Some(&r) => r,
                None => {
                    rope.push(DBuf::from_slice(&be, &crate::rope::cos_sin_table(spec, opts.max_pos))?);
                    rope_keys.insert(key, rope.len() - 1);
                    rope.len() - 1
                }
            };
            layers.push(Layer {
                spec: spec.clone(),
                in_ln: vec("input_layernorm.weight")?,
                q: lin("self_attn.q_proj")?,
                k: lin("self_attn.k_proj")?,
                v: if spec.k_eq_v { None } else { Some(lin("self_attn.v_proj")?) },
                o: lin("self_attn.o_proj")?,
                q_norm: vec("self_attn.q_norm.weight")?,
                k_norm: vec("self_attn.k_norm.weight")?,
                post_attn_ln: vec("post_attention_layernorm.weight")?,
                pre_ff_ln: vec("pre_feedforward_layernorm.weight")?,
                gate: lin("mlp.gate_proj")?,
                up: lin("mlp.up_proj")?,
                down: lin("mlp.down_proj")?,
                post_ff_ln1: vec("post_feedforward_layernorm_1.weight")?,
                post_ff_ln: vec("post_feedforward_layernorm.weight")?,
                moe,
                layer_scalar: vec("layer_scalar")?,
                rope: ri,
            });
        }
        let (mut kpool, mut vpool) = (vec![], vec![]);
        for spec in &cfg.layers {
            let pages = if spec.kind == LayerKind::Sliding { opts.sliding_pages + 1 } else { opts.num_pages };
            let bytes = pages * spec.n_kv_heads * PAGE * spec.head_dim * 2;
            kpool.push(DBuf::zeros(&be, bytes)?);
            vpool.push(DBuf::zeros(&be, bytes)?);
        }
        let smap = DBuf::zeros(&be, opts.num_pages * 4)?;
        let r = opts.max_rows;
        let max_q = cfg.layers.iter().map(|s| cfg.num_heads * s.head_dim).max().unwrap();
        let max_kv = cfg.layers.iter().map(|s| s.n_kv_heads * s.head_dim).max().unwrap();
        let dff = cfg.intermediate_size;
        let k = cfg.top_k_experts;
        let z = |bytes: usize| DBuf::zeros(&be, bytes);
        let bt_cap = r * 8;
        let max_k = h.max(max_q).max(dff);
        let maxld = cfg.layers.iter().map(|s| cfg.num_heads * s.head_dim + 2 * s.n_kv_heads * s.head_dim).max().unwrap();
        let max_part = cfg.layers.iter().map(|s| dllm_kernels::attn::decode_part_bytes(16.min(r), cfg.num_heads, 32, s.head_dim)).max().unwrap();
        let tl = || z((r + 8) * 4);
        let ws = Ws {
            rows: r,
            x: z(r * h * 2)?,
            qraw: z(r * max_q * 2)?,
            kraw: z(r * max_kv * 2)?,
            vraw: z(r * max_kv * 2)?,
            q: z(r * max_q * 2)?,
            ao: z(r * max_q * 2)?,
            o: z(r * h * 2)?,
            mid: z(r * h * 2)?,
            act: z(r * dff * 2)?,
            m: z(r * h * 2)?,
            rn: z(r * h * 2)?,
            topk_ids: z(r * k * 4)?,
            topk_w: z(r * k * 4)?,
            counts: z(ne * 4)?,
            offsets: z((ne + 1) * 4)?,
            cursor: z(ne * 4)?,
            pair_of_slot: z(r * k * 4)?,
            slot_of_pair: z(r * k * 4)?,
            moe_o: z(r * k * h * 2)?,
            hidden: z(r * h * 2)?,
            tokens: z(r * 4)?,
            pos: z(r * 4)?,
            slot: z(r * 4)?,
            bt_off: z(r * 4)?,
            kv_hi: z(r * 4)?,
            kv_hi_sl: z(r * 4)?,
            bt: z(bt_cap * 4)?,
            bt_cap,
            aq: z(r * max_k)?,
            a_s: z(r * 4)?,
            mxq: z(r * h / 2)?,
            mxs: z(r * h / 16)?,
            mhq: z(r * k * im / 2)?,
            mhs: z(r * k * im / 16)?,
            qkv: z(16 * maxld * 2)?,
            fq: z(r * h)?,
            fs: z(r * 4)?,
            rlog: z(r * 128 * 4)?,
            ticket: z(16)?,
            part: z(max_part)?,
            attn_cnt: z(4096 * 4)?,
            grp: z((257 + 128 * 16) * 4)?,
            head: dllm_kernels::head::HeadScratch::new(&be)?,
            head_idx: z(16 * 4)?,
            head_val: z(16 * 4)?,
            tile_row0: [tl()?, tl()?],
            tile_n: [tl()?, tl()?],
            ntiles: [AtomicUsize::new(0), AtomicUsize::new(0)],
        };
        let embed_scale = crate::rope::bf16_round((cfg.hidden_size as f32).sqrt());
        let soft_idx = DBuf::zeros(&gpu.be, opts.max_rows * 8)?;
        Ok(GpuModel {
            gpu,
            cfg: cfg.clone(),
            opts,
            embed_w,
            embed_s,
            final_norm,
            layers,
            rope,
            kpool,
            vpool,
            smap,
            ws,
            embed_scale,
            soft_idx,
            soft_src: Default::default(),
            soft_n: Default::default(),
            graphs: Default::default(),
            graph_busy: Default::default(),
            force_splits: Default::default(),
        })
    }

    /// run the whole decoder over `b`; the final-norm hidden rows stay on the device (`hidden_ptr`)
    pub fn forward(&mut self, b: &RowBatch) -> Result<()> {
        let r = b.rows();
        ensure!(r > 0 && r <= self.ws.rows, "rows {r} exceeds max_rows {}", self.ws.rows);
        // metadata upload
        if b.block_table.len() > self.ws.bt_cap {
            self.ws.bt_cap = b.block_table.len() * 2;
            self.ws.bt = DBuf::zeros(&self.gpu.be, self.ws.bt_cap * 4)?;
            self.graphs_clear();
        }
        let ws = &self.ws;
        ws.tokens.upload_slice(&b.tokens)?;
        ws.pos.upload_slice(&b.pos)?;
        ws.slot.upload_slice(&b.slot)?;
        ws.bt_off.upload_slice(&b.bt_off)?;
        ws.kv_hi.upload_slice(&b.kv_hi)?;
        ws.kv_hi_sl.upload_slice(&b.kv_hi_sliding)?;
        ws.bt.upload_slice(&b.block_table)?;
        match &b.soft {
            Some(s) if !s.row_idx.is_empty() => {
                ensure!(s.row_idx.len() == s.src_idx.len() && s.row_idx.len() <= r, "soft injection: {} rows / {} src idx (batch rows {r})", s.row_idx.len(), s.src_idx.len());
                ensure!(s.row_idx.iter().all(|&x| x >= 0 && (x as usize) < r), "soft injection row index out of range");
                self.soft_idx.upload_slice(&s.row_idx)?;
                self.gpu.be.upload(self.soft_idx.at(self.opts.max_rows * 4), unsafe { std::slice::from_raw_parts(s.src_idx.as_ptr() as *const u8, s.src_idx.len() * 4) })?;
                self.soft_src.store(s.src, Ordering::Relaxed);
                self.soft_n.store(s.row_idx.len(), Ordering::Relaxed);
            }
            _ => self.soft_n.store(0, Ordering::Relaxed),
        }
        let max_kv = b.kv_hi.iter().chain(&b.kv_hi_sliding).copied().max().unwrap_or(0) as usize + 1;
        let mut splits = if b.decode { self.decode_splits(max_kv) } else { 1 };
        if b.decode && self.graphs_enabled() {
            let f = self.force_splits.load(Ordering::Relaxed);
            if f > 0 {
                splits = graphs::quant_splits(f);
            }
        }
        let m = DevMeta { tokens: ws.tokens.ptr, pos: ws.pos.ptr, slot: ws.slot.ptr, bt_off: ws.bt_off.ptr, kv_hi: ws.kv_hi.ptr, kv_hi_sl: ws.kv_hi_sl.ptr, bt: ws.bt.ptr };
        self.prepare_tiles(&b.bt_off, b.decode)?;
        self.forward_dev(r, b.decode, splits, m)
    }

    /// KV splits of a decode forward whose longest row sees `max_kv` positions (changes only at multiples of 1024)
    pub fn decode_splits(&self, max_kv: usize) -> usize {
        let s = ((max_kv + 1023) / 1024).clamp(1, 8);
        if self.graphs_enabled() { graphs::quant_splits(s) } else { s }
    }

    /// the forward over `r` rows whose metadata already lives in device buffers `m` (built by device kernels, e.g. the speculative verify batch);
    /// `forward` is `upload metadata + forward_dev`. Prefill batches need `forward` (host tile table).
    pub fn forward_dev(&self, r: usize, decode: bool, splits: usize, m: DevMeta) -> Result<()> {
        ensure!(r > 0 && r <= self.ws.rows, "rows {r} exceeds max_rows {}", self.ws.rows);
        if decode {
            if let Some(res) = self.try_graph(r, splits, &m) {
                return res;
            }
        }
        self.forward_rows(r, decode, splits, m)
    }

    /// device pointer of the final-norm hidden rows [rows, H] of the last forward
    pub fn hidden_ptr(&self) -> u64 {
        self.ws.hidden.ptr
    }

    fn gather_hidden(&self, rows: &[usize]) -> Result<DBuf> {
        let idx: Vec<i32> = rows.iter().map(|&r| r as i32).collect();
        let ib = DBuf::from_slice(&self.gpu.be, &idx)?;
        let out = DBuf::zeros(&self.gpu.be, rows.len() * self.cfg.hidden_size * 2)?;
        self.gpu.gather_rows(out.ptr, self.ws.hidden.ptr, ib.ptr, rows.len(), self.cfg.hidden_size)?;
        Ok(out)
    }

    /// full-vocab fp8 scan of `m` (<= 16) hidden rows -> softcapped fp32 logits
    fn scan(&self, out: u64, hid: u64, m: usize) -> Result<()> {
        let (v, h) = (self.cfg.vocab_size, self.cfg.hidden_size);
        self.gpu.lm_head_logits_fp8(out, hid, self.embed_w.ptr, self.embed_s.ptr, m, v, h, self.softcap())
    }

    /// softcapped full-vocab logits of the given hidden rows
    pub fn logits_rows(&self, rows: &[usize]) -> Result<Vec<Vec<f32>>> {
        let mut res = vec![];
        for chunk in rows.chunks(16) {
            let hid = self.gather_hidden(chunk)?;
            let v = self.cfg.vocab_size;
            let out = DBuf::zeros(&self.gpu.be, chunk.len() * v * 4)?;
            self.scan(out.ptr, hid.ptr, chunk.len())?;
            self.gpu.be.sync()?;
            let all: Vec<f32> = out.to_vec(chunk.len() * v)?;
            res.extend(all.chunks(v).map(|c| c.to_vec()));
        }
        Ok(res)
    }

    /// softcapped logits of the option-letter tokens for each hidden row (letter readout)
    pub fn letter_logits(&self, rows: &[usize], label_ids: &[i32]) -> Result<Vec<Vec<f32>>> {
        let hid = self.gather_hidden(rows)?;
        let lb = DBuf::from_slice(&self.gpu.be, label_ids)?;
        let out = DBuf::zeros(&self.gpu.be, rows.len() * label_ids.len() * 4)?;
        self.gpu.letter_logits_fp8(out.ptr, hid.ptr, self.embed_w.ptr, self.embed_s.ptr, lb.ptr, rows.len(), label_ids.len(), self.cfg.hidden_size, self.softcap())?;
        self.gpu.be.sync()?;
        let all: Vec<f32> = out.to_vec(rows.len() * label_ids.len())?;
        Ok(all.chunks(label_ids.len()).map(|c| c.to_vec()).collect())
    }

    /// greedy token per row (optionally excluding two ids, e.g. end-of-thought / EOS), full fp32 logits + argmax kernel
    pub fn argmax_rows(&self, rows: &[usize], ex0: i32, ex1: i32) -> Result<Vec<i32>> {
        let mut res = vec![];
        for chunk in rows.chunks(16) {
            let hid = self.gather_hidden(chunk)?;
            let v = self.cfg.vocab_size;
            let lg = DBuf::zeros(&self.gpu.be, chunk.len() * v * 4)?;
            self.scan(lg.ptr, hid.ptr, chunk.len())?;
            let idx = DBuf::zeros(&self.gpu.be, chunk.len() * 4)?;
            let val = DBuf::zeros(&self.gpu.be, chunk.len() * 4)?;
            self.gpu.argmax_rows(idx.ptr, val.ptr, lg.ptr, chunk.len(), v, ex0, ex1)?;
            self.gpu.be.sync()?;
            res.extend(idx.to_vec::<i32>(chunk.len())?);
        }
        Ok(res)
    }

    /// greedy token of the first `t` (<= 16) hidden rows of the last forward through the fused lm_head scan + argmax (bf16 activations for t <= 8,
    /// per-token fp8 above); the decode step's fast path
    pub fn greedy_first_rows(&self, t: usize, ex0: i32, ex1: i32) -> Result<Vec<i32>> {
        let ws = &self.ws;
        dllm_kernels::head::lm_head_argmax(&self.gpu.be, &ws.head, ws.head_idx.ptr, ws.head_val.ptr, ws.hidden.ptr, self.embed_w.ptr, self.embed_s.ptr, t, self.cfg.vocab_size, self.cfg.hidden_size, ex0, ex1, self.softcap(), t > 8)?;
        self.gpu.be.sync()?;
        ws.head_idx.to_vec::<i32>(t)
    }

    // ---- accessors for the speculative decoder and the draft head

    /// K / V pool device address of layer `layer` ([page][n_kv][16][hd] bf16)
    pub fn kv_pool_ptrs(&self, layer: usize) -> (u64, u64) {
        (self.kpool[layer].ptr, self.vpool[layer].ptr)
    }
    /// the fp8 tied table (e4m3 bytes, per-row scale) and the embedding scale
    pub fn embed_desc(&self) -> (u64, u64, f32) {
        (self.embed_w.ptr, self.embed_s.ptr, self.embed_scale)
    }
    pub fn softcap(&self) -> f32 {
        self.cfg.final_logit_softcap.unwrap_or(0.0)
    }
    /// while set, `forward_dev` never routes through the per-shape CUDA-graph store (needed when the caller records its own graph around forward_dev)
    pub fn suspend_graphs(&self, on: bool) {
        self.graph_busy.store(on, Ordering::Relaxed);
    }
    /// page map of layer `layer`'s KV pool for the kernels: the sliding-window pool map for sliding layers, 0 (identity) for full-attention layers
    pub fn kv_pmap(&self, layer: usize) -> u64 {
        if self.cfg.layers[layer].kind == LayerKind::Sliding {
            self.smap.ptr
        } else {
            0
        }
    }
    pub fn sliding_pages(&self) -> usize {
        self.opts.sliding_pages
    }
    /// attention window of the sliding layers (positions incl. self)
    pub fn sliding_window(&self) -> usize {
        self.cfg.layers.iter().filter(|s| s.kind == LayerKind::Sliding).filter_map(|s| s.window).max().unwrap_or(0)
    }
    /// write `vals` into the sliding page map at full-page ids `first..first+vals.len()` (stream-ordered before the next forward). Sliding page ids are 1..=sliding_pages, 0 = none.
    pub fn set_sliding_map(&self, first: usize, vals: &[i32]) -> Result<()> {
        ensure!(first + vals.len() <= self.opts.num_pages, "sliding map update out of range");
        ensure!(vals.iter().all(|&v| v >= 0 && v as usize <= self.opts.sliding_pages), "sliding page id out of range");
        self.gpu.be.upload(self.smap.at(first * 4), unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) })
    }
}