metrale_model_arch/kimi_k3/
dense_cuda.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: CUDA SiTU-GLU MLP for Kimi K3 (a dense layer's MLP or a LatentMoE layer's shared expert) on the resident BF16/FP32 weights, with FP32 activations.
4//!
5//! Decode uses it only with `K3_CUDA_DENSE=1`. Only the input row and the output
6//! row cross the host boundary; the weights are read where the loader put them.
7//!
8//! Owner: model-arch, Kimi K3.
9//! Invariants:
10//! - Every scratch buffer a launch allocates gets a `free` call when it
11//!   returns, on success and on error (`Scratch` on drop); a failed `free` is
12//!   ignored.
13
14use anyhow::{Context, Result, ensure};
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
16use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
17use metrale_model_weights::weights::WeightDtype;
18
19use super::bound::K3BoundLayer;
20
21fn dtype_tag(dtype: WeightDtype) -> Result<u32> {
22    match dtype {
23        WeightDtype::FP32 => Ok(0),
24        WeightDtype::BF16 => Ok(1),
25        other => anyhow::bail!("K3 dense GPU: unsupported resident weight dtype {other:?}"),
26    }
27}
28
29struct Scratch<'a> {
30    gpu: &'a dyn GpuBackend,
31    pointers: Vec<DevicePtr>,
32}
33impl Scratch<'_> {
34    fn alloc(&mut self, floats: usize) -> Result<DevicePtr> {
35        let bytes = floats
36            .checked_mul(4)
37            .context("K3 dense scratch size overflow")?;
38        let p = self.gpu.alloc(bytes)?;
39        self.pointers.push(p);
40        Ok(p)
41    }
42}
43impl Drop for Scratch<'_> {
44    fn drop(&mut self) {
45        for p in self.pointers.drain(..).rev() {
46            let _ = self.gpu.free(p);
47        }
48    }
49}
50
51/// 2026-09-25: Finds the layer's resident `gate`/`up`/`down` weights (under
52/// `.mlp.` for a dense layer, `.block_sparse_moe.shared_experts.` for a
53/// LatentMoE layer), checks each holds `hidden * inter` values, and runs
54/// [`launch_dense_matrices`].
55#[allow(clippy::too_many_arguments)]
56pub fn launch_dense_mlp(
57    layer: &K3BoundLayer,
58    gpu: &dyn GpuBackend,
59    x: &[f32],
60    hidden: usize,
61    inter: usize,
62    beta: f32,
63    linear_beta: f32,
64    stream: u64,
65) -> Result<Vec<f32>> {
66    ensure!(
67        hidden > 0 && inter > 0 && x.len() == hidden,
68        "K3 dense GPU: invalid activation geometry"
69    );
70    ensure!(
71        beta.is_finite() && beta > 0.0 && linear_beta.is_finite() && linear_beta > 0.0,
72        "K3 dense GPU: invalid SiTU coefficients"
73    );
74    let elements = hidden
75        .checked_mul(inter)
76        .context("K3 dense weight size overflow")?;
77    ensure!(
78        elements <= u32::MAX as usize,
79        "K3 dense GPU: matrix exceeds 32-bit kernel indexing"
80    );
81    let prefix = match layer.spec.mlp {
82        metrale_model_weights::kimi_k3_host::MlpKind::Dense => ".mlp.",
83        metrale_model_weights::kimi_k3_host::MlpKind::LatentMoe => {
84            ".block_sparse_moe.shared_experts."
85        }
86    };
87    let find = |role: &str| -> Result<_> {
88        let suffix = format!("{prefix}{role}_proj.weight");
89        let (weight, meta) = layer
90            .weights
91            .iter()
92            .zip(&layer.weight_meta)
93            .find(|(_, meta)| meta.name.ends_with(&suffix))
94            .with_context(|| format!("K3 dense GPU: resident {suffix} missing"))?;
95        ensure!(
96            meta.numel == elements,
97            "K3 dense GPU: {} has {} values, expected {elements}",
98            meta.name,
99            meta.numel
100        );
101        Ok((weight.weight, dtype_tag(meta.dtype)?))
102    };
103    let (gate, gate_dtype) = find("gate")?;
104    let (up, up_dtype) = find("up")?;
105    let (down, down_dtype) = find("down")?;
106    ensure!(
107        gate_dtype == up_dtype,
108        "K3 dense GPU: gate/up dtypes differ"
109    );
110    launch_dense_matrices(
111        gpu,
112        x,
113        gate,
114        up,
115        down,
116        gate_dtype,
117        down_dtype,
118        hidden,
119        inter,
120        beta,
121        linear_beta,
122        stream,
123    )
124}
125
126/// 2026-09-25: Runs the gate/up SiTU kernel and the down kernel on explicit
127/// device weights. Dtype tags: 0 = FP32, 1 = BF16.
128#[allow(clippy::too_many_arguments)]
129pub fn launch_dense_matrices(
130    gpu: &dyn GpuBackend,
131    x: &[f32],
132    gate: DevicePtr,
133    up: DevicePtr,
134    down: DevicePtr,
135    gate_dtype: u32,
136    down_dtype: u32,
137    hidden: usize,
138    inter: usize,
139    beta: f32,
140    linear_beta: f32,
141    stream: u64,
142) -> Result<Vec<f32>> {
143    ensure!(
144        hidden > 0 && inter > 0 && x.len() == hidden,
145        "K3 dense GPU: invalid dimensions"
146    );
147    ensure!(
148        hidden
149            .checked_mul(inter)
150            .is_some_and(|n| n <= u32::MAX as usize),
151        "K3 dense GPU: index overflow"
152    );
153    ensure!(
154        gate_dtype <= 1 && down_dtype <= 1,
155        "K3 dense GPU: invalid dtype tag"
156    );
157    ensure!(
158        beta.is_finite() && beta > 0.0 && linear_beta.is_finite() && linear_beta > 0.0,
159        "K3 dense GPU: invalid activation coefficients"
160    );
161    let gate_kernel = gpu.kernel("dense_f32io", "k3_dense_gate_up_situ_f32io")?;
162    let down_kernel = gpu.kernel("dense_f32io", "k3_dense_down_f32io")?;
163    let mut scratch = Scratch {
164        gpu,
165        pointers: Vec::new(),
166    };
167    let input = scratch.alloc(hidden)?;
168    let mid = scratch.alloc(inter)?;
169    let output = scratch.alloc(hidden)?;
170    gpu.copy_h2d(
171        &x.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<_>>(),
172        input,
173    )?;
174    KernelLaunch::new(gpu, gate_kernel)
175        .grid([div_ceil(inter as u32, 4), 1, 1])
176        .block([128, 1, 1])
177        .arg_ptr(input)
178        .arg_ptr(gate)
179        .arg_ptr(up)
180        .arg_ptr(mid)
181        .arg_u32(inter as u32)
182        .arg_u32(hidden as u32)
183        .arg_u32(gate_dtype)
184        .arg_f32(beta)
185        .arg_f32(linear_beta)
186        .launch(stream)?;
187    KernelLaunch::new(gpu, down_kernel)
188        .grid([div_ceil(hidden as u32, 4), 1, 1])
189        .block([128, 1, 1])
190        .arg_ptr(mid)
191        .arg_ptr(down)
192        .arg_ptr(output)
193        .arg_u32(hidden as u32)
194        .arg_u32(inter as u32)
195        .arg_u32(down_dtype)
196        .launch(stream)?;
197    gpu.synchronize(stream)?;
198    let mut raw = vec![0; hidden * 4];
199    gpu.copy_d2h(output, &mut raw)?;
200    let result: Vec<f32> = raw
201        .chunks_exact(4)
202        .map(|b| f32::from_le_bytes(b.try_into().unwrap()))
203        .collect();
204    ensure!(
205        result.iter().all(|v| v.is_finite()),
206        "K3 dense GPU: nonfinite output"
207    );
208    Ok(result)
209}
metrale_model_arch/kimi_k3/dense_cuda.rs

metrale_model_arch/kimi_k3/
dense_cuda.rs
