metrale_model_arch/glm5next_kda/decode.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The KDA recurrent path: single-token `decode` and K-row `decode_k`.
4//!
5//! Owner: model-arch (GLM-5.3-Flash KDA).
6//! Invariants:
7//! - The recurrent state advances one row at a time, in row order (`stateful_row`).
8
9use super::*;
10
11impl Glm5NextKdaLayer {
12 /// 2026-09-25: The stateful half of one KDA token: the conv window update (SiLU and L2 fused),
13 /// then one recurrent step, both on row `row` of the workspace, updating `state` in place.
14 ///
15 /// The state after row `t + 1` depends on the state after row `t`, so [`Self::decode_k`] calls
16 /// this once per row, in order, and batches only the projections around it. The chunked
17 /// [`Self::prefill`] computes the same recurrence in a different order, so its results are not
18 /// guaranteed to match this path bit for bit.
19 ///
20 /// `q`/`k` reach `kda_recurrent` already L2-normalised by the conv, which is the input that
21 /// kernel expects; normalising them again would change the bf16-rounded values.
22 fn stateful_row(
23 &self,
24 gpu: &dyn GpuBackend,
25 row: usize,
26 state: &KdaSeqState,
27 ws: &Glm5NextKdaWorkspace,
28 stream: u64,
29 ) -> Result<()> {
30 let c = &self.cfg;
31 let qkv = c.qkv_dim();
32 let cd = c.conv_dim();
33
34 ops::conv1d_update_l2norm(
35 gpu,
36 self.kernels.conv_decode,
37 state.conv,
38 ws.qkv_proj.offset(row * cd * 2),
39 &self.weights.conv,
40 ws.conv_out.offset(row * cd * 2),
41 c.conv_dim() as u32,
42 c.conv_kernel as u32,
43 1,
44 c.qk_channels() as u32,
45 c.head_dim as u32,
46 c.l2_eps,
47 stream,
48 )?;
49
50 let d = c.head_dim;
51 // 2026-09-25: 1R+1W when the target has the shared-memory kernel. Each block owns `vpb`
52 // V columns and grid.y covers the rest; the request is `3 * d` floats plus
53 // `vpb * (d + 1)` floats of column scratch.
54 let vpb = KDA_V_PER_BLOCK.min(d);
55 let smem_smem = (3 * d + vpb * (d + 1)) * 4;
56 if self.kernels.recurrent_smem.0 != 0
57 && d.is_multiple_of(vpb)
58 && smem_smem <= KDA_SMEM_BUDGET
59 && !kda_no_smem()
60 {
61 KernelLaunch::new(gpu, self.kernels.recurrent_smem)
62 .grid([c.heads as u32, (d / vpb) as u32, 1])
63 .block([vpb as u32, 1, 1])
64 .shared_mem(smem_smem as u32)
65 .arg_ptr(ws.conv_out.offset(row * cd * 2))
66 .arg_ptr(ws.conv_out.offset(row * cd * 2 + qkv * 2))
67 .arg_ptr(ws.conv_out.offset(row * cd * 2 + qkv * 4))
68 .arg_ptr(ws.gate.offset(row * qkv * 4))
69 .arg_ptr(ws.beta.offset(row * c.heads * 4))
70 .arg_ptr(state.recurrent)
71 .arg_ptr(ws.core.offset(row * qkv * 4))
72 .arg_u32(c.heads as u32)
73 .arg_u32(d as u32)
74 .arg_f32(1.0 / (d as f32).sqrt())
75 .arg_u32(vpb as u32)
76 .launch(stream)?;
77 } else {
78 KernelLaunch::new(gpu, self.kernels.recurrent)
79 .grid([c.heads as u32, 1, 1])
80 .block([BLOCK.min(d as u32), 1, 1])
81 .shared_mem((3 * d * 4) as u32)
82 .arg_ptr(ws.conv_out.offset(row * cd * 2))
83 .arg_ptr(ws.conv_out.offset(row * cd * 2 + qkv * 2))
84 .arg_ptr(ws.conv_out.offset(row * cd * 2 + qkv * 4))
85 .arg_ptr(ws.gate.offset(row * qkv * 4))
86 .arg_ptr(ws.beta.offset(row * c.heads * 4))
87 .arg_ptr(state.recurrent)
88 .arg_ptr(ws.core.offset(row * qkv * 4))
89 .arg_u32(c.heads as u32)
90 .arg_u32(d as u32)
91 .arg_f32(1.0 / (d as f32).sqrt())
92 .launch(stream)?;
93 }
94
95 Ok(())
96 }
97
98 /// 2026-09-25: Single-token decode, carrying both states. The result lands in `ws.final_out`;
99 /// `state` is updated in place.
100 pub fn decode(
101 &self,
102 gpu: &dyn GpuBackend,
103 hidden: DevicePtr,
104 state: &KdaSeqState,
105 ws: &Glm5NextKdaWorkspace,
106 stream: u64,
107 ) -> Result<()> {
108 self.front_end(gpu, hidden, 1, ws, stream)?;
109 self.stateful_row(gpu, 0, state, ws, stream)?;
110 self.back_end(gpu, 1, ws, stream)
111 }
112
113 /// 2026-09-25: `k` tokens of one sequence: the projections batched, the recurrence one row at a
114 /// time. Used for the speculative verify, and for prefill sub-chunks that do not take the
115 /// chunked arm (`glm5next_layer/steps/forward.rs`).
116 ///
117 /// `front_end` and `back_end` run once over all `k` rows. For `2 <= k <=
118 /// DENSE_GEMV_BATCHM_MAX_M` on a target with `dense_gemv_bf16_batchm`, the output matches `k`
119 /// serial [`Self::decode`] calls bit for bit: the batched GEMV gives each row the bits of the
120 /// M = 1 GEMV, the pack, gate, sigmoid and `o_norm` kernels treat rows independently, and
121 /// `stateful_row` walks the state one row at a time as `decode` does.
122 ///
123 /// `snapshots[t]`, `(h_dst, conv_dst)`, receives the state after row `t`; rows past
124 /// `snapshots.len()` take no snapshot. Refuses `k == 0` or `k` above the workspace.
125 pub fn decode_k(
126 &self,
127 gpu: &dyn GpuBackend,
128 hidden: DevicePtr,
129 k: usize,
130 state: &KdaSeqState,
131 ws: &Glm5NextKdaWorkspace,
132 snapshots: &[(DevicePtr, DevicePtr)],
133 stream: u64,
134 ) -> Result<()> {
135 if k == 0 || k > ws.max_tokens {
136 bail!(
137 "KDA decode_k of {k} tokens does not fit a workspace built for {}",
138 ws.max_tokens
139 );
140 }
141 use crate::glm5next_layer::profile;
142 let c = &self.cfg;
143 let (h_bytes, conv_bytes) = (c.recurrent_state_elems() * 4, c.conv_state_elems() * 4);
144 // 2026-09-25: Three profile buckets (front, recurrence, back), each closed where its span
145 // ends.
146 let t_front = profile::start();
147 self.front_end(gpu, hidden, k, ws, stream)?;
148 profile::end(profile::KDA_FRONT, t_front, gpu, stream);
149 let t_recur = profile::start();
150 for row in 0..k {
151 self.stateful_row(gpu, row, state, ws, stream)?;
152 if let Some((h_dst, conv_dst)) = snapshots.get(row) {
153 gpu.copy_d2d_async(state.recurrent, *h_dst, h_bytes, stream)?;
154 gpu.copy_d2d_async(state.conv, *conv_dst, conv_bytes, stream)?;
155 }
156 }
157 profile::end(profile::KDA_RECUR, t_recur, gpu, stream);
158 let t_back = profile::start();
159 let r = self.back_end(gpu, k, ws, stream);
160 profile::end(profile::KDA_BACK, t_back, gpu, stream);
161 r
162 }
163}