metrale_model_layers/layers/ops/
moe_atomic_c4.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the two C=4 routed MoE decode kernels of `moe_decode_atomic_c4.cu`: `moe_decode_atomic_c4_silu_down_accum` and `moe_decode_atomic_c4_finalize`.
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7//!
8//! The SiLU+down kernel adds each routed expert's weighted output into `routed_accum` with FP32
9//! `atomicAdd`. The multi-sequence SSM decode takes this path only for n == 4 with
10//! `METRALE_MOE_ATOMIC_C4_DECODE=1` (`qwen3_ssm/trait_decode_multi_seq.rs`).
11
12#![allow(unused_imports)]
13
14use anyhow::Result;
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
16use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
17
18use crate::weight_map::QuantizedWeight;
19
20#[allow(clippy::too_many_arguments)]
21pub fn moe_decode_atomic_c4_silu_down_accum(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    gate_out: DevicePtr,
25    up_out: DevicePtr,
26    packed_ptrs: DevicePtr,
27    scale_ptrs: DevicePtr,
28    scale2_vals: DevicePtr,
29    expert_indices: DevicePtr,
30    expert_weights: DevicePtr,
31    routed_accum: DevicePtr,
32    sh_gate_in: DevicePtr,
33    sh_up_in: DevicePtr,
34    sh_down: &QuantizedWeight,
35    sh_down_out: DevicePtr,
36    hidden: u32,
37    inter: u32,
38    top_k: u32,
39    num_tokens: u32,
40    stream: u64,
41) -> Result<()> {
42    let smem_bytes = (inter as usize * std::mem::size_of::<f32>()) as u32;
43    KernelLaunch::new(gpu, kernel)
44        .grid([div_ceil(hidden, 8), num_tokens * (top_k + 1), 1])
45        .block([128, 1, 1])
46        .shared_mem(smem_bytes)
47        .arg_ptr(gate_out)
48        .arg_ptr(up_out)
49        .arg_ptr(packed_ptrs)
50        .arg_ptr(scale_ptrs)
51        .arg_ptr(scale2_vals)
52        .arg_ptr(expert_indices)
53        .arg_ptr(expert_weights)
54        .arg_ptr(routed_accum)
55        .arg_ptr(sh_gate_in)
56        .arg_ptr(sh_up_in)
57        .arg_ptr(sh_down.weight)
58        .arg_ptr(sh_down.weight_scale)
59        .arg_f32(sh_down.weight_scale_2)
60        .arg_ptr(sh_down_out)
61        .arg_u32(hidden)
62        .arg_u32(inter)
63        .arg_u32(top_k)
64        .arg_u32(num_tokens)
65        .launch(stream)
66}
67
68#[allow(clippy::too_many_arguments)]
69pub fn moe_decode_atomic_c4_finalize(
70    gpu: &dyn GpuBackend,
71    kernel: KernelHandle,
72    output: DevicePtr,
73    routed_accum: DevicePtr,
74    shared_out: DevicePtr,
75    input: DevicePtr,
76    gate_weight: DevicePtr,
77    hidden: u32,
78    num_tokens: u32,
79    include_shared: bool,
80    stream: u64,
81) -> Result<()> {
82    KernelLaunch::new(gpu, kernel)
83        .grid([div_ceil(hidden, 256), num_tokens, 1])
84        .block([256, 1, 1])
85        .arg_ptr(output)
86        .arg_ptr(routed_accum)
87        .arg_ptr(shared_out)
88        .arg_ptr(input)
89        .arg_ptr(gate_weight)
90        .arg_u32(hidden)
91        .arg_u32(num_tokens)
92        .arg_u32(if include_shared { 1 } else { 0 })
93        .launch(stream)
94}
metrale_model_layers/layers/ops/moe_atomic_c4.rs

metrale_model_layers/layers/ops/
moe_atomic_c4.rs
