spark_model/kimi_k3/
kda_cuda.rs

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

spark_model/kimi_k3/
kda_cuda.rs
