metrale_model_layers/layers/ops/moe_grouped_a/
fused_gate_up.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: Launchers for the fused gate+up grouped MoE prefill GEMMs over NVFP4 pointer
4//! tables (moved out of `moe_grouped_a.rs` unchanged; the grid and the pointer-table layout are
5//! described there).
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use metrale_gpu_runtime::kernel_args::KernelLaunch;
13
14use super::super::*;
15
16/// 2026-09-25: [`moe_w4a16_fused_gate_up_n128`] for the kernels with a 64-wide K step.
17#[allow(clippy::too_many_arguments)]
18pub fn moe_w4a16_fused_gate_up_k64_n128(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    a: DevicePtr,
22    gate_packed_ptrs: DevicePtr,
23    gate_scale_ptrs: DevicePtr,
24    gate_scale2_vals: DevicePtr,
25    up_packed_ptrs: DevicePtr,
26    up_scale_ptrs: DevicePtr,
27    up_scale2_vals: DevicePtr,
28    c_gate: DevicePtr,
29    c_up: DevicePtr,
30    expert_offsets: DevicePtr,
31    sorted_token_ids: DevicePtr,
32    num_experts: u32,
33    n_out: u32,
34    k: u32,
35    max_m_tiles: u32,
36    stream: u64,
37) -> Result<()> {
38    check_e4m3_grouped(
39        gpu,
40        kernel,
41        a,
42        expert_offsets,
43        sorted_token_ids,
44        num_experts,
45        k,
46        stream,
47    )?;
48    KernelLaunch::new(gpu, kernel)
49        .grid([
50            n_tile_blocks(gpu, kernel, 2 * n_out)?,
51            max_m_tiles,
52            num_experts,
53        ])
54        .block([128, 1, 1])
55        .arg_ptr(a)
56        .arg_ptr(gate_packed_ptrs)
57        .arg_ptr(gate_scale_ptrs)
58        .arg_ptr(gate_scale2_vals)
59        .arg_ptr(up_packed_ptrs)
60        .arg_ptr(up_scale_ptrs)
61        .arg_ptr(up_scale2_vals)
62        .arg_ptr(c_gate)
63        .arg_ptr(c_up)
64        .arg_ptr(expert_offsets)
65        .arg_ptr(sorted_token_ids)
66        .arg_u32(num_experts)
67        .arg_u32(n_out)
68        .arg_u32(k)
69        .launch(stream)
70}
71
72/// 2026-09-25: [`moe_w4a16_fused_gate_up_k64_n128`] with a 128-row M tile and a 256-thread
73/// block. `max_m_tiles_m128` counts 128-row tiles: the call site halves the 64-row count
74/// (`forward_prefill_routed.rs`).
75#[allow(clippy::too_many_arguments)]
76pub fn moe_w4a16_fused_gate_up_k64_m128(
77    gpu: &dyn GpuBackend,
78    kernel: KernelHandle,
79    a: DevicePtr,
80    gate_packed_ptrs: DevicePtr,
81    gate_scale_ptrs: DevicePtr,
82    gate_scale2_vals: DevicePtr,
83    up_packed_ptrs: DevicePtr,
84    up_scale_ptrs: DevicePtr,
85    up_scale2_vals: DevicePtr,
86    c_gate: DevicePtr,
87    c_up: DevicePtr,
88    expert_offsets: DevicePtr,
89    sorted_token_ids: DevicePtr,
90    num_experts: u32,
91    n_out: u32,
92    k: u32,
93    max_m_tiles_m128: u32,
94    stream: u64,
95) -> Result<()> {
96    check_e4m3_grouped(
97        gpu,
98        kernel,
99        a,
100        expert_offsets,
101        sorted_token_ids,
102        num_experts,
103        k,
104        stream,
105    )?;
106    KernelLaunch::new(gpu, kernel)
107        .grid([
108            n_tile_blocks(gpu, kernel, 2 * n_out)?,
109            max_m_tiles_m128,
110            num_experts,
111        ])
112        .block([256, 1, 1])
113        .arg_ptr(a)
114        .arg_ptr(gate_packed_ptrs)
115        .arg_ptr(gate_scale_ptrs)
116        .arg_ptr(gate_scale2_vals)
117        .arg_ptr(up_packed_ptrs)
118        .arg_ptr(up_scale_ptrs)
119        .arg_ptr(up_scale2_vals)
120        .arg_ptr(c_gate)
121        .arg_ptr(c_up)
122        .arg_ptr(expert_offsets)
123        .arg_ptr(sorted_token_ids)
124        .arg_u32(num_experts)
125        .arg_u32(n_out)
126        .arg_u32(k)
127        .launch(stream)
128}
129
130/// 2026-09-25: Gate and up grouped GEMMs in one launch: the grid spans `2 * n_out` columns, the
131/// first `n_out` for gate and the rest for up.
132#[allow(clippy::too_many_arguments)]
133pub fn moe_w4a16_fused_gate_up_n128(
134    gpu: &dyn GpuBackend,
135    kernel: KernelHandle,
136    a: DevicePtr,
137    gate_packed_ptrs: DevicePtr,
138    gate_scale_ptrs: DevicePtr,
139    gate_scale2_vals: DevicePtr,
140    up_packed_ptrs: DevicePtr,
141    up_scale_ptrs: DevicePtr,
142    up_scale2_vals: DevicePtr,
143    c_gate: DevicePtr,
144    c_up: DevicePtr,
145    expert_offsets: DevicePtr,
146    sorted_token_ids: DevicePtr,
147    num_experts: u32,
148    n_out: u32,
149    k: u32,
150    max_m_tiles: u32,
151    stream: u64,
152) -> Result<()> {
153    check_e4m3_grouped(
154        gpu,
155        kernel,
156        a,
157        expert_offsets,
158        sorted_token_ids,
159        num_experts,
160        k,
161        stream,
162    )?;
163    KernelLaunch::new(gpu, kernel)
164        .grid([
165            n_tile_blocks(gpu, kernel, 2 * n_out)?,
166            max_m_tiles,
167            num_experts,
168        ])
169        .block([128, 1, 1])
170        .arg_ptr(a)
171        .arg_ptr(gate_packed_ptrs)
172        .arg_ptr(gate_scale_ptrs)
173        .arg_ptr(gate_scale2_vals)
174        .arg_ptr(up_packed_ptrs)
175        .arg_ptr(up_scale_ptrs)
176        .arg_ptr(up_scale2_vals)
177        .arg_ptr(c_gate)
178        .arg_ptr(c_up)
179        .arg_ptr(expert_offsets)
180        .arg_ptr(sorted_token_ids)
181        .arg_u32(num_experts)
182        .arg_u32(n_out)
183        .arg_u32(k)
184        .launch(stream)
185}
metrale_model_layers/layers/ops/moe_grouped_a/fused_gate_up.rs

metrale_model_layers/layers/ops/moe_grouped_a/
fused_gate_up.rs
