metrale_model_arch/kimi_k3/
kda_cuda.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host launch of the Kimi K3 KDA decode kernels (PTX module [`MODULE`]) for one token.
4//!
5//! Two launches per token: [`CONV_ENTRY`] updates the short-conv state and
6//! applies SiLU, then [`RECURRENT_ENTRY`] L2-normalises q and k and runs the
7//! gated delta-rule step. The CPU reference is
8//! [`kda_decode_token`](metrale_model_weights::kimi_k3_host::kda_decode_token).
9//! Decode uses [`launch_k3_kda_decode_token_on_device`] unless `K3_CUDA_KDA=0`.
10//!
11//! Owner: model-arch, Kimi K3.
12//! Invariants:
13//! - Every temporary device buffer a launch allocates gets a `free` call
14//!   before it returns, on success and on error; a failed `free` is ignored.
15//! - [`KdaDeviceState::alloc_and_upload`], when it fails, calls `free` on what
16//!   it allocated, ignoring a failed `free`.
17
18use anyhow::{Context, Result, bail, ensure};
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
20use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
21use metrale_model_weights::kimi_k3_host::{KDA_L2_EPS, KdaConfig, KdaState};
22
23/// 2026-09-25: Built from `kernels/gb10/kimi-k3/bf16/kda_decode.cu`.
24pub const MODULE: &str = "kda_decode";
25pub const CONV_ENTRY: &str = "k3_kda_conv_update_f32";
26pub const RECURRENT_ENTRY: &str = "k3_kda_recurrent_step_f32";
27const CONV_BLOCK: u32 = 128;
28
29#[derive(Clone, Copy, Debug)]
30pub struct K3KdaDecodeKernels {
31    pub conv: KernelHandle,
32    pub recurrent: KernelHandle,
33}
34
35impl K3KdaDecodeKernels {
36    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
37        Ok(Self {
38            conv: gpu.kernel(MODULE, CONV_ENTRY)?,
39            recurrent: gpu.kernel(MODULE, RECURRENT_ENTRY)?,
40        })
41    }
42}
43
44fn f32_bytes(v: &[f32]) -> Vec<u8> {
45    v.iter().flat_map(|x| x.to_le_bytes()).collect()
46}
47
48fn bytes_f32(b: &[u8]) -> Vec<f32> {
49    b.chunks_exact(4)
50        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
51        .collect()
52}
53
54fn up(gpu: &dyn GpuBackend, v: &[f32], hold: &mut Vec<DevicePtr>) -> Result<DevicePtr> {
55    let b = f32_bytes(v);
56    let p = gpu.alloc(b.len().max(1))?;
57    hold.push(p);
58    gpu.copy_h2d(&b, p)?;
59    Ok(p)
60}
61
62/// 2026-09-25: One sequence's KDA conv and recurrent state on the device.
63/// [`launch_k3_kda_decode_token_on_device`] updates it in place and copies
64/// only the output back; `download` copies it to a host [`KdaState`].
65///
66/// [`launch_k3_kda_decode_token`] instead uploads and downloads the state
67/// every token; at [`KdaConfig::production`] the recurrent state alone is
68/// 96 x 128 x 128 FP32, 6.3 MB per layer.
69pub struct KdaDeviceState {
70    pub conv: DevicePtr,
71    pub recurrent: DevicePtr,
72    conv_elems: usize,
73    recurrent_elems: usize,
74}
75
76impl KdaDeviceState {
77    pub fn alloc_and_upload(gpu: &dyn GpuBackend, host: &KdaState) -> Result<Self> {
78        let conv_b = f32_bytes(&host.conv);
79        let rec_b = f32_bytes(&host.recurrent);
80        let conv = gpu.alloc(conv_b.len().max(1))?;
81        let recurrent = match gpu.alloc(rec_b.len().max(1)) {
82            Ok(ptr) => ptr,
83            Err(error) => {
84                let _ = gpu.free(conv);
85                return Err(error);
86            }
87        };
88        let upload = gpu
89            .copy_h2d(&conv_b, conv)
90            .and_then(|()| gpu.copy_h2d(&rec_b, recurrent));
91        if let Err(error) = upload {
92            let _ = gpu.free(conv);
93            let _ = gpu.free(recurrent);
94            return Err(error);
95        }
96        Ok(Self {
97            conv,
98            recurrent,
99            conv_elems: host.conv.len(),
100            recurrent_elems: host.recurrent.len(),
101        })
102    }
103
104    pub fn download(&self, gpu: &dyn GpuBackend, host: &mut KdaState) -> Result<()> {
105        self.validate_geometry(host.conv.len(), host.recurrent.len())?;
106        let mut conv_b = vec![0u8; host.conv.len() * 4];
107        let mut rec_b = vec![0u8; host.recurrent.len() * 4];
108        gpu.copy_d2h(self.conv, &mut conv_b)?;
109        gpu.copy_d2h(self.recurrent, &mut rec_b)?;
110        host.conv = bytes_f32(&conv_b);
111        host.recurrent = bytes_f32(&rec_b);
112        Ok(())
113    }
114
115    fn validate_geometry(&self, conv_elems: usize, recurrent_elems: usize) -> Result<()> {
116        ensure!(
117            self.conv_elems == conv_elems && self.recurrent_elems == recurrent_elems,
118            "k3 kda: resident state rank"
119        );
120        Ok(())
121    }
122
123    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
124        let conv = gpu.free(self.conv);
125        let recurrent = gpu.free(self.recurrent);
126        conv.and(recurrent)
127    }
128}
129
130/// 2026-09-25: One decode token against a host [`KdaState`]: uploads the
131/// state, runs both kernels, and downloads the updated state.
132pub fn launch_k3_kda_decode_token(
133    gpu: &dyn GpuBackend,
134    kernels: &K3KdaDecodeKernels,
135    x_qkv: &[f32],
136    conv_w: &[f32],
137    gate: &[f32],
138    beta: &[f32],
139    cfg: &KdaConfig,
140    state: &mut KdaState,
141    stream: u64,
142) -> Result<Vec<f32>> {
143    launch_k3_kda_decode_inner(
144        gpu,
145        kernels,
146        x_qkv,
147        conv_w,
148        gate,
149        beta,
150        cfg,
151        Some(state),
152        None,
153        stream,
154    )
155}
156
157/// 2026-09-25: One decode token against a [`KdaDeviceState`]; only the output
158/// is copied back to the host.
159pub fn launch_k3_kda_decode_token_on_device(
160    gpu: &dyn GpuBackend,
161    kernels: &K3KdaDecodeKernels,
162    x_qkv: &[f32],
163    conv_w: &[f32],
164    gate: &[f32],
165    beta: &[f32],
166    cfg: &KdaConfig,
167    device: &KdaDeviceState,
168    stream: u64,
169) -> Result<Vec<f32>> {
170    launch_k3_kda_decode_inner(
171        gpu,
172        kernels,
173        x_qkv,
174        conv_w,
175        gate,
176        beta,
177        cfg,
178        None,
179        Some(device),
180        stream,
181    )
182}
183
184#[allow(clippy::too_many_arguments)]
185fn launch_k3_kda_decode_inner(
186    gpu: &dyn GpuBackend,
187    kernels: &K3KdaDecodeKernels,
188    x_qkv: &[f32],
189    conv_w: &[f32],
190    gate: &[f32],
191    beta: &[f32],
192    cfg: &KdaConfig,
193    mut host: Option<&mut KdaState>,
194    device: Option<&KdaDeviceState>,
195    stream: u64,
196) -> Result<Vec<f32>> {
197    let (h, d, k) = (cfg.heads, cfg.head_dim, cfg.conv_kernel);
198    let c = cfg.conv_dim();
199    if x_qkv.len() != c {
200        bail!("k3 kda: x_qkv {} != conv_dim {c}", x_qkv.len());
201    }
202    if conv_w.len() != cfg.conv_elems() {
203        bail!("k3 kda: conv_w {} != {}", conv_w.len(), cfg.conv_elems());
204    }
205    if gate.len() != cfg.qkv_dim() || beta.len() != h {
206        bail!("k3 kda: gate/beta rank");
207    }
208    if let Some(state) = host.as_ref()
209        && (state.conv.len() != cfg.conv_elems() || state.recurrent.len() != cfg.recurrent_elems())
210    {
211        bail!("k3 kda: state rank");
212    }
213    if k == 0 || d == 0 {
214        bail!("k3 kda: D and conv_kernel must be > 0");
215    }
216    if let Some(state) = device {
217        state.validate_geometry(cfg.conv_elems(), cfg.recurrent_elems())?;
218    }
219
220    let mut hold = Vec::new();
221    let run = (|| {
222        let dx = up(gpu, x_qkv, &mut hold)?;
223        let dw = up(gpu, conv_w, &mut hold)?;
224        let dgate = up(gpu, gate, &mut hold)?;
225        let dbeta = up(gpu, beta, &mut hold)?;
226        let (dconv, drec, pull_state) = if let Some(d) = device {
227            (d.conv, d.recurrent, false)
228        } else {
229            let s = host.as_ref().context("k3 kda: host or device state")?;
230            (
231                up(gpu, &s.conv, &mut hold)?,
232                up(gpu, &s.recurrent, &mut hold)?,
233                true,
234            )
235        };
236        let dy = gpu.alloc((c * 4).max(1))?;
237        hold.push(dy);
238        let dout = gpu.alloc((cfg.qkv_dim() * 4).max(1))?;
239        hold.push(dout);
240
241        KernelLaunch::new(gpu, kernels.conv)
242            .grid([div_ceil(c as u32, CONV_BLOCK), 1, 1])
243            .block([CONV_BLOCK, 1, 1])
244            .arg_ptr(dx)
245            .arg_ptr(dw)
246            .arg_ptr(dconv)
247            .arg_ptr(dy)
248            .arg_u32(c as u32)
249            .arg_u32(k as u32)
250            .launch(stream)
251            .context("k3_kda_conv_update_f32")?;
252
253        let rec_block = (d as u32).min(128);
254        KernelLaunch::new(gpu, kernels.recurrent)
255            .grid([h as u32, 1, 1])
256            .block([rec_block, 1, 1])
257            .shared_mem((3 * d * 4) as u32)
258            .arg_ptr(dy)
259            .arg_ptr(dgate)
260            .arg_ptr(dbeta)
261            .arg_ptr(drec)
262            .arg_ptr(dout)
263            .arg_u32(h as u32)
264            .arg_u32(d as u32)
265            .arg_f32(KDA_L2_EPS)
266            .launch(stream)
267            .context("k3_kda_recurrent_step_f32")?;
268        gpu.synchronize(stream)?;
269
270        let mut out_b = vec![0u8; cfg.qkv_dim() * 4];
271        gpu.copy_d2h(dout, &mut out_b)?;
272        if pull_state {
273            let state = host.as_mut().context("k3 kda: host state for D2H")?;
274            let mut conv_b = vec![0u8; cfg.conv_elems() * 4];
275            let mut rec_b = vec![0u8; cfg.recurrent_elems() * 4];
276            gpu.copy_d2h(dconv, &mut conv_b)?;
277            gpu.copy_d2h(drec, &mut rec_b)?;
278            state.conv = bytes_f32(&conv_b);
279            state.recurrent = bytes_f32(&rec_b);
280        }
281        Ok(bytes_f32(&out_b))
282    })();
283    for p in hold {
284        let _ = gpu.free(p);
285    }
286    run
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use metrale_gpu_runtime::gpu::mock::{MockArg, MockGpuBackend};
293    use metrale_model_weights::kimi_k3_host::kda_decode_token;
294
295    #[test]
296    fn resolve_looks_up_k3_entries() {
297        let gpu = MockGpuBackend::new();
298        let _ = K3KdaDecodeKernels::resolve(&gpu).unwrap();
299        assert_eq!(
300            gpu.kernel_lookups_snapshot(),
301            vec![
302                (MODULE.to_string(), CONV_ENTRY.to_string()),
303                (MODULE.to_string(), RECURRENT_ENTRY.to_string()),
304            ]
305        );
306    }
307
308    #[test]
309    fn mock_launch_contract_twin_geometry() {
310        let gpu = MockGpuBackend::new();
311        let k = K3KdaDecodeKernels::resolve(&gpu).unwrap();
312        let cfg = KdaConfig::twin_0_40b();
313        let mut state = KdaState::new(&cfg);
314        let x = vec![0.1f32; cfg.conv_dim()];
315        let w = vec![0.2f32; cfg.conv_elems()];
316        let gate = vec![-0.5f32; cfg.qkv_dim()];
317        let beta = vec![0.25f32; cfg.heads];
318        let _ = launch_k3_kda_decode_token(&gpu, &k, &x, &w, &gate, &beta, &cfg, &mut state, 3)
319            .unwrap();
320        let launches = gpu.launches_snapshot();
321        assert_eq!(launches.len(), 2, "conv then recurrent");
322        let c = cfg.conv_dim() as u32;
323        assert_eq!(launches[0].grid, [div_ceil(c, CONV_BLOCK), 1, 1]);
324        assert_eq!(launches[0].block, [CONV_BLOCK, 1, 1]);
325        assert_eq!(launches[0].shared_mem, 0);
326        assert_eq!(launches[0].stream, 3);
327        assert_eq!(launches[0].args.len(), 6);
328        assert_eq!(
329            launches[0].args[4],
330            MockArg::Bytes(c.to_le_bytes().to_vec())
331        );
332        assert_eq!(
333            launches[0].args[5],
334            MockArg::Bytes((cfg.conv_kernel as u32).to_le_bytes().to_vec())
335        );
336
337        let rec = &launches[1];
338        assert_eq!(rec.grid, [cfg.heads as u32, 1, 1]);
339        assert_eq!(rec.block, [cfg.head_dim as u32, 1, 1]);
340        assert_eq!(rec.shared_mem, (3 * cfg.head_dim * 4) as u32);
341        assert_eq!(rec.args.len(), 8);
342        assert_eq!(
343            rec.args[7],
344            MockArg::Bytes(KDA_L2_EPS.to_le_bytes().to_vec())
345        );
346    }
347
348    #[test]
349    fn cpu_oracle_is_the_compare_target() {
350        let cfg = KdaConfig::twin_0_40b();
351        let mut state = KdaState::new(&cfg);
352        let x = vec![0.1f32; cfg.conv_dim()];
353        let w = vec![0.2f32; cfg.conv_elems()];
354        let gate = vec![-0.5f32; cfg.qkv_dim()];
355        let beta = vec![0.25f32; cfg.heads];
356        let y = kda_decode_token(&x, &w, &gate, &beta, &cfg, &mut state);
357        assert_eq!(y.len(), cfg.qkv_dim());
358        assert!(y.iter().any(|v| v.abs() > 1e-8));
359    }
360
361    #[test]
362    fn device_resident_state_skips_conv_recurrent_d2h() {
363        let gpu = MockGpuBackend::new();
364        let k = K3KdaDecodeKernels::resolve(&gpu).unwrap();
365        let cfg = KdaConfig::twin_0_40b();
366        let host = KdaState::new(&cfg);
367        let x = vec![0.1f32; cfg.conv_dim()];
368        let w = vec![0.2f32; cfg.conv_elems()];
369        let gate = vec![-0.5f32; cfg.qkv_dim()];
370        let beta = vec![0.25f32; cfg.heads];
371        let device = KdaDeviceState::alloc_and_upload(&gpu, &host).unwrap();
372        let before = gpu.d2h_blocking_count();
373        let _ =
374            launch_k3_kda_decode_token_on_device(&gpu, &k, &x, &w, &gate, &beta, &cfg, &device, 0)
375                .unwrap();
376        let pulled = gpu.d2h_blocking_count() - before;
377        assert_eq!(
378            pulled, 1,
379            "resident decode must D2H the output only, not conv/recurrent (got {pulled})"
380        );
381        device.free(&gpu).unwrap();
382    }
383}
metrale_model_arch/kimi_k3/kda_cuda.rs

metrale_model_arch/kimi_k3/
kda_cuda.rs
