spark_model/kimi_k3/
moe_cuda.rs

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

spark_model/kimi_k3/
moe_cuda.rs
