metrale_model_layers/layers/ops/
moe_grouped_a.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the grouped (per-expert) MoE prefill GEMMs over NVFP4 pointer tables, the token gather and `moe_silu_mul`; the batched routers are in `moe_grouped_a/topk.rs`.
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7//!
8//! 2026-09-29: The pointer-table GEMMs launch `(n tiles, max_m_tiles, num_experts)`, the N tile
9//! being the one the resolved kernel publishes (`n_tile_blocks`): `blockIdx.z` is the
10//! expert and `blockIdx.y` a row tile within its rows (`expert_offsets[e]..expert_offsets[e+1]`
11//! of the expert-sorted rows, `sorted_token_ids`). `max_m_tiles` must cover the busiest expert
12//! in the kernel's own M tile.
13
14#![allow(unused_imports)]
15
16use anyhow::Result;
17use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
18use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
19
20use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
21
22use super::*;
23
24// 2026-09-25: `ops.rs` loads this file through `#[path]`, so its child needs a `#[path]` too.
25#[path = "moe_grouped_a/topk.rs"]
26mod topk;
27// 2026-09-25: Re-exported so the routers are reached as `ops::moe_topk_*`.
28pub use topk::{moe_topk_sigmoid_batched, moe_topk_softmax_batched, moe_topk_sqrtsoftplus_batched};
29// 2026-09-30: The fused gate+up grouped GEMMs, split out to keep this file under the size cap.
30#[path = "moe_grouped_a/fused_gate_up.rs"]
31mod fused_gate_up;
32pub use fused_gate_up::{
33    moe_w4a16_fused_gate_up_k64_m128, moe_w4a16_fused_gate_up_k64_n128,
34    moe_w4a16_fused_gate_up_n128,
35};
36
37/// 2026-09-25: Launcher for `moe_w4a16_grouped_gemm`, whose weights are contiguous per expert
38/// rather than in pointer tables. No production code calls it (@human-review: its grid
39/// `(num_experts, 1, 1)` does not match the kernel, which takes the expert from `blockIdx.z`).
40pub fn moe_w4a16_grouped_gemm(
41    gpu: &dyn GpuBackend,
42    kernel: KernelHandle,
43    a: DevicePtr,
44    b_packed: DevicePtr,
45    b_scale: DevicePtr,
46    scale2: f32,
47    c: DevicePtr,
48    expert_offsets: DevicePtr,
49    num_experts: u32,
50    n: u32,
51    k: u32,
52    stream: u64,
53) -> Result<()> {
54    check_e4m3_grouped(
55        gpu,
56        kernel,
57        a,
58        expert_offsets,
59        DevicePtr::NULL,
60        num_experts,
61        k,
62        stream,
63    )?;
64    KernelLaunch::new(gpu, kernel)
65        .grid([num_experts, 1, 1])
66        .block([256, 1, 1])
67        .arg_ptr(a)
68        .arg_ptr(b_packed)
69        .arg_ptr(b_scale)
70        .arg_f32(scale2)
71        .arg_ptr(c)
72        .arg_ptr(expert_offsets)
73        .arg_u32(num_experts)
74        .arg_u32(n)
75        .arg_u32(k)
76        .launch(stream)
77}
78
79/// 2026-09-25: `moe_w4a16_grouped_gemm_ptrtable` with a 256-row M tile and a 512-thread block.
80/// `max_m_tiles` counts 256-row tiles: `MoeLayer::launch_grouped_gemm` divides the 64-row count
81/// by 4. Loaded only under `METRALE_MOE_GROUPED_M256=1` (`moe/init.rs`); the measurement behind
82/// that is on `launch_grouped_gemm`.
83#[allow(clippy::too_many_arguments)]
84pub fn moe_w4a16_grouped_gemm_ptrtable_m256(
85    gpu: &dyn GpuBackend,
86    kernel: KernelHandle,
87    a: DevicePtr,
88    b_packed_ptrs: DevicePtr,
89    b_scale_ptrs: DevicePtr,
90    scale2_vals: DevicePtr,
91    c: DevicePtr,
92    expert_offsets: DevicePtr,
93    sorted_token_ids: DevicePtr,
94    num_experts: u32,
95    n_out: u32,
96    k: u32,
97    max_m_tiles: u32,
98    stream: u64,
99) -> Result<()> {
100    check_e4m3_grouped(
101        gpu,
102        kernel,
103        a,
104        expert_offsets,
105        sorted_token_ids,
106        num_experts,
107        k,
108        stream,
109    )?;
110    KernelLaunch::new(gpu, kernel)
111        .grid([n_tile_blocks(gpu, kernel, n_out)?, max_m_tiles, num_experts])
112        .block([512, 1, 1])
113        .arg_ptr(a)
114        .arg_ptr(b_packed_ptrs)
115        .arg_ptr(b_scale_ptrs)
116        .arg_ptr(scale2_vals)
117        .arg_ptr(c)
118        .arg_ptr(expert_offsets)
119        .arg_ptr(sorted_token_ids)
120        .arg_u32(num_experts)
121        .arg_u32(n_out)
122        .arg_u32(k)
123        .launch(stream)
124}
125
126/// 2026-09-29: Pointer-table grouped W4A16 GEMM, 128 threads; one launch covers all
127/// experts.
128#[allow(clippy::too_many_arguments)]
129pub fn moe_w4a16_grouped_gemm_ptrtable(
130    gpu: &dyn GpuBackend,
131    kernel: KernelHandle,
132    a: DevicePtr,
133    b_packed_ptrs: DevicePtr,
134    b_scale_ptrs: DevicePtr,
135    scale2_vals: DevicePtr,
136    c: DevicePtr,
137    expert_offsets: DevicePtr,
138    sorted_token_ids: DevicePtr,
139    num_experts: u32,
140    n_out: u32,
141    k: u32,
142    max_m_tiles: u32,
143    stream: u64,
144) -> Result<()> {
145    check_e4m3_grouped(
146        gpu,
147        kernel,
148        a,
149        expert_offsets,
150        sorted_token_ids,
151        num_experts,
152        k,
153        stream,
154    )?;
155    KernelLaunch::new(gpu, kernel)
156        .grid([n_tile_blocks(gpu, kernel, n_out)?, max_m_tiles, num_experts])
157        .block([128, 1, 1])
158        .arg_ptr(a)
159        .arg_ptr(b_packed_ptrs)
160        .arg_ptr(b_scale_ptrs)
161        .arg_ptr(scale2_vals)
162        .arg_ptr(c)
163        .arg_ptr(expert_offsets)
164        .arg_ptr(sorted_token_ids)
165        .arg_u32(num_experts)
166        .arg_u32(n_out)
167        .arg_u32(k)
168        .launch(stream)
169}
170
171/// 2026-09-29: Pointer-table grouped W4A16 GEMM, 128 threads. The name is historical: the
172/// grid follows the kernel's published N tile, 64 or 128.
173#[allow(clippy::too_many_arguments)]
174pub fn moe_w4a16_grouped_gemm_ptrtable_n128(
175    gpu: &dyn GpuBackend,
176    kernel: KernelHandle,
177    a: DevicePtr,
178    b_packed_ptrs: DevicePtr,
179    b_scale_ptrs: DevicePtr,
180    scale2_vals: DevicePtr,
181    c: DevicePtr,
182    expert_offsets: DevicePtr,
183    sorted_token_ids: DevicePtr,
184    num_experts: u32,
185    n_out: u32,
186    k: u32,
187    max_m_tiles: u32,
188    stream: u64,
189) -> Result<()> {
190    check_e4m3_grouped(
191        gpu,
192        kernel,
193        a,
194        expert_offsets,
195        sorted_token_ids,
196        num_experts,
197        k,
198        stream,
199    )?;
200    KernelLaunch::new(gpu, kernel)
201        .grid([n_tile_blocks(gpu, kernel, n_out)?, max_m_tiles, num_experts])
202        .block([128, 1, 1])
203        .arg_ptr(a)
204        .arg_ptr(b_packed_ptrs)
205        .arg_ptr(b_scale_ptrs)
206        .arg_ptr(scale2_vals)
207        .arg_ptr(c)
208        .arg_ptr(expert_offsets)
209        .arg_ptr(sorted_token_ids)
210        .arg_u32(num_experts)
211        .arg_u32(n_out)
212        .arg_u32(k)
213        .launch(stream)
214}
215
216/// 2026-09-25: Pointer-table grouped GEMM with FP8 E4M3 activations and transposed NVFP4
217/// weights; `a_fp8` must already be FP8. Same launch shape as
218/// [`moe_w4a16_grouped_gemm_ptrtable_n128`].
219#[allow(clippy::too_many_arguments)]
220pub fn moe_fp8_grouped_gemm_ptrtable_n128(
221    gpu: &dyn GpuBackend,
222    kernel: KernelHandle,
223    a_fp8: DevicePtr,
224    b_packed_ptrs: DevicePtr,
225    b_scale_ptrs: DevicePtr,
226    scale2_vals: DevicePtr,
227    c: DevicePtr,
228    expert_offsets: DevicePtr,
229    sorted_token_ids: DevicePtr,
230    num_experts: u32,
231    n_out: u32,
232    k: u32,
233    max_m_tiles: u32,
234    stream: u64,
235) -> Result<()> {
236    KernelLaunch::new(gpu, kernel)
237        .grid([n_tile_blocks(gpu, kernel, n_out)?, max_m_tiles, num_experts])
238        .block([128, 1, 1])
239        .arg_ptr(a_fp8)
240        .arg_ptr(b_packed_ptrs)
241        .arg_ptr(b_scale_ptrs)
242        .arg_ptr(scale2_vals)
243        .arg_ptr(c)
244        .arg_ptr(expert_offsets)
245        .arg_ptr(sorted_token_ids)
246        .arg_u32(num_experts)
247        .arg_u32(n_out)
248        .arg_u32(k)
249        .launch(stream)
250}
251
252/// 2026-09-25: Same launch shape as [`moe_w4a16_grouped_gemm_ptrtable_n128`], for the kernels
253/// with a 64-wide K step.
254#[allow(clippy::too_many_arguments)]
255pub fn moe_w4a16_grouped_gemm_ptrtable_k64_n128(
256    gpu: &dyn GpuBackend,
257    kernel: KernelHandle,
258    a: DevicePtr,
259    b_packed_ptrs: DevicePtr,
260    b_scale_ptrs: DevicePtr,
261    scale2_vals: DevicePtr,
262    c: DevicePtr,
263    expert_offsets: DevicePtr,
264    sorted_token_ids: DevicePtr,
265    num_experts: u32,
266    n_out: u32,
267    k: u32,
268    max_m_tiles: u32,
269    stream: u64,
270) -> Result<()> {
271    check_e4m3_grouped(
272        gpu,
273        kernel,
274        a,
275        expert_offsets,
276        sorted_token_ids,
277        num_experts,
278        k,
279        stream,
280    )?;
281    KernelLaunch::new(gpu, kernel)
282        .grid([n_tile_blocks(gpu, kernel, n_out)?, max_m_tiles, num_experts])
283        .block([128, 1, 1])
284        .arg_ptr(a)
285        .arg_ptr(b_packed_ptrs)
286        .arg_ptr(b_scale_ptrs)
287        .arg_ptr(scale2_vals)
288        .arg_ptr(c)
289        .arg_ptr(expert_offsets)
290        .arg_ptr(sorted_token_ids)
291        .arg_u32(num_experts)
292        .arg_u32(n_out)
293        .arg_u32(k)
294        .launch(stream)
295}
296
297/// 2026-09-25: Gather token rows into expert-sorted order:
298/// `permuted[i] = hidden[sorted_token_ids[i]]`, `permuted` being `[total_expanded, hidden]`. One
299/// block per output row. No code calls it.
300#[allow(clippy::too_many_arguments, dead_code)]
301pub fn moe_permute_tokens(
302    gpu: &dyn GpuBackend,
303    kernel: KernelHandle,
304    hidden_states: DevicePtr,
305    permuted: DevicePtr,
306    sorted_token_ids: DevicePtr,
307    hidden: u32,
308    total_expanded: u32,
309    stream: u64,
310) -> Result<()> {
311    let threads = hidden.clamp(1, 256);
312    KernelLaunch::new(gpu, kernel)
313        .grid([total_expanded, 1, 1])
314        .block([threads, 1, 1])
315        .arg_ptr(hidden_states)
316        .arg_ptr(permuted)
317        .arg_ptr(sorted_token_ids)
318        .arg_u32(hidden)
319        .arg_u32(total_expanded)
320        .launch(stream)
321}
322
323/// 2026-09-25: `output[i] = silu(gate[i]) * up[i]` over `total_elements` BF16 values.
324pub fn moe_silu_mul(
325    gpu: &dyn GpuBackend,
326    kernel: KernelHandle,
327    gate: DevicePtr,
328    up: DevicePtr,
329    output: DevicePtr,
330    total_elements: u32,
331    stream: u64,
332) -> Result<()> {
333    KernelLaunch::new(gpu, kernel)
334        .grid([div_ceil(total_elements, 256), 1, 1])
335        .block([256, 1, 1])
336        .arg_ptr(gate)
337        .arg_ptr(up)
338        .arg_ptr(output)
339        .arg_u32(total_elements)
340        .launch(stream)
341}
metrale_model_layers/layers/ops/moe_grouped_a.rs

metrale_model_layers/layers/ops/
moe_grouped_a.rs
