cherenkov/qwen4_exp/gpu/
sampling.rs

1//! Final hyper-connection read, logits projection, and greedy argmax.
2
3use super::*;
4
5impl Gpu<'_> {
6    /// Final mixer read of `nb` rows of `hyper` (injecting `pending`
7    /// first), LM head into `logits` `[nb][vocab]`, greedy argmax of each
8    /// row into ids[ids_out + b].
9    #[allow(clippy::too_many_arguments)]
10    pub(super) fn head_b(
11        &self,
12        enc: &Enc,
13        mixer: &Hc,
14        nb: usize,
15        hyper: &Buf,
16        hyper_off: usize,
17        pending: Option<&Buf>,
18        logits: &Buf,
19        ids_out: usize,
20    ) {
21        let s = &self.scratch;
22        let vocab = self.p.cfg.vocab_size;
23
24        self.hc_read_b(enc, mixer, nb, hyper, hyper_off, pending, &s.hc, false);
25
26        if !self.skips("lmhead") {
27            self.qmv_h(enc, &self.lm_head, logits, nb, &s.hc.h1);
28        }
29
30        let n = vocab as u32;
31        let np = ARGMAX_TGS as u32;
32
33        for b in 0..nb {
34            self.dispatch(
35                enc,
36                &self.pipes.argmax_partial,
37                |e| {
38                    self.bind(e, 0, logits, b * vocab * 4);
39                    self.bind(e, 1, &s.partials, 0);
40                    set_bytes(e, 2, &n);
41                },
42                ARGMAX_TGS,
43                256,
44                true,
45            );
46
47            let step = (ids_out + b - 1) as u32;
48
49            self.dispatch(
50                enc,
51                &self.pipes.argmax_final,
52                |e| {
53                    self.bind(e, 0, &s.partials, 0);
54                    self.bind(e, 1, &s.ids, 0);
55                    set_bytes(e, 2, &np);
56                    set_bytes(e, 3, &step);
57                },
58                1,
59                1024,
60                true,
61            );
62        }
63    }
64}
cherenkov/qwen4_exp/gpu/sampling.rs

cherenkov/qwen4_exp/gpu/
sampling.rs
