spark_model/kimi_k3/
mla_cuda.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Host launch for K3 CUDA gated-NoPE MLA decode (`mla_decode` PTX module).
4//!
5//! Two kernels, one token: maybe_rope (NoPE skips rotate), then SDPA + gate.
6//! CPU oracle: [`avarok_core::kimi_k3::mla_decode_token`]. BoundLayer serve
7//! FullAttention default is this launch (`K3_CUDA_MLA=0` keeps CPU).
8
9use anyhow::{Context, Result, bail, ensure};
10use avarok_core::kimi_k3::{MlaConfig, MlaKv};
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/mla_decode.cu`.
15pub const MODULE: &str = "mla_decode";
16pub const ROPE_ENTRY: &str = "k3_mla_maybe_rope_f32";
17pub const SDPA_ENTRY: &str = "k3_mla_sdpa_gate_f32";
18const ROPE_BLOCK: u32 = 128;
19const SDPA_BLOCK: u32 = 32;
20
21#[derive(Clone, Copy, Debug)]
22pub struct K3MlaDecodeKernels {
23    pub rope: KernelHandle,
24    pub sdpa: KernelHandle,
25}
26
27impl K3MlaDecodeKernels {
28    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
29        Ok(Self {
30            rope: gpu.kernel(MODULE, ROPE_ENTRY)?,
31            sdpa: gpu.kernel(MODULE, SDPA_ENTRY)?,
32        })
33    }
34}
35
36fn f32_bytes(v: &[f32]) -> Vec<u8> {
37    v.iter().flat_map(|x| x.to_le_bytes()).collect()
38}
39
40fn bytes_f32(b: &[u8]) -> Vec<f32> {
41    b.chunks_exact(4)
42        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
43        .collect()
44}
45
46fn up(gpu: &dyn GpuBackend, v: &[f32], hold: &mut Vec<DevicePtr>) -> Result<DevicePtr> {
47    let b = f32_bytes(v);
48    let p = gpu.alloc(b.len().max(1))?;
49    hold.push(p);
50    gpu.copy_h2d(&b, p)?;
51    Ok(p)
52}
53
54/// Device-resident MLA KV. Append is one-row D2D/H2D; SDPA reads the buffer.
55/// Host `MlaKv` remains the CPU oracle. Do not re-upload `[0..T]` each token.
56pub struct MlaDeviceKv {
57    pub k: DevicePtr,
58    pub v: DevicePtr,
59    pub seq_len: usize,
60    pub cap: usize,
61    k_row: usize,
62    v_row: usize,
63}
64
65impl MlaDeviceKv {
66    pub fn k_row(&self) -> usize {
67        self.k_row
68    }
69
70    pub fn v_row(&self) -> usize {
71        self.v_row
72    }
73
74    pub fn alloc(gpu: &dyn GpuBackend, cap: usize, k_row: usize, v_row: usize) -> Result<Self> {
75        ensure!(k_row > 0 && v_row > 0, "k3 mla: resident row rank");
76        let cap = cap.max(1);
77        let k = gpu.alloc((cap * k_row * 4).max(1))?;
78        let v = match gpu.alloc((cap * v_row * 4).max(1)) {
79            Ok(ptr) => ptr,
80            Err(error) => {
81                let _ = gpu.free(k);
82                return Err(error);
83            }
84        };
85        Ok(Self {
86            k,
87            v,
88            seq_len: 0,
89            cap,
90            k_row,
91            v_row,
92        })
93    }
94
95    pub fn alloc_and_upload(
96        gpu: &dyn GpuBackend,
97        host: &MlaKv,
98        cap: usize,
99        k_row: usize,
100        v_row: usize,
101    ) -> Result<Self> {
102        ensure!(
103            host.k.len() == host.seq_len.saturating_mul(k_row)
104                && host.v.len() == host.seq_len.saturating_mul(v_row),
105            "k3 mla: host KV rank"
106        );
107        ensure!(
108            cap >= host.seq_len.max(1),
109            "k3 mla: device KV cap {cap} < seq {}",
110            host.seq_len
111        );
112        let mut kv = Self::alloc(gpu, cap, k_row, v_row)?;
113        if host.seq_len == 0 {
114            return Ok(kv);
115        }
116        let upload = gpu
117            .copy_h2d(&f32_bytes(&host.k), kv.k)
118            .and_then(|()| gpu.copy_h2d(&f32_bytes(&host.v), kv.v));
119        if let Err(error) = upload {
120            let _ = kv.free(gpu);
121            return Err(error);
122        }
123        kv.seq_len = host.seq_len;
124        Ok(kv)
125    }
126
127    pub fn validate_cfg(&self, cfg: &MlaConfig) -> Result<()> {
128        let k_row = cfg.heads * cfg.qk_head_dim();
129        let v_row = cfg.heads * cfg.v_head_dim;
130        ensure!(
131            self.k_row == k_row && self.v_row == v_row,
132            "k3 mla: resident KV rank"
133        );
134        Ok(())
135    }
136
137    pub fn download(&self, gpu: &dyn GpuBackend, host: &mut MlaKv) -> Result<()> {
138        let n = self.seq_len;
139        let mut kb = vec![0u8; n.saturating_mul(self.k_row) * 4];
140        let mut vb = vec![0u8; n.saturating_mul(self.v_row) * 4];
141        if n > 0 {
142            gpu.copy_d2h(self.k, &mut kb)?;
143            gpu.copy_d2h(self.v, &mut vb)?;
144        }
145        host.k = bytes_f32(&kb);
146        host.v = bytes_f32(&vb);
147        host.seq_len = n;
148        Ok(())
149    }
150
151    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
152        let k = gpu.free(self.k);
153        let v = gpu.free(self.v);
154        k.and(v)
155    }
156
157    fn append_from_device(
158        &mut self,
159        gpu: &dyn GpuBackend,
160        k_new: DevicePtr,
161        v_host: &[f32],
162    ) -> Result<()> {
163        if self.seq_len >= self.cap {
164            bail!("k3 mla: device KV cap {} full", self.cap);
165        }
166        ensure!(
167            v_host.len() == self.v_row,
168            "k3 mla: V row rank {} != {}",
169            v_host.len(),
170            self.v_row
171        );
172        let k_off = DevicePtr(self.k.0 + (self.seq_len * self.k_row * 4) as u64);
173        gpu.copy_d2d(k_new, k_off, self.k_row * 4)?;
174        let v_off = DevicePtr(self.v.0 + (self.seq_len * self.v_row * 4) as u64);
175        gpu.copy_h2d(&f32_bytes(v_host), v_off)?;
176        self.seq_len += 1;
177        Ok(())
178    }
179}
180
181/// One-token CUDA gated-NoPE MLA. Updates `q`/`k` (rope) and `kv` (append).
182#[allow(clippy::too_many_arguments)]
183pub fn launch_k3_mla_decode_token(
184    gpu: &dyn GpuBackend,
185    kernels: &K3MlaDecodeKernels,
186    q: &mut [f32],
187    k: &mut [f32],
188    v: &[f32],
189    g: &[f32],
190    kv: &mut MlaKv,
191    cfg: &MlaConfig,
192    pos: usize,
193    theta: f32,
194    stream: u64,
195) -> Result<Vec<f32>> {
196    let (h, dq, dv) = (cfg.heads, cfg.qk_head_dim(), cfg.v_head_dim);
197    if q.len() != h * dq || k.len() != h * dq {
198        bail!("k3 mla: q/k rank");
199    }
200    if v.len() != h * dv || g.len() != h * dv {
201        bail!("k3 mla: v/g rank");
202    }
203    if dq == 0 {
204        bail!("k3 mla: dq must be > 0");
205    }
206
207    let mut hold = Vec::new();
208    let run = (|| {
209        let dq_ptr = up(gpu, q, &mut hold)?;
210        let dk_new = up(gpu, k, &mut hold)?;
211        KernelLaunch::new(gpu, kernels.rope)
212            .grid([div_ceil(h as u32, ROPE_BLOCK), 1, 1])
213            .block([ROPE_BLOCK, 1, 1])
214            .arg_ptr(dq_ptr)
215            .arg_ptr(dk_new)
216            .arg_u32(h as u32)
217            .arg_u32(cfg.qk_nope_head_dim as u32)
218            .arg_u32(cfg.qk_rope_head_dim as u32)
219            .arg_u32(pos as u32)
220            .arg_f32(theta)
221            .arg_u32(u32::from(cfg.mla_use_nope))
222            .launch(stream)
223            .context("k3_mla_maybe_rope_f32")?;
224        gpu.synchronize(stream)?;
225
226        let mut qb = vec![0u8; q.len() * 4];
227        let mut kb = vec![0u8; k.len() * 4];
228        gpu.copy_d2h(dq_ptr, &mut qb)?;
229        gpu.copy_d2h(dk_new, &mut kb)?;
230        q.copy_from_slice(&bytes_f32(&qb));
231        k.copy_from_slice(&bytes_f32(&kb));
232        kv.append(k, v);
233
234        let dq2 = up(gpu, q, &mut hold)?;
235        let dk = up(gpu, &kv.k, &mut hold)?;
236        let dv_ptr = up(gpu, &kv.v, &mut hold)?;
237        let dg = up(gpu, g, &mut hold)?;
238        let dout = gpu.alloc((h * dv * 4).max(1))?;
239        hold.push(dout);
240        KernelLaunch::new(gpu, kernels.sdpa)
241            .grid([h as u32, 1, 1])
242            .block([SDPA_BLOCK, 1, 1])
243            .arg_ptr(dq2)
244            .arg_ptr(dk)
245            .arg_ptr(dv_ptr)
246            .arg_ptr(dg)
247            .arg_ptr(dout)
248            .arg_u32(kv.seq_len as u32)
249            .arg_u32(h as u32)
250            .arg_u32(dq as u32)
251            .arg_u32(dv as u32)
252            .arg_u32(u32::from(cfg.mla_use_output_gate))
253            .launch(stream)
254            .context("k3_mla_sdpa_gate_f32")?;
255        gpu.synchronize(stream)?;
256
257        let mut out_b = vec![0u8; h * dv * 4];
258        gpu.copy_d2h(dout, &mut out_b)?;
259        Ok(bytes_f32(&out_b))
260    })();
261    for p in hold {
262        let _ = gpu.free(p);
263    }
264    run
265}
266
267/// Device-resident KV: rope stays on device, append is one row, one D2H (output).
268#[allow(clippy::too_many_arguments)]
269pub fn launch_k3_mla_decode_token_on_device(
270    gpu: &dyn GpuBackend,
271    kernels: &K3MlaDecodeKernels,
272    q: &[f32],
273    k: &[f32],
274    v: &[f32],
275    g: &[f32],
276    kv: &mut MlaDeviceKv,
277    cfg: &MlaConfig,
278    pos: usize,
279    theta: f32,
280    stream: u64,
281) -> Result<Vec<f32>> {
282    let (h, dq, dv) = (cfg.heads, cfg.qk_head_dim(), cfg.v_head_dim);
283    if q.len() != h * dq || k.len() != h * dq || v.len() != h * dv || g.len() != h * dv {
284        bail!("k3 mla: q/k/v/g rank");
285    }
286    let mut hold = Vec::new();
287    let run = (|| {
288        let dq_ptr = up(gpu, q, &mut hold)?;
289        let dk_new = up(gpu, k, &mut hold)?;
290        KernelLaunch::new(gpu, kernels.rope)
291            .grid([div_ceil(h as u32, ROPE_BLOCK), 1, 1])
292            .block([ROPE_BLOCK, 1, 1])
293            .arg_ptr(dq_ptr)
294            .arg_ptr(dk_new)
295            .arg_u32(h as u32)
296            .arg_u32(cfg.qk_nope_head_dim as u32)
297            .arg_u32(cfg.qk_rope_head_dim as u32)
298            .arg_u32(pos as u32)
299            .arg_f32(theta)
300            .arg_u32(u32::from(cfg.mla_use_nope))
301            .launch(stream)
302            .context("k3_mla_maybe_rope_f32")?;
303        kv.append_from_device(gpu, dk_new, v)?;
304        let dg = up(gpu, g, &mut hold)?;
305        let dout = gpu.alloc((h * dv * 4).max(1))?;
306        hold.push(dout);
307        KernelLaunch::new(gpu, kernels.sdpa)
308            .grid([h as u32, 1, 1])
309            .block([SDPA_BLOCK, 1, 1])
310            .arg_ptr(dq_ptr)
311            .arg_ptr(kv.k)
312            .arg_ptr(kv.v)
313            .arg_ptr(dg)
314            .arg_ptr(dout)
315            .arg_u32(kv.seq_len as u32)
316            .arg_u32(h as u32)
317            .arg_u32(dq as u32)
318            .arg_u32(dv as u32)
319            .arg_u32(u32::from(cfg.mla_use_output_gate))
320            .launch(stream)
321            .context("k3_mla_sdpa_gate_f32")?;
322        gpu.synchronize(stream)?;
323        let mut out_b = vec![0u8; h * dv * 4];
324        gpu.copy_d2h(dout, &mut out_b)?;
325        Ok(bytes_f32(&out_b))
326    })();
327    for p in hold {
328        let _ = gpu.free(p);
329    }
330    run
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336    use spark_runtime::gpu::mock::{MockArg, MockGpuBackend};
337
338    #[test]
339    fn resolve_looks_up_k3_entries() {
340        let gpu = MockGpuBackend::new();
341        let _ = K3MlaDecodeKernels::resolve(&gpu).unwrap();
342        assert_eq!(
343            gpu.kernel_lookups_snapshot(),
344            vec![
345                (MODULE.to_string(), ROPE_ENTRY.to_string()),
346                (MODULE.to_string(), SDPA_ENTRY.to_string()),
347            ]
348        );
349    }
350
351    #[test]
352    fn mock_launch_contract_twin_geometry() {
353        let gpu = MockGpuBackend::new();
354        let kernels = K3MlaDecodeKernels::resolve(&gpu).unwrap();
355        let cfg = MlaConfig::twin_0_40b();
356        let mut q = vec![0.1f32; cfg.heads * cfg.qk_head_dim()];
357        let mut k = vec![0.2f32; cfg.heads * cfg.qk_head_dim()];
358        let v = vec![0.3f32; cfg.heads * cfg.v_head_dim];
359        let g = vec![0.0f32; cfg.heads * cfg.v_head_dim];
360        let mut kv = MlaKv::default();
361        let _ = launch_k3_mla_decode_token(
362            &gpu, &kernels, &mut q, &mut k, &v, &g, &mut kv, &cfg, 3, 10000.0, 3,
363        )
364        .unwrap();
365        assert_eq!(kv.seq_len, 1);
366        let launches = gpu.launches_snapshot();
367        assert_eq!(launches.len(), 2, "rope then sdpa_gate");
368        assert_eq!(
369            launches[0].grid,
370            [div_ceil(cfg.heads as u32, ROPE_BLOCK), 1, 1]
371        );
372        assert_eq!(launches[0].block, [ROPE_BLOCK, 1, 1]);
373        assert_eq!(launches[0].args.len(), 8);
374        assert_eq!(
375            launches[0].args[7],
376            MockArg::Bytes(1u32.to_le_bytes().to_vec()),
377            "twin is NoPE"
378        );
379        let sdpa = &launches[1];
380        assert_eq!(sdpa.grid, [cfg.heads as u32, 1, 1]);
381        assert_eq!(sdpa.block, [SDPA_BLOCK, 1, 1]);
382        assert_eq!(sdpa.shared_mem, 0);
383        assert_eq!(sdpa.stream, 3);
384        assert_eq!(sdpa.args.len(), 10);
385        assert_eq!(
386            sdpa.args[9],
387            MockArg::Bytes(1u32.to_le_bytes().to_vec()),
388            "twin uses output gate"
389        );
390    }
391
392    #[test]
393    fn device_kv_second_token_d2h_is_output_only() {
394        let gpu = MockGpuBackend::new();
395        let kernels = K3MlaDecodeKernels::resolve(&gpu).unwrap();
396        let cfg = MlaConfig::twin_0_40b();
397        let q = vec![0.1f32; cfg.heads * cfg.qk_head_dim()];
398        let k = vec![0.2f32; cfg.heads * cfg.qk_head_dim()];
399        let v = vec![0.3f32; cfg.heads * cfg.v_head_dim];
400        let g = vec![0.0f32; cfg.heads * cfg.v_head_dim];
401        let mut kv = MlaDeviceKv::alloc(
402            &gpu,
403            8,
404            cfg.heads * cfg.qk_head_dim(),
405            cfg.heads * cfg.v_head_dim,
406        )
407        .unwrap();
408        let _ = launch_k3_mla_decode_token_on_device(
409            &gpu, &kernels, &q, &k, &v, &g, &mut kv, &cfg, 0, 10000.0, 0,
410        )
411        .unwrap();
412        let before = gpu.d2h_blocking_count();
413        let _ = launch_k3_mla_decode_token_on_device(
414            &gpu, &kernels, &q, &k, &v, &g, &mut kv, &cfg, 1, 10000.0, 0,
415        )
416        .unwrap();
417        let pulled = gpu.d2h_blocking_count() - before;
418        assert_eq!(
419            pulled, 1,
420            "resident MLA must not re-download KV; D2H is output only (got {pulled})"
421        );
422        assert_eq!(kv.seq_len, 2);
423    }
424
425    #[test]
426    fn host_kv_wrapper_h2d_grows_with_history_known_bad() {
427        // Oracle: the old serving wrapper re-uploads [0..T] each token.
428        // This known-bad case must keep failing so the resident path is not
429        // confused with it.
430        let gpu = MockGpuBackend::new();
431        let kernels = K3MlaDecodeKernels::resolve(&gpu).unwrap();
432        let cfg = MlaConfig::twin_0_40b();
433        let mut q = vec![0.1f32; cfg.heads * cfg.qk_head_dim()];
434        let mut k = vec![0.2f32; cfg.heads * cfg.qk_head_dim()];
435        let v = vec![0.3f32; cfg.heads * cfg.v_head_dim];
436        let g = vec![0.0f32; cfg.heads * cfg.v_head_dim];
437        let mut kv = MlaKv::default();
438        let _ = launch_k3_mla_decode_token(
439            &gpu, &kernels, &mut q, &mut k, &v, &g, &mut kv, &cfg, 0, 10000.0, 0,
440        )
441        .unwrap();
442        let after_t0 = gpu.h2d_bytes();
443        let _ = launch_k3_mla_decode_token(
444            &gpu, &kernels, &mut q, &mut k, &v, &g, &mut kv, &cfg, 1, 10000.0, 0,
445        )
446        .unwrap();
447        let t1 = gpu.h2d_bytes() - after_t0;
448        assert!(
449            t1 > after_t0,
450            "host-KV wrapper must re-upload history (t0={after_t0} t1={t1})"
451        );
452    }
453
454    #[test]
455    fn resident_kv_h2d_does_not_grow_with_history() {
456        let gpu = MockGpuBackend::new();
457        let kernels = K3MlaDecodeKernels::resolve(&gpu).unwrap();
458        let cfg = MlaConfig::twin_0_40b();
459        let q = vec![0.1f32; cfg.heads * cfg.qk_head_dim()];
460        let k = vec![0.2f32; cfg.heads * cfg.qk_head_dim()];
461        let v = vec![0.3f32; cfg.heads * cfg.v_head_dim];
462        let g = vec![0.0f32; cfg.heads * cfg.v_head_dim];
463        let mut kv = MlaDeviceKv::alloc(
464            &gpu,
465            8,
466            cfg.heads * cfg.qk_head_dim(),
467            cfg.heads * cfg.v_head_dim,
468        )
469        .unwrap();
470        let _ = launch_k3_mla_decode_token_on_device(
471            &gpu, &kernels, &q, &k, &v, &g, &mut kv, &cfg, 0, 10000.0, 0,
472        )
473        .unwrap();
474        let t0 = gpu.h2d_bytes();
475        let _ = launch_k3_mla_decode_token_on_device(
476            &gpu, &kernels, &q, &k, &v, &g, &mut kv, &cfg, 1, 10000.0, 0,
477        )
478        .unwrap();
479        let t1 = gpu.h2d_bytes() - t0;
480        assert_eq!(
481            t1, t0,
482            "resident MLA token-append H2D must stay O(1) in T (t0={t0} t1={t1})"
483        );
484    }
485
486    #[test]
487    fn failed_second_buffer_allocation_frees_k() {
488        let gpu = MockGpuBackend::new();
489        // K fits the cap; V is larger so the second alloc fails.
490        gpu.set_max_allocation_bytes(8 * 4);
491        assert!(MlaDeviceKv::alloc(&gpu, 1, 8, 16).is_err());
492        assert_eq!(gpu.alloc_count(), 0);
493    }
494}
spark_model/kimi_k3/mla_cuda.rs

spark_model/kimi_k3/
mla_cuda.rs
