metrale_model_layers/layers/ops/
fp8_moe_batch_a.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the transposed-weight decode MoE kernels: NVFP4
4//! for two and three tokens, and FP8 for one token.
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::layers::moe;
16use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
17
18use super::*;
19
20/// 2026-09-25: NVFP4 fused gate+up GEMV over transposed weights, for two tokens.
21#[allow(clippy::too_many_arguments)]
22pub fn moe_expert_gate_up_shared_batch2_t(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    input: DevicePtr,
26    gate_packed_t_ptrs: DevicePtr,
27    gate_scale_t_ptrs: DevicePtr,
28    gate_scale2_vals: DevicePtr,
29    gate_out: DevicePtr,
30    up_packed_t_ptrs: DevicePtr,
31    up_scale_t_ptrs: DevicePtr,
32    up_scale2_vals: DevicePtr,
33    up_out: DevicePtr,
34    expert_indices: DevicePtr,
35    sh_gate_t: &QuantizedWeight,
36    sh_gate_out: DevicePtr,
37    sh_up_t: &QuantizedWeight,
38    sh_up_out: DevicePtr,
39    n: u32,
40    k: u32,
41    top_k: u32,
42    stream: u64,
43) -> Result<()> {
44    KernelLaunch::new(gpu, kernel)
45        .grid([div_ceil(n, T_BLOCK), 2 * (top_k + 1), 2])
46        .block([T_BLOCK, 1, 1])
47        .arg_ptr(input)
48        .arg_ptr(gate_packed_t_ptrs)
49        .arg_ptr(gate_scale_t_ptrs)
50        .arg_ptr(gate_scale2_vals)
51        .arg_ptr(gate_out)
52        .arg_ptr(up_packed_t_ptrs)
53        .arg_ptr(up_scale_t_ptrs)
54        .arg_ptr(up_scale2_vals)
55        .arg_ptr(up_out)
56        .arg_ptr(expert_indices)
57        .arg_ptr(sh_gate_t.weight)
58        .arg_ptr(sh_gate_t.weight_scale)
59        .arg_f32(sh_gate_t.weight_scale_2)
60        .arg_ptr(sh_gate_out)
61        .arg_ptr(sh_up_t.weight)
62        .arg_ptr(sh_up_t.weight_scale)
63        .arg_f32(sh_up_t.weight_scale_2)
64        .arg_ptr(sh_up_out)
65        .arg_u32(n)
66        .arg_u32(k)
67        .arg_u32(top_k)
68        .launch(stream)
69}
70
71/// 2026-09-25: NVFP4 fused SiLU+down GEMV over transposed weights, for two tokens.
72#[allow(clippy::too_many_arguments)]
73pub fn moe_expert_silu_down_shared_batch2_t(
74    gpu: &dyn GpuBackend,
75    kernel: KernelHandle,
76    gate_out: DevicePtr,
77    up_out: DevicePtr,
78    packed_t_ptrs: DevicePtr,
79    scale_t_ptrs: DevicePtr,
80    scale2_vals: DevicePtr,
81    output: DevicePtr,
82    expert_indices: DevicePtr,
83    sh_gate_in: DevicePtr,
84    sh_up_in: DevicePtr,
85    sh_down_t: &QuantizedWeight,
86    sh_down_out: DevicePtr,
87    n: u32,
88    k: u32,
89    top_k: u32,
90    stream: u64,
91) -> Result<()> {
92    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
93    KernelLaunch::new(gpu, kernel)
94        .grid([div_ceil(n, T_BLOCK), 2 * (top_k + 1), 1])
95        .block([T_BLOCK, 1, 1])
96        .shared_mem(smem_bytes)
97        .arg_ptr(gate_out)
98        .arg_ptr(up_out)
99        .arg_ptr(packed_t_ptrs)
100        .arg_ptr(scale_t_ptrs)
101        .arg_ptr(scale2_vals)
102        .arg_ptr(output)
103        .arg_ptr(expert_indices)
104        .arg_ptr(sh_gate_in)
105        .arg_ptr(sh_up_in)
106        .arg_ptr(sh_down_t.weight)
107        .arg_ptr(sh_down_t.weight_scale)
108        .arg_f32(sh_down_t.weight_scale_2)
109        .arg_ptr(sh_down_out)
110        .arg_u32(n)
111        .arg_u32(k)
112        .arg_u32(top_k)
113        .launch(stream)
114}
115
116/// 2026-09-25: NVFP4 fused gate+up GEMV over transposed weights, for three tokens.
117#[allow(clippy::too_many_arguments)]
118pub fn moe_expert_gate_up_shared_batch3_t(
119    gpu: &dyn GpuBackend,
120    kernel: KernelHandle,
121    input: DevicePtr,
122    gate_packed_t_ptrs: DevicePtr,
123    gate_scale_t_ptrs: DevicePtr,
124    gate_scale2_vals: DevicePtr,
125    gate_out: DevicePtr,
126    up_packed_t_ptrs: DevicePtr,
127    up_scale_t_ptrs: DevicePtr,
128    up_scale2_vals: DevicePtr,
129    up_out: DevicePtr,
130    expert_indices: DevicePtr,
131    sh_gate_t: &QuantizedWeight,
132    sh_gate_out: DevicePtr,
133    sh_up_t: &QuantizedWeight,
134    sh_up_out: DevicePtr,
135    n: u32,
136    k: u32,
137    top_k: u32,
138    stream: u64,
139) -> Result<()> {
140    KernelLaunch::new(gpu, kernel)
141        .grid([div_ceil(n, T_BLOCK), 3 * (top_k + 1), 2])
142        .block([T_BLOCK, 1, 1])
143        .arg_ptr(input)
144        .arg_ptr(gate_packed_t_ptrs)
145        .arg_ptr(gate_scale_t_ptrs)
146        .arg_ptr(gate_scale2_vals)
147        .arg_ptr(gate_out)
148        .arg_ptr(up_packed_t_ptrs)
149        .arg_ptr(up_scale_t_ptrs)
150        .arg_ptr(up_scale2_vals)
151        .arg_ptr(up_out)
152        .arg_ptr(expert_indices)
153        .arg_ptr(sh_gate_t.weight)
154        .arg_ptr(sh_gate_t.weight_scale)
155        .arg_f32(sh_gate_t.weight_scale_2)
156        .arg_ptr(sh_gate_out)
157        .arg_ptr(sh_up_t.weight)
158        .arg_ptr(sh_up_t.weight_scale)
159        .arg_f32(sh_up_t.weight_scale_2)
160        .arg_ptr(sh_up_out)
161        .arg_u32(n)
162        .arg_u32(k)
163        .arg_u32(top_k)
164        .launch(stream)
165}
166
167/// 2026-09-25: NVFP4 fused SiLU+down GEMV over transposed weights, for three tokens.
168#[allow(clippy::too_many_arguments)]
169pub fn moe_expert_silu_down_shared_batch3_t(
170    gpu: &dyn GpuBackend,
171    kernel: KernelHandle,
172    gate_out: DevicePtr,
173    up_out: DevicePtr,
174    packed_t_ptrs: DevicePtr,
175    scale_t_ptrs: DevicePtr,
176    scale2_vals: DevicePtr,
177    output: DevicePtr,
178    expert_indices: DevicePtr,
179    sh_gate_in: DevicePtr,
180    sh_up_in: DevicePtr,
181    sh_down_t: &QuantizedWeight,
182    sh_down_out: DevicePtr,
183    n: u32,
184    k: u32,
185    top_k: u32,
186    stream: u64,
187) -> Result<()> {
188    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
189    KernelLaunch::new(gpu, kernel)
190        .grid([div_ceil(n, T_BLOCK), 3 * (top_k + 1), 1])
191        .block([T_BLOCK, 1, 1])
192        .shared_mem(smem_bytes)
193        .arg_ptr(gate_out)
194        .arg_ptr(up_out)
195        .arg_ptr(packed_t_ptrs)
196        .arg_ptr(scale_t_ptrs)
197        .arg_ptr(scale2_vals)
198        .arg_ptr(output)
199        .arg_ptr(expert_indices)
200        .arg_ptr(sh_gate_in)
201        .arg_ptr(sh_up_in)
202        .arg_ptr(sh_down_t.weight)
203        .arg_ptr(sh_down_t.weight_scale)
204        .arg_f32(sh_down_t.weight_scale_2)
205        .arg_ptr(sh_down_out)
206        .arg_u32(n)
207        .arg_u32(k)
208        .arg_u32(top_k)
209        .launch(stream)
210}
211
212/// 2026-09-25: FP8 fused gate+up GEMV over transposed weights, for one token.
213#[allow(clippy::too_many_arguments)]
214pub fn moe_expert_gate_up_shared_fp8_t(
215    gpu: &dyn GpuBackend,
216    kernel: KernelHandle,
217    input: DevicePtr,
218    gate_weight_t_ptrs: DevicePtr,
219    gate_block_scale_t_ptrs: DevicePtr,
220    gate_out: DevicePtr,
221    up_weight_t_ptrs: DevicePtr,
222    up_block_scale_t_ptrs: DevicePtr,
223    up_out: DevicePtr,
224    expert_indices: DevicePtr,
225    sh_gate_t: &Fp8Weight,
226    sh_gate_out: DevicePtr,
227    sh_up_t: &Fp8Weight,
228    sh_up_out: DevicePtr,
229    n: u32,
230    k: u32,
231    top_k: u32,
232    stream: u64,
233) -> Result<()> {
234    KernelLaunch::new(gpu, kernel)
235        .grid([div_ceil(n, T_BLOCK), top_k + 1, 2])
236        .block([T_BLOCK, 1, 1])
237        .arg_ptr(input)
238        .arg_ptr(gate_weight_t_ptrs)
239        .arg_ptr(gate_block_scale_t_ptrs)
240        .arg_ptr(gate_out)
241        .arg_ptr(up_weight_t_ptrs)
242        .arg_ptr(up_block_scale_t_ptrs)
243        .arg_ptr(up_out)
244        .arg_ptr(expert_indices)
245        .arg_ptr(sh_gate_t.weight)
246        .arg_ptr(sh_gate_t.row_scale)
247        .arg_ptr(sh_gate_out)
248        .arg_ptr(sh_up_t.weight)
249        .arg_ptr(sh_up_t.row_scale)
250        .arg_ptr(sh_up_out)
251        .arg_u32(n)
252        .arg_u32(k)
253        .arg_u32(top_k)
254        .launch(stream)
255}
256
257/// 2026-09-25: FP8 fused SiLU+down GEMV over transposed weights, for one token.
258#[allow(clippy::too_many_arguments)]
259pub fn moe_expert_silu_down_shared_fp8_t(
260    gpu: &dyn GpuBackend,
261    kernel: KernelHandle,
262    gate_out: DevicePtr,
263    up_out: DevicePtr,
264    weight_t_ptrs: DevicePtr,
265    block_scale_t_ptrs: DevicePtr,
266    output: DevicePtr,
267    expert_indices: DevicePtr,
268    sh_gate_in: DevicePtr,
269    sh_up_in: DevicePtr,
270    sh_down_t: &Fp8Weight,
271    sh_down_out: DevicePtr,
272    n: u32,
273    k: u32,
274    top_k: u32,
275    stream: u64,
276) -> Result<()> {
277    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
278    KernelLaunch::new(gpu, kernel)
279        .grid([div_ceil(n, T_BLOCK), top_k + 1, 1])
280        .block([T_BLOCK, 1, 1])
281        .shared_mem(smem_bytes)
282        .arg_ptr(gate_out)
283        .arg_ptr(up_out)
284        .arg_ptr(weight_t_ptrs)
285        .arg_ptr(block_scale_t_ptrs)
286        .arg_ptr(output)
287        .arg_ptr(expert_indices)
288        .arg_ptr(sh_gate_in)
289        .arg_ptr(sh_up_in)
290        .arg_ptr(sh_down_t.weight)
291        .arg_ptr(sh_down_t.row_scale)
292        .arg_ptr(sh_down_out)
293        .arg_u32(n)
294        .arg_u32(k)
295        .arg_u32(top_k)
296        .launch(stream)
297}
metrale_model_layers/layers/ops/fp8_moe_batch_a.rs

metrale_model_layers/layers/ops/
fp8_moe_batch_a.rs
