metrale_model_arch/kimi_k3/
mla_cuda.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host launch of the Kimi K3 MLA decode kernels (PTX module [`MODULE`]) for one token.
4//!
5//! Two launches per token: [`ROPE_ENTRY`] rotates the rope slice of q and k
6//! (it returns without rotating when `mla_use_nope` is set), then
7//! [`SDPA_ENTRY`] attends over the cached KV and applies the sigmoid output
8//! gate when `mla_use_output_gate` is set. The CPU reference is
9//! [`mla_decode_token`](metrale_model_weights::kimi_k3_host::mla_decode_token).
10//! Decode uses [`launch_k3_mla_decode_token_on_device`] unless `K3_CUDA_MLA=0`.
11//!
12//! Owner: model-arch, Kimi K3.
13//! Invariants:
14//! - Every temporary device buffer a launch allocates gets a `free` call
15//!   before it returns, on success and on error; a failed `free` is ignored.
16//! - [`MlaDeviceKv::alloc`] and [`MlaDeviceKv::alloc_and_upload`], when they
17//!   fail, call `free` on what they allocated, ignoring a failed `free`.
18
19use anyhow::{Context, Result, bail, ensure};
20use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
21use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
22use metrale_model_weights::kimi_k3_host::{MlaConfig, MlaKv};
23
24/// 2026-09-25: Built from `kernels/gb10/kimi-k3/bf16/mla_decode.cu`.
25pub const MODULE: &str = "mla_decode";
26pub const ROPE_ENTRY: &str = "k3_mla_maybe_rope_f32";
27pub const SDPA_ENTRY: &str = "k3_mla_sdpa_gate_f32";
28const ROPE_BLOCK: u32 = 128;
29const SDPA_BLOCK: u32 = 32;
30
31#[derive(Clone, Copy, Debug)]
32pub struct K3MlaDecodeKernels {
33    pub rope: KernelHandle,
34    pub sdpa: KernelHandle,
35}
36
37impl K3MlaDecodeKernels {
38    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
39        Ok(Self {
40            rope: gpu.kernel(MODULE, ROPE_ENTRY)?,
41            sdpa: gpu.kernel(MODULE, SDPA_ENTRY)?,
42        })
43    }
44}
45
46fn f32_bytes(v: &[f32]) -> Vec<u8> {
47    v.iter().flat_map(|x| x.to_le_bytes()).collect()
48}
49
50fn bytes_f32(b: &[u8]) -> Vec<f32> {
51    b.chunks_exact(4)
52        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
53        .collect()
54}
55
56fn up(gpu: &dyn GpuBackend, v: &[f32], hold: &mut Vec<DevicePtr>) -> Result<DevicePtr> {
57    let b = f32_bytes(v);
58    let p = gpu.alloc(b.len().max(1))?;
59    hold.push(p);
60    gpu.copy_h2d(&b, p)?;
61    Ok(p)
62}
63
64/// 2026-09-25: One sequence's MLA K and V on the device, allocated for `cap`
65/// tokens. A decode token appends one row (K device to device after rope, V
66/// from the host) and attention reads the buffers in place; `download` copies
67/// the history to a host [`MlaKv`].
68pub struct MlaDeviceKv {
69    pub k: DevicePtr,
70    pub v: DevicePtr,
71    pub seq_len: usize,
72    pub cap: usize,
73    k_row: usize,
74    v_row: usize,
75}
76
77impl MlaDeviceKv {
78    pub fn k_row(&self) -> usize {
79        self.k_row
80    }
81
82    pub fn v_row(&self) -> usize {
83        self.v_row
84    }
85
86    pub fn alloc(gpu: &dyn GpuBackend, cap: usize, k_row: usize, v_row: usize) -> Result<Self> {
87        ensure!(k_row > 0 && v_row > 0, "k3 mla: resident row rank");
88        let cap = cap.max(1);
89        let k = gpu.alloc((cap * k_row * 4).max(1))?;
90        let v = match gpu.alloc((cap * v_row * 4).max(1)) {
91            Ok(ptr) => ptr,
92            Err(error) => {
93                let _ = gpu.free(k);
94                return Err(error);
95            }
96        };
97        Ok(Self {
98            k,
99            v,
100            seq_len: 0,
101            cap,
102            k_row,
103            v_row,
104        })
105    }
106
107    pub fn alloc_and_upload(
108        gpu: &dyn GpuBackend,
109        host: &MlaKv,
110        cap: usize,
111        k_row: usize,
112        v_row: usize,
113    ) -> Result<Self> {
114        ensure!(
115            host.k.len() == host.seq_len.saturating_mul(k_row)
116                && host.v.len() == host.seq_len.saturating_mul(v_row),
117            "k3 mla: host KV rank"
118        );
119        ensure!(
120            cap >= host.seq_len.max(1),
121            "k3 mla: device KV cap {cap} < seq {}",
122            host.seq_len
123        );
124        let mut kv = Self::alloc(gpu, cap, k_row, v_row)?;
125        if host.seq_len == 0 {
126            return Ok(kv);
127        }
128        let upload = gpu
129            .copy_h2d(&f32_bytes(&host.k), kv.k)
130            .and_then(|()| gpu.copy_h2d(&f32_bytes(&host.v), kv.v));
131        if let Err(error) = upload {
132            let _ = kv.free(gpu);
133            return Err(error);
134        }
135        kv.seq_len = host.seq_len;
136        Ok(kv)
137    }
138
139    pub fn validate_cfg(&self, cfg: &MlaConfig) -> Result<()> {
140        let k_row = cfg.heads * cfg.qk_head_dim();
141        let v_row = cfg.heads * cfg.v_head_dim;
142        ensure!(
143            self.k_row == k_row && self.v_row == v_row,
144            "k3 mla: resident KV rank"
145        );
146        Ok(())
147    }
148
149    pub fn download(&self, gpu: &dyn GpuBackend, host: &mut MlaKv) -> Result<()> {
150        let n = self.seq_len;
151        let mut kb = vec![0u8; n.saturating_mul(self.k_row) * 4];
152        let mut vb = vec![0u8; n.saturating_mul(self.v_row) * 4];
153        if n > 0 {
154            gpu.copy_d2h(self.k, &mut kb)?;
155            gpu.copy_d2h(self.v, &mut vb)?;
156        }
157        host.k = bytes_f32(&kb);
158        host.v = bytes_f32(&vb);
159        host.seq_len = n;
160        Ok(())
161    }
162
163    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
164        let k = gpu.free(self.k);
165        let v = gpu.free(self.v);
166        k.and(v)
167    }
168
169    fn append_from_device(
170        &mut self,
171        gpu: &dyn GpuBackend,
172        k_new: DevicePtr,
173        v_host: &[f32],
174    ) -> Result<()> {
175        if self.seq_len >= self.cap {
176            bail!("k3 mla: device KV cap {} full", self.cap);
177        }
178        ensure!(
179            v_host.len() == self.v_row,
180            "k3 mla: V row rank {} != {}",
181            v_host.len(),
182            self.v_row
183        );
184        let k_off = DevicePtr(self.k.0 + (self.seq_len * self.k_row * 4) as u64);
185        gpu.copy_d2d(k_new, k_off, self.k_row * 4)?;
186        let v_off = DevicePtr(self.v.0 + (self.seq_len * self.v_row * 4) as u64);
187        gpu.copy_h2d(&f32_bytes(v_host), v_off)?;
188        self.seq_len += 1;
189        Ok(())
190    }
191}
192
193/// 2026-09-25: One decode token against a host [`MlaKv`]: writes the roped
194/// `q`/`k` back, appends to `kv`, and uploads the whole KV history.
195#[allow(clippy::too_many_arguments)]
196pub fn launch_k3_mla_decode_token(
197    gpu: &dyn GpuBackend,
198    kernels: &K3MlaDecodeKernels,
199    q: &mut [f32],
200    k: &mut [f32],
201    v: &[f32],
202    g: &[f32],
203    kv: &mut MlaKv,
204    cfg: &MlaConfig,
205    pos: usize,
206    theta: f32,
207    stream: u64,
208) -> Result<Vec<f32>> {
209    let (h, dq, dv) = (cfg.heads, cfg.qk_head_dim(), cfg.v_head_dim);
210    if q.len() != h * dq || k.len() != h * dq {
211        bail!("k3 mla: q/k rank");
212    }
213    if v.len() != h * dv || g.len() != h * dv {
214        bail!("k3 mla: v/g rank");
215    }
216    if dq == 0 {
217        bail!("k3 mla: dq must be > 0");
218    }
219
220    let mut hold = Vec::new();
221    let run = (|| {
222        let dq_ptr = up(gpu, q, &mut hold)?;
223        let dk_new = up(gpu, k, &mut hold)?;
224        KernelLaunch::new(gpu, kernels.rope)
225            .grid([div_ceil(h as u32, ROPE_BLOCK), 1, 1])
226            .block([ROPE_BLOCK, 1, 1])
227            .arg_ptr(dq_ptr)
228            .arg_ptr(dk_new)
229            .arg_u32(h as u32)
230            .arg_u32(cfg.qk_nope_head_dim as u32)
231            .arg_u32(cfg.qk_rope_head_dim as u32)
232            .arg_u32(pos as u32)
233            .arg_f32(theta)
234            .arg_u32(u32::from(cfg.mla_use_nope))
235            .launch(stream)
236            .context("k3_mla_maybe_rope_f32")?;
237        gpu.synchronize(stream)?;
238
239        let mut qb = vec![0u8; q.len() * 4];
240        let mut kb = vec![0u8; k.len() * 4];
241        gpu.copy_d2h(dq_ptr, &mut qb)?;
242        gpu.copy_d2h(dk_new, &mut kb)?;
243        q.copy_from_slice(&bytes_f32(&qb));
244        k.copy_from_slice(&bytes_f32(&kb));
245        kv.append(k, v);
246
247        let dq2 = up(gpu, q, &mut hold)?;
248        let dk = up(gpu, &kv.k, &mut hold)?;
249        let dv_ptr = up(gpu, &kv.v, &mut hold)?;
250        let dg = up(gpu, g, &mut hold)?;
251        let dout = gpu.alloc((h * dv * 4).max(1))?;
252        hold.push(dout);
253        KernelLaunch::new(gpu, kernels.sdpa)
254            .grid([h as u32, 1, 1])
255            .block([SDPA_BLOCK, 1, 1])
256            .arg_ptr(dq2)
257            .arg_ptr(dk)
258            .arg_ptr(dv_ptr)
259            .arg_ptr(dg)
260            .arg_ptr(dout)
261            .arg_u32(kv.seq_len as u32)
262            .arg_u32(h as u32)
263            .arg_u32(dq as u32)
264            .arg_u32(dv as u32)
265            .arg_u32(u32::from(cfg.mla_use_output_gate))
266            .launch(stream)
267            .context("k3_mla_sdpa_gate_f32")?;
268        gpu.synchronize(stream)?;
269
270        let mut out_b = vec![0u8; h * dv * 4];
271        gpu.copy_d2h(dout, &mut out_b)?;
272        Ok(bytes_f32(&out_b))
273    })();
274    for p in hold {
275        let _ = gpu.free(p);
276    }
277    run
278}
279
280/// 2026-09-25: One decode token against an [`MlaDeviceKv`]: rope stays on the
281/// device, the append is one row, and only the output is copied back.
282#[allow(clippy::too_many_arguments)]
283pub fn launch_k3_mla_decode_token_on_device(
284    gpu: &dyn GpuBackend,
285    kernels: &K3MlaDecodeKernels,
286    q: &[f32],
287    k: &[f32],
288    v: &[f32],
289    g: &[f32],
290    kv: &mut MlaDeviceKv,
291    cfg: &MlaConfig,
292    pos: usize,
293    theta: f32,
294    stream: u64,
295) -> Result<Vec<f32>> {
296    let (h, dq, dv) = (cfg.heads, cfg.qk_head_dim(), cfg.v_head_dim);
297    if q.len() != h * dq || k.len() != h * dq || v.len() != h * dv || g.len() != h * dv {
298        bail!("k3 mla: q/k/v/g rank");
299    }
300    let mut hold = Vec::new();
301    let run = (|| {
302        let dq_ptr = up(gpu, q, &mut hold)?;
303        let dk_new = up(gpu, k, &mut hold)?;
304        KernelLaunch::new(gpu, kernels.rope)
305            .grid([div_ceil(h as u32, ROPE_BLOCK), 1, 1])
306            .block([ROPE_BLOCK, 1, 1])
307            .arg_ptr(dq_ptr)
308            .arg_ptr(dk_new)
309            .arg_u32(h as u32)
310            .arg_u32(cfg.qk_nope_head_dim as u32)
311            .arg_u32(cfg.qk_rope_head_dim as u32)
312            .arg_u32(pos as u32)
313            .arg_f32(theta)
314            .arg_u32(u32::from(cfg.mla_use_nope))
315            .launch(stream)
316            .context("k3_mla_maybe_rope_f32")?;
317        kv.append_from_device(gpu, dk_new, v)?;
318        let dg = up(gpu, g, &mut hold)?;
319        let dout = gpu.alloc((h * dv * 4).max(1))?;
320        hold.push(dout);
321        KernelLaunch::new(gpu, kernels.sdpa)
322            .grid([h as u32, 1, 1])
323            .block([SDPA_BLOCK, 1, 1])
324            .arg_ptr(dq_ptr)
325            .arg_ptr(kv.k)
326            .arg_ptr(kv.v)
327            .arg_ptr(dg)
328            .arg_ptr(dout)
329            .arg_u32(kv.seq_len as u32)
330            .arg_u32(h as u32)
331            .arg_u32(dq as u32)
332            .arg_u32(dv as u32)
333            .arg_u32(u32::from(cfg.mla_use_output_gate))
334            .launch(stream)
335            .context("k3_mla_sdpa_gate_f32")?;
336        gpu.synchronize(stream)?;
337        let mut out_b = vec![0u8; h * dv * 4];
338        gpu.copy_d2h(dout, &mut out_b)?;
339        Ok(bytes_f32(&out_b))
340    })();
341    for p in hold {
342        let _ = gpu.free(p);
343    }
344    run
345}
346
347#[cfg(test)]
348#[path = "mla_cuda_tests.rs"]
349mod tests;
metrale_model_arch/kimi_k3/mla_cuda.rs

metrale_model_arch/kimi_k3/
mla_cuda.rs
