spark_model/layers/ops/moe_grouped_a/
topk.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Which experts, not the multiply.
4//!
5//! Split from `moe_grouped_a.rs` on the 500-line cap. The seam is the one the
6//! file already drew a banner for: everything here answers "which experts does
7//! this token go to", and what stays behind is the grouped GEMM that runs once
8//! that is known. The three variants differ only in the normaliser the router
9//! was trained with — softmax, sigmoid, sqrt-softplus — and choosing the wrong
10//! one is silent, so they are worth reading side by side.
11
12#![allow(unused_imports)]
13
14use anyhow::Result;
15use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
16use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
17
18use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
19
20use super::super::*;
21
22// ── Grouped MoE prefill ops ─────────────────────────────────────
23
24/// Batched top-K softmax: N tokens in parallel.
25///
26/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
27#[allow(clippy::too_many_arguments)]
28pub fn moe_topk_softmax_batched(
29    gpu: &dyn GpuBackend,
30    kernel: KernelHandle,
31    gate_logits: DevicePtr,
32    expert_indices: DevicePtr,
33    expert_weights: DevicePtr,
34    num_experts: u32,
35    top_k: u32,
36    normalize: bool,
37    num_tokens: u32,
38    stream: u64,
39) -> Result<()> {
40    KernelLaunch::new(gpu, kernel)
41        .grid([num_tokens, 1, 1])
42        .block([256, 1, 1])
43        .arg_ptr(gate_logits)
44        .arg_ptr(expert_indices)
45        .arg_ptr(expert_weights)
46        .arg_u32(num_experts)
47        .arg_u32(top_k)
48        .arg_u32(if normalize { 1 } else { 0 })
49        .launch(stream)
50}
51
52/// Batched sigmoid + correction-bias top-K MoE routing.
53///
54/// Kernel: `moe_topk_sigmoid_batched(gate_logits, bias, expert_indices,
55///         expert_weights, num_experts, top_k, normalize, scaling_factor)`
56/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
57#[allow(clippy::too_many_arguments)]
58pub fn moe_topk_sigmoid_batched(
59    gpu: &dyn GpuBackend,
60    kernel: KernelHandle,
61    gate_logits: DevicePtr,
62    bias: DevicePtr,
63    expert_indices: DevicePtr,
64    expert_weights: DevicePtr,
65    num_experts: u32,
66    top_k: u32,
67    normalize: bool,
68    scaling_factor: f32,
69    num_tokens: u32,
70    stream: u64,
71) -> Result<()> {
72    KernelLaunch::new(gpu, kernel)
73        .grid([num_tokens, 1, 1])
74        .block([256, 1, 1])
75        .arg_ptr(gate_logits)
76        .arg_ptr(bias)
77        .arg_ptr(expert_indices)
78        .arg_ptr(expert_weights)
79        .arg_u32(num_experts)
80        .arg_u32(top_k)
81        .arg_u32(if normalize { 1 } else { 0 })
82        .arg_f32(scaling_factor)
83        .launch(stream)
84}
85
86/// Batched sqrtsoftplus + correction-bias routing (DeepSeek-V4 prefill).
87///
88/// Same I/O as [`moe_topk_sigmoid_batched`] but scores experts with
89/// `sqrt(log(1+exp(logits)))` (matching the single-token decode path), so
90/// V4 prefill and decode route identically. Grid (N) / Block (256).
91#[allow(clippy::too_many_arguments)]
92pub fn moe_topk_sqrtsoftplus_batched(
93    gpu: &dyn GpuBackend,
94    kernel: KernelHandle,
95    gate_logits: DevicePtr,
96    bias: DevicePtr,
97    expert_indices: DevicePtr,
98    expert_weights: DevicePtr,
99    num_experts: u32,
100    top_k: u32,
101    normalize: bool,
102    scaling_factor: f32,
103    num_tokens: u32,
104    stream: u64,
105) -> Result<()> {
106    KernelLaunch::new(gpu, kernel)
107        .grid([num_tokens, 1, 1])
108        .block([256, 1, 1])
109        .arg_ptr(gate_logits)
110        .arg_ptr(bias)
111        .arg_ptr(expert_indices)
112        .arg_ptr(expert_weights)
113        .arg_u32(num_experts)
114        .arg_u32(top_k)
115        .arg_u32(if normalize { 1 } else { 0 })
116        .arg_f32(scaling_factor)
117        .launch(stream)
118}
spark_model/layers/ops/moe_grouped_a/topk.rs

spark_model/layers/ops/moe_grouped_a/
topk.rs
