metrale_model_layers/layers/ops/
moe_gate.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the MoE routers (top-k softmax, softmax with correction bias, sigmoid, sqrt-softplus, hash) and the zero-expert blend, plus the sigmoid router's compile-time bounds.
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
13
14use crate::layers::moe;
15use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
16
17use super::*;
18
19/// 2026-09-25: Largest `top_k` the sigmoid routing kernel holds: `#define MAX_TOP_K` in
20/// `kernels/gb10/common/moe_topk_sigmoid.cu`, which sizes its shared `s_top_vals` /
21/// `s_top_idxs`. The kernel clamps a larger `top_k` to it, and `NemotronMoeLayer` construction
22/// refuses one (`nemotron_moe.rs`). `crates/model-engine/tests/moe_topk_sigmoid_bounds.rs` pins
23/// this mirror to the define and fails if a directory other than `common/` holds a copy of the
24/// kernel.
25pub const MOE_TOPK_SIGMOID_MAX_TOP_K: usize = 32;
26
27/// 2026-09-25: Largest `num_experts` the sigmoid routing kernel holds, from `#define MAX_EXPERTS`
28/// in the same file. Beyond it the kernel considers only the first `MAX_EXPERTS` experts
29/// (`actual_n` is a `min`).
30pub const MOE_TOPK_SIGMOID_MAX_EXPERTS: usize = 512;
31
32/// 2026-09-25: Top-k experts of one token's gate logits, with softmax weights, in one block.
33/// The kernel handle decides the logit type (`moe_topk_softmax` BF16, `moe_topk_softmax_f32`).
34pub fn moe_topk_softmax(
35    gpu: &dyn GpuBackend,
36    kernel: KernelHandle,
37    gate_logits: DevicePtr,
38    expert_indices: DevicePtr,
39    expert_weights: DevicePtr,
40    num_experts: u32,
41    top_k: u32,
42    normalize: bool,
43    stream: u64,
44) -> Result<()> {
45    KernelLaunch::new(gpu, kernel)
46        .grid([1, 1, 1])
47        .block([256, 1, 1])
48        .arg_ptr(gate_logits)
49        .arg_ptr(expert_indices)
50        .arg_ptr(expert_weights)
51        .arg_u32(num_experts)
52        .arg_u32(top_k)
53        .arg_u32(if normalize { 1 } else { 0 })
54        .launch(stream)
55}
56
57/// 2026-09-25: Softmax router with a correction bias and the zero-expert fold, one token. Logits
58/// cover `num_logits` experts, of which ids `>= num_routed` are identity experts: a selected one
59/// adds its weight to `zero_accum` and its slot becomes expert 0 with weight 0.
60#[allow(clippy::too_many_arguments)]
61pub fn moe_topk_softmax_bias(
62    gpu: &dyn GpuBackend,
63    kernel: KernelHandle,
64    gate_logits: DevicePtr,
65    bias: DevicePtr,
66    expert_indices: DevicePtr,
67    expert_weights: DevicePtr,
68    zero_accum: DevicePtr,
69    num_logits: u32,
70    num_routed: u32,
71    top_k: u32,
72    normalize: bool,
73    scaling_factor: f32,
74    stream: u64,
75) -> Result<()> {
76    KernelLaunch::new(gpu, kernel)
77        .grid([1, 1, 1])
78        .block([256, 1, 1])
79        .arg_ptr(gate_logits)
80        .arg_ptr(bias)
81        .arg_ptr(expert_indices)
82        .arg_ptr(expert_weights)
83        .arg_ptr(zero_accum)
84        .arg_u32(num_logits)
85        .arg_u32(num_routed)
86        .arg_u32(top_k)
87        .arg_u32(normalize as u32)
88        .arg_f32(scaling_factor)
89        .launch(stream)
90}
91
92/// 2026-09-25: [`moe_topk_softmax_bias`] for `n` tokens, one block per token.
93#[allow(clippy::too_many_arguments)]
94pub fn moe_topk_softmax_bias_batched(
95    gpu: &dyn GpuBackend,
96    kernel: KernelHandle,
97    gate_logits: DevicePtr,
98    bias: DevicePtr,
99    expert_indices: DevicePtr,
100    expert_weights: DevicePtr,
101    zero_accum: DevicePtr,
102    num_logits: u32,
103    num_routed: u32,
104    top_k: u32,
105    normalize: bool,
106    scaling_factor: f32,
107    n: u32,
108    stream: u64,
109) -> Result<()> {
110    KernelLaunch::new(gpu, kernel)
111        .grid([n, 1, 1])
112        .block([256, 1, 1])
113        .arg_ptr(gate_logits)
114        .arg_ptr(bias)
115        .arg_ptr(expert_indices)
116        .arg_ptr(expert_weights)
117        .arg_ptr(zero_accum)
118        .arg_u32(num_logits)
119        .arg_u32(num_routed)
120        .arg_u32(top_k)
121        .arg_u32(normalize as u32)
122        .arg_f32(scaling_factor)
123        .launch(stream)
124}
125
126/// 2026-09-25: `out[t, :] += zero_accum[t] * x[t, :]`, the identity-expert blend.
127pub fn moe_zero_expert_add(
128    gpu: &dyn GpuBackend,
129    kernel: KernelHandle,
130    out: DevicePtr,
131    x: DevicePtr,
132    zero_accum: DevicePtr,
133    n: u32,
134    h: u32,
135    stream: u64,
136) -> Result<()> {
137    use metrale_gpu_runtime::kernel_args::div_ceil;
138    KernelLaunch::new(gpu, kernel)
139        .grid([div_ceil(n * h, 256), 1, 1])
140        .block([256, 1, 1])
141        .arg_ptr(out)
142        .arg_ptr(x)
143        .arg_ptr(zero_accum)
144        .arg_u32(n)
145        .arg_u32(h)
146        .launch(stream)
147}
148
149/// 2026-09-25: Sigmoid router, one token. The bias is added to the scores for selection only;
150/// the weights come from the pre-bias sigmoid scores. `top_k` and `num_experts` are bounded by
151/// [`MOE_TOPK_SIGMOID_MAX_TOP_K`] and [`MOE_TOPK_SIGMOID_MAX_EXPERTS`].
152pub fn moe_topk_sigmoid(
153    gpu: &dyn GpuBackend,
154    kernel: KernelHandle,
155    gate_logits: DevicePtr,
156    bias: DevicePtr,
157    expert_indices: DevicePtr,
158    expert_weights: DevicePtr,
159    num_experts: u32,
160    top_k: u32,
161    normalize: bool,
162    scaling_factor: f32,
163    stream: u64,
164) -> Result<()> {
165    KernelLaunch::new(gpu, kernel)
166        .grid([1, 1, 1])
167        .block([256, 1, 1])
168        .arg_ptr(gate_logits)
169        .arg_ptr(bias)
170        .arg_ptr(expert_indices)
171        .arg_ptr(expert_weights)
172        .arg_u32(num_experts)
173        .arg_u32(top_k)
174        .arg_u32(if normalize { 1 } else { 0 })
175        .arg_f32(scaling_factor)
176        .launch(stream)
177}
178
179/// 2026-09-25: Sqrt-softplus router (`sqrt(log(1 + exp(logit)))`), one token. As in
180/// [`moe_topk_sigmoid`], the bias steers selection only and the weights come from the pre-bias
181/// scores.
182#[allow(clippy::too_many_arguments)]
183pub fn moe_topk_sqrtsoftplus(
184    gpu: &dyn GpuBackend,
185    kernel: KernelHandle,
186    gate_logits: DevicePtr,
187    bias: DevicePtr,
188    expert_indices: DevicePtr,
189    expert_weights: DevicePtr,
190    num_experts: u32,
191    top_k: u32,
192    normalize: bool,
193    scaling_factor: f32,
194    stream: u64,
195) -> Result<()> {
196    KernelLaunch::new(gpu, kernel)
197        .grid([1, 1, 1])
198        .block([256, 1, 1])
199        .arg_ptr(gate_logits)
200        .arg_ptr(bias)
201        .arg_ptr(expert_indices)
202        .arg_ptr(expert_weights)
203        .arg_u32(num_experts)
204        .arg_u32(top_k)
205        .arg_u32(if normalize { 1 } else { 0 })
206        .arg_f32(scaling_factor)
207        .launch(stream)
208}
209
210/// 2026-09-25: Hash router, one token: the experts are the static `tid2eid[token_id]` row, and
211/// the gate's sqrt-softplus scores weight them. The token id is read on the device from
212/// `token_id_ptr`.
213#[allow(clippy::too_many_arguments)]
214pub fn moe_hash_route(
215    gpu: &dyn GpuBackend,
216    kernel: KernelHandle,
217    gate_logits: DevicePtr,
218    tid2eid: DevicePtr,
219    token_id_ptr: DevicePtr,
220    expert_indices: DevicePtr,
221    expert_weights: DevicePtr,
222    num_experts: u32,
223    top_k: u32,
224    normalize: bool,
225    scaling_factor: f32,
226    stream: u64,
227) -> Result<()> {
228    KernelLaunch::new(gpu, kernel)
229        .grid([1, 1, 1])
230        .block([256, 1, 1])
231        .arg_ptr(gate_logits)
232        .arg_ptr(tid2eid)
233        .arg_ptr(token_id_ptr)
234        .arg_ptr(expert_indices)
235        .arg_ptr(expert_weights)
236        .arg_u32(num_experts)
237        .arg_u32(top_k)
238        .arg_u32(if normalize { 1 } else { 0 })
239        .arg_f32(scaling_factor)
240        .launch(stream)
241}
242
243/// 2026-09-25: [`moe_hash_route`] for `n` tokens, one block per token, reading `token_ids[n]`.
244#[allow(clippy::too_many_arguments)]
245pub fn moe_hash_route_batched(
246    gpu: &dyn GpuBackend,
247    kernel: KernelHandle,
248    gate_logits: DevicePtr,
249    tid2eid: DevicePtr,
250    token_ids: DevicePtr,
251    expert_indices: DevicePtr,
252    expert_weights: DevicePtr,
253    num_experts: u32,
254    top_k: u32,
255    normalize: bool,
256    scaling_factor: f32,
257    n: u32,
258    stream: u64,
259) -> Result<()> {
260    KernelLaunch::new(gpu, kernel)
261        .grid([n, 1, 1])
262        .block([256, 1, 1])
263        .arg_ptr(gate_logits)
264        .arg_ptr(tid2eid)
265        .arg_ptr(token_ids)
266        .arg_ptr(expert_indices)
267        .arg_ptr(expert_weights)
268        .arg_u32(num_experts)
269        .arg_u32(top_k)
270        .arg_u32(if normalize { 1 } else { 0 })
271        .arg_f32(scaling_factor)
272        .launch(stream)
273}
metrale_model_layers/layers/ops/moe_gate.rs

metrale_model_layers/layers/ops/
moe_gate.rs
