File size: 8,297 Bytes
19886eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | //! 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(())
}
}
|