engine/crates/model/src/gpu_model/forward.rs

//! The decoder forward over one row batch (child module of `gpu_model`: sees the private fields): embedding + soft injection -> 30 layers -> final norm -> `hidden`.
//! Per layer: fused norm + fp8 quant producers, FP8 MMA GEMV (<= 16 rows) / GEMM, mma attention, NVFP4 W4A4 block-scaled MMA MoE. In decode steps the dense
//! MLP runs on a second stream, overlapping the router + MoE chain.
use super::*;
use dllm_kernels::attn as fattn;
use dllm_kernels::linear as flin;
use dllm_kernels::moe::MoeDecodeScratch;

impl GpuModel {
    /// host-side prep of a prefill forward: attention tiles (consecutive rows sharing a block table)
    pub(super) fn prepare_tiles(&self, bt_off: &[i32], decode: bool) -> Result<()> {
        if decode {
            return Ok(());
        }
        let ws = &self.ws;
        for (i, hd) in [256usize, 512].iter().enumerate() {
            let (r0, n) = fattn::make_tiles(bt_off, fattn::prefill_rows_per_tile(*hd));
            ws.tile_row0[i].upload_slice(&r0)?;
            ws.tile_n[i].upload_slice(&n)?;
            ws.ntiles[i].store(r0.len(), Ordering::Relaxed);
        }
        Ok(())
    }

    pub(super) fn forward_rows(&self, r: usize, decode: bool, splits: usize, m: DevMeta) -> Result<()> {
        let gpu = self.gpu.clone();
        let be = &gpu.be;
        let cfg = &self.cfg;
        let (h, eps, dff) = (cfg.hidden_size, cfg.rms_eps, cfg.intermediate_size);
        let (k_top, im, ne) = (cfg.top_k_experts, cfg.moe_intermediate_size, cfg.num_experts);
        let ws = &self.ws;
        gpu.embed_gather_fp8(ws.x.ptr, self.embed_w.ptr, self.embed_s.ptr, m.tokens, r, h, self.embed_scale)?;
        let n_soft = self.soft_n.swap(0, Ordering::Relaxed);
        if n_soft > 0 {
            gpu.inject_soft(ws.x.ptr, self.soft_src.load(Ordering::Relaxed), self.soft_idx.ptr, self.soft_idx.at(self.opts.max_rows * 4), n_soft, h)?;
        }
        let small = r <= 16;
        let (aq, a_s) = (ws.aq.ptr, ws.a_s.ptr);
        flin::quant_rows(be, aq, a_s, ws.x.ptr, self.layers[0].in_ln.ptr, eps, r, h, true)?;
        let nl = self.layers.len();
        for (i, l) in self.layers.iter().enumerate() {
            let sp = &l.spec;
            let moe = &l.moe;
            let fe = &moe.experts;
            let (nh, nkv, hd) = (cfg.num_heads, sp.n_kv_heads, sp.head_dim);
            let (nq, nk) = (nh * hd, nkv * hd);
            let nv = if l.v.is_some() { nk } else { 0 };
            let hi = if sp.kind == LayerKind::Sliding { m.kv_hi_sl } else { m.kv_hi };
            let window = if sp.kind == LayerKind::Sliding { sp.window.unwrap_or(0) } else { 0 };
            let pm = self.kv_pmap(i);
            // ---- qkv
            if small {
                let ld = nq + nk + nv;
                let mut segs = vec![l.q.seg(), l.k.seg()];
                if let Some(v) = &l.v {
                    segs.push(v.seg());
                }
                flin::gemv_pqd(be, ws.qkv.ptr, ld, aq, a_s, r, h, &segs, false, None)?;
                let vr = if l.v.is_some() { ws.qkv.at((nq + nk) * 2) } else { ws.qkv.at(nq * 2) };
                fattn::qkv_prep_ld(&gpu, ws.q.ptr, self.kpool[i].ptr, self.vpool[i].ptr, ws.qkv.ptr, ws.qkv.at(nq * 2), vr, l.q_norm.ptr, l.k_norm.ptr, self.rope[l.rope].ptr, m.pos, m.slot, r, nh, nkv, hd, eps, ld, ld, ld, pm)?;
            } else {
                let (q, k) = (l.q.seg(), l.k.seg());
                flin::gemm(be, ws.qraw.ptr, aq, a_s, q.w, q.s, r, nq, h)?;
                flin::gemm(be, ws.kraw.ptr, aq, a_s, k.w, k.s, r, nk, h)?;
                let vsrc = if let Some(v) = &l.v {
                    let v = v.seg();
                    flin::gemm(be, ws.vraw.ptr, aq, a_s, v.w, v.s, r, nv, h)?;
                    ws.vraw.ptr
                } else {
                    ws.kraw.ptr
                };
                fattn::qkv_prep(&gpu, ws.q.ptr, self.kpool[i].ptr, self.vpool[i].ptr, ws.qraw.ptr, ws.kraw.ptr, vsrc, l.q_norm.ptr, l.k_norm.ptr, self.rope[l.rope].ptr, m.pos, m.slot, r, nh, nkv, hd, eps, pm)?;
            }
            // ---- attention
            if decode {
                let sp_n = fattn::choose_splits(splits.max(1) * 1024, r, nkv, if r == 1 { 1 } else { 5 });
                let (part, cnt) = if sp_n > 1 { (ws.part.ptr, ws.attn_cnt.ptr) } else { (0, 0) };
                fattn::attn_decode(&gpu, ws.ao.ptr, part, ws.q.ptr, self.kpool[i].ptr, self.vpool[i].ptr, m.bt, m.bt_off, m.pos, hi, r, nh, nkv, hd, window, sp_n, cnt, pm)?;
            } else {
                let ti = if hd == 512 { 1 } else { 0 };
                let nt = ws.ntiles[ti].load(Ordering::Relaxed);
                ensure!(nt > 0, "prefill needs the host tile table (call GpuModel::forward, not forward_dev, for prefill batches)");
                fattn::attn_prefill(&gpu, ws.ao.ptr, ws.q.ptr, self.kpool[i].ptr, self.vpool[i].ptr, m.bt, m.bt_off, m.pos, hi, ws.tile_row0[ti].ptr, ws.tile_n[ti].ptr, nt, nh, nkv, hd, window, pm)?;
            }
            // ---- o_proj
            flin::quant_rows(be, aq, a_s, ws.ao.ptr, 0, eps, r, nq, false)?;
            let o = l.o.seg();
            if small {
                flin::gemv_pqd(be, ws.o.ptr, h, aq, a_s, r, nq, &[o], false, None)?;
            } else {
                flin::gemm(be, ws.o.ptr, aq, a_s, o.w, o.s, r, h, nq)?;
            }
            // ---- sandwich + producers of the dense / expert / router inputs
            gpu.post_attn(ws.mid.ptr, ws.fq.ptr, ws.fs.ptr, ws.mxq.ptr, ws.mxs.ptr, ws.rn.ptr, ws.x.ptr, ws.o.ptr, l.post_attn_ln.ptr, l.pre_ff_ln.ptr, moe.pre_ff_ln2.ptr, moe.router_scale.ptr, fe.act_g, r, h, eps)?;
            // ---- dense MLP (own stream branch in decode: overlaps the router + MoE chain)
            if small {
                be.branch_fork()?;
                be.branch_use(true);
            }
            let (g, u, d) = (l.gate.seg(), l.up.seg(), l.down.seg());
            if small {
                flin::gemv_pqd(be, ws.act.ptr, dff, ws.fq.ptr, ws.fs.ptr, r, h, &[g, u], true, None)?;
                flin::quant_rows(be, aq, a_s, ws.act.ptr, 0, eps, r, dff, false)?;
                flin::gemv_pqd(be, ws.m.ptr, h, aq, a_s, r, dff, &[d], false, None)?;
            } else {
                flin::gemm_glu(be, ws.act.ptr, ws.fq.ptr, ws.fs.ptr, g.w, g.s, u.w, u.s, r, dff, h)?;
                flin::quant_rows(be, aq, a_s, ws.act.ptr, 0, eps, r, dff, false)?;
                flin::gemm(be, ws.m.ptr, aq, a_s, d.w, d.s, r, h, dff)?;
            }
            if small {
                be.branch_use(false);
            }
            // ---- MoE
            let proj = moe.router_proj.ptr;
            if small {
                let s = MoeDecodeScratch { logits: ws.rlog.ptr, ids: ws.topk_ids.ptr, wts: ws.topk_w.ptr, ticket: ws.ticket.ptr, hq: ws.mhq.ptr, hs: ws.mhs.ptr, opair: ws.moe_o.ptr, grp: ws.grp.ptr };
                gpu.moe_router_decode(&s, ws.rn.ptr, proj, moe.per_expert_scale.ptr, r, h)?;
                gpu.moe_glu_decode(fe, &s, ws.mxq.ptr, ws.mxs.ptr, r, h, im)?;
                gpu.moe_down_decode(fe, &s, r, h, im)?;
            } else {
                gpu.moe_router_logits_prefill(ws.rlog.ptr, ws.rn.ptr, proj, r, h)?;
                gpu.moe_topk_prefill(ws.topk_ids.ptr, ws.topk_w.ptr, ws.rlog.ptr, moe.per_expert_scale.ptr, r)?;
                gpu.moe_route(ws.counts.ptr, ws.offsets.ptr, ws.cursor.ptr, ws.pair_of_slot.ptr, ws.slot_of_pair.ptr, ws.topk_ids.ptr, r, k_top, ne)?;
                gpu.moe_glu_prefill(fe, ws.mhq.ptr, ws.mhs.ptr, ws.mxq.ptr, ws.mxs.ptr, ws.pair_of_slot.ptr, ws.offsets.ptr, ws.counts.ptr, r, h, im, ne)?;
                gpu.moe_down_prefill(fe, ws.moe_o.ptr, ws.mhq.ptr, ws.mhs.ptr, ws.pair_of_slot.ptr, ws.offsets.ptr, ws.counts.ptr, r, h, im, ne)?;
            }
            if small {
                be.branch_join()?;
            }
            // ---- combine + layer end (+ next layer's input norm / fp8 quant)
            let next = if i + 1 < nl { Some((self.layers[i + 1].in_ln.ptr, aq, a_s)) } else { None };
            gpu.layer_final(ws.x.ptr, ws.m.ptr, ws.moe_o.ptr, ws.topk_w.ptr, ws.mid.ptr, l.post_ff_ln1.ptr, moe.post_ff_ln2.ptr, l.post_ff_ln.ptr, l.layer_scalar.ptr, r, h, eps, next)?;
        }
        gpu.rmsnorm(ws.hidden.ptr, ws.x.ptr, self.final_norm.ptr, r, h, eps)?;
        Ok(())
    }
}