spark_model/layers/ops/
moe_prefill.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/// Fused gate+up expert GEMV for N-token prefill with shared expert.
17///
18/// Grid: (ceil(inter/8), num_tokens*(top_k+1), 2)  Block: (128, 1, 1)
19#[allow(clippy::too_many_arguments)]
20pub fn moe_expert_gate_up_shared_prefill(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    input: DevicePtr, // [num_tokens, H] BF16
24    gate_packed_ptrs: DevicePtr,
25    gate_scale_ptrs: DevicePtr,
26    gate_scale2_vals: DevicePtr,
27    gate_out: DevicePtr, // [num_tokens*top_k, inter] BF16
28    up_packed_ptrs: DevicePtr,
29    up_scale_ptrs: DevicePtr,
30    up_scale2_vals: DevicePtr,
31    up_out: DevicePtr,         // [num_tokens*top_k, inter] BF16
32    expert_indices: DevicePtr, // [num_tokens*top_k] u32
33    sh_gate: &QuantizedWeight,
34    sh_gate_out: DevicePtr, // [num_tokens, inter] BF16
35    sh_up: &QuantizedWeight,
36    sh_up_out: DevicePtr, // [num_tokens, inter] BF16
37    n: u32,
38    k: u32,
39    top_k: u32,
40    num_tokens: u32,
41    stream: u64,
42) -> Result<()> {
43    KernelLaunch::new(gpu, kernel)
44        .grid([div_ceil(n, 8), num_tokens * (top_k + 1), 2])
45        .block([128, 1, 1])
46        .arg_ptr(input)
47        .arg_ptr(gate_packed_ptrs)
48        .arg_ptr(gate_scale_ptrs)
49        .arg_ptr(gate_scale2_vals)
50        .arg_ptr(gate_out)
51        .arg_ptr(up_packed_ptrs)
52        .arg_ptr(up_scale_ptrs)
53        .arg_ptr(up_scale2_vals)
54        .arg_ptr(up_out)
55        .arg_ptr(expert_indices)
56        .arg_ptr(sh_gate.weight)
57        .arg_ptr(sh_gate.weight_scale)
58        .arg_f32(sh_gate.weight_scale_2)
59        .arg_ptr(sh_gate_out)
60        .arg_ptr(sh_up.weight)
61        .arg_ptr(sh_up.weight_scale)
62        .arg_f32(sh_up.weight_scale_2)
63        .arg_ptr(sh_up_out)
64        .arg_u32(n)
65        .arg_u32(k)
66        .arg_u32(top_k)
67        .arg_u32(num_tokens)
68        .launch(stream)
69}
70
71/// Fused SiLU+down expert GEMV for N-token prefill with shared expert.
72///
73/// Grid: (ceil(hidden/64), num_tokens*(top_k+1), 1)  Block: (128, 1, 1)
74#[allow(clippy::too_many_arguments)]
75pub fn moe_expert_silu_down_shared_prefill(
76    gpu: &dyn GpuBackend,
77    kernel: KernelHandle,
78    gate_out: DevicePtr, // [num_tokens*top_k, inter] BF16
79    up_out: DevicePtr,   // [num_tokens*top_k, inter] BF16
80    packed_ptrs: DevicePtr,
81    scale_ptrs: DevicePtr,
82    scale2_vals: DevicePtr,
83    output: DevicePtr,         // [num_tokens*top_k, H] BF16
84    expert_indices: DevicePtr, // [num_tokens*top_k] u32
85    sh_gate_in: DevicePtr,     // [num_tokens, inter] BF16
86    sh_up_in: DevicePtr,       // [num_tokens, inter] BF16
87    sh_down: &QuantizedWeight,
88    sh_down_out: DevicePtr, // [num_tokens, H] BF16
89    n: u32,
90    k: u32,
91    top_k: u32,
92    num_tokens: u32,
93    stream: u64,
94) -> Result<()> {
95    KernelLaunch::new(gpu, kernel)
96        .grid([div_ceil(n, 8), num_tokens * (top_k + 1), 1])
97        .block([128, 1, 1])
98        .arg_ptr(gate_out)
99        .arg_ptr(up_out)
100        .arg_ptr(packed_ptrs)
101        .arg_ptr(scale_ptrs)
102        .arg_ptr(scale2_vals)
103        .arg_ptr(output)
104        .arg_ptr(expert_indices)
105        .arg_ptr(sh_gate_in)
106        .arg_ptr(sh_up_in)
107        .arg_ptr(sh_down.weight)
108        .arg_ptr(sh_down.weight_scale)
109        .arg_f32(sh_down.weight_scale_2)
110        .arg_ptr(sh_down_out)
111        .arg_u32(n)
112        .arg_u32(k)
113        .arg_u32(top_k)
114        .arg_u32(num_tokens)
115        .launch(stream)
116}
117
118/// Fused weighted sum + sigmoid blend for N-token prefill.
119///
120/// Grid: (ceil(hidden/256), num_tokens, 1)  Block: (256, 1, 1)
121#[allow(clippy::too_many_arguments)]
122pub fn moe_weighted_sum_blend_prefill(
123    gpu: &dyn GpuBackend,
124    kernel: KernelHandle,
125    output: DevicePtr,         // [num_tokens, hidden] BF16
126    expert_out: DevicePtr,     // [num_tokens*top_k, hidden] BF16
127    expert_weights: DevicePtr, // [num_tokens*top_k] f32
128    shared_out: DevicePtr,     // [num_tokens, hidden] BF16
129    input: DevicePtr,          // [num_tokens, K] BF16
130    gate_weight: DevicePtr,    // [1, K] BF16 (shared)
131    hidden: u32,
132    top_k: u32,
133    k: u32,
134    num_tokens: u32,
135    stream: u64,
136) -> Result<()> {
137    KernelLaunch::new(gpu, kernel)
138        .grid([div_ceil(hidden, 256), num_tokens, 1])
139        .block([256, 1, 1])
140        .arg_ptr(output)
141        .arg_ptr(expert_out)
142        .arg_ptr(expert_weights)
143        .arg_ptr(shared_out)
144        .arg_ptr(input)
145        .arg_ptr(gate_weight)
146        .arg_u32(hidden)
147        .arg_u32(top_k)
148        .arg_u32(k)
149        .arg_u32(num_tokens)
150        .launch(stream)
151}
152
153/// W4A16 dual GEMV: two projections sharing the same BF16 input, one launch.
154///
155/// blockIdx.z selects projection 0 vs 1. Both N dimensions must be equal.
156///
157/// Grid: (ceil(N/4), 1, 2)  Block: (256, 1, 1)
158#[allow(clippy::too_many_arguments)]
159pub fn w4a16_gemv_dual(
160    gpu: &dyn GpuBackend,
161    kernel: KernelHandle,
162    input: DevicePtr,
163    weight1: &QuantizedWeight,
164    output1: DevicePtr,
165    weight2: &QuantizedWeight,
166    output2: DevicePtr,
167    n: u32,
168    k: u32,
169    stream: u64,
170) -> Result<()> {
171    KernelLaunch::new(gpu, kernel)
172        .grid([w4a16_gemv_grid_x(n), 1, 2])
173        .block([256, 1, 1])
174        .arg_ptr(input)
175        .arg_ptr(weight1.weight)
176        .arg_ptr(weight1.weight_scale)
177        .arg_f32(weight1.weight_scale_2)
178        .arg_ptr(output1)
179        .arg_ptr(weight2.weight)
180        .arg_ptr(weight2.weight_scale)
181        .arg_f32(weight2.weight_scale_2)
182        .arg_ptr(output2)
183        .arg_u32(n)
184        .arg_u32(k)
185        .launch(stream)
186}
187
188/// Single-warp-per-output variant of `w4a16_gemv_dual` (8 outputs/block → N/8
189/// grid). Bit-identical output (see w4a16_gemv_fused.cu). Default ON via
190/// `ModelLevers::gemv_sw`; kill with `ATLAS_NO_GEMV_SW=1`.
191#[allow(clippy::too_many_arguments)]
192pub fn w4a16_gemv_dual_sw(
193    gpu: &dyn GpuBackend,
194    kernel: KernelHandle,
195    input: DevicePtr,
196    weight1: &QuantizedWeight,
197    output1: DevicePtr,
198    weight2: &QuantizedWeight,
199    output2: DevicePtr,
200    n: u32,
201    k: u32,
202    stream: u64,
203) -> Result<()> {
204    KernelLaunch::new(gpu, kernel)
205        .grid([w4a16_gemv_sw_grid_x(n), 1, 2])
206        .block([256, 1, 1])
207        .arg_ptr(input)
208        .arg_ptr(weight1.weight)
209        .arg_ptr(weight1.weight_scale)
210        .arg_f32(weight1.weight_scale_2)
211        .arg_ptr(output1)
212        .arg_ptr(weight2.weight)
213        .arg_ptr(weight2.weight_scale)
214        .arg_f32(weight2.weight_scale_2)
215        .arg_ptr(output2)
216        .arg_u32(n)
217        .arg_u32(k)
218        .launch(stream)
219}
220
221/// Single-warp-per-output variant of `w4a16_gemv_silu_input` (N/8 grid).
222/// Bit-identical. Default ON via `ModelLevers::gemv_sw`.
223#[allow(clippy::too_many_arguments)]
224pub fn w4a16_gemv_silu_input_sw(
225    gpu: &dyn GpuBackend,
226    kernel: KernelHandle,
227    gate_out: DevicePtr,
228    up_out: DevicePtr,
229    weight: &QuantizedWeight,
230    output: DevicePtr,
231    n: u32,
232    k: u32,
233    stream: u64,
234) -> Result<()> {
235    KernelLaunch::new(gpu, kernel)
236        .grid([w4a16_gemv_sw_grid_x(n), 1, 1])
237        .block([256, 1, 1])
238        .arg_ptr(gate_out)
239        .arg_ptr(up_out)
240        .arg_ptr(weight.weight)
241        .arg_ptr(weight.weight_scale)
242        .arg_f32(weight.weight_scale_2)
243        .arg_ptr(output)
244        .arg_u32(n)
245        .arg_u32(k)
246        .launch(stream)
247}
248
249/// W4A16 GEMV with fused SiLU input: silu(gate)*up as activation, GEMV with down weights.
250///
251/// Reads `gate_out[K]` and `up_out[K]` BF16, computes silu(gate)*up per element
252/// inline, then multiplies by dequanted NVFP4 weights. Eliminates silu_mul kernel.
253///
254/// Grid: (ceil(N/4), 1, 1)  Block: (256, 1, 1)
255#[allow(clippy::too_many_arguments)]
256pub fn w4a16_gemv_silu_input(
257    gpu: &dyn GpuBackend,
258    kernel: KernelHandle,
259    gate_out: DevicePtr,
260    up_out: DevicePtr,
261    weight: &QuantizedWeight,
262    output: DevicePtr,
263    n: u32,
264    k: u32,
265    stream: u64,
266) -> Result<()> {
267    KernelLaunch::new(gpu, kernel)
268        .grid([div_ceil(n, 4), 1, 1])
269        .block([256, 1, 1])
270        .arg_ptr(gate_out)
271        .arg_ptr(up_out)
272        .arg_ptr(weight.weight)
273        .arg_ptr(weight.weight_scale)
274        .arg_f32(weight.weight_scale_2)
275        .arg_ptr(output)
276        .arg_u32(n)
277        .arg_u32(k)
278        .launch(stream)
279}
280
281/// W8A16 (FP8 E4M3) dual GEMV: two projections sharing the same BF16 input,
282/// one launch. blockIdx.z selects projection 0 (gate) vs 1 (up). Both N must be
283/// equal. Mirrors `w4a16_gemv_dual` but takes RAW DevicePtrs (FP8 weights are
284/// `fp8w.weight` / `fp8w.row_scale`, no QuantizedWeight wrapper, no scale2 f32).
285///
286/// Grid: (ceil(N/4), 1, 2)  Block: (256, 1, 1)
287#[allow(clippy::too_many_arguments)]
288pub fn w8a16_gemv_dual(
289    gpu: &dyn GpuBackend,
290    kernel: KernelHandle,
291    input: DevicePtr,
292    weight1: DevicePtr,
293    row_scale1: DevicePtr,
294    output1: DevicePtr,
295    weight2: DevicePtr,
296    row_scale2: DevicePtr,
297    output2: DevicePtr,
298    n: u32,
299    k: u32,
300    stream: u64,
301) -> Result<()> {
302    KernelLaunch::new(gpu, kernel)
303        .grid([div_ceil(n, 4), 1, 2])
304        .block([256, 1, 1])
305        .arg_ptr(input)
306        .arg_ptr(weight1)
307        .arg_ptr(row_scale1)
308        .arg_ptr(output1)
309        .arg_ptr(weight2)
310        .arg_ptr(row_scale2)
311        .arg_ptr(output2)
312        .arg_u32(n)
313        .arg_u32(k)
314        .launch(stream)
315}
316
317/// W8A16 (FP8 E4M3) GEMV with fused SiLU input: silu(gate)*up as activation,
318/// GEMV with FP8 down weights. Reads `gate_out[K]` and `up_out[K]` BF16, computes
319/// silu(gate)*up per element inline, then multiplies by dequanted FP8 down
320/// weights. Eliminates the separate silu_mul kernel + down GEMV. Mirrors
321/// `w4a16_gemv_silu_input` but with RAW DevicePtrs (no scale2 f32).
322///
323/// Grid: (ceil(N/4), 1, 1)  Block: (256, 1, 1)
324#[allow(clippy::too_many_arguments)]
325pub fn w8a16_gemv_silu_input(
326    gpu: &dyn GpuBackend,
327    kernel: KernelHandle,
328    gate_out: DevicePtr,
329    up_out: DevicePtr,
330    weight: DevicePtr,
331    block_scale: DevicePtr,
332    output: DevicePtr,
333    n: u32,
334    k: u32,
335    stream: u64,
336) -> Result<()> {
337    KernelLaunch::new(gpu, kernel)
338        .grid([div_ceil(n, 4), 1, 1])
339        .block([256, 1, 1])
340        .arg_ptr(gate_out)
341        .arg_ptr(up_out)
342        .arg_ptr(weight)
343        .arg_ptr(block_scale)
344        .arg_ptr(output)
345        .arg_u32(n)
346        .arg_u32(k)
347        .launch(stream)
348}
349
350/// Sigmoid-gated blend reading gate scalar from device memory.
351///
352/// `output[i] += sigmoid(bf16_to_f32(*gate_ptr)) * src[i]`
353///
354/// Kernel: `bf16_sigmoid_blend_device(output, src, gate_ptr, n)`
355/// Grid: (ceil(n/256), 1, 1)  Block: (256, 1, 1)
356pub fn sigmoid_blend_device(
357    gpu: &dyn GpuBackend,
358    kernel: KernelHandle,
359    output: DevicePtr,
360    src: DevicePtr,
361    gate_ptr: DevicePtr,
362    num_elements: u32,
363    stream: u64,
364) -> Result<()> {
365    KernelLaunch::new(gpu, kernel)
366        .grid([div_ceil(num_elements, 256), 1, 1])
367        .block([256, 1, 1])
368        .arg_ptr(output)
369        .arg_ptr(src)
370        .arg_ptr(gate_ptr)
371        .arg_u32(num_elements)
372        .launch(stream)
373}
374
375// ── MoE grouped GEMM (future) ──────────────────────────────────
spark_model/layers/ops/moe_prefill.rs

spark_model/layers/ops/
moe_prefill.rs
