metrale_model_layers/layers/ops/
fp8_moe_batch_b.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the FP8 transposed-weight decode MoE kernels, for
4//! two and three tokens.
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::layers::moe;
16use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
17
18use super::*;
19
20/// 2026-09-25: FP8 fused gate+up GEMV over transposed weights, for two tokens.
21#[allow(clippy::too_many_arguments)]
22pub fn moe_expert_gate_up_shared_fp8_batch2_t(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    input: DevicePtr,
26    gate_weight_t_ptrs: DevicePtr,
27    gate_block_scale_t_ptrs: DevicePtr,
28    gate_out: DevicePtr,
29    up_weight_t_ptrs: DevicePtr,
30    up_block_scale_t_ptrs: DevicePtr,
31    up_out: DevicePtr,
32    expert_indices: DevicePtr,
33    sh_gate_t: &Fp8Weight,
34    sh_gate_out: DevicePtr,
35    sh_up_t: &Fp8Weight,
36    sh_up_out: DevicePtr,
37    n: u32,
38    k: u32,
39    top_k: u32,
40    stream: u64,
41) -> Result<()> {
42    KernelLaunch::new(gpu, kernel)
43        .grid([div_ceil(n, T_BLOCK), 2 * (top_k + 1), 2])
44        .block([T_BLOCK, 1, 1])
45        .arg_ptr(input)
46        .arg_ptr(gate_weight_t_ptrs)
47        .arg_ptr(gate_block_scale_t_ptrs)
48        .arg_ptr(gate_out)
49        .arg_ptr(up_weight_t_ptrs)
50        .arg_ptr(up_block_scale_t_ptrs)
51        .arg_ptr(up_out)
52        .arg_ptr(expert_indices)
53        .arg_ptr(sh_gate_t.weight)
54        .arg_ptr(sh_gate_t.row_scale)
55        .arg_ptr(sh_gate_out)
56        .arg_ptr(sh_up_t.weight)
57        .arg_ptr(sh_up_t.row_scale)
58        .arg_ptr(sh_up_out)
59        .arg_u32(n)
60        .arg_u32(k)
61        .arg_u32(top_k)
62        .launch(stream)
63}
64
65/// 2026-09-25: FP8 fused SiLU+down GEMV over transposed weights, for two tokens.
66#[allow(clippy::too_many_arguments)]
67pub fn moe_expert_silu_down_shared_fp8_batch2_t(
68    gpu: &dyn GpuBackend,
69    kernel: KernelHandle,
70    gate_out: DevicePtr,
71    up_out: DevicePtr,
72    weight_t_ptrs: DevicePtr,
73    block_scale_t_ptrs: DevicePtr,
74    output: DevicePtr,
75    expert_indices: DevicePtr,
76    sh_gate_in: DevicePtr,
77    sh_up_in: DevicePtr,
78    sh_down_t: &Fp8Weight,
79    sh_down_out: DevicePtr,
80    n: u32,
81    k: u32,
82    top_k: u32,
83    stream: u64,
84) -> Result<()> {
85    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
86    KernelLaunch::new(gpu, kernel)
87        .grid([div_ceil(n, T_BLOCK), 2 * (top_k + 1), 1])
88        .block([T_BLOCK, 1, 1])
89        .shared_mem(smem_bytes)
90        .arg_ptr(gate_out)
91        .arg_ptr(up_out)
92        .arg_ptr(weight_t_ptrs)
93        .arg_ptr(block_scale_t_ptrs)
94        .arg_ptr(output)
95        .arg_ptr(expert_indices)
96        .arg_ptr(sh_gate_in)
97        .arg_ptr(sh_up_in)
98        .arg_ptr(sh_down_t.weight)
99        .arg_ptr(sh_down_t.row_scale)
100        .arg_ptr(sh_down_out)
101        .arg_u32(n)
102        .arg_u32(k)
103        .arg_u32(top_k)
104        .launch(stream)
105}
106
107/// 2026-09-25: FP8 fused gate+up GEMV over transposed weights, for three tokens.
108#[allow(clippy::too_many_arguments)]
109pub fn moe_expert_gate_up_shared_fp8_batch3_t(
110    gpu: &dyn GpuBackend,
111    kernel: KernelHandle,
112    input: DevicePtr,
113    gate_weight_t_ptrs: DevicePtr,
114    gate_block_scale_t_ptrs: DevicePtr,
115    gate_out: DevicePtr,
116    up_weight_t_ptrs: DevicePtr,
117    up_block_scale_t_ptrs: DevicePtr,
118    up_out: DevicePtr,
119    expert_indices: DevicePtr,
120    sh_gate_t: &Fp8Weight,
121    sh_gate_out: DevicePtr,
122    sh_up_t: &Fp8Weight,
123    sh_up_out: DevicePtr,
124    n: u32,
125    k: u32,
126    top_k: u32,
127    stream: u64,
128) -> Result<()> {
129    KernelLaunch::new(gpu, kernel)
130        .grid([div_ceil(n, T_BLOCK), 3 * (top_k + 1), 2])
131        .block([T_BLOCK, 1, 1])
132        .arg_ptr(input)
133        .arg_ptr(gate_weight_t_ptrs)
134        .arg_ptr(gate_block_scale_t_ptrs)
135        .arg_ptr(gate_out)
136        .arg_ptr(up_weight_t_ptrs)
137        .arg_ptr(up_block_scale_t_ptrs)
138        .arg_ptr(up_out)
139        .arg_ptr(expert_indices)
140        .arg_ptr(sh_gate_t.weight)
141        .arg_ptr(sh_gate_t.row_scale)
142        .arg_ptr(sh_gate_out)
143        .arg_ptr(sh_up_t.weight)
144        .arg_ptr(sh_up_t.row_scale)
145        .arg_ptr(sh_up_out)
146        .arg_u32(n)
147        .arg_u32(k)
148        .arg_u32(top_k)
149        .launch(stream)
150}
151
152/// 2026-09-25: FP8 fused SiLU+down GEMV over transposed weights, for three tokens.
153#[allow(clippy::too_many_arguments)]
154pub fn moe_expert_silu_down_shared_fp8_batch3_t(
155    gpu: &dyn GpuBackend,
156    kernel: KernelHandle,
157    gate_out: DevicePtr,
158    up_out: DevicePtr,
159    weight_t_ptrs: DevicePtr,
160    block_scale_t_ptrs: DevicePtr,
161    output: DevicePtr,
162    expert_indices: DevicePtr,
163    sh_gate_in: DevicePtr,
164    sh_up_in: DevicePtr,
165    sh_down_t: &Fp8Weight,
166    sh_down_out: DevicePtr,
167    n: u32,
168    k: u32,
169    top_k: u32,
170    stream: u64,
171) -> Result<()> {
172    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
173    KernelLaunch::new(gpu, kernel)
174        .grid([div_ceil(n, T_BLOCK), 3 * (top_k + 1), 1])
175        .block([T_BLOCK, 1, 1])
176        .shared_mem(smem_bytes)
177        .arg_ptr(gate_out)
178        .arg_ptr(up_out)
179        .arg_ptr(weight_t_ptrs)
180        .arg_ptr(block_scale_t_ptrs)
181        .arg_ptr(output)
182        .arg_ptr(expert_indices)
183        .arg_ptr(sh_gate_in)
184        .arg_ptr(sh_up_in)
185        .arg_ptr(sh_down_t.weight)
186        .arg_ptr(sh_down_t.row_scale)
187        .arg_ptr(sh_down_out)
188        .arg_u32(n)
189        .arg_u32(k)
190        .arg_u32(top_k)
191        .launch(stream)
192}
metrale_model_layers/layers/ops/fp8_moe_batch_b.rs

metrale_model_layers/layers/ops/
fp8_moe_batch_b.rs
