cherenkov/qwen4_exp/gpu/
mtp.rs

1//! Draft-head execution and folded/chained MTP inputs.
2
3use super::*;
4
5impl Gpu<'_> {
6    /// MTP draft head after a commit of `next.len()` rows: row r pairs the
7    /// trunk's wide residual at batch_pos + r with `next[r]`, the token that
8    /// follows it. Fills the head's KV cache for those positions and returns
9    /// `chain` greedy drafts for the positions after the last committed
10    /// token (each chained draft re-enters the head on its own residual).
11    pub fn mtp_draft(&mut self, next: &[u32], chain: usize) -> Result<Vec<u32>> {
12        let n = next.len();
13
14        anyhow::ensure!(
15            n >= 1 && n <= self.batch_nb,
16            "MTP rows must match the committed rows"
17        );
18        anyhow::ensure!(
19            self.mtp_len == self.batch_pos,
20            "MTP cache out of step: {} vs {}",
21            self.mtp_len,
22            self.batch_pos
23        );
24
25        let t0 = std::time::Instant::now();
26        let mut drafts = Vec::with_capacity(chain);
27        let mut token_rows: Vec<u32> = next.to_vec();
28        let mut src_hyper_row = 0usize;
29        let mut first = true;
30        let mut c0 = 0;
31
32        if let Some(out) = self.folded_mtp.take() {
33            // The trunk step already ran the head over its rows with the
34            // trunk's predictions, which are `next` for every accepted row.
35            self.mtp_len = self.batch_pos + n;
36            src_hyper_row = n - 1;
37            first = false;
38            token_rows = vec![out[n - 1]];
39
40            if chain > 0 {
41                drafts.push(out[n - 1]);
42            }
43
44            c0 = 1;
45        }
46
47        for c in c0..chain.max(1) {
48            let rows = if first { n } else { 1 };
49            let base_pos = if first {
50                self.batch_pos
51            } else {
52                self.batch_pos + n - 1 + c
53            };
54
55            anyhow::ensure!(base_pos + rows <= self.max_t, "context capacity exceeded");
56
57            let out = self.mtp_pass(&token_rows, rows, base_pos, first, src_hyper_row)?;
58            let d = out[rows - 1];
59
60            if first {
61                self.mtp_len = self.batch_pos + n;
62                src_hyper_row = n - 1;
63            } else {
64                src_hyper_row = 0;
65            }
66
67            if c < chain {
68                drafts.push(d);
69            }
70
71            token_rows = vec![d];
72            first = false;
73        }
74
75        self.mtp_ms.push(t0.elapsed().as_secs_f64() * 1e3);
76
77        Ok(drafts)
78    }
79
80    /// One MTP pass over `rows` rows: hidden input from the trunk's wide
81    /// residual rows 0.. (`from_trunk`) or from the head's own residual row
82    /// `src_row`; tokens in `tokens`; positions base_pos... Returns the
83    /// argmax per row (logits stay in scratch.mtp_logits).
84    pub(super) fn mtp_pass(
85        &mut self,
86        tokens: &[u32],
87        rows: usize,
88        base_pos: usize,
89        from_trunk: bool,
90        src_row: usize,
91    ) -> Result<Vec<u32>> {
92        let hh = self.p.cfg.hc_hidden();
93        let phase_clock = self.phase_clock();
94        let mtp = self.mtp.as_ref().context("no MTP head")?;
95
96        unsafe {
97            let ids = self.scratch.ids.contents().cast::<u32>().as_ptr();
98
99            for (i, &t) in tokens.iter().enumerate() {
100                ids.add(IDS_MTP_IN + i).write(t);
101            }
102        }
103
104        self.dispatch_count.set(0);
105
106        self.step_no += 1;
107        let seq = self.event_base + 1;
108        self.event_base += 4;
109        let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
110        let mut enc = self.phase_encoder(&cb, None, 4 * self.layers.len())?;
111        let (src, off) = if from_trunk {
112            (&self.scratch.hyper, 0)
113        } else {
114            (&self.scratch.mtp_hyper, src_row * hh * 4)
115        };
116
117        self.encode_mtp_prelude(&enc, mtp, rows, IDS_MTP_IN, src, off);
118
119        let slot_row = self.layers.len();
120
121        self.encode_block(
122            &cb,
123            &mut enc,
124            &mtp.layer,
125            slot_row,
126            base_pos,
127            rows,
128            0,
129            &self.scratch.mtp_hyper,
130            None,
131            None,
132            seq,
133        )?;
134
135        let logits = if from_trunk {
136            &self.scratch.mtp_logits
137        } else {
138            &self.scratch.mtp_logits2
139        };
140
141        self.head_b(
142            &enc,
143            &mtp.mixer,
144            rows,
145            &self.scratch.mtp_hyper,
146            0,
147            Some(&self.scratch.moe_out),
148            logits,
149            IDS_MTP_OUT,
150        );
151        enc.endEncoding();
152        cb.commit();
153
154        let record_layer = mtp.layer.moe.record_layer;
155        let mut predicted = std::collections::VecDeque::new();
156        let (mut io_s, mut turn_s) = (0.0, 0.0);
157
158        self.service_block(
159            record_layer,
160            slot_row,
161            rows,
162            seq,
163            &mut predicted,
164            None,
165            &mut io_s,
166            &mut turn_s,
167        )?;
168        cb.waitUntilCompleted();
169        self.collect_phases(&[(slot_row, record_layer)], phase_clock);
170
171        Ok(self.read_u32(&self.scratch.ids, IDS_MTP_OUT + rows)[IDS_MTP_OUT..].to_vec())
172    }
173
174    /// The MTP head's input stream for `rows` rows: tokens from
175    /// scratch.ids[ids_in..], residual rows from `src` (byte offset `off`),
176    /// folded into scratch.mtp_hyper.
177    pub(super) fn encode_mtp_prelude(
178        &self,
179        enc: &Enc,
180        mtp: &Mtp,
181        rows: usize,
182        ids_in: usize,
183        src: &Buf,
184        off: usize,
185    ) {
186        let c = &self.p.cfg;
187        let h = c.hidden_size as u32;
188        let hc = c.hc_count;
189        let hh = c.hc_hidden();
190        let s = &self.scratch;
191        let (i0, nbu) = (ids_in as u32, rows as u32);
192
193        self.dispatch(
194            enc,
195            &self.pipes.embed_rows,
196            |e| {
197                self.bind(e, 0, &s.ids, 0);
198                self.bind(e, 1, &self.dense, self.embed.w);
199                self.bind(e, 2, &self.dense, self.embed.s);
200                self.bind(e, 3, &self.dense, self.embed.b);
201                self.bind(e, 4, &s.e, 0);
202                set_bytes(e, 5, &h);
203                set_bytes(e, 6, &i0);
204                set_bytes(e, 7, &nbu);
205            },
206            rows * h as usize,
207            256,
208            false,
209        );
210        // e_norm = rmsnorm(e) * (1 + enorm)  ->  fe = fc_e(e_norm)
211        self.group_norm_b(enc, &s.e, 0, mtp.enorm, &s.hc.mixed, h, 1, 1.0, rows);
212        self.prep_h(enc, &s.hc.mixed, 0, h, rows, &s.hc.h1);
213        self.qmv_h(enc, &mtp.fc_e, &s.fe, rows, &s.hc.h1);
214        // Per stream: fh[b*hc+g] = fc_h(rmsnorm(src[b][g]) * (1 + hnorm[g]))
215        self.group_norm_b(
216            enc,
217            src,
218            off,
219            mtp.hnorm,
220            &s.hc.normed,
221            h,
222            hc as u32,
223            1.0,
224            rows,
225        );
226        self.prep_h(enc, &s.hc.normed, 0, h, rows * hc, &s.hc.h1);
227        self.qmv_h(enc, &mtp.fc_h, &s.fh, rows * hc, &s.hc.h1);
228
229        let gp = self.group_params(false, 0.0);
230
231        self.dispatch(
232            enc,
233            &self.pipes.mtp_fold,
234            |e| {
235                self.bind(e, 0, &s.fe, 0);
236                self.bind(e, 1, &s.fh, 0);
237                self.bind(e, 2, &s.mtp_hyper, 0);
238                set_bytes(e, 3, &gp);
239                set_bytes(e, 4, &nbu);
240            },
241            rows * hh,
242            256,
243            false,
244        );
245    }
246
247    /// A second draft chained on the MTP residual left by `prefill_chunk`.
248    pub fn mtp_chain(&mut self, draft: u32) -> Result<u32> {
249        let out = self.mtp_pass(&[draft], 1, self.pos, false, 0)?;
250
251        Ok(out[0])
252    }
253}
cherenkov/qwen4_exp/gpu/mtp.rs

cherenkov/qwen4_exp/gpu/
mtp.rs
