cherenkov/qwen4_exp/gpu/
ple.rs

1//! N-gram prefetch/gather and PLE gating/convolution dispatch.
2
3use super::*;
4
5impl Gpu<'_> {
6    /// The 16 n-gram row ids of the token at index `t` of `self.tokens`.
7    pub(super) fn ngram_ids_at(&self, pl: &Ple, t: usize) -> Vec<u64> {
8        CpuModel::ngram_ids_from(
9            &self.p.cfg,
10            &pl.multipliers,
11            &pl.head_offsets,
12            &pl.head_sizes,
13            &self.tokens[..=t],
14        )
15    }
16
17    /// Start pulling the n-gram rows of tokens t0..t0+n into the page
18    /// cache on a thread pool, so the block's gather (which faults the
19    /// mapping one row at a time) finds them warm. The ids only depend
20    /// on the tokens, which are known when a step or chunk starts.
21    pub(super) fn ngram_prefetch_start(&self, t0: usize, n: usize) {
22        let Some(pl) = self.layers.iter().find_map(|l| l.ple.as_ref()) else {
23            return;
24        };
25
26        self.ngram_prefetch_join();
27
28        let mut ids: Vec<u64> = (t0..t0 + n)
29            .flat_map(|t| self.ngram_ids_at(pl, t))
30            .collect();
31
32        ids.sort_unstable();
33        ids.dedup();
34
35        let row_bytes = self.p.manifest.ngram.row_bytes;
36        let file = self.ngram_file.try_clone().expect("dup ngram fd");
37        let h = std::thread::spawn(move || prefetch_ngram_rows(&file, &ids, row_bytes));
38        *self.ngram_prefetch.borrow_mut() = Some(h);
39    }
40
41    pub(super) fn ngram_prefetch_join(&self) {
42        if let Some(h) = self.ngram_prefetch.borrow_mut().take() {
43            let _ = h.join();
44        }
45    }
46
47    /// PLE block over `nb` rows at positions base_pos..: n-gram rows are
48    /// gathered on the CPU from `self.tokens` (which must hold the batch).
49    pub(super) fn ple_b(&self, enc: &Enc, pl: &Ple, nb: usize, base_pos: usize) -> Result<()> {
50        let c = &self.p.cfg;
51        let s = &self.scratch;
52        let h = c.hidden_size as u32;
53        let hh = c.hc_hidden();
54        let dim = self.p.manifest.ngram.dim;
55        let e = unsafe {
56            std::slice::from_raw_parts_mut(
57                pl.e.contents().cast::<f32>().as_ptr(),
58                MAX_NB * c.ple_embed_dim,
59            )
60        };
61        let t_gather = std::time::Instant::now();
62
63        self.ngram_prefetch_join();
64
65        for b in 0..nb {
66            let ids = self.ngram_ids_at(pl, base_pos + b);
67
68            anyhow::ensure!(
69                ids.len() * dim == c.ple_embed_dim,
70                "n-gram head layout mismatch"
71            );
72
73            let eb = &mut e[b * c.ple_embed_dim..(b + 1) * c.ple_embed_dim];
74
75            for (hi, &id) in ids.iter().enumerate() {
76                self.p.ngram_row(id, &mut eb[hi * dim..(hi + 1) * dim]);
77            }
78        }
79
80        self.ngram_gather_s
81            .set(self.ngram_gather_s.get() + t_gather.elapsed().as_secs_f64());
82        self.prep_h(enc, &pl.e, 0, c.ple_embed_dim as u32, nb, &s.hc.h1);
83        self.qmv_h(enc, &pl.key, &s.ple_key, nb, &s.hc.h1);
84        self.qmv_h(enc, &pl.value, &s.ple_value, nb, &s.hc.h1);
85
86        let groups = c.hc_count as u32;
87
88        self.group_norm_b(
89            enc,
90            &s.ple_key,
91            0,
92            pl.norm_key,
93            &s.ple_keyn,
94            h,
95            groups,
96            0.0,
97            nb,
98        );
99        self.group_norm_b(
100            enc,
101            &s.hyper,
102            0,
103            pl.norm_query,
104            &s.ple_query,
105            h,
106            groups,
107            0.0,
108            nb,
109        );
110
111        let gp = self.group_params(true, 0.0);
112
113        self.dispatch(
114            enc,
115            &self.pipes.ple_gate_b,
116            |e| {
117                self.bind(e, 0, &s.ple_keyn, 0);
118                self.bind(e, 1, &s.ple_query, 0);
119                self.bind(e, 2, &s.ple_value, 0);
120                self.bind(e, 3, &s.ple_gated, 0);
121                set_bytes(e, 4, &gp);
122            },
123            nb * c.hc_count,
124            256,
125            true,
126        );
127        self.group_norm_b(
128            enc,
129            &s.ple_gated,
130            0,
131            pl.norm_conv,
132            &s.ple_gvn,
133            h,
134            groups,
135            0.0,
136            nb,
137        );
138
139        let cp = PleConvParams {
140            channels: hh as u32,
141            ksize: pl.kernel,
142            dilation: pl.dilation,
143            span: pl.span,
144            filled: base_pos as u32,
145            nb: nb as u32,
146        };
147
148        self.dispatch(
149            enc,
150            &self.pipes.ple_conv_b,
151            |e| {
152                self.bind(e, 0, &s.ple_gated, 0);
153                self.bind(e, 1, &s.ple_gvn, 0);
154                self.bind(e, 2, &self.dense, pl.conv.0);
155                self.bind(e, 3, &pl.hist, 0);
156                self.bind(e, 4, &s.ple_out, 0);
157                set_bytes(e, 5, &cp);
158            },
159            hh,
160            256,
161            false,
162        );
163        self.add(enc, &s.hyper, &s.ple_out, nb * hh);
164
165        Ok(())
166    }
167}
168
169/// Touch each selected row once; workers claim disjoint indices from the queue.
170fn prefetch_ngram_rows(file: &std::fs::File, ids: &[u64], row_bytes: u64) {
171    use std::os::unix::fs::FileExt as _;
172
173    use std::sync::atomic::{AtomicUsize, Ordering};
174
175    let next = AtomicUsize::new(0);
176    let workers = 16.min(ids.len().max(1));
177
178    std::thread::scope(|s| {
179        for _ in 0..workers {
180            s.spawn(|| {
181                let mut buf = [0u8; 256];
182
183                loop {
184                    let i = next.fetch_add(1, Ordering::Relaxed);
185
186                    if i >= ids.len() {
187                        break;
188                    }
189
190                    let _ = file.read_at(&mut buf[..row_bytes as usize], ids[i] * row_bytes);
191                }
192            });
193        }
194    });
195}
cherenkov/qwen4_exp/gpu/ple.rs

cherenkov/qwen4_exp/gpu/
ple.rs
