spark_model/kimi_k3/
dense_cuda.rs

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

spark_model/kimi_k3/
dense_cuda.rs
