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

spark_model/layers/ops/
fp8_moe_batch_b.rs
