cherenkov/qwen4_exp/gpu/
hyperconnection.rs

1//! Grouped normalization and gated residual mixing/injection.
2
3use super::*;
4
5impl Gpu<'_> {
6    pub(super) fn group_params(&self, eps: bool, shift: f32) -> GroupParams {
7        let c = &self.p.cfg;
8
9        GroupParams {
10            n: c.hidden_size as u32,
11            groups: c.hc_count as u32,
12            eps: if eps { c.rms_norm_eps as f32 } else { 0.0 },
13            shift,
14        }
15    }
16
17    /// Grouped RMSNorm of `nb` rows (each `groups` slices of `n`), weight
18    /// shifted by `shift` (1.0 for raw HF norm weights).
19    #[allow(clippy::too_many_arguments)]
20    pub(super) fn group_norm_b(
21        &self,
22        enc: &Enc,
23        x: &Buf,
24        x_off: usize,
25        w: T,
26        y: &Buf,
27        n: u32,
28        groups: u32,
29        shift: f32,
30        nb: usize,
31    ) {
32        let p = GroupParams {
33            n,
34            groups,
35            eps: self.p.cfg.rms_norm_eps as f32,
36            shift,
37        };
38
39        self.dispatch(
40            enc,
41            &self.pipes.group_norm_b,
42            |e| {
43                self.bind(e, 0, x, x_off);
44                self.bind(e, 1, &self.dense, w.0);
45                self.bind(e, 2, y, 0);
46                set_bytes(e, 3, &p);
47            },
48            nb * groups as usize,
49            256,
50            true,
51        );
52    }
53
54    /// Gated residual read of `nb` rows of `hyper` for one block: the block
55    /// input lands in `bufs.mixed` and, prepped, in `bufs.h1`; injection
56    /// logits in `bufs.inj`. `pending` is a block output not yet injected
57    /// into the stream; the fused norm applies it first (in place).
58    #[allow(clippy::too_many_arguments)]
59    pub(super) fn hc_read_b(
60        &self,
61        enc: &Enc,
62        hc: &Hc,
63        nb: usize,
64        hyper: &Buf,
65        hyper_off: usize,
66        pending: Option<&Buf>,
67        bufs: &HcBufs,
68        inject: bool,
69    ) {
70        let c = &self.p.cfg;
71        let h = c.hidden_size as u32;
72        let nbu = nb as u32;
73        let np = NormPrepParams {
74            n: h,
75            groups: c.hc_count as u32,
76            eps: c.rms_norm_eps as f32,
77            inject: pending.is_some() as u32,
78        };
79        let out = pending.unwrap_or(&self.scratch.mix_out);
80
81        self.dispatch(
82            enc,
83            &self.pipes.norm_prep_b,
84            |e| {
85                self.bind(e, 0, hyper, hyper_off);
86                self.bind(e, 1, &self.dense, hc.norm.0);
87                self.bind(e, 2, &bufs.normed, 0);
88                self.bind(e, 3, &bufs.h1.xe, 0);
89                self.bind(e, 4, &bufs.h1.xo, 0);
90                self.bind(e, 5, &bufs.h1.xsum, 0);
91                self.bind(e, 6, out, 0);
92                self.bind(e, 7, &bufs.inj, 0);
93                set_bytes(e, 8, &np);
94            },
95            nb * c.hc_count,
96            256,
97            true,
98        );
99
100        // Bottleneck down-projection with silu(./hc).
101        let div = c.hc_count as f32;
102        let qp = QmvParams {
103            out_dim: hc.down.out,
104            in_dim: hc.down.inp,
105        };
106
107        self.dispatch(
108            enc,
109            &self.pipes.qmv_silu_b[nb - 1],
110            |e| {
111                self.bind(e, 0, &self.dense, hc.down.w);
112                self.bind(e, 1, &self.dense, hc.down.s);
113                self.bind(e, 2, &self.dense, hc.down.b);
114                self.bind(e, 3, &bufs.h1.xe, 0);
115                self.bind(e, 4, &bufs.h1.xo, 0);
116                self.bind(e, 5, &bufs.h1.xsum, 0);
117                self.bind(e, 6, &bufs.d, 0);
118                set_bytes(e, 7, &qp);
119                set_bytes(e, 8, &div);
120            },
121            (hc.down.out as usize).div_ceil(4),
122            128,
123            true,
124        );
125
126        // Injection logits from the normed stream.
127        if inject && let Some(inj) = &hc.inject {
128            self.qmv_h(enc, inj, &bufs.inj, nb, &bufs.h1);
129        }
130
131        // Up-projection back to the stream width, then the mix.
132        self.prep_h(enc, &bufs.d, 0, hc.down.out, nb, &bufs.h2);
133        self.qmv_h(enc, &hc.up, &bufs.u, nb, &bufs.h2);
134
135        let gp = self.group_params(false, 0.0);
136
137        self.dispatch(
138            enc,
139            &self.pipes.hc_mix_b,
140            |e| {
141                self.bind(e, 0, &bufs.u, 0);
142                self.bind(e, 1, &bufs.normed, 0);
143                self.bind(e, 2, &bufs.mixed, 0);
144                set_bytes(e, 3, &gp);
145                set_bytes(e, 4, &nbu);
146            },
147            nb * h as usize,
148            256,
149            false,
150        );
151        self.prep_h(enc, &bufs.mixed, 0, h, nb, &bufs.h1);
152    }
153
154    pub(super) fn inject_b(&self, enc: &Enc, hyper: &Buf, out: &Buf, nb: usize) {
155        let c = &self.p.cfg;
156        let gp = self.group_params(false, 0.0);
157        let nbu = nb as u32;
158
159        self.dispatch(
160            enc,
161            &self.pipes.inject_b,
162            |e| {
163                self.bind(e, 0, hyper, 0);
164                self.bind(e, 1, out, 0);
165                self.bind(e, 2, &self.scratch.hc.inj, 0);
166                set_bytes(e, 3, &gp);
167                set_bytes(e, 4, &nbu);
168            },
169            nb * c.hc_hidden(),
170            256,
171            false,
172        );
173    }
174}
cherenkov/qwen4_exp/gpu/hyperconnection.rs

cherenkov/qwen4_exp/gpu/
hyperconnection.rs
