cherenkov/qwen4_exp/gpu/prefill/
projection.rs

1//! Dense prefill projections with narrow-output fallback.
2
3use super::*;
4
5impl Gpu<'_> {
6    /// `y[b] = W x[b]` for nb rows with the simdgroup-matrix GEMM: the
7    /// x buffer must hold rows padded to the token tile (32).
8    pub(super) fn qmm_from(&self, enc: &Enc, wb: &Buf, q: &Q, x: &Buf, y: &Buf, nb: usize) {
9        let p = QmvParams {
10            out_dim: q.out,
11            in_dim: q.inp,
12        };
13
14        if q.out < 128 {
15            // The GEMM tiles clamp-load and store 8-row fragments; narrow
16            // outputs go through the plain per-output kernel.
17            let nbu = nb as u32;
18
19            self.dispatch(
20                enc,
21                &self.pipes.qmv_small_b,
22                |e| {
23                    self.bind(e, 0, wb, q.w);
24                    self.bind(e, 1, wb, q.s);
25                    self.bind(e, 2, wb, q.b);
26                    self.bind(e, 3, x, 0);
27                    self.bind(e, 4, y, 0);
28                    set_bytes(e, 5, &p);
29                    set_bytes(e, 6, &nbu);
30                },
31                (nb * q.out as usize).div_ceil(4),
32                128,
33                true,
34            );
35
36            return;
37        }
38
39        self.qmm_tiled(
40            enc,
41            wb,
42            q,
43            x,
44            y,
45            nb,
46            [&self.pipes.qmm_n8, &self.pipes.qmm_n16, &self.pipes.qmm_w],
47        );
48    }
49
50    /// Low-bit expert records share the tiled dispatch and output layout,
51    /// but use their own unpacking kernels. Dense projections stay Q4.
52    #[allow(clippy::too_many_arguments)]
53    pub(super) fn expert_qmm_from(
54        &self,
55        enc: &Enc,
56        wb: &Buf,
57        q: &Q,
58        x: &Buf,
59        y: &Buf,
60        nb: usize,
61        bits: u32,
62    ) {
63        if bits == 4 {
64            self.qmm_from(enc, wb, q, x, y, nb);
65        } else {
66            let pipes = &self.pipes.expert_qmm[(bits - 2) as usize];
67
68            self.qmm_tiled(enc, wb, q, x, y, nb, [&pipes[0], &pipes[1], &pipes[2]]);
69        }
70    }
71
72    #[allow(clippy::too_many_arguments)]
73    fn qmm_tiled(&self, enc: &Enc, wb: &Buf, q: &Q, x: &Buf, y: &Buf, nb: usize, pipes: [&Pso; 3]) {
74        let p = QmvParams {
75            out_dim: q.out,
76            in_dim: q.inp,
77        };
78        let (pipe, tile) = if nb <= 8 {
79            (pipes[0], 8)
80        } else if nb <= 16 {
81            (pipes[1], 16)
82        } else {
83            (pipes[2], 32)
84        };
85        let ntt = nb.div_ceil(tile).max(1);
86        let nttu = ntt as u32;
87
88        self.dispatch(
89            enc,
90            pipe,
91            |e| {
92                self.bind(e, 0, wb, q.w);
93                self.bind(e, 1, wb, q.s);
94                self.bind(e, 2, wb, q.b);
95                self.bind(e, 3, x, 0);
96                self.bind(e, 4, y, 0);
97                set_bytes(e, 5, &p);
98                set_bytes(e, 6, &nttu);
99            },
100            (q.out as usize).div_ceil(128) * ntt,
101            128,
102            true,
103        );
104    }
105
106    pub(super) fn qmm(&self, enc: &Enc, q: &Q, x: &Buf, y: &Buf, nb: usize) {
107        self.qmm_from(enc, &self.dense, q, x, y, nb);
108    }
109}
cherenkov/qwen4_exp/gpu/prefill/projection.rs

cherenkov/qwen4_exp/gpu/prefill/
projection.rs
