metrale_model_layers/layers/ops/
moe_prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the token-major MoE prefill kernels
4//! (`kernels/gb10/common/moe_prefill.cu`), the fused dual and SiLU-input W4A16
5//! and W8A16 GEMVs, and the device-gated sigmoid blend.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::layers::moe;
17use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
18
19use super::*;
20
21/// 2026-09-25: Gate and up GEMVs for `num_tokens` tokens, routed and shared
22/// experts in one launch. Row `y < num_tokens * top_k` is routed (token
23/// `y / top_k`, expert `expert_indices[y]`) and writes `gate_out`/`up_out`
24/// (`[num_tokens * top_k, n]` BF16); the remaining `num_tokens` rows run the
25/// shared expert and write `sh_gate_out`/`sh_up_out` (`[num_tokens, n]` BF16).
26/// `blockIdx.z` picks gate (0) or up (1). `input` is `[num_tokens, k]` BF16.
27#[allow(clippy::too_many_arguments)]
28pub fn moe_expert_gate_up_shared_prefill(
29    gpu: &dyn GpuBackend,
30    kernel: KernelHandle,
31    input: DevicePtr,
32    gate_packed_ptrs: DevicePtr,
33    gate_scale_ptrs: DevicePtr,
34    gate_scale2_vals: DevicePtr,
35    gate_out: DevicePtr,
36    up_packed_ptrs: DevicePtr,
37    up_scale_ptrs: DevicePtr,
38    up_scale2_vals: DevicePtr,
39    up_out: DevicePtr,
40    expert_indices: DevicePtr,
41    sh_gate: &QuantizedWeight,
42    sh_gate_out: DevicePtr,
43    sh_up: &QuantizedWeight,
44    sh_up_out: DevicePtr,
45    n: u32,
46    k: u32,
47    top_k: u32,
48    num_tokens: u32,
49    stream: u64,
50) -> Result<()> {
51    KernelLaunch::new(gpu, kernel)
52        .grid([div_ceil(n, 8), num_tokens * (top_k + 1), 2])
53        .block([128, 1, 1])
54        .arg_ptr(input)
55        .arg_ptr(gate_packed_ptrs)
56        .arg_ptr(gate_scale_ptrs)
57        .arg_ptr(gate_scale2_vals)
58        .arg_ptr(gate_out)
59        .arg_ptr(up_packed_ptrs)
60        .arg_ptr(up_scale_ptrs)
61        .arg_ptr(up_scale2_vals)
62        .arg_ptr(up_out)
63        .arg_ptr(expert_indices)
64        .arg_ptr(sh_gate.weight)
65        .arg_ptr(sh_gate.weight_scale)
66        .arg_f32(sh_gate.weight_scale_2)
67        .arg_ptr(sh_gate_out)
68        .arg_ptr(sh_up.weight)
69        .arg_ptr(sh_up.weight_scale)
70        .arg_f32(sh_up.weight_scale_2)
71        .arg_ptr(sh_up_out)
72        .arg_u32(n)
73        .arg_u32(k)
74        .arg_u32(top_k)
75        .arg_u32(num_tokens)
76        .launch(stream)
77}
78
79/// 2026-09-25: SiLU(gate) * up, then the down GEMV, for the routed and shared
80/// rows of [`moe_expert_gate_up_shared_prefill`], with the same row split.
81/// Routed rows read `gate_out`/`up_out` and write `output`
82/// (`[num_tokens * top_k, n]` BF16); shared rows read `sh_gate_in`/`sh_up_in`
83/// and write `sh_down_out` (`[num_tokens, n]` BF16).
84#[allow(clippy::too_many_arguments)]
85pub fn moe_expert_silu_down_shared_prefill(
86    gpu: &dyn GpuBackend,
87    kernel: KernelHandle,
88    gate_out: DevicePtr,
89    up_out: DevicePtr,
90    packed_ptrs: DevicePtr,
91    scale_ptrs: DevicePtr,
92    scale2_vals: DevicePtr,
93    output: DevicePtr,
94    expert_indices: DevicePtr,
95    sh_gate_in: DevicePtr,
96    sh_up_in: DevicePtr,
97    sh_down: &QuantizedWeight,
98    sh_down_out: DevicePtr,
99    n: u32,
100    k: u32,
101    top_k: u32,
102    num_tokens: u32,
103    stream: u64,
104) -> Result<()> {
105    KernelLaunch::new(gpu, kernel)
106        .grid([div_ceil(n, 8), num_tokens * (top_k + 1), 1])
107        .block([128, 1, 1])
108        .arg_ptr(gate_out)
109        .arg_ptr(up_out)
110        .arg_ptr(packed_ptrs)
111        .arg_ptr(scale_ptrs)
112        .arg_ptr(scale2_vals)
113        .arg_ptr(output)
114        .arg_ptr(expert_indices)
115        .arg_ptr(sh_gate_in)
116        .arg_ptr(sh_up_in)
117        .arg_ptr(sh_down.weight)
118        .arg_ptr(sh_down.weight_scale)
119        .arg_f32(sh_down.weight_scale_2)
120        .arg_ptr(sh_down_out)
121        .arg_u32(n)
122        .arg_u32(k)
123        .arg_u32(top_k)
124        .arg_u32(num_tokens)
125        .launch(stream)
126}
127
128/// 2026-09-25: Per token `t`: `output[t] = sum_e expert_weights[t*top_k+e] *
129/// expert_out[t*top_k+e] + sigmoid(dot(input[t], gate_weight)) * shared_out[t]`.
130/// `expert_out` is `[num_tokens * top_k, hidden]` BF16, `expert_weights`
131/// `[num_tokens * top_k]` f32, `input` `[num_tokens, k]` BF16 and `gate_weight`
132/// the shared-expert gate `[1, k]` BF16.
133#[allow(clippy::too_many_arguments)]
134pub fn moe_weighted_sum_blend_prefill(
135    gpu: &dyn GpuBackend,
136    kernel: KernelHandle,
137    output: DevicePtr,
138    expert_out: DevicePtr,
139    expert_weights: DevicePtr,
140    shared_out: DevicePtr,
141    input: DevicePtr,
142    gate_weight: DevicePtr,
143    hidden: u32,
144    top_k: u32,
145    k: u32,
146    num_tokens: u32,
147    stream: u64,
148) -> Result<()> {
149    KernelLaunch::new(gpu, kernel)
150        .grid([div_ceil(hidden, 256), num_tokens, 1])
151        .block([256, 1, 1])
152        .arg_ptr(output)
153        .arg_ptr(expert_out)
154        .arg_ptr(expert_weights)
155        .arg_ptr(shared_out)
156        .arg_ptr(input)
157        .arg_ptr(gate_weight)
158        .arg_u32(hidden)
159        .arg_u32(top_k)
160        .arg_u32(k)
161        .arg_u32(num_tokens)
162        .launch(stream)
163}
164
165/// 2026-09-25: Two W4A16 GEMVs over the same BF16 input in one launch;
166/// `blockIdx.z` picks `weight1` or `weight2`. Both projections have `n` outputs.
167#[allow(clippy::too_many_arguments)]
168pub fn w4a16_gemv_dual(
169    gpu: &dyn GpuBackend,
170    kernel: KernelHandle,
171    input: DevicePtr,
172    weight1: &QuantizedWeight,
173    output1: DevicePtr,
174    weight2: &QuantizedWeight,
175    output2: DevicePtr,
176    n: u32,
177    k: u32,
178    stream: u64,
179) -> Result<()> {
180    KernelLaunch::new(gpu, kernel)
181        .grid([w4a16_gemv_grid_x(n), 1, 2])
182        .block([256, 1, 1])
183        .arg_ptr(input)
184        .arg_ptr(weight1.weight)
185        .arg_ptr(weight1.weight_scale)
186        .arg_f32(weight1.weight_scale_2)
187        .arg_ptr(output1)
188        .arg_ptr(weight2.weight)
189        .arg_ptr(weight2.weight_scale)
190        .arg_f32(weight2.weight_scale_2)
191        .arg_ptr(output2)
192        .arg_u32(n)
193        .arg_u32(k)
194        .launch(stream)
195}
196
197/// 2026-09-25: One warp per output variant of [`w4a16_gemv_dual`], 8 outputs per
198/// block. It shares the per-lane partial `w4a16_dual_partial` with the base
199/// kernel (`w4a16_gemv_fused.cu`). Callers pick it through [`use_gemv_sw`] with
200/// `ModelLevers::gemv_sw`, which is on unless `METRALE_NO_GEMV_SW=1`.
201#[allow(clippy::too_many_arguments)]
202pub fn w4a16_gemv_dual_sw(
203    gpu: &dyn GpuBackend,
204    kernel: KernelHandle,
205    input: DevicePtr,
206    weight1: &QuantizedWeight,
207    output1: DevicePtr,
208    weight2: &QuantizedWeight,
209    output2: DevicePtr,
210    n: u32,
211    k: u32,
212    stream: u64,
213) -> Result<()> {
214    KernelLaunch::new(gpu, kernel)
215        .grid([w4a16_gemv_sw_grid_x(n), 1, 2])
216        .block([256, 1, 1])
217        .arg_ptr(input)
218        .arg_ptr(weight1.weight)
219        .arg_ptr(weight1.weight_scale)
220        .arg_f32(weight1.weight_scale_2)
221        .arg_ptr(output1)
222        .arg_ptr(weight2.weight)
223        .arg_ptr(weight2.weight_scale)
224        .arg_f32(weight2.weight_scale_2)
225        .arg_ptr(output2)
226        .arg_u32(n)
227        .arg_u32(k)
228        .launch(stream)
229}
230
231/// 2026-09-25: One warp per output variant of [`w4a16_gemv_silu_input`], 8
232/// outputs per block, picked the same way as [`w4a16_gemv_dual_sw`].
233#[allow(clippy::too_many_arguments)]
234pub fn w4a16_gemv_silu_input_sw(
235    gpu: &dyn GpuBackend,
236    kernel: KernelHandle,
237    gate_out: DevicePtr,
238    up_out: DevicePtr,
239    weight: &QuantizedWeight,
240    output: DevicePtr,
241    n: u32,
242    k: u32,
243    stream: u64,
244) -> Result<()> {
245    KernelLaunch::new(gpu, kernel)
246        .grid([w4a16_gemv_sw_grid_x(n), 1, 1])
247        .block([256, 1, 1])
248        .arg_ptr(gate_out)
249        .arg_ptr(up_out)
250        .arg_ptr(weight.weight)
251        .arg_ptr(weight.weight_scale)
252        .arg_f32(weight.weight_scale_2)
253        .arg_ptr(output)
254        .arg_u32(n)
255        .arg_u32(k)
256        .launch(stream)
257}
258
259/// 2026-09-25: W4A16 GEMV whose activation is `silu(gate_out) * up_out`,
260/// computed per element from the two `[k]` BF16 vectors inside the kernel.
261#[allow(clippy::too_many_arguments)]
262pub fn w4a16_gemv_silu_input(
263    gpu: &dyn GpuBackend,
264    kernel: KernelHandle,
265    gate_out: DevicePtr,
266    up_out: DevicePtr,
267    weight: &QuantizedWeight,
268    output: DevicePtr,
269    n: u32,
270    k: u32,
271    stream: u64,
272) -> Result<()> {
273    KernelLaunch::new(gpu, kernel)
274        .grid([div_ceil(n, 4), 1, 1])
275        .block([256, 1, 1])
276        .arg_ptr(gate_out)
277        .arg_ptr(up_out)
278        .arg_ptr(weight.weight)
279        .arg_ptr(weight.weight_scale)
280        .arg_f32(weight.weight_scale_2)
281        .arg_ptr(output)
282        .arg_u32(n)
283        .arg_u32(k)
284        .launch(stream)
285}
286
287/// 2026-09-25: Two FP8 E4M3 GEMVs over the same BF16 input in one launch;
288/// `blockIdx.z` picks projection 1 or 2. Both have `n` outputs. Each weight is
289/// `[n, k]` bytes with an FP32 `[n/128, k/128]` block-scale buffer.
290#[allow(clippy::too_many_arguments)]
291pub fn w8a16_gemv_dual(
292    gpu: &dyn GpuBackend,
293    kernel: KernelHandle,
294    input: DevicePtr,
295    weight1: DevicePtr,
296    row_scale1: DevicePtr,
297    output1: DevicePtr,
298    weight2: DevicePtr,
299    row_scale2: DevicePtr,
300    output2: DevicePtr,
301    n: u32,
302    k: u32,
303    stream: u64,
304) -> Result<()> {
305    KernelLaunch::new(gpu, kernel)
306        .grid([div_ceil(n, 4), 1, 2])
307        .block([256, 1, 1])
308        .arg_ptr(input)
309        .arg_ptr(weight1)
310        .arg_ptr(row_scale1)
311        .arg_ptr(output1)
312        .arg_ptr(weight2)
313        .arg_ptr(row_scale2)
314        .arg_ptr(output2)
315        .arg_u32(n)
316        .arg_u32(k)
317        .launch(stream)
318}
319
320/// 2026-09-25: FP8 E4M3 GEMV whose activation is `silu(gate_out) * up_out`,
321/// computed per element from the two `[k]` BF16 vectors inside the kernel.
322/// `block_scale` is FP32 `[n/128, k/128]`.
323#[allow(clippy::too_many_arguments)]
324pub fn w8a16_gemv_silu_input(
325    gpu: &dyn GpuBackend,
326    kernel: KernelHandle,
327    gate_out: DevicePtr,
328    up_out: DevicePtr,
329    weight: DevicePtr,
330    block_scale: DevicePtr,
331    output: DevicePtr,
332    n: u32,
333    k: u32,
334    stream: u64,
335) -> Result<()> {
336    KernelLaunch::new(gpu, kernel)
337        .grid([div_ceil(n, 4), 1, 1])
338        .block([256, 1, 1])
339        .arg_ptr(gate_out)
340        .arg_ptr(up_out)
341        .arg_ptr(weight)
342        .arg_ptr(block_scale)
343        .arg_ptr(output)
344        .arg_u32(n)
345        .arg_u32(k)
346        .launch(stream)
347}
348
349/// 2026-09-25: `output[i] += sigmoid(*gate_ptr) * src[i]`, where `gate_ptr` is
350/// one BF16 scalar in device memory (`bf16_sigmoid_blend_device` in
351/// `kernels/gb10/common/residual_add.cu`).
352pub fn sigmoid_blend_device(
353    gpu: &dyn GpuBackend,
354    kernel: KernelHandle,
355    output: DevicePtr,
356    src: DevicePtr,
357    gate_ptr: DevicePtr,
358    num_elements: u32,
359    stream: u64,
360) -> Result<()> {
361    KernelLaunch::new(gpu, kernel)
362        .grid([div_ceil(num_elements, 256), 1, 1])
363        .block([256, 1, 1])
364        .arg_ptr(output)
365        .arg_ptr(src)
366        .arg_ptr(gate_ptr)
367        .arg_u32(num_elements)
368        .launch(stream)
369}
metrale_model_layers/layers/ops/moe_prefill.rs

metrale_model_layers/layers/ops/
moe_prefill.rs
