spark_model/layers/ops/
fp8_moe_batch_a.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/// NVFP4 fused gate+up GEMV (transposed). K=2 batch.
17#[allow(clippy::too_many_arguments)]
18pub fn moe_expert_gate_up_shared_batch2_t(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    input: DevicePtr,
22    gate_packed_t_ptrs: DevicePtr,
23    gate_scale_t_ptrs: DevicePtr,
24    gate_scale2_vals: DevicePtr,
25    gate_out: DevicePtr,
26    up_packed_t_ptrs: DevicePtr,
27    up_scale_t_ptrs: DevicePtr,
28    up_scale2_vals: DevicePtr,
29    up_out: DevicePtr,
30    expert_indices: DevicePtr,
31    sh_gate_t: &QuantizedWeight,
32    sh_gate_out: DevicePtr,
33    sh_up_t: &QuantizedWeight,
34    sh_up_out: DevicePtr,
35    n: u32,
36    k: u32,
37    top_k: u32,
38    stream: u64,
39) -> Result<()> {
40    KernelLaunch::new(gpu, kernel)
41        .grid([div_ceil(n, T_BLOCK), 2 * (top_k + 1), 2])
42        .block([T_BLOCK, 1, 1])
43        .arg_ptr(input)
44        .arg_ptr(gate_packed_t_ptrs)
45        .arg_ptr(gate_scale_t_ptrs)
46        .arg_ptr(gate_scale2_vals)
47        .arg_ptr(gate_out)
48        .arg_ptr(up_packed_t_ptrs)
49        .arg_ptr(up_scale_t_ptrs)
50        .arg_ptr(up_scale2_vals)
51        .arg_ptr(up_out)
52        .arg_ptr(expert_indices)
53        .arg_ptr(sh_gate_t.weight)
54        .arg_ptr(sh_gate_t.weight_scale)
55        .arg_f32(sh_gate_t.weight_scale_2)
56        .arg_ptr(sh_gate_out)
57        .arg_ptr(sh_up_t.weight)
58        .arg_ptr(sh_up_t.weight_scale)
59        .arg_f32(sh_up_t.weight_scale_2)
60        .arg_ptr(sh_up_out)
61        .arg_u32(n)
62        .arg_u32(k)
63        .arg_u32(top_k)
64        .launch(stream)
65}
66
67/// NVFP4 fused SiLU+down GEMV (transposed). K=2 batch.
68#[allow(clippy::too_many_arguments)]
69pub fn moe_expert_silu_down_shared_batch2_t(
70    gpu: &dyn GpuBackend,
71    kernel: KernelHandle,
72    gate_out: DevicePtr,
73    up_out: DevicePtr,
74    packed_t_ptrs: DevicePtr,
75    scale_t_ptrs: DevicePtr,
76    scale2_vals: DevicePtr,
77    output: DevicePtr,
78    expert_indices: DevicePtr,
79    sh_gate_in: DevicePtr,
80    sh_up_in: DevicePtr,
81    sh_down_t: &QuantizedWeight,
82    sh_down_out: DevicePtr,
83    n: u32,
84    k: u32,
85    top_k: u32,
86    stream: u64,
87) -> Result<()> {
88    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
89    KernelLaunch::new(gpu, kernel)
90        .grid([div_ceil(n, T_BLOCK), 2 * (top_k + 1), 1])
91        .block([T_BLOCK, 1, 1])
92        .shared_mem(smem_bytes)
93        .arg_ptr(gate_out)
94        .arg_ptr(up_out)
95        .arg_ptr(packed_t_ptrs)
96        .arg_ptr(scale_t_ptrs)
97        .arg_ptr(scale2_vals)
98        .arg_ptr(output)
99        .arg_ptr(expert_indices)
100        .arg_ptr(sh_gate_in)
101        .arg_ptr(sh_up_in)
102        .arg_ptr(sh_down_t.weight)
103        .arg_ptr(sh_down_t.weight_scale)
104        .arg_f32(sh_down_t.weight_scale_2)
105        .arg_ptr(sh_down_out)
106        .arg_u32(n)
107        .arg_u32(k)
108        .arg_u32(top_k)
109        .launch(stream)
110}
111
112/// NVFP4 fused gate+up GEMV (transposed). K=3 batch.
113#[allow(clippy::too_many_arguments)]
114pub fn moe_expert_gate_up_shared_batch3_t(
115    gpu: &dyn GpuBackend,
116    kernel: KernelHandle,
117    input: DevicePtr,
118    gate_packed_t_ptrs: DevicePtr,
119    gate_scale_t_ptrs: DevicePtr,
120    gate_scale2_vals: DevicePtr,
121    gate_out: DevicePtr,
122    up_packed_t_ptrs: DevicePtr,
123    up_scale_t_ptrs: DevicePtr,
124    up_scale2_vals: DevicePtr,
125    up_out: DevicePtr,
126    expert_indices: DevicePtr,
127    sh_gate_t: &QuantizedWeight,
128    sh_gate_out: DevicePtr,
129    sh_up_t: &QuantizedWeight,
130    sh_up_out: DevicePtr,
131    n: u32,
132    k: u32,
133    top_k: u32,
134    stream: u64,
135) -> Result<()> {
136    KernelLaunch::new(gpu, kernel)
137        .grid([div_ceil(n, T_BLOCK), 3 * (top_k + 1), 2])
138        .block([T_BLOCK, 1, 1])
139        .arg_ptr(input)
140        .arg_ptr(gate_packed_t_ptrs)
141        .arg_ptr(gate_scale_t_ptrs)
142        .arg_ptr(gate_scale2_vals)
143        .arg_ptr(gate_out)
144        .arg_ptr(up_packed_t_ptrs)
145        .arg_ptr(up_scale_t_ptrs)
146        .arg_ptr(up_scale2_vals)
147        .arg_ptr(up_out)
148        .arg_ptr(expert_indices)
149        .arg_ptr(sh_gate_t.weight)
150        .arg_ptr(sh_gate_t.weight_scale)
151        .arg_f32(sh_gate_t.weight_scale_2)
152        .arg_ptr(sh_gate_out)
153        .arg_ptr(sh_up_t.weight)
154        .arg_ptr(sh_up_t.weight_scale)
155        .arg_f32(sh_up_t.weight_scale_2)
156        .arg_ptr(sh_up_out)
157        .arg_u32(n)
158        .arg_u32(k)
159        .arg_u32(top_k)
160        .launch(stream)
161}
162
163/// NVFP4 fused SiLU+down GEMV (transposed). K=3 batch.
164#[allow(clippy::too_many_arguments)]
165pub fn moe_expert_silu_down_shared_batch3_t(
166    gpu: &dyn GpuBackend,
167    kernel: KernelHandle,
168    gate_out: DevicePtr,
169    up_out: DevicePtr,
170    packed_t_ptrs: DevicePtr,
171    scale_t_ptrs: DevicePtr,
172    scale2_vals: DevicePtr,
173    output: DevicePtr,
174    expert_indices: DevicePtr,
175    sh_gate_in: DevicePtr,
176    sh_up_in: DevicePtr,
177    sh_down_t: &QuantizedWeight,
178    sh_down_out: DevicePtr,
179    n: u32,
180    k: u32,
181    top_k: u32,
182    stream: u64,
183) -> Result<()> {
184    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
185    KernelLaunch::new(gpu, kernel)
186        .grid([div_ceil(n, T_BLOCK), 3 * (top_k + 1), 1])
187        .block([T_BLOCK, 1, 1])
188        .shared_mem(smem_bytes)
189        .arg_ptr(gate_out)
190        .arg_ptr(up_out)
191        .arg_ptr(packed_t_ptrs)
192        .arg_ptr(scale_t_ptrs)
193        .arg_ptr(scale2_vals)
194        .arg_ptr(output)
195        .arg_ptr(expert_indices)
196        .arg_ptr(sh_gate_in)
197        .arg_ptr(sh_up_in)
198        .arg_ptr(sh_down_t.weight)
199        .arg_ptr(sh_down_t.weight_scale)
200        .arg_f32(sh_down_t.weight_scale_2)
201        .arg_ptr(sh_down_out)
202        .arg_u32(n)
203        .arg_u32(k)
204        .arg_u32(top_k)
205        .launch(stream)
206}
207
208/// FP8 fused gate+up GEMV (transposed weight). Single-token decode.
209#[allow(clippy::too_many_arguments)]
210pub fn moe_expert_gate_up_shared_fp8_t(
211    gpu: &dyn GpuBackend,
212    kernel: KernelHandle,
213    input: DevicePtr,
214    gate_weight_t_ptrs: DevicePtr,
215    gate_block_scale_t_ptrs: DevicePtr,
216    gate_out: DevicePtr,
217    up_weight_t_ptrs: DevicePtr,
218    up_block_scale_t_ptrs: DevicePtr,
219    up_out: DevicePtr,
220    expert_indices: DevicePtr,
221    sh_gate_t: &Fp8Weight,
222    sh_gate_out: DevicePtr,
223    sh_up_t: &Fp8Weight,
224    sh_up_out: DevicePtr,
225    n: u32,
226    k: u32,
227    top_k: u32,
228    stream: u64,
229) -> Result<()> {
230    KernelLaunch::new(gpu, kernel)
231        .grid([div_ceil(n, T_BLOCK), top_k + 1, 2])
232        .block([T_BLOCK, 1, 1])
233        .arg_ptr(input)
234        .arg_ptr(gate_weight_t_ptrs)
235        .arg_ptr(gate_block_scale_t_ptrs)
236        .arg_ptr(gate_out)
237        .arg_ptr(up_weight_t_ptrs)
238        .arg_ptr(up_block_scale_t_ptrs)
239        .arg_ptr(up_out)
240        .arg_ptr(expert_indices)
241        .arg_ptr(sh_gate_t.weight)
242        .arg_ptr(sh_gate_t.row_scale)
243        .arg_ptr(sh_gate_out)
244        .arg_ptr(sh_up_t.weight)
245        .arg_ptr(sh_up_t.row_scale)
246        .arg_ptr(sh_up_out)
247        .arg_u32(n)
248        .arg_u32(k)
249        .arg_u32(top_k)
250        .launch(stream)
251}
252
253/// FP8 fused SiLU+down GEMV (transposed weight). Single-token decode.
254#[allow(clippy::too_many_arguments)]
255pub fn moe_expert_silu_down_shared_fp8_t(
256    gpu: &dyn GpuBackend,
257    kernel: KernelHandle,
258    gate_out: DevicePtr,
259    up_out: DevicePtr,
260    weight_t_ptrs: DevicePtr,
261    block_scale_t_ptrs: DevicePtr,
262    output: DevicePtr,
263    expert_indices: DevicePtr,
264    sh_gate_in: DevicePtr,
265    sh_up_in: DevicePtr,
266    sh_down_t: &Fp8Weight,
267    sh_down_out: DevicePtr,
268    n: u32,
269    k: u32,
270    top_k: u32,
271    stream: u64,
272) -> Result<()> {
273    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
274    KernelLaunch::new(gpu, kernel)
275        .grid([div_ceil(n, T_BLOCK), top_k + 1, 1])
276        .block([T_BLOCK, 1, 1])
277        .shared_mem(smem_bytes)
278        .arg_ptr(gate_out)
279        .arg_ptr(up_out)
280        .arg_ptr(weight_t_ptrs)
281        .arg_ptr(block_scale_t_ptrs)
282        .arg_ptr(output)
283        .arg_ptr(expert_indices)
284        .arg_ptr(sh_gate_in)
285        .arg_ptr(sh_up_in)
286        .arg_ptr(sh_down_t.weight)
287        .arg_ptr(sh_down_t.row_scale)
288        .arg_ptr(sh_down_out)
289        .arg_u32(n)
290        .arg_u32(k)
291        .arg_u32(top_k)
292        .launch(stream)
293}
spark_model/layers/ops/fp8_moe_batch_a.rs

spark_model/layers/ops/
fp8_moe_batch_a.rs
