metrale_model_arch/kimi_k3/
moe_cuda.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host launch of the Kimi K3 packed LatentMoE experts on the MXFP4 E8M0 pointer-table GEMM, [`PTRTABLE_E8M0`] in PTX module [`MODULE`].
4//!
5//! A failed kernel lookup or a missing expert projection is an error: packed
6//! experts have no host fallback.
7//!
8//! Owner: model-arch, Kimi K3.
9//! Invariants:
10//! - Every device buffer a launch allocates gets a `free` call before it
11//!   returns, on success and on error; a failed `free` is ignored.
12
13use std::collections::HashMap;
14
15use anyhow::{Context, Result, bail, ensure};
16use half::bf16;
17use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
18use metrale_model_weights::kimi_k3_host::{LatentMoeConfig, situ_glu_vec};
19
20use metrale_model_layers::layers::ops::moe_w4a16_grouped_gemm_ptrtable;
21use metrale_model_layers::weight_map::QuantizedWeight;
22
23/// 2026-09-25: The gb10 and b300 kimi-k3 `KERNEL.toml` files build this module
24/// from `deepseek-v4-flash/nvfp4/moe_w4a16_grouped_gemm.cu` (`[sources] use`,
25/// renamed by `[modules]`).
26pub const MODULE: &str = "moe_w4a16";
27pub const PTRTABLE_E8M0: &str = "moe_w4a16_grouped_gemm_ptrtable_e8m0";
28pub const E8M0_ENTRY: &str = PTRTABLE_E8M0;
29
30#[derive(Clone, Copy, Debug)]
31pub struct K3MoeGemmKernels {
32    pub ptrtable: KernelHandle,
33}
34
35impl K3MoeGemmKernels {
36    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
37        Ok(Self {
38            ptrtable: gpu.kernel(MODULE, PTRTABLE_E8M0).with_context(|| {
39                format!(
40                    "K3 MXFP4: {MODULE}::{PTRTABLE_E8M0} missing; \
41                     packed experts cannot silently run host F32"
42                )
43            })?,
44        })
45    }
46}
47
48/// 2026-09-25: One grouped E8M0 GEMM, `C[M, N] = A[M, K] @ B_packed[N, K/2]`,
49/// with one pointer-table entry per expert in `packed`. The launch covers one
50/// M tile per expert; `gemm_rows` gives every expert exactly one row.
51#[allow(clippy::too_many_arguments)]
52pub fn launch_k3_moe_e8m0_ptrtable(
53    gpu: &dyn GpuBackend,
54    kernel: KernelHandle,
55    a: DevicePtr,
56    packed: &[QuantizedWeight],
57    c: DevicePtr,
58    expert_offsets: DevicePtr,
59    sorted_token_ids: DevicePtr,
60    num_experts: u32,
61    n_out: u32,
62    k: u32,
63    stream: u64,
64) -> Result<()> {
65    ensure!(
66        packed.len() == num_experts as usize && num_experts > 0,
67        "K3 MXFP4: ptrtable length {} != num_experts {num_experts}",
68        packed.len()
69    );
70    let mut hold = Vec::new();
71    let run = (|| {
72        let (packed_ptrs, scale_ptrs, scale2_vals) = upload_ptr_table(gpu, packed, &mut hold)?;
73        moe_w4a16_grouped_gemm_ptrtable(
74            gpu,
75            kernel,
76            a,
77            packed_ptrs,
78            scale_ptrs,
79            scale2_vals,
80            c,
81            expert_offsets,
82            sorted_token_ids,
83            num_experts,
84            n_out,
85            k,
86            1,
87            stream,
88        )
89        .context("moe_w4a16_grouped_gemm_ptrtable_e8m0")
90    })();
91    for p in hold {
92        let _ = gpu.free(p);
93    }
94    run
95}
96
97/// 2026-09-25: Computes what
98/// [`mix_routed_experts`](metrale_model_weights::kimi_k3_host::mix_routed_experts)
99/// computes, `sum_e mix_w[e] * w2(SiTU(w1 x, w3 x))` over `ids`, from the
100/// packed experts, with BF16 GEMM inputs and outputs. Empty `ids` returns zeros
101/// without a launch.
102#[allow(clippy::too_many_arguments)]
103pub fn launch_k3_latent_moe_experts(
104    gpu: &dyn GpuBackend,
105    kernels: &K3MoeGemmKernels,
106    packed: &[(String, QuantizedWeight)],
107    latent: &[f32],
108    ids: &[usize],
109    mix_w: &[f32],
110    cfg: &LatentMoeConfig,
111    stream: u64,
112) -> Result<Vec<f32>> {
113    ensure!(
114        latent.len() == cfg.latent,
115        "K3 MXFP4: latent {} != {}",
116        latent.len(),
117        cfg.latent
118    );
119    ensure!(ids.len() == mix_w.len(), "K3 MXFP4: ids/weights rank");
120    if ids.is_empty() {
121        return Ok(vec![0.0; cfg.latent]);
122    }
123    let table = index_packed(packed)?;
124    let w1 = gather_proj(&table, ids, 0, "w1")?;
125    let w2 = gather_proj(&table, ids, 1, "w2")?;
126    let w3 = gather_proj(&table, ids, 2, "w3")?;
127    let m = ids.len();
128    // 2026-09-25: Every row picks its own weight pointer, so gate (w1) and up
129    // (w3) run as one launch of 2m rows: rows 0..m are gate, m..2m are up, and
130    // each row reads the same latent.
131    let gate_up_weights: Vec<_> = w1.into_iter().chain(w3).collect();
132    let gate_up_input: Vec<f32> = latent
133        .iter()
134        .copied()
135        .cycle()
136        .take(2 * m * cfg.latent)
137        .collect();
138    let gate_up = gemm_rows(
139        gpu,
140        kernels,
141        &gate_up_input,
142        2 * m,
143        &gate_up_weights,
144        cfg.expert_hidden as u32,
145        cfg.latent as u32,
146        stream,
147    )?;
148    let (gate, up) = gate_up.split_at(m * cfg.expert_hidden);
149    let mut mid = Vec::with_capacity(m * cfg.expert_hidden);
150    for e in 0..m {
151        let g = &gate[e * cfg.expert_hidden..(e + 1) * cfg.expert_hidden];
152        let u = &up[e * cfg.expert_hidden..(e + 1) * cfg.expert_hidden];
153        mid.extend(situ_glu_vec(g, u, cfg.situ_beta, cfg.situ_linear_beta));
154    }
155    let down_rows = gemm_rows(
156        gpu,
157        kernels,
158        &mid,
159        m,
160        &w2,
161        cfg.latent as u32,
162        cfg.expert_hidden as u32,
163        stream,
164    )?;
165    let mut mixed = vec![0.0f32; cfg.latent];
166    for (e, &w) in mix_w.iter().enumerate() {
167        let y = &down_rows[e * cfg.latent..(e + 1) * cfg.latent];
168        for (acc, yy) in mixed.iter_mut().zip(y) {
169            *acc += w * *yy;
170        }
171    }
172    Ok(mixed)
173}
174
175fn gemm_rows(
176    gpu: &dyn GpuBackend,
177    kernels: &K3MoeGemmKernels,
178    a_f32: &[f32],
179    m: usize,
180    packed: &[QuantizedWeight],
181    n_out: u32,
182    k: u32,
183    stream: u64,
184) -> Result<Vec<f32>> {
185    ensure!(
186        a_f32.len() == m * k as usize,
187        "K3 MXFP4: A {} vs {m}x{k}",
188        a_f32.len()
189    );
190    let mut hold = Vec::new();
191    let run = (|| {
192        let a = up_bf16(gpu, a_f32, &mut hold)?;
193        let c = gpu.alloc((m * n_out as usize * 2).max(1))?;
194        hold.push(c);
195        let off: Vec<i32> = (0..=m as i32).collect();
196        let ids: Vec<i32> = (0..m as i32).collect();
197        let offsets = up_i32(gpu, &off, &mut hold)?;
198        let sorted = up_i32(gpu, &ids, &mut hold)?;
199        launch_k3_moe_e8m0_ptrtable(
200            gpu,
201            kernels.ptrtable,
202            a,
203            packed,
204            c,
205            offsets,
206            sorted,
207            m as u32,
208            n_out,
209            k,
210            stream,
211        )?;
212        gpu.synchronize(stream)?;
213        let mut raw = vec![0u8; m * n_out as usize * 2];
214        gpu.copy_d2h(c, &mut raw)?;
215        Ok(bf16_to_f32(&raw))
216    })();
217    for p in hold {
218        let _ = gpu.free(p);
219    }
220    run
221}
222
223fn upload_ptr_table(
224    gpu: &dyn GpuBackend,
225    packed: &[QuantizedWeight],
226    hold: &mut Vec<DevicePtr>,
227) -> Result<(DevicePtr, DevicePtr, DevicePtr)> {
228    let n = packed.len();
229    let packed_bytes: Vec<u8> = packed
230        .iter()
231        .flat_map(|w| w.weight.0.to_le_bytes())
232        .collect();
233    let scale_bytes: Vec<u8> = packed
234        .iter()
235        .flat_map(|w| w.weight_scale.0.to_le_bytes())
236        .collect();
237    let scale2_bytes: Vec<u8> = packed
238        .iter()
239        .flat_map(|w| w.weight_scale_2.to_le_bytes())
240        .collect();
241    let packed_ptrs = gpu.alloc((n * 8).max(1))?;
242    hold.push(packed_ptrs);
243    gpu.copy_h2d(&packed_bytes, packed_ptrs)?;
244    let scale_ptrs = gpu.alloc((n * 8).max(1))?;
245    hold.push(scale_ptrs);
246    gpu.copy_h2d(&scale_bytes, scale_ptrs)?;
247    let scale2_vals = gpu.alloc((n * 4).max(1))?;
248    hold.push(scale2_vals);
249    gpu.copy_h2d(&scale2_bytes, scale2_vals)?;
250    Ok((packed_ptrs, scale_ptrs, scale2_vals))
251}
252
253fn index_packed(
254    packed: &[(String, QuantizedWeight)],
255) -> Result<HashMap<(usize, usize), QuantizedWeight>> {
256    let mut t = HashMap::new();
257    for (prefix, w) in packed {
258        let (id, proj) = parse_expert_proj(prefix)?;
259        if t.insert((id, proj), *w).is_some() {
260            bail!("K3 MXFP4: duplicate packed {prefix}");
261        }
262    }
263    Ok(t)
264}
265
266fn parse_expert_proj(prefix: &str) -> Result<(usize, usize)> {
267    let (head, proj) = prefix
268        .rsplit_once('.')
269        .with_context(|| format!("K3 MXFP4: expert prefix {prefix}"))?;
270    let proj_i = match proj {
271        "w1" => 0,
272        "w2" => 1,
273        "w3" => 2,
274        _ => bail!("K3 MXFP4: expected w1|w2|w3 in {prefix}"),
275    };
276    let id_s = head
277        .rsplit_once(".experts.")
278        .map(|(_, id)| id)
279        .with_context(|| format!("K3 MXFP4: experts.id in {prefix}"))?;
280    let id: usize = id_s
281        .parse()
282        .with_context(|| format!("K3 MXFP4: expert id {id_s}"))?;
283    Ok((id, proj_i))
284}
285
286fn gather_proj(
287    table: &HashMap<(usize, usize), QuantizedWeight>,
288    ids: &[usize],
289    proj: usize,
290    name: &str,
291) -> Result<Vec<QuantizedWeight>> {
292    ids.iter()
293        .map(|&id| {
294            table.get(&(id, proj)).copied().with_context(|| {
295                format!(
296                    "K3 MXFP4: expert {id} {name} missing; packed experts cannot silently run host F32"
297                )
298            })
299        })
300        .collect()
301}
302
303fn up_bf16(gpu: &dyn GpuBackend, v: &[f32], hold: &mut Vec<DevicePtr>) -> Result<DevicePtr> {
304    let b: Vec<u8> = v
305        .iter()
306        .flat_map(|&f| bf16::from_f32(f).to_le_bytes())
307        .collect();
308    let p = gpu.alloc(b.len().max(1))?;
309    hold.push(p);
310    gpu.copy_h2d(&b, p)?;
311    Ok(p)
312}
313
314fn up_i32(gpu: &dyn GpuBackend, v: &[i32], hold: &mut Vec<DevicePtr>) -> Result<DevicePtr> {
315    let b: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
316    let p = gpu.alloc(b.len().max(1))?;
317    hold.push(p);
318    gpu.copy_h2d(&b, p)?;
319    Ok(p)
320}
321
322fn bf16_to_f32(raw: &[u8]) -> Vec<f32> {
323    raw.chunks_exact(2)
324        .map(|b| bf16::from_le_bytes([b[0], b[1]]).to_f32())
325        .collect()
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use metrale_gpu_runtime::gpu::mock::{MockArg, MockGpuBackend};
332    use metrale_gpu_runtime::kernel_args::div_ceil;
333
334    fn dummy_qw(gpu: &MockGpuBackend) -> QuantizedWeight {
335        QuantizedWeight {
336            weight: gpu.alloc(16).unwrap(),
337            weight_scale: gpu.alloc(1).unwrap(),
338            weight_scale_2: 1.0,
339            input_scale: DevicePtr::NULL,
340            weight_scale_2_vec: DevicePtr::NULL,
341            act: metrale_config::Nvfp4Act::Unstamped,
342        }
343    }
344
345    #[test]
346    fn resolve_looks_up_e8m0_ptrtable() {
347        let gpu = MockGpuBackend::new();
348        let _ = K3MoeGemmKernels::resolve(&gpu).unwrap();
349        assert_eq!(
350            gpu.kernel_lookups_snapshot(),
351            vec![(MODULE.to_string(), PTRTABLE_E8M0.to_string())]
352        );
353    }
354
355    #[test]
356    fn deny_kernel_resolve_bails_not_silent_cpu() {
357        let gpu = MockGpuBackend::new();
358        gpu.deny_kernel(MODULE, PTRTABLE_E8M0);
359        let err = K3MoeGemmKernels::resolve(&gpu).unwrap_err().to_string();
360        assert!(
361            err.contains(PTRTABLE_E8M0) && err.contains("cannot silently run host F32"),
362            "{err}"
363        );
364        assert_eq!(gpu.launch_count(), 0);
365    }
366
367    #[test]
368    fn mock_launch_contract_one_expert() {
369        let gpu = MockGpuBackend::new();
370        let k = K3MoeGemmKernels::resolve(&gpu).unwrap();
371        // 2026-09-29: The deepseek-v4-flash tree's `_ptrtable_e8m0` publishes a 64-column tile.
372        gpu.set_kernel_n_tile(k.ptrtable, 64);
373        let packed = dummy_qw(&gpu);
374        let a = gpu.alloc(64).unwrap();
375        let c = gpu.alloc(128).unwrap();
376        let off = gpu.alloc(8).unwrap();
377        let ids = gpu.alloc(4).unwrap();
378        let n_out = 64u32;
379        let kk = 32u32;
380        launch_k3_moe_e8m0_ptrtable(&gpu, k.ptrtable, a, &[packed], c, off, ids, 1, n_out, kk, 3)
381            .unwrap();
382        let launches = gpu.launches_snapshot();
383        assert_eq!(launches.len(), 1);
384        assert_eq!(launches[0].grid, [div_ceil(n_out, 64), 1, 1]);
385        assert_eq!(launches[0].block, [128, 1, 1]);
386        assert_eq!(launches[0].stream, 3);
387        assert_eq!(launches[0].args.len(), 10);
388        assert_eq!(
389            launches[0].args[7],
390            MockArg::Bytes(1u32.to_le_bytes().to_vec())
391        );
392        assert_eq!(
393            launches[0].args[8],
394            MockArg::Bytes(n_out.to_le_bytes().to_vec())
395        );
396        assert_eq!(
397            launches[0].args[9],
398            MockArg::Bytes(kk.to_le_bytes().to_vec())
399        );
400    }
401
402    #[test]
403    fn parse_k3_expert_prefix() {
404        let p = "language_model.model.layers.12.block_sparse_moe.experts.7.w1";
405        assert_eq!(parse_expert_proj(p).unwrap(), (7, 0));
406        assert_eq!(
407            parse_expert_proj("model.layers.1.block_sparse_moe.experts.0.w3").unwrap(),
408            (0, 2)
409        );
410    }
411
412    #[test]
413    fn empty_ids_does_not_launch() {
414        let gpu = MockGpuBackend::new();
415        let k = K3MoeGemmKernels::resolve(&gpu).unwrap();
416        let cfg = LatentMoeConfig {
417            hidden: 2,
418            latent: 2,
419            expert_hidden: 2,
420            n_routed: 1,
421            top_k: 1,
422            n_shared: 0,
423            situ_beta: 4.0,
424            situ_linear_beta: 25.0,
425            use_norm: false,
426            renormalize: true,
427        };
428        let y =
429            launch_k3_latent_moe_experts(&gpu, &k, &[], &[1.0, 0.0], &[], &[], &cfg, 0).unwrap();
430        assert_eq!(y, vec![0.0, 0.0]);
431        assert_eq!(gpu.launch_count(), 0);
432    }
433
434    #[test]
435    fn missing_packed_proj_bails_not_cpu() {
436        let gpu = MockGpuBackend::new();
437        let k = K3MoeGemmKernels::resolve(&gpu).unwrap();
438        let cfg = LatentMoeConfig {
439            hidden: 2,
440            latent: 2,
441            expert_hidden: 2,
442            n_routed: 1,
443            top_k: 1,
444            n_shared: 0,
445            situ_beta: 4.0,
446            situ_linear_beta: 25.0,
447            use_norm: false,
448            renormalize: true,
449        };
450        let packed = [(
451            "model.layers.1.block_sparse_moe.experts.0.w1".to_string(),
452            dummy_qw(&gpu),
453        )];
454        let before = gpu.launch_count();
455        let err =
456            launch_k3_latent_moe_experts(&gpu, &k, &packed, &[1.0, 0.0], &[0], &[1.0], &cfg, 0)
457                .unwrap_err()
458                .to_string();
459        assert!(
460            err.contains("w2") && err.contains("cannot silently run host F32"),
461            "{err}"
462        );
463        assert_eq!(gpu.launch_count(), before);
464    }
465}
metrale_model_arch/kimi_k3/moe_cuda.rs

metrale_model_arch/kimi_k3/
moe_cuda.rs
