metrale_model_arch/glm5next_kda/
prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The KDA chunked prefill (`prefill`, `prefill_with_pad_fill`).
4//!
5//! Owner: model-arch (GLM-5.3-Flash KDA).
6//! Invariants:
7//! - A call covers `1..=ws.max_tokens()` tokens; any other count is refused before a launch.
8
9use super::*;
10
11impl Glm5NextKdaLayer {
12    /// 2026-09-25: Chunked prefill over `t` tokens, starting from the carried state.
13    pub fn prefill(
14        &self,
15        gpu: &dyn GpuBackend,
16        hidden: DevicePtr,
17        t: usize,
18        state: &KdaSeqState,
19        ws: &Glm5NextKdaWorkspace,
20        stream: u64,
21    ) -> Result<()> {
22        self.prefill_with_pad_fill(gpu, hidden, t, state, ws, 0.0, stream)
23    }
24
25    /// 2026-09-25: [`Self::prefill`] with the padded q/k/v tails filled with `pad_fill`.
26    ///
27    /// [`Self::prefill`] passes zero. The layer harness passes other values to show that
28    /// `kda_chunk_*` ignores positions at or past `t` in-kernel, which a zeroed tail cannot show.
29    #[allow(clippy::too_many_arguments)]
30    pub fn prefill_with_pad_fill(
31        &self,
32        gpu: &dyn GpuBackend,
33        hidden: DevicePtr,
34        t: usize,
35        state: &KdaSeqState,
36        ws: &Glm5NextKdaWorkspace,
37        pad_fill: f32,
38        stream: u64,
39    ) -> Result<()> {
40        let c = &self.cfg;
41        let (qkv, cd, hd) = (c.qkv_dim(), c.conv_dim(), c.head_dim);
42        if t == 0 || t > ws.max_tokens {
43            bail!(
44                "prefill of {t} tokens does not fit a workspace built for {}",
45                ws.max_tokens
46            );
47        }
48        let nchunks = t.div_ceil(c.chunk);
49        let tp = nchunks * c.chunk;
50        // 2026-09-25: The same three profile buckets as `decode_k`.
51        let t_front = crate::glm5next_layer::profile::start();
52        self.front_end(gpu, hidden, t, ws, stream)?;
53        crate::glm5next_layer::profile::end(
54            crate::glm5next_layer::profile::KDA_FRONT,
55            t_front,
56            gpu,
57            stream,
58        );
59        let t_recur = crate::glm5next_layer::profile::start();
60
61        // 2026-09-25: The prefill conv is conv + SiLU only; L2 is a separate launch over q|k.
62        KernelLaunch::new(gpu, self.kernels.conv_prefill)
63            .grid([div_ceil(cd as u32, 256), 1, 1])
64            .block([256, 1, 1])
65            .arg_ptr(state.conv)
66            .arg_ptr(ws.qkv_proj)
67            .arg_ptr(self.weights.conv.weight)
68            .arg_ptr(DevicePtr::NULL)
69            .arg_ptr(ws.conv_out)
70            .arg_u32(cd as u32)
71            .arg_u32(c.conv_kernel as u32)
72            .arg_u32(t as u32)
73            .arg_u32(cd as u32)
74            .arg_u32(cd as u32)
75            .launch(stream)?;
76        KernelLaunch::new(gpu, self.kernels.l2)
77            .grid([(c.qk_channels() / hd) as u32, t as u32, 1])
78            .block([hd as u32, 1, 1])
79            .arg_ptr(ws.conv_out)
80            .arg_u32(hd as u32)
81            .arg_f32(c.l2_eps)
82            .arg_u32(cd as u32)
83            .launch(stream)?;
84
85        // 2026-09-25: Zero the gate and beta pad tails, so only the q/k/v tails carry `pad_fill`.
86        for (buf, real, padded) in [
87            (ws.gate, t * qkv, tp * qkv),
88            (ws.beta, t * c.heads, tp * c.heads),
89        ] {
90            if padded == real {
91                continue;
92            }
93            KernelLaunch::new(gpu, self.kernels.fill)
94                .grid([div_ceil((padded - real) as u32, 256), 1, 1])
95                .block([256, 1, 1])
96                .arg_ptr(buf.offset(real * 4))
97                .arg_u32((padded - real) as u32)
98                .arg_f32(0.0)
99                .launch(stream)?;
100        }
101        for p in [ws.q_f32, ws.k_f32, ws.v_f32] {
102            KernelLaunch::new(gpu, self.kernels.fill)
103                .grid([div_ceil((tp * qkv) as u32, 256), 1, 1])
104                .block([256, 1, 1])
105                .arg_ptr(p)
106                .arg_u32((tp * qkv) as u32)
107                .arg_f32(pad_fill)
108                .launch(stream)?;
109        }
110        KernelLaunch::new(gpu, self.kernels.split_widen)
111            .grid([div_ceil(qkv as u32, 256), t as u32, 1])
112            .block([256, 1, 1])
113            .arg_ptr(ws.conv_out)
114            .arg_ptr(ws.q_f32)
115            .arg_ptr(ws.k_f32)
116            .arg_ptr(ws.v_f32)
117            .arg_u32(t as u32)
118            .arg_u32(qkv as u32)
119            .launch(stream)?;
120
121        KernelLaunch::new(gpu, self.kernels.chunk_prepare)
122            .grid([nchunks as u32, c.heads as u32, 1])
123            .block([BLOCK, 1, 1])
124            .shared_mem(c.smem_prepare() as u32)
125            .arg_ptr(ws.k_f32)
126            .arg_ptr(ws.v_f32)
127            .arg_ptr(ws.gate)
128            .arg_ptr(ws.beta)
129            .arg_ptr(ws.chunk_gc)
130            .arg_ptr(ws.chunk_u)
131            .arg_ptr(ws.chunk_w)
132            .arg_u32(c.heads as u32)
133            .arg_u32(hd as u32)
134            .arg_u32(c.chunk as u32)
135            .arg_u32(t as u32)
136            .launch(stream)?;
137        KernelLaunch::new(gpu, self.kernels.chunk_scan)
138            .grid([c.heads as u32, 1, 1])
139            .block([BLOCK, 1, 1])
140            .shared_mem(c.smem_scan() as u32)
141            .arg_ptr(ws.q_f32)
142            .arg_ptr(ws.k_f32)
143            .arg_ptr(ws.chunk_gc)
144            .arg_ptr(ws.chunk_u)
145            .arg_ptr(ws.chunk_w)
146            .arg_ptr(state.recurrent)
147            .arg_ptr(ws.core)
148            .arg_u32(c.heads as u32)
149            .arg_u32(hd as u32)
150            .arg_u32(c.chunk as u32)
151            .arg_u32(nchunks as u32)
152            .arg_u32(t as u32)
153            .arg_f32(1.0 / (hd as f32).sqrt())
154            .launch(stream)?;
155        crate::glm5next_layer::profile::end(
156            crate::glm5next_layer::profile::KDA_RECUR,
157            t_recur,
158            gpu,
159            stream,
160        );
161
162        let t_back = crate::glm5next_layer::profile::start();
163        let r = self.back_end(gpu, t, ws, stream);
164        crate::glm5next_layer::profile::end(
165            crate::glm5next_layer::profile::KDA_BACK,
166            t_back,
167            gpu,
168            stream,
169        );
170        r
171    }
172}
metrale_model_arch/glm5next_kda/prefill.rs

metrale_model_arch/glm5next_kda/
prefill.rs
