spark_model/layers/ops/
moe_expert_more.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 SiLU+down expert GEMV, wide variant (16 outputs/block for small K).
17///
18/// Same semantics as `moe_expert_gemv_silu_down` but 4x more outputs per block
19/// with sub-warp reduction. Optimal for K<=512 where the narrow kernel has
20/// insufficient inner loop iterations for memory latency hiding.
21///
22/// Fused weighted sum + sigmoid blend + gate scalar GEMV.
23///
24/// Computes gate_scalar = dot(input, gate_weight) inline, then:
25/// `output[j] = sum_e weights[e] * expert_out[e,j] + sigmoid(gate_scalar) * shared_out[j]`
26///
27/// Each block independently computes the gate scalar dot product (redundant but
28/// only 8KB per block for K=2048 — negligible). Eliminates the separate dense_gemv
29/// kernel for the shared expert gate scalar (saves 48 graph nodes).
30///
31/// Grid: (ceil(hidden/256), 1, 1)  Block: (256, 1, 1)
32#[allow(clippy::too_many_arguments)]
33pub fn moe_weighted_sum_blend(
34    gpu: &dyn GpuBackend,
35    kernel: KernelHandle,
36    output: DevicePtr,
37    expert_out: DevicePtr,
38    expert_weights: DevicePtr,
39    shared_out: DevicePtr,
40    input: DevicePtr,
41    gate_weight: DevicePtr,
42    hidden: u32,
43    top_k: u32,
44    k: u32,
45    stream: u64,
46) -> Result<()> {
47    KernelLaunch::new(gpu, kernel)
48        .grid([div_ceil(hidden, 256), 1, 1])
49        .block([256, 1, 1])
50        .arg_ptr(output)
51        .arg_ptr(expert_out)
52        .arg_ptr(expert_weights)
53        .arg_ptr(shared_out)
54        .arg_ptr(input)
55        .arg_ptr(gate_weight)
56        .arg_u32(hidden)
57        .arg_u32(top_k)
58        .arg_u32(k)
59        .launch(stream)
60}
61
62// ═══════════════════════════════════════════════════════════════════
63// K=2 batch MoE variants — process 2 tokens in single kernel launches
64// ═══════════════════════════════════════════════════════════════════
65
66/// Fused gate+up expert GEMV for K=2 tokens with shared expert.
67///
68/// Expands blockIdx.y from (top_k+1) to 2*(top_k+1) to process both tokens.
69/// Token index = blockIdx.y / (top_k+1), expert slot = blockIdx.y % (top_k+1).
70/// Shared expert blocks use direct weight pointers; routed use pointer table.
71///
72/// Grid: (ceil(N/8), 2*(top_k+1), 2)  Block: (128, 1, 1)
73#[allow(clippy::too_many_arguments)]
74pub fn moe_expert_gate_up_shared_batch2(
75    gpu: &dyn GpuBackend,
76    kernel: KernelHandle,
77    input: DevicePtr, // [2, H] BF16
78    gate_packed_ptrs: DevicePtr,
79    gate_scale_ptrs: DevicePtr,
80    gate_scale2_vals: DevicePtr,
81    gate_out: DevicePtr, // [2*top_k, inter] BF16
82    up_packed_ptrs: DevicePtr,
83    up_scale_ptrs: DevicePtr,
84    up_scale2_vals: DevicePtr,
85    up_out: DevicePtr,         // [2*top_k, inter] BF16
86    expert_indices: DevicePtr, // [2*top_k] u32
87    sh_gate: &QuantizedWeight,
88    sh_gate_out: DevicePtr, // [2, inter] BF16
89    sh_up: &QuantizedWeight,
90    sh_up_out: DevicePtr, // [2, inter] BF16
91    n: u32,
92    k: u32,
93    top_k: u32,
94    block_size: u32,
95    stream: u64,
96) -> Result<()> {
97    KernelLaunch::new(gpu, kernel)
98        .grid([div_ceil(n, 8), 2 * (top_k + 1), 2])
99        .block([block_size, 1, 1])
100        .arg_ptr(input)
101        .arg_ptr(gate_packed_ptrs)
102        .arg_ptr(gate_scale_ptrs)
103        .arg_ptr(gate_scale2_vals)
104        .arg_ptr(gate_out)
105        .arg_ptr(up_packed_ptrs)
106        .arg_ptr(up_scale_ptrs)
107        .arg_ptr(up_scale2_vals)
108        .arg_ptr(up_out)
109        .arg_ptr(expert_indices)
110        .arg_ptr(sh_gate.weight)
111        .arg_ptr(sh_gate.weight_scale)
112        .arg_f32(sh_gate.weight_scale_2)
113        .arg_ptr(sh_gate_out)
114        .arg_ptr(sh_up.weight)
115        .arg_ptr(sh_up.weight_scale)
116        .arg_f32(sh_up.weight_scale_2)
117        .arg_ptr(sh_up_out)
118        .arg_u32(n)
119        .arg_u32(k)
120        .arg_u32(top_k)
121        .launch(stream)
122}
123
124/// Fused SiLU+down expert GEMV for K=2 tokens with shared expert.
125///
126/// Grid: (ceil(N/8), 2*(top_k+1), 1)  Block: (block_size, 1, 1)
127#[allow(clippy::too_many_arguments)]
128pub fn moe_expert_silu_down_shared_batch2(
129    gpu: &dyn GpuBackend,
130    kernel: KernelHandle,
131    gate_out: DevicePtr, // [2*top_k, inter] BF16
132    up_out: DevicePtr,   // [2*top_k, inter] BF16
133    packed_ptrs: DevicePtr,
134    scale_ptrs: DevicePtr,
135    scale2_vals: DevicePtr,
136    output: DevicePtr,         // [2*top_k, H] BF16
137    expert_indices: DevicePtr, // [2*top_k] u32
138    sh_gate_in: DevicePtr,     // [2, inter] BF16
139    sh_up_in: DevicePtr,       // [2, inter] BF16
140    sh_down: &QuantizedWeight,
141    sh_down_out: DevicePtr, // [2, H] BF16
142    n: u32,
143    k: u32,
144    top_k: u32,
145    block_size: u32,
146    stream: u64,
147) -> Result<()> {
148    // s_act is extern shared: K floats (issue #85 -- static 1024 overflowed
149    // for Mistral-Small-4's expert_hidden_dim=2048).
150    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
151    KernelLaunch::new(gpu, kernel)
152        .grid([div_ceil(n, 8), 2 * (top_k + 1), 1])
153        .block([block_size, 1, 1])
154        .shared_mem(smem_bytes)
155        .arg_ptr(gate_out)
156        .arg_ptr(up_out)
157        .arg_ptr(packed_ptrs)
158        .arg_ptr(scale_ptrs)
159        .arg_ptr(scale2_vals)
160        .arg_ptr(output)
161        .arg_ptr(expert_indices)
162        .arg_ptr(sh_gate_in)
163        .arg_ptr(sh_up_in)
164        .arg_ptr(sh_down.weight)
165        .arg_ptr(sh_down.weight_scale)
166        .arg_f32(sh_down.weight_scale_2)
167        .arg_ptr(sh_down_out)
168        .arg_u32(n)
169        .arg_u32(k)
170        .arg_u32(top_k)
171        .launch(stream)
172}
173
174/// Fused weighted sum + sigmoid blend for K=2 tokens.
175///
176/// blockIdx.y = token index (0 or 1). Each block computes gate scalar
177/// independently from per-token input and shared gate weight.
178///
179/// Grid: (ceil(hidden/256), 2, 1)  Block: (256, 1, 1)
180#[allow(clippy::too_many_arguments)]
181pub fn moe_weighted_sum_blend_batch2(
182    gpu: &dyn GpuBackend,
183    kernel: KernelHandle,
184    output: DevicePtr,         // [2, hidden] BF16
185    expert_out: DevicePtr,     // [2*top_k, hidden] BF16
186    expert_weights: DevicePtr, // [2*top_k] f32
187    shared_out: DevicePtr,     // [2, hidden] BF16
188    input: DevicePtr,          // [2, K] BF16
189    gate_weight: DevicePtr,    // [1, K] BF16 (shared)
190    hidden: u32,
191    top_k: u32,
192    k: u32,
193    stream: u64,
194) -> Result<()> {
195    KernelLaunch::new(gpu, kernel)
196        .grid([div_ceil(hidden, 256), 2, 1])
197        .block([256, 1, 1])
198        .arg_ptr(output)
199        .arg_ptr(expert_out)
200        .arg_ptr(expert_weights)
201        .arg_ptr(shared_out)
202        .arg_ptr(input)
203        .arg_ptr(gate_weight)
204        .arg_u32(hidden)
205        .arg_u32(top_k)
206        .arg_u32(k)
207        .launch(stream)
208}
209
210/// Fused gate+up expert GEMV for K=3 tokens with shared expert.
211///
212/// Grid: (ceil(N/8), 3*(top_k+1), 2)  Block: (128, 1, 1)
213#[allow(clippy::too_many_arguments)]
214pub fn moe_expert_gate_up_shared_batch3(
215    gpu: &dyn GpuBackend,
216    kernel: KernelHandle,
217    input: DevicePtr, // [3, H] BF16
218    gate_packed_ptrs: DevicePtr,
219    gate_scale_ptrs: DevicePtr,
220    gate_scale2_vals: DevicePtr,
221    gate_out: DevicePtr, // [3*top_k, inter] BF16
222    up_packed_ptrs: DevicePtr,
223    up_scale_ptrs: DevicePtr,
224    up_scale2_vals: DevicePtr,
225    up_out: DevicePtr,         // [3*top_k, inter] BF16
226    expert_indices: DevicePtr, // [3*top_k] u32
227    sh_gate: &QuantizedWeight,
228    sh_gate_out: DevicePtr, // [3, inter] BF16
229    sh_up: &QuantizedWeight,
230    sh_up_out: DevicePtr, // [3, inter] BF16
231    n: u32,
232    k: u32,
233    top_k: u32,
234    stream: u64,
235) -> Result<()> {
236    KernelLaunch::new(gpu, kernel)
237        .grid([div_ceil(n, 8), 3 * (top_k + 1), 2])
238        .block([128, 1, 1])
239        .arg_ptr(input)
240        .arg_ptr(gate_packed_ptrs)
241        .arg_ptr(gate_scale_ptrs)
242        .arg_ptr(gate_scale2_vals)
243        .arg_ptr(gate_out)
244        .arg_ptr(up_packed_ptrs)
245        .arg_ptr(up_scale_ptrs)
246        .arg_ptr(up_scale2_vals)
247        .arg_ptr(up_out)
248        .arg_ptr(expert_indices)
249        .arg_ptr(sh_gate.weight)
250        .arg_ptr(sh_gate.weight_scale)
251        .arg_f32(sh_gate.weight_scale_2)
252        .arg_ptr(sh_gate_out)
253        .arg_ptr(sh_up.weight)
254        .arg_ptr(sh_up.weight_scale)
255        .arg_f32(sh_up.weight_scale_2)
256        .arg_ptr(sh_up_out)
257        .arg_u32(n)
258        .arg_u32(k)
259        .arg_u32(top_k)
260        .launch(stream)
261}
262
263/// Fused SiLU+down expert GEMV for K=3 tokens with shared expert.
264///
265/// Grid: (ceil(N/8), 3*(top_k+1), 1)  Block: (128, 1, 1)
266#[allow(clippy::too_many_arguments)]
267pub fn moe_expert_silu_down_shared_batch3(
268    gpu: &dyn GpuBackend,
269    kernel: KernelHandle,
270    gate_out: DevicePtr, // [3*top_k, inter] BF16
271    up_out: DevicePtr,   // [3*top_k, inter] BF16
272    packed_ptrs: DevicePtr,
273    scale_ptrs: DevicePtr,
274    scale2_vals: DevicePtr,
275    output: DevicePtr,         // [3*top_k, H] BF16
276    expert_indices: DevicePtr, // [3*top_k] u32
277    sh_gate_in: DevicePtr,     // [3, inter] BF16
278    sh_up_in: DevicePtr,       // [3, inter] BF16
279    sh_down: &QuantizedWeight,
280    sh_down_out: DevicePtr, // [3, H] BF16
281    n: u32,
282    k: u32,
283    top_k: u32,
284    stream: u64,
285) -> Result<()> {
286    // s_act is extern shared: K floats (issue #85 -- static 1024 overflowed
287    // for expert inter dims > 1024).
288    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
289    KernelLaunch::new(gpu, kernel)
290        .grid([div_ceil(n, 8), 3 * (top_k + 1), 1])
291        .block([128, 1, 1])
292        .shared_mem(smem_bytes)
293        .arg_ptr(gate_out)
294        .arg_ptr(up_out)
295        .arg_ptr(packed_ptrs)
296        .arg_ptr(scale_ptrs)
297        .arg_ptr(scale2_vals)
298        .arg_ptr(output)
299        .arg_ptr(expert_indices)
300        .arg_ptr(sh_gate_in)
301        .arg_ptr(sh_up_in)
302        .arg_ptr(sh_down.weight)
303        .arg_ptr(sh_down.weight_scale)
304        .arg_f32(sh_down.weight_scale_2)
305        .arg_ptr(sh_down_out)
306        .arg_u32(n)
307        .arg_u32(k)
308        .arg_u32(top_k)
309        .launch(stream)
310}
311
312/// Fused weighted sum + sigmoid blend for K=3 tokens.
313///
314/// Grid: (ceil(hidden/256), 3, 1)  Block: (256, 1, 1)
315#[allow(clippy::too_many_arguments)]
316pub fn moe_weighted_sum_blend_batch3(
317    gpu: &dyn GpuBackend,
318    kernel: KernelHandle,
319    output: DevicePtr,         // [3, hidden] BF16
320    expert_out: DevicePtr,     // [3*top_k, hidden] BF16
321    expert_weights: DevicePtr, // [3*top_k] f32
322    shared_out: DevicePtr,     // [3, hidden] BF16
323    input: DevicePtr,          // [3, K] BF16
324    gate_weight: DevicePtr,    // [1, K] BF16 (shared)
325    hidden: u32,
326    top_k: u32,
327    k: u32,
328    stream: u64,
329) -> Result<()> {
330    KernelLaunch::new(gpu, kernel)
331        .grid([div_ceil(hidden, 256), 3, 1])
332        .block([256, 1, 1])
333        .arg_ptr(output)
334        .arg_ptr(expert_out)
335        .arg_ptr(expert_weights)
336        .arg_ptr(shared_out)
337        .arg_ptr(input)
338        .arg_ptr(gate_weight)
339        .arg_u32(hidden)
340        .arg_u32(top_k)
341        .arg_u32(k)
342        .launch(stream)
343}
344
345// ── MoE prefill (N-token batch) ──────────────────────────────────
spark_model/layers/ops/moe_expert_more.rs

spark_model/layers/ops/
moe_expert_more.rs
