spark_model/layers/ops/
moe_gate.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// Largest `top_k` the sigmoid routing kernels can hold.
17///
18/// `moe_topk_sigmoid` stages the running top-K in `__shared__ float
19/// s_top_vals[MAX_TOP_K]` / `s_top_idxs[MAX_TOP_K]`, so a larger `top_k` walks
20/// off the end of those arrays and into the shared block that follows them.
21/// Nothing on the device can report that; the arrays are sized at compile time
22/// and the selection loop was bounded by the expert count, not by the array.
23///
24/// The authoritative value is `#define MAX_TOP_K` in
25/// `kernels/gb10/common/moe_topk_sigmoid.cu`. This mirror is pinned to it by
26/// `tests/moe_topk_sigmoid_bounds.rs`, which also fails if a model directory
27/// reintroduces a shadow copy of that kernel with a different cap — the drift
28/// that had the two Nemotron shadows capped at 24 against a common file of 32.
29pub const MOE_TOPK_SIGMOID_MAX_TOP_K: usize = 32;
30
31/// Largest `num_experts` the sigmoid routing kernels can hold, from
32/// `#define MAX_EXPERTS` in the same file. Beyond it the kernel silently
33/// considers only the first `MAX_EXPERTS` experts (`actual_n` is a `min`), so
34/// routing stays memory-safe but stops matching the checkpoint.
35pub const MOE_TOPK_SIGMOID_MAX_EXPERTS: usize = 512;
36
37/// GPU-side MoE top-K softmax.
38///
39/// Finds top-K experts from BF16 gate logits, computes softmax weights.
40///
41/// Kernel: `moe_topk_softmax(gate_logits, expert_indices, expert_weights,
42///          num_experts, top_k, normalize)`
43/// Grid: (1, 1, 1)  Block: (256, 1, 1)
44pub fn moe_topk_softmax(
45    gpu: &dyn GpuBackend,
46    kernel: KernelHandle,
47    gate_logits: DevicePtr,
48    expert_indices: DevicePtr,
49    expert_weights: DevicePtr,
50    num_experts: u32,
51    top_k: u32,
52    normalize: bool,
53    stream: u64,
54) -> Result<()> {
55    KernelLaunch::new(gpu, kernel)
56        .grid([1, 1, 1])
57        .block([256, 1, 1])
58        .arg_ptr(gate_logits)
59        .arg_ptr(expert_indices)
60        .arg_ptr(expert_weights)
61        .arg_u32(num_experts)
62        .arg_u32(top_k)
63        .arg_u32(if normalize { 1 } else { 0 })
64        .launch(stream)
65}
66
67/// GPU-side MoE top-K sigmoid routing (Nemotron-H).
68///
69/// Uses sigmoid scoring (not softmax). Bias affects expert selection only,
70/// not their weights. Weights come from pre-bias sigmoid scores.
71///
72/// Kernel: `moe_topk_sigmoid(gate_logits, bias, expert_indices, expert_weights,
73///          num_experts, top_k, normalize, scaling_factor)`
74/// Grid: (1, 1, 1)  Block: (256, 1, 1)
75/// LongCat softmax + e_score_correction_bias router with the zero-expert
76/// fold (single token). `zero_accum` receives the token's summed
77/// zero-expert weight; folded slots are rewritten (expert 0, weight 0).
78#[allow(clippy::too_many_arguments)]
79pub fn moe_topk_softmax_bias(
80    gpu: &dyn GpuBackend,
81    kernel: KernelHandle,
82    gate_logits: DevicePtr,
83    bias: DevicePtr,
84    expert_indices: DevicePtr,
85    expert_weights: DevicePtr,
86    zero_accum: DevicePtr,
87    num_logits: u32, // routed + zero
88    num_routed: u32,
89    top_k: u32,
90    normalize: bool,
91    scaling_factor: f32,
92    stream: u64,
93) -> Result<()> {
94    KernelLaunch::new(gpu, kernel)
95        .grid([1, 1, 1])
96        .block([256, 1, 1])
97        .arg_ptr(gate_logits)
98        .arg_ptr(bias)
99        .arg_ptr(expert_indices)
100        .arg_ptr(expert_weights)
101        .arg_ptr(zero_accum)
102        .arg_u32(num_logits)
103        .arg_u32(num_routed)
104        .arg_u32(top_k)
105        .arg_u32(normalize as u32)
106        .arg_f32(scaling_factor)
107        .launch(stream)
108}
109
110/// Batched twin: one block per token.
111#[allow(clippy::too_many_arguments)]
112pub fn moe_topk_softmax_bias_batched(
113    gpu: &dyn GpuBackend,
114    kernel: KernelHandle,
115    gate_logits: DevicePtr,
116    bias: DevicePtr,
117    expert_indices: DevicePtr,
118    expert_weights: DevicePtr,
119    zero_accum: DevicePtr,
120    num_logits: u32,
121    num_routed: u32,
122    top_k: u32,
123    normalize: bool,
124    scaling_factor: f32,
125    n: u32,
126    stream: u64,
127) -> Result<()> {
128    KernelLaunch::new(gpu, kernel)
129        .grid([n, 1, 1])
130        .block([256, 1, 1])
131        .arg_ptr(gate_logits)
132        .arg_ptr(bias)
133        .arg_ptr(expert_indices)
134        .arg_ptr(expert_weights)
135        .arg_ptr(zero_accum)
136        .arg_u32(num_logits)
137        .arg_u32(num_routed)
138        .arg_u32(top_k)
139        .arg_u32(normalize as u32)
140        .arg_f32(scaling_factor)
141        .launch(stream)
142}
143
144/// `out[t, :] += zero_accum[t] * x[t, :]` — the identity-expert blend.
145pub fn moe_zero_expert_add(
146    gpu: &dyn GpuBackend,
147    kernel: KernelHandle,
148    out: DevicePtr,
149    x: DevicePtr,
150    zero_accum: DevicePtr,
151    n: u32,
152    h: u32,
153    stream: u64,
154) -> Result<()> {
155    use spark_runtime::kernel_args::div_ceil;
156    KernelLaunch::new(gpu, kernel)
157        .grid([div_ceil(n * h, 256), 1, 1])
158        .block([256, 1, 1])
159        .arg_ptr(out)
160        .arg_ptr(x)
161        .arg_ptr(zero_accum)
162        .arg_u32(n)
163        .arg_u32(h)
164        .launch(stream)
165}
166
167pub fn moe_topk_sigmoid(
168    gpu: &dyn GpuBackend,
169    kernel: KernelHandle,
170    gate_logits: DevicePtr,
171    bias: DevicePtr,
172    expert_indices: DevicePtr,
173    expert_weights: DevicePtr,
174    num_experts: u32,
175    top_k: u32,
176    normalize: bool,
177    scaling_factor: f32,
178    stream: u64,
179) -> Result<()> {
180    KernelLaunch::new(gpu, kernel)
181        .grid([1, 1, 1])
182        .block([256, 1, 1])
183        .arg_ptr(gate_logits)
184        .arg_ptr(bias)
185        .arg_ptr(expert_indices)
186        .arg_ptr(expert_weights)
187        .arg_u32(num_experts)
188        .arg_u32(top_k)
189        .arg_u32(if normalize { 1 } else { 0 })
190        .arg_f32(scaling_factor)
191        .launch(stream)
192}
193
194/// GPU-side MoE top-K sqrtsoftplus routing (DeepSeek-V4).
195///
196/// Uses sqrtsoftplus scoring (not sigmoid/softmax). Bias affects expert
197/// selection only, not their weights. Weights come from pre-bias
198/// sqrtsoftplus scores.
199///
200/// Kernel: `moe_topk_sqrtsoftplus(gate_logits, bias, expert_indices, expert_weights,
201///          num_experts, top_k, normalize, scaling_factor)`
202/// Grid: (1, 1, 1)  Block: (256, 1, 1)
203#[allow(clippy::too_many_arguments)]
204pub fn moe_topk_sqrtsoftplus(
205    gpu: &dyn GpuBackend,
206    kernel: KernelHandle,
207    gate_logits: DevicePtr,
208    bias: DevicePtr,
209    expert_indices: DevicePtr,
210    expert_weights: DevicePtr,
211    num_experts: u32,
212    top_k: u32,
213    normalize: bool,
214    scaling_factor: f32,
215    stream: u64,
216) -> Result<()> {
217    KernelLaunch::new(gpu, kernel)
218        .grid([1, 1, 1])
219        .block([256, 1, 1])
220        .arg_ptr(gate_logits)
221        .arg_ptr(bias)
222        .arg_ptr(expert_indices)
223        .arg_ptr(expert_weights)
224        .arg_u32(num_experts)
225        .arg_u32(top_k)
226        .arg_u32(if normalize { 1 } else { 0 })
227        .arg_f32(scaling_factor)
228        .launch(stream)
229}
230
231/// GPU-side MoE hash routing (DeepSeek-V4 hash_moe layers).
232///
233/// Expert selection is a static `tid2eid[token_id]` lookup (frozen table);
234/// the learned gate still supplies the sqrtsoftplus scores that weight the
235/// selected experts. Mirrors [`moe_topk_sqrtsoftplus`] but with static
236/// selection instead of top-K.
237///
238/// Kernel: `moe_hash_route(gate_logits, tid2eid, token_id_ptr, expert_indices,
239///          expert_weights, num_experts, top_k, normalize, scaling_factor)`
240/// Grid: (1, 1, 1)  Block: (256, 1, 1)
241#[allow(clippy::too_many_arguments)]
242pub fn moe_hash_route(
243    gpu: &dyn GpuBackend,
244    kernel: KernelHandle,
245    gate_logits: DevicePtr,
246    tid2eid: DevicePtr,
247    token_id_ptr: DevicePtr,
248    expert_indices: DevicePtr,
249    expert_weights: DevicePtr,
250    num_experts: u32,
251    top_k: u32,
252    normalize: bool,
253    scaling_factor: f32,
254    stream: u64,
255) -> Result<()> {
256    KernelLaunch::new(gpu, kernel)
257        .grid([1, 1, 1])
258        .block([256, 1, 1])
259        .arg_ptr(gate_logits)
260        .arg_ptr(tid2eid)
261        .arg_ptr(token_id_ptr)
262        .arg_ptr(expert_indices)
263        .arg_ptr(expert_weights)
264        .arg_u32(num_experts)
265        .arg_u32(top_k)
266        .arg_u32(if normalize { 1 } else { 0 })
267        .arg_f32(scaling_factor)
268        .launch(stream)
269}
270
271/// Batched GPU-side MoE hash routing (DeepSeek-V4 hash_moe layers, prefill).
272///
273/// One block per token; reads `token_ids[N]` and the static `tid2eid` table.
274/// Grid: (N, 1, 1)  Block: (256, 1, 1)
275#[allow(clippy::too_many_arguments)]
276pub fn moe_hash_route_batched(
277    gpu: &dyn GpuBackend,
278    kernel: KernelHandle,
279    gate_logits: DevicePtr,
280    tid2eid: DevicePtr,
281    token_ids: DevicePtr,
282    expert_indices: DevicePtr,
283    expert_weights: DevicePtr,
284    num_experts: u32,
285    top_k: u32,
286    normalize: bool,
287    scaling_factor: f32,
288    n: u32,
289    stream: u64,
290) -> Result<()> {
291    KernelLaunch::new(gpu, kernel)
292        .grid([n, 1, 1])
293        .block([256, 1, 1])
294        .arg_ptr(gate_logits)
295        .arg_ptr(tid2eid)
296        .arg_ptr(token_ids)
297        .arg_ptr(expert_indices)
298        .arg_ptr(expert_weights)
299        .arg_u32(num_experts)
300        .arg_u32(top_k)
301        .arg_u32(if normalize { 1 } else { 0 })
302        .arg_f32(scaling_factor)
303        .launch(stream)
304}
305
306// ── Batched MoE Expert GEMV ──────────────────────────────────
spark_model/layers/ops/moe_gate.rs

spark_model/layers/ops/
moe_gate.rs
