cherenkov/qwen4_exp/gpu/
prefill.rs

1//! Prefill engine: a prompt chunk of up to a few thousand tokens runs
2//! through the model layer by layer, so each layer's experts are streamed
3//! once for the whole chunk instead of once per few tokens.
4//!
5//! Dense projections use simdgroup-matrix GEMM kernels;
6//! attention is the GEMM form (S = QK^T, causal softmax with the QSA
7//! visibility mask, O = PV) in sub-chunks of queries; the DeltaNet scan
8//! and the PLE conv walk the chunk sequentially per channel. Routed
9//! experts are handled per expert: its tokens' rows are gathered, run
10//! through gate/up/act/down GEMMs and scatter-added back with their
11//! routing weights; records stream through a ring of pool slots, a fetch
12//! thread filling batches ahead of the GPU with two events per batch.
13//! Every side effect (KV cache, index keys, recurrent state, PLE ring,
14//! MTP cache) is left exactly as the row-batched engine would have left
15//! it, so decode continues from the result.
16
17use super::*;
18
19pub(super) mod allocation;
20mod attention;
21mod deltanet;
22mod experts;
23mod hyperconnection;
24mod ple;
25mod projection;
26mod stats;
27pub use stats::PrefillStats;
28
29/// Slots of the expert stream ring, and experts per fetch batch.
30pub(super) const RING: usize = 64;
31const GROUP: usize = 8;
32/// Query rows per attention sub-chunk (bounds the S matrix).
33const QS: usize = 256;
34
35/// Timing of one prefill chunk.
36#[derive(Clone, Copy, Default)]
37pub struct ChunkStats {
38    pub tokens: usize,
39    pub secs: f64,
40    /// Expert records read from disk (ring + pool).
41    pub fetched: usize,
42    /// Bytes read in each record's stored precision.
43    pub fetched_bytes: usize,
44    /// Seconds the fetch thread waited for ring slots.
45    pub wait_s: f64,
46    /// GPU seconds: DeltaNet blocks, attention blocks, expert streams, MTP.
47    pub gpu_delta_s: f64,
48    pub gpu_attn_s: f64,
49    pub gpu_experts_s: f64,
50    pub gpu_mtp_s: f64,
51    /// CPU seconds gathering n-gram rows (page faults included).
52    pub ngram_s: f64,
53}
54
55#[derive(Clone, Copy)]
56struct MoeRef {
57    record_layer: usize,
58    sg: Q,
59    su: Q,
60    sd: Q,
61    gate: T,
62}
63
64pub(super) struct PrefillScratch<T = Buf> {
65    allocated_bytes: usize,
66    rows: usize,
67    ids: T,
68    e: T,
69    hyper: T,
70    normed: T,
71    d: T,
72    u: T,
73    mixed: T,
74    inj: T,
75    mix_out: T,
76    moe_out: T,
77    qg: T,
78    k: T,
79    v: T,
80    attn_out: T,
81    iqk: T,
82    iq: T,
83    bscore: T,
84    vis: T,
85    nvis: T,
86    vmask: T,
87    ag_qh: T,
88    ag_kh: T,
89    ag_vt: T,
90    ag_s: T,
91    ag_p: T,
92    ag_o: T,
93    qkv: T,
94    z: T,
95    a: T,
96    b: T,
97    kqn: T,
98    gbuf: T,
99    delta_y: T,
100    router: T,
101    topk_idx: T,
102    topk_w: T,
103    csr_rows: T,
104    csr_w: T,
105    xg: T,
106    ge: T,
107    ue: T,
108    hg: T,
109    ye: T,
110    mtp_hyper: T,
111    fe: T,
112    fh: T,
113    logits_all: Option<T>,
114}
115
116pub const MAX_PREFILL_ROWS: usize = 4096;
117
118impl<'a> Gpu<'a> {
119    /// One decoder block over t rows of `hyper`: everything up to the
120    /// router in one command buffer, then the expert stream. Returns
121    /// (records fetched, bytes fetched, ring wait seconds, block GPU seconds, expert
122    /// stream GPU seconds).
123    #[allow(clippy::too_many_arguments)]
124    fn pf_layer(
125        &mut self,
126        li: Option<usize>,
127        base: usize,
128        t: usize,
129        pf: &PrefillScratch,
130        hyper: &Buf,
131        pending: Option<&Buf>,
132    ) -> Result<(usize, usize, f64, f64, f64)> {
133        let (moe, block_s) = {
134            let layer = match li {
135                Some(li) => &self.layers[li],
136                None => &self.mtp.as_ref().context("no MTP head")?.layer,
137            };
138            let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
139            let enc = cb.computeCommandEncoder().context("encoder")?;
140            let mut pending = pending;
141
142            if let Some(pl) = &layer.ple {
143                if let Some(out) = pending.take() {
144                    self.pf_inject(&enc, hyper, out, &pf.inj, t);
145                }
146
147                self.pf_ple(&enc, pl, base, t, pf)?;
148            }
149
150            self.pf_hc_read(&enc, &layer.attn_hc, t, pf, hyper, pending, true);
151
152            if !self.skips("mixer") {
153                match &layer.mix {
154                    Mix::Attn(a) => self.pf_attention(&enc, a, base, t, pf),
155                    Mix::Delta(d) => self.pf_deltanet(&enc, d, t, pf),
156                }
157            }
158
159            self.pf_hc_read(&enc, &layer.mlp_hc, t, pf, hyper, Some(&pf.mix_out), true);
160            self.router_b(
161                &enc,
162                &layer.moe,
163                t,
164                &pf.mixed,
165                &pf.router,
166                &pf.topk_idx,
167                &pf.topk_w,
168                self.p.cfg.num_experts_per_tok,
169            );
170            enc.endEncoding();
171            cb.commit();
172            cb.waitUntilCompleted();
173
174            (
175                MoeRef {
176                    record_layer: layer.moe.record_layer,
177                    sg: layer.moe.sg,
178                    su: layer.moe.su,
179                    sd: layer.moe.sd,
180                    gate: layer.moe.shared_gate,
181                },
182                cb.GPUEndTime() - cb.GPUStartTime(),
183            )
184        };
185        let (fetched, fetched_bytes, wait_s, experts_s) = self.pf_experts(moe, t, pf)?;
186
187        Ok((fetched, fetched_bytes, wait_s, block_s, experts_s))
188    }
189
190    /// Run one prompt chunk at positions pos.. through the trunk and the
191    /// MTP head, committing everything. `next_after` is the token that
192    /// follows the chunk when known (the next chunk's first token); the
193    /// trunk's own prediction pairs with the last row otherwise. With
194    /// `all_logits` every row's logits are kept for checks. Returns the
195    /// greedy token after the chunk and the MTP's first draft.
196    pub fn prefill_chunk(
197        &mut self,
198        tokens: &[u32],
199        next_after: Option<u32>,
200        all_logits: bool,
201    ) -> Result<(u32, u32)> {
202        let t = tokens.len();
203
204        anyhow::ensure!(t >= 1, "empty chunk");
205        anyhow::ensure!(self.pos + t <= self.max_t, "context capacity exceeded");
206
207        let t0 = std::time::Instant::now();
208
209        if self
210            .pf
211            .as_ref()
212            .is_none_or(|pf| pf.rows < t || (all_logits && pf.logits_all.is_none()))
213        {
214            self.pf = None;
215            self.pf = Some(self.pf_alloc(t, all_logits)?);
216        }
217
218        let pf = self.pf.take().unwrap();
219        let c = &self.p.cfg;
220        let h = c.hidden_size as u32;
221        let hh = c.hc_hidden();
222        let pos = self.pos;
223
224        self.tokens.truncate(pos);
225        self.tokens.extend_from_slice(tokens);
226        self.ngram_prefetch_start(pos, t);
227
228        let write_ids = |buf: &Buf, ids: &[u32]| unsafe {
229            std::ptr::copy_nonoverlapping(
230                ids.as_ptr(),
231                buf.contents().cast::<u32>().as_ptr(),
232                ids.len(),
233            );
234        };
235
236        write_ids(&pf.ids, tokens);
237
238        let mut st = ChunkStats {
239            tokens: t,
240            ..Default::default()
241        };
242        let ngram0 = self.ngram_gather_s.get();
243
244        // Embedding into the replicated stream.
245        {
246            let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
247            let enc = cb.computeCommandEncoder().context("encoder")?;
248            let (i0, nbu) = (0u32, t as u32);
249
250            self.dispatch(
251                &enc,
252                &self.pipes.embed_rows,
253                |e| {
254                    self.bind(e, 0, &pf.ids, 0);
255                    self.bind(e, 1, &self.dense, self.embed.w);
256                    self.bind(e, 2, &self.dense, self.embed.s);
257                    self.bind(e, 3, &self.dense, self.embed.b);
258                    self.bind(e, 4, &pf.e, 0);
259                    set_bytes(e, 5, &h);
260                    set_bytes(e, 6, &i0);
261                    set_bytes(e, 7, &nbu);
262                },
263                t * h as usize,
264                256,
265                false,
266            );
267
268            let gp = self.group_params(false, 0.0);
269
270            self.dispatch(
271                &enc,
272                &self.pipes.replicate_b,
273                |e| {
274                    self.bind(e, 0, &pf.e, 0);
275                    self.bind(e, 1, &pf.hyper, 0);
276                    set_bytes(e, 2, &gp);
277                    set_bytes(e, 3, &nbu);
278                },
279                t * hh,
280                256,
281                false,
282            );
283            enc.endEncoding();
284            cb.commit();
285            cb.waitUntilCompleted();
286        }
287
288        let n_layers = self.layers.len().min(layer_cap());
289
290        for li in 0..n_layers {
291            let pending = if li > 0 { Some(&pf.moe_out) } else { None };
292            let (f, bytes, w, block_s, experts_s) =
293                self.pf_layer(Some(li), pos, t, &pf, &pf.hyper, pending)?;
294            st.fetched += f;
295            st.fetched_bytes += bytes;
296            st.wait_s += w;
297            st.gpu_experts_s += experts_s;
298
299            if matches!(self.layers[li].mix, Mix::Delta(_)) {
300                st.gpu_delta_s += block_s;
301            } else {
302                st.gpu_attn_s += block_s;
303            }
304        }
305
306        // Final injection, then the LM head on the last row (all rows when
307        // checking).
308        let cur = {
309            let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
310            let enc = cb.computeCommandEncoder().context("encoder")?;
311
312            self.pf_inject(&enc, &pf.hyper, &pf.moe_out, &pf.inj, t);
313
314            if let Some(la) = &pf.logits_all {
315                self.pf_hc_read(&enc, &self.final_mixer, t, &pf, &pf.hyper, None, false);
316                self.qmm(&enc, &self.lm_head, &pf.mixed, la, t);
317            }
318
319            self.head_b(
320                &enc,
321                &self.final_mixer,
322                1,
323                &pf.hyper,
324                (t - 1) * hh * 4,
325                None,
326                &self.scratch.logits,
327                IDS_OUT,
328            );
329            enc.endEncoding();
330            cb.commit();
331            cb.waitUntilCompleted();
332
333            self.read_u32(&self.scratch.ids, IDS_OUT + 1)[IDS_OUT]
334        };
335        // MTP head over the chunk: row r pairs the trunk residual at pos+r
336        // with the token after it.
337        let mut draft = cur;
338
339        if let Some(mtp) = self.mtp.as_ref() {
340            let (enorm, hnorm, fc_e, fc_h) = (mtp.enorm, mtp.hnorm, mtp.fc_e, mtp.fc_h);
341            let mut next: Vec<u32> = tokens[1..].to_vec();
342
343            next.push(next_after.unwrap_or(cur));
344            write_ids(&pf.ids, &next);
345
346            {
347                let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
348                let enc = cb.computeCommandEncoder().context("encoder")?;
349                let (i0, nbu) = (0u32, t as u32);
350
351                self.dispatch(
352                    &enc,
353                    &self.pipes.embed_rows,
354                    |e| {
355                        self.bind(e, 0, &pf.ids, 0);
356                        self.bind(e, 1, &self.dense, self.embed.w);
357                        self.bind(e, 2, &self.dense, self.embed.s);
358                        self.bind(e, 3, &self.dense, self.embed.b);
359                        self.bind(e, 4, &pf.e, 0);
360                        set_bytes(e, 5, &h);
361                        set_bytes(e, 6, &i0);
362                        set_bytes(e, 7, &nbu);
363                    },
364                    t * h as usize,
365                    256,
366                    false,
367                );
368                self.group_norm_b(&enc, &pf.e, 0, enorm, &pf.mixed, h, 1, 1.0, t);
369                self.qmm(&enc, &fc_e, &pf.mixed, &pf.fe, t);
370                self.group_norm_b(
371                    &enc,
372                    &pf.hyper,
373                    0,
374                    hnorm,
375                    &pf.normed,
376                    h,
377                    c.hc_count as u32,
378                    1.0,
379                    t,
380                );
381                self.qmm(&enc, &fc_h, &pf.normed, &pf.fh, t * c.hc_count);
382
383                let gp = self.group_params(false, 0.0);
384
385                self.dispatch(
386                    &enc,
387                    &self.pipes.mtp_fold,
388                    |e| {
389                        self.bind(e, 0, &pf.fe, 0);
390                        self.bind(e, 1, &pf.fh, 0);
391                        self.bind(e, 2, &pf.mtp_hyper, 0);
392                        set_bytes(e, 3, &gp);
393                        set_bytes(e, 4, &nbu);
394                    },
395                    t * hh,
396                    256,
397                    false,
398                );
399                enc.endEncoding();
400                cb.commit();
401                cb.waitUntilCompleted();
402            }
403
404            let (f, bytes, w, block_s, experts_s) =
405                self.pf_layer(None, pos, t, &pf, &pf.mtp_hyper, None)?;
406            st.fetched += f;
407            st.fetched_bytes += bytes;
408            st.wait_s += w;
409            st.gpu_mtp_s += block_s + experts_s;
410            let mtp = self.mtp.as_ref().unwrap();
411            let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
412            let enc = cb.computeCommandEncoder().context("encoder")?;
413
414            self.pf_inject(&enc, &pf.mtp_hyper, &pf.moe_out, &pf.inj, t);
415            self.head_b(
416                &enc,
417                &mtp.mixer,
418                1,
419                &pf.mtp_hyper,
420                (t - 1) * hh * 4,
421                None,
422                &self.scratch.mtp_logits,
423                IDS_MTP_OUT,
424            );
425
426            // Keep the last row's residual where chained drafts re-enter.
427            let n = hh as u32;
428
429            self.dispatch(
430                &enc,
431                &self.pipes.copy_f32,
432                |e| {
433                    self.bind(e, 0, &pf.mtp_hyper, (t - 1) * hh * 4);
434                    self.bind(e, 1, &self.scratch.mtp_hyper, 0);
435                    set_bytes(e, 2, &n);
436                },
437                hh,
438                256,
439                false,
440            );
441            enc.endEncoding();
442            cb.commit();
443            cb.waitUntilCompleted();
444
445            draft = self.read_u32(&self.scratch.ids, IDS_MTP_OUT + 1)[IDS_MTP_OUT];
446            self.mtp_len = pos + t;
447        }
448
449        self.pos = pos + t;
450        self.batch_pos = self.pos;
451        self.batch_nb = 0;
452
453        self.join_pending()?;
454
455        st.secs = t0.elapsed().as_secs_f64();
456        st.ngram_s = self.ngram_gather_s.get() - ngram0;
457
458        self.activity.prefill.record(st);
459        self.prefill_stats.push(st);
460
461        self.pf = Some(pf);
462
463        Ok((cur, draft))
464    }
465
466    /// Release the prefill scratch (a few hundred MB to GB).
467    pub fn prefill_release(&mut self) {
468        self.pf = None;
469    }
470
471    /// Debug: blocks selected for row `r` of the last prefill chunk (its
472    /// last attention layer), from the block bitmask.
473    pub fn debug_engine_blocks(&self, r: usize) -> Vec<u32> {
474        let c = &self.p.cfg;
475        let max_blocks = self.max_t / c.indexer_compress_ratio + 1;
476        let words = max_blocks.div_ceil(32);
477        let pf = self.pf.as_ref().expect("prefill scratch released");
478        let m = self.read_u32(&pf.vmask, (r + 1) * words);
479
480        (0..max_blocks as u32)
481            .filter(|&j| (m[r * words + (j / 32) as usize] >> (j % 32)) & 1 == 1)
482            .collect()
483    }
484
485    /// Debug: visible tokens of row `r` of the last row-batched step (its
486    /// last attention layer).
487    pub fn debug_row_vis(&self, r: usize) -> Vec<u32> {
488        let c = &self.p.cfg;
489        let stride = c.indexer_budget + c.indexer_compress_ratio;
490        let n = self.read_u32(&self.scratch.nvis, r + 1)[r] as usize;
491
492        self.read_u32(&self.scratch.vis, (r + 1) * stride)[r * stride..r * stride + n].to_vec()
493    }
494
495    /// Row `r` of the last prefill chunk's logits (all-logits mode).
496    pub fn pf_logits_row(&self, r: usize) -> &[f32] {
497        let v = self.p.cfg.vocab_size;
498        let la = self
499            .pf
500            .as_ref()
501            .and_then(|pf| pf.logits_all.as_ref())
502            .expect("prefill logits not kept");
503        let ptr = la.contents().cast::<f32>();
504
505        unsafe { std::slice::from_raw_parts(ptr.as_ptr().add(r * v), v) }
506    }
507}
508
509#[cfg(test)]
510#[path = "../../../tests/unit/qwen4_exp/gpu/prefill.rs"]
511mod tests;
cherenkov/qwen4_exp/gpu/prefill.rs

cherenkov/qwen4_exp/gpu/
prefill.rs
