metrale_model_layers/layers/ops/moe_grouped_a/
topk.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the batched MoE routers (softmax, sigmoid, sqrt-softplus): one block per token, each choosing that token's experts before the grouped GEMMs run.
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::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
15
16use super::super::*;
17
18/// 2026-09-25: The top-k softmax router for `num_tokens` tokens. The handle picks the
19/// tie-break: `moe_topk_softmax_batched` keeps the lower lane on an equal logit,
20/// `moe_topk_softmax_rows` (2026-09-26) the lower expert index, as `moe_topk_softmax` does.
21#[allow(clippy::too_many_arguments)]
22pub fn moe_topk_softmax_batched(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    gate_logits: DevicePtr,
26    expert_indices: DevicePtr,
27    expert_weights: DevicePtr,
28    num_experts: u32,
29    top_k: u32,
30    normalize: bool,
31    num_tokens: u32,
32    stream: u64,
33) -> Result<()> {
34    KernelLaunch::new(gpu, kernel)
35        .grid([num_tokens, 1, 1])
36        .block([256, 1, 1])
37        .arg_ptr(gate_logits)
38        .arg_ptr(expert_indices)
39        .arg_ptr(expert_weights)
40        .arg_u32(num_experts)
41        .arg_u32(top_k)
42        .arg_u32(if normalize { 1 } else { 0 })
43        .launch(stream)
44}
45
46/// 2026-09-25: The sigmoid router with a selection-only bias for `num_tokens` tokens.
47#[allow(clippy::too_many_arguments)]
48pub fn moe_topk_sigmoid_batched(
49    gpu: &dyn GpuBackend,
50    kernel: KernelHandle,
51    gate_logits: DevicePtr,
52    bias: DevicePtr,
53    expert_indices: DevicePtr,
54    expert_weights: DevicePtr,
55    num_experts: u32,
56    top_k: u32,
57    normalize: bool,
58    scaling_factor: f32,
59    num_tokens: u32,
60    stream: u64,
61) -> Result<()> {
62    KernelLaunch::new(gpu, kernel)
63        .grid([num_tokens, 1, 1])
64        .block([256, 1, 1])
65        .arg_ptr(gate_logits)
66        .arg_ptr(bias)
67        .arg_ptr(expert_indices)
68        .arg_ptr(expert_weights)
69        .arg_u32(num_experts)
70        .arg_u32(top_k)
71        .arg_u32(if normalize { 1 } else { 0 })
72        .arg_f32(scaling_factor)
73        .launch(stream)
74}
75
76/// 2026-09-25: Same arguments as [`moe_topk_sigmoid_batched`], scoring experts with
77/// `sqrt(log(1 + exp(logit)))`, the score of the single-token `moe_topk_sqrtsoftplus`.
78#[allow(clippy::too_many_arguments)]
79pub fn moe_topk_sqrtsoftplus_batched(
80    gpu: &dyn GpuBackend,
81    kernel: KernelHandle,
82    gate_logits: DevicePtr,
83    bias: DevicePtr,
84    expert_indices: DevicePtr,
85    expert_weights: DevicePtr,
86    num_experts: u32,
87    top_k: u32,
88    normalize: bool,
89    scaling_factor: f32,
90    num_tokens: u32,
91    stream: u64,
92) -> Result<()> {
93    KernelLaunch::new(gpu, kernel)
94        .grid([num_tokens, 1, 1])
95        .block([256, 1, 1])
96        .arg_ptr(gate_logits)
97        .arg_ptr(bias)
98        .arg_ptr(expert_indices)
99        .arg_ptr(expert_weights)
100        .arg_u32(num_experts)
101        .arg_u32(top_k)
102        .arg_u32(if normalize { 1 } else { 0 })
103        .arg_f32(scaling_factor)
104        .launch(stream)
105}
metrale_model_layers/layers/ops/moe_grouped_a/topk.rs

metrale_model_layers/layers/ops/moe_grouped_a/
topk.rs
