metrale_model_layers/layers/ops/
moe_expert.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the single-token routed MoE expert GEMVs over NVFP4 pointer tables (`moe_expert_gemv.cu`, `moe_expert_gemv_fused.cu`); the shared-expert variants are in `moe_expert_shared.rs`.
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7//!
8//! Every kernel here reads `expert_indices[slot]` on the device (`blockIdx.y` is the slot) and
9//! looks the expert's weights up in device-side pointer tables, so no routing reaches the host.
10
11#![allow(unused_imports)]
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
16
17use crate::layers::moe;
18use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
19
20use super::*;
21#[path = "moe_expert_shared.rs"]
22mod shared;
23pub use shared::{
24    moe_expert_gate_up_shared, moe_expert_gate_up_shared_bf16,
25    moe_expert_gate_up_shared_bf16_batch2, moe_expert_gate_up_shared_fp8,
26    moe_expert_silu_down_shared, moe_expert_silu_down_shared_bf16,
27    moe_expert_silu_down_shared_bf16_batch2, moe_expert_silu_down_shared_fp8,
28};
29
30/// 2026-09-25: All top_k expert W4A16 GEMVs in one launch, output `[top_k, N]`. `input_stride`
31/// is 0 when every expert reads the same input row (gate/up) and K when slot e reads row e
32/// (down). A null weight pointer (a remote expert under EP) writes zeros.
33pub fn moe_expert_gemv(
34    gpu: &dyn GpuBackend,
35    kernel: KernelHandle,
36    input: DevicePtr,
37    packed_ptrs: DevicePtr,
38    scale_ptrs: DevicePtr,
39    scale2_vals: DevicePtr,
40    output: DevicePtr,
41    expert_indices: DevicePtr,
42    n: u32,
43    k: u32,
44    top_k: u32,
45    input_stride: u32,
46    stream: u64,
47) -> Result<()> {
48    KernelLaunch::new(gpu, kernel)
49        .grid([div_ceil(n, 4), top_k, 1])
50        .block([128, 1, 1])
51        .arg_ptr(input)
52        .arg_ptr(packed_ptrs)
53        .arg_ptr(scale_ptrs)
54        .arg_ptr(scale2_vals)
55        .arg_ptr(output)
56        .arg_ptr(expert_indices)
57        .arg_u32(n)
58        .arg_u32(k)
59        .arg_u32(top_k)
60        .arg_u32(input_stride)
61        .launch(stream)
62}
63
64/// 2026-09-25: Gate and up expert GEMVs in one launch; `blockIdx.z` 0 is gate, 1 is up.
65#[allow(clippy::too_many_arguments)]
66pub fn moe_expert_gemv_gate_up(
67    gpu: &dyn GpuBackend,
68    kernel: KernelHandle,
69    input: DevicePtr,
70    gate_packed_ptrs: DevicePtr,
71    gate_scale_ptrs: DevicePtr,
72    gate_scale2_vals: DevicePtr,
73    gate_out: DevicePtr,
74    up_packed_ptrs: DevicePtr,
75    up_scale_ptrs: DevicePtr,
76    up_scale2_vals: DevicePtr,
77    up_out: DevicePtr,
78    expert_indices: DevicePtr,
79    n: u32,
80    k: u32,
81    top_k: u32,
82    stream: u64,
83) -> Result<()> {
84    KernelLaunch::new(gpu, kernel)
85        .grid([div_ceil(n, 4), top_k, 2])
86        .block([128, 1, 1])
87        .arg_ptr(input)
88        .arg_ptr(gate_packed_ptrs)
89        .arg_ptr(gate_scale_ptrs)
90        .arg_ptr(gate_scale2_vals)
91        .arg_ptr(gate_out)
92        .arg_ptr(up_packed_ptrs)
93        .arg_ptr(up_scale_ptrs)
94        .arg_ptr(up_scale2_vals)
95        .arg_ptr(up_out)
96        .arg_ptr(expert_indices)
97        .arg_u32(n)
98        .arg_u32(k)
99        .arg_u32(top_k)
100        .launch(stream)
101}
102
103/// 2026-09-25: [`moe_expert_gemv_gate_up`] with two adjacent output rows per 32-thread group,
104/// 8 outputs per block.
105#[allow(clippy::too_many_arguments)]
106pub fn moe_expert_gemv_gate_up_2x(
107    gpu: &dyn GpuBackend,
108    kernel: KernelHandle,
109    input: DevicePtr,
110    gate_packed_ptrs: DevicePtr,
111    gate_scale_ptrs: DevicePtr,
112    gate_scale2_vals: DevicePtr,
113    gate_out: DevicePtr,
114    up_packed_ptrs: DevicePtr,
115    up_scale_ptrs: DevicePtr,
116    up_scale2_vals: DevicePtr,
117    up_out: DevicePtr,
118    expert_indices: DevicePtr,
119    n: u32,
120    k: u32,
121    top_k: u32,
122    stream: u64,
123) -> Result<()> {
124    KernelLaunch::new(gpu, kernel)
125        .grid([div_ceil(n, 8), top_k, 2])
126        .block([128, 1, 1])
127        .arg_ptr(input)
128        .arg_ptr(gate_packed_ptrs)
129        .arg_ptr(gate_scale_ptrs)
130        .arg_ptr(gate_scale2_vals)
131        .arg_ptr(gate_out)
132        .arg_ptr(up_packed_ptrs)
133        .arg_ptr(up_scale_ptrs)
134        .arg_ptr(up_scale2_vals)
135        .arg_ptr(up_out)
136        .arg_ptr(expert_indices)
137        .arg_u32(n)
138        .arg_u32(k)
139        .arg_u32(top_k)
140        .launch(stream)
141}
142
143/// 2026-09-25: Down expert GEMV whose activation `silu(gate) * up` is computed inline from
144/// `gate_out` and `up_out`.
145#[allow(clippy::too_many_arguments)]
146pub fn moe_expert_gemv_silu_down(
147    gpu: &dyn GpuBackend,
148    kernel: KernelHandle,
149    gate_out: DevicePtr,
150    up_out: DevicePtr,
151    packed_ptrs: DevicePtr,
152    scale_ptrs: DevicePtr,
153    scale2_vals: DevicePtr,
154    output: DevicePtr,
155    expert_indices: DevicePtr,
156    n: u32,
157    k: u32,
158    top_k: u32,
159    stream: u64,
160) -> Result<()> {
161    KernelLaunch::new(gpu, kernel)
162        .grid([div_ceil(n, 4), top_k, 1])
163        .block([128, 1, 1])
164        .arg_ptr(gate_out)
165        .arg_ptr(up_out)
166        .arg_ptr(packed_ptrs)
167        .arg_ptr(scale_ptrs)
168        .arg_ptr(scale2_vals)
169        .arg_ptr(output)
170        .arg_ptr(expert_indices)
171        .arg_u32(n)
172        .arg_u32(k)
173        .arg_u32(top_k)
174        .launch(stream)
175}
176
177/// 2026-09-25: [`moe_expert_gemv_silu_down`] with two adjacent output rows per 32-thread group,
178/// 8 outputs per block.
179#[allow(clippy::too_many_arguments)]
180pub fn moe_expert_gemv_silu_down_2x(
181    gpu: &dyn GpuBackend,
182    kernel: KernelHandle,
183    gate_out: DevicePtr,
184    up_out: DevicePtr,
185    packed_ptrs: DevicePtr,
186    scale_ptrs: DevicePtr,
187    scale2_vals: DevicePtr,
188    output: DevicePtr,
189    expert_indices: DevicePtr,
190    n: u32,
191    k: u32,
192    top_k: u32,
193    stream: u64,
194) -> Result<()> {
195    KernelLaunch::new(gpu, kernel)
196        .grid([div_ceil(n, 8), top_k, 1])
197        .block([128, 1, 1])
198        .arg_ptr(gate_out)
199        .arg_ptr(up_out)
200        .arg_ptr(packed_ptrs)
201        .arg_ptr(scale_ptrs)
202        .arg_ptr(scale2_vals)
203        .arg_ptr(output)
204        .arg_ptr(expert_indices)
205        .arg_u32(n)
206        .arg_u32(k)
207        .arg_u32(top_k)
208        .launch(stream)
209}
metrale_model_layers/layers/ops/moe_expert.rs

metrale_model_layers/layers/ops/
moe_expert.rs
