spark_model/layers/moe_v41/
forward.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! `MoeV41::forward`: the routed experts from the cache plus the shared
5//! expert, and the `launch_n` elementwise helper it runs on.
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::KernelLaunch;
10use spark_runtime::weights::expert_stream::{ExpertLru, ExpertSource};
11
12use super::{MoeV41, MoeV41LayerWeights, MoeV41Timing, RouterWeights, prefetch_k, shared_side};
13use crate::layers::ops::{
14    self, Q2K_MMQ_SMEM, Q3K_MMQ_SMEM, Q8_1_BLOCK_BYTES, kquant_mmq_gemm, kquant_mmvq_w,
15    kquant_q8_1_rows,
16};
17
18/// `ATLAS_DS41_PREFILL_GEMV=1`, read once: prefill groups of more than eight
19/// rows run through the decode GEMV arm in chunks of eight (see `forward`).
20fn m1_loop() -> bool {
21    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22    *V.get_or_init(|| std::env::var("ATLAS_DS41_M1_LOOP").is_ok_and(|v| v == "1"))
23}
24
25fn prefill_gemv() -> bool {
26    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
27    *V.get_or_init(|| std::env::var("ATLAS_DS41_PREFILL_GEMV").is_ok_and(|v| v == "1"))
28}
29
30impl MoeV41 {
31    pub(super) fn launch_n(
32        &self,
33        gpu: &dyn GpuBackend,
34        k: KernelHandle,
35        n: usize,
36        stream: u64,
37        f: impl FnOnce(KernelLaunch) -> KernelLaunch,
38    ) -> Result<()> {
39        if n == 0 {
40            return Ok(());
41        }
42        f(KernelLaunch::new(gpu, k)
43            .grid([(n as u32).div_ceil(256), 1, 1])
44            .block([256, 1, 1]))
45        .launch(stream)
46    }
47
48    /// One layer's MoE for `m` tokens: routed experts from the cache, plus the
49    /// shared expert. Returns the bf16 `[m, dim]` output and the routing.
50    #[allow(clippy::too_many_arguments)]
51    pub fn forward<S: ExpertSource + ?Sized>(
52        &self,
53        gpu: &dyn GpuBackend,
54        w: &MoeV41LayerWeights,
55        lru: &mut ExpertLru,
56        src: &S,
57        x: DevicePtr,
58        m: usize,
59        reader_threads: usize,
60        next: Option<&RouterWeights>,
61        stream: u64,
62    ) -> Result<(DevicePtr, Vec<f32>, Vec<usize>)> {
63        let c = &self.cfg;
64        let t0 = std::time::Instant::now();
65        // ATLAS_DS41_M1_LOOP=1 (diagnostic): run the single token through the
66        // per-expert loop below (one-row groups on the decode GEMV) instead of
67        // the pointer-table arm, to tell the two apart.
68        let m1_arm = m == 1 && !m1_loop();
69        // At decode the shared expert needs only `x`: it runs on the side
70        // stream from here, under the router launch, the read-back and the
71        // host's selection, and joins the main stream before the tail.
72        let side = m1_arm && shared_side();
73        // the token's q8_1 rows once, for the routed experts and both shared
74        // projections (phase 7 A2; the side stream reads them after `ev_in`)
75        if m1_arm {
76            self.ffn_input_q8_m1(gpu, x, stream)?;
77        }
78        if side {
79            gpu.record_event(self.ev_in, stream)?;
80        }
81        // the router, and at decode the next layer's router on the same
82        // input (its read-back rides the same drain), then the selection
83        self.route_launch(gpu, w, x, m, stream)?;
84        if side {
85            gpu.stream_wait_event(self.side, self.ev_in)?;
86            self.shared_expert_body(gpu, w, x, 1, self.a_q8, self.sh_q8, true, self.side)?;
87            gpu.record_event(self.ev_out, self.side)?;
88        }
89        // ATLAS_DS41_DEVICE_ROUTE=1: the whole selection, the plan and the
90        // pointer table on the device, one read-back of the header, the host
91        // fetching misses on its flag (`device_route.rs`)
92        let dev_route = m1_arm && self.device_route_ok(w);
93        let predict = next.filter(|_| m == 1 && !dev_route && prefetch_k() > 0 && lru.has_pool());
94        if let Some(nw) = predict {
95            self.predict_launch(gpu, nw, x, stream)?;
96        }
97        let (weights, indices, plan, slots, before, after, t1) = if dev_route {
98            let before = lru.stats();
99            let (weights, indices) =
100                self.route_device_m1(gpu, w, lru, src, reader_threads, stream)?;
101            let after = lru.stats();
102            let t1 = std::time::Instant::now();
103            (weights, indices, Vec::new(), Vec::new(), before, after, t1)
104        } else {
105            let (weights, indices) = self.route_select(gpu, w, m, stream)?;
106            let predicted = match predict {
107                Some(nw) => self.predict_select(gpu, nw, prefetch_k(), stream)?,
108                None => Vec::new(),
109            };
110            let t1 = std::time::Instant::now();
111            // this token batch's experts, gathered once; the predicted ones of
112            // the next layer start reading in the background
113            lru.begin_token();
114            let before = lru.stats();
115            let keys: Vec<(u32, u32)> = indices.iter().map(|&e| (w.layer, e as u32)).collect();
116            let slots = lru.fetch_many_on(gpu, stream, src, &keys, &predicted, reader_threads)?;
117            let after = lru.stats();
118            let (rows_host, w_host, plan) = self.plan(&indices, &weights, m);
119            gpu.copy_h2d_async(&rows_host, self.rows_dev, stream)?;
120            gpu.copy_h2d_async(&w_host, self.weight_dev, stream)?;
121            (weights, indices, plan, slots, before, after, t1)
122        };
123        let t2 = std::time::Instant::now();
124        gpu.memset_async(self.acc, 0, m * c.dim * 4, stream)?;
125        if m1_arm {
126            // the single-token arm: the token is every expert's activation, so
127            // quantise it once and run each projection as one launch over the
128            // experts (pointer table), the routing weight folded in at the
129            // SwiGLU and the expert rows summed into `acc` in plan order, the
130            // same order and the same per-row math as the loop below
131            let ne = if dev_route {
132                c.topk
133            } else {
134                self.upload_expert_table(gpu, &plan, &slots, stream)?
135            };
136            self.routed_m1(gpu, x, ne, true, stream)?;
137        }
138        // ATLAS_DS41_PREFILL_GEMV=1: every group through the decode GEMV arm in
139        // chunks of eight rows, so prefill quantises activations exactly as the
140        // single-token step does (q8 block 32) instead of the MMQ tile's D2S6
141        // block 64 for Q2_K. A diagnostic: it makes prefill and decode
142        // numerically the same path at the cost of the tensor-core arm.
143        let prefill_gemv = prefill_gemv();
144        for &(a0, off, r) in plan.iter().filter(|_| !m1_arm) {
145            let slot = slots[a0];
146            let rows_ptr = DevicePtr(self.rows_dev.0 + (off * 4) as u64);
147            let w_ptr = DevicePtr(self.weight_dev.0 + (off * 4) as u64);
148            KernelLaunch::new(gpu, self.k.gather)
149                .grid([r as u32, 1, 1])
150                .block([256, 1, 1])
151                .arg_ptr(x)
152                .arg_ptr(rows_ptr)
153                .arg_ptr(self.a_rows)
154                .arg_u32(c.dim as u32)
155                .launch(stream)?;
156            if prefill_gemv && r > 8 {
157                // chunks of <= 8 rows, each an offset view into the group's buffers
158                let q8_row = (c.dim / 32) * Q8_1_BLOCK_BYTES;
159                let mut r0 = 0usize;
160                while r0 < r {
161                    let rr = (r - r0).min(8);
162                    let a = DevicePtr(self.a_rows.0 + (r0 * c.dim * 2) as u64);
163                    let q = DevicePtr(self.a_q8.0 + (r0 * q8_row) as u64);
164                    let go = DevicePtr(self.gate_out.0 + (r0 * c.inter * 2) as u64);
165                    let uo = DevicePtr(self.up_out.0 + (r0 * c.inter * 2) as u64);
166                    kquant_q8_1_rows(gpu, self.k.q8_rows, a, q, rr as u32, c.dim as u32, stream)?;
167                    kquant_mmvq_w(
168                        gpu,
169                        self.k.mmvq_q2k,
170                        slot.gate,
171                        q,
172                        go,
173                        c.inter as u32,
174                        c.dim as u32,
175                        rr as u32,
176                        stream,
177                    )?;
178                    kquant_mmvq_w(
179                        gpu,
180                        self.k.mmvq_q2k,
181                        slot.up,
182                        q,
183                        uo,
184                        c.inter as u32,
185                        c.dim as u32,
186                        rr as u32,
187                        stream,
188                    )?;
189                    r0 += rr;
190                }
191            } else if r <= 8 {
192                // the decode GEMV: plain q8_1 rows, one weight read shared by the rows
193                kquant_q8_1_rows(
194                    gpu,
195                    self.k.q8_rows,
196                    self.a_rows,
197                    self.a_q8,
198                    r as u32,
199                    c.dim as u32,
200                    stream,
201                )?;
202                kquant_mmvq_w(
203                    gpu,
204                    self.k.mmvq_q2k,
205                    slot.gate,
206                    self.a_q8,
207                    self.gate_out,
208                    c.inter as u32,
209                    c.dim as u32,
210                    r as u32,
211                    stream,
212                )?;
213                kquant_mmvq_w(
214                    gpu,
215                    self.k.mmvq_q2k,
216                    slot.up,
217                    self.a_q8,
218                    self.up_out,
219                    c.inter as u32,
220                    c.dim as u32,
221                    r as u32,
222                    stream,
223                )?;
224            } else {
225                // the prefill MMQ: tensor cores on the raw blocks, D2S6 activations for Q2_K
226                ops::quantize_act_q8_1(
227                    gpu,
228                    self.k.quant_d2s6,
229                    self.a_rows,
230                    self.a_q8,
231                    r as u32,
232                    c.dim as u32,
233                    stream,
234                )?;
235                kquant_mmq_gemm(
236                    gpu,
237                    self.k.mmq_q2k_nc,
238                    self.k.mmq_q2k_wc,
239                    self.a_q8,
240                    slot.gate,
241                    self.gate_out,
242                    r as u32,
243                    c.inter as u32,
244                    c.dim as u32,
245                    Q2K_MMQ_SMEM,
246                    stream,
247                )?;
248                kquant_mmq_gemm(
249                    gpu,
250                    self.k.mmq_q2k_nc,
251                    self.k.mmq_q2k_wc,
252                    self.a_q8,
253                    slot.up,
254                    self.up_out,
255                    r as u32,
256                    c.inter as u32,
257                    c.dim as u32,
258                    Q2K_MMQ_SMEM,
259                    stream,
260                )?;
261            }
262            self.launch_n(gpu, self.k.swiglu, r * c.inter, stream, |l| {
263                l.arg_ptr(self.gate_out)
264                    .arg_ptr(self.up_out)
265                    .arg_ptr(w_ptr)
266                    .arg_ptr(self.h)
267                    .arg_u32(r as u32)
268                    .arg_u32(c.inter as u32)
269                    .arg_f32(c.swiglu_limit)
270            })?;
271            if prefill_gemv && r > 8 {
272                let q8_row = (c.inter / 32) * Q8_1_BLOCK_BYTES;
273                let mut r0 = 0usize;
274                while r0 < r {
275                    let rr = (r - r0).min(8);
276                    let h = DevicePtr(self.h.0 + (r0 * c.inter * 2) as u64);
277                    let q = DevicePtr(self.h_q8.0 + (r0 * q8_row) as u64);
278                    let d = DevicePtr(self.down_out.0 + (r0 * c.dim * 2) as u64);
279                    kquant_q8_1_rows(gpu, self.k.q8_rows, h, q, rr as u32, c.inter as u32, stream)?;
280                    kquant_mmvq_w(
281                        gpu,
282                        self.k.mmvq_q3k,
283                        slot.down,
284                        q,
285                        d,
286                        c.dim as u32,
287                        c.inter as u32,
288                        rr as u32,
289                        stream,
290                    )?;
291                    r0 += rr;
292                }
293            } else if r <= 8 {
294                kquant_q8_1_rows(
295                    gpu,
296                    self.k.q8_rows,
297                    self.h,
298                    self.h_q8,
299                    r as u32,
300                    c.inter as u32,
301                    stream,
302                )?;
303                kquant_mmvq_w(
304                    gpu,
305                    self.k.mmvq_q3k,
306                    slot.down,
307                    self.h_q8,
308                    self.down_out,
309                    c.dim as u32,
310                    c.inter as u32,
311                    r as u32,
312                    stream,
313                )?;
314            } else {
315                ops::quantize_act_q8_1(
316                    gpu,
317                    self.k.quant_d4,
318                    self.h,
319                    self.h_q8,
320                    r as u32,
321                    c.inter as u32,
322                    stream,
323                )?;
324                kquant_mmq_gemm(
325                    gpu,
326                    self.k.mmq_q3k_nc,
327                    self.k.mmq_q3k_wc,
328                    self.h_q8,
329                    slot.down,
330                    self.down_out,
331                    r as u32,
332                    c.dim as u32,
333                    c.inter as u32,
334                    Q3K_MMQ_SMEM,
335                    stream,
336                )?;
337            }
338            KernelLaunch::new(gpu, self.k.scatter_add)
339                .grid([r as u32, 1, 1])
340                .block([256, 1, 1])
341                .arg_ptr(self.acc)
342                .arg_ptr(self.down_out)
343                .arg_ptr(rows_ptr)
344                .arg_u32(c.dim as u32)
345                .launch(stream)?;
346        }
347        if side {
348            // `sd` from the side stream, then the same accumulate + finish
349            gpu.stream_wait_event(stream, self.ev_out)?;
350            self.shared_expert_tail(gpu, m, stream)?;
351        } else {
352            self.shared_expert(gpu, w, x, m, m1_arm, stream)?;
353        }
354        if self.timing_sync {
355            // ATLAS_DS41_DIAG=1: make compute_ms the GPU time, not the launch time
356            gpu.synchronize(stream)?;
357        }
358        self.last.set(MoeV41Timing {
359            route_ms: (t1 - t0).as_secs_f64() * 1e3,
360            fetch_ms: (t2 - t1).as_secs_f64() * 1e3,
361            compute_ms: t2.elapsed().as_secs_f64() * 1e3,
362            hits: after.hits - before.hits,
363            misses: after.misses - before.misses,
364            bytes_read: after.bytes_read - before.bytes_read,
365        });
366        Ok((self.out, weights, indices))
367    }
368
369    /// The shared expert added into `acc`, then `acc` finished into `out`.
370    /// Device work only.
371    pub(super) fn shared_expert(
372        &self,
373        gpu: &dyn GpuBackend,
374        w: &MoeV41LayerWeights,
375        x: DevicePtr,
376        m: usize,
377        x_pre: bool,
378        stream: u64,
379    ) -> Result<()> {
380        self.shared_expert_body(gpu, w, x, m, self.a_q8, self.h_q8, x_pre, stream)?;
381        self.shared_expert_tail(gpu, m, stream)
382    }
383
384    /// `acc += sd`, then `out = bf16(acc)`: the shared expert's join.
385    pub(super) fn shared_expert_tail(
386        &self,
387        gpu: &dyn GpuBackend,
388        m: usize,
389        stream: u64,
390    ) -> Result<()> {
391        let c = &self.cfg;
392        self.launch_n(gpu, self.k.accumulate, m * c.dim, stream, |l| {
393            l.arg_ptr(self.acc)
394                .arg_ptr(self.sd)
395                .arg_u32((m * c.dim) as u32)
396        })?;
397        self.launch_n(gpu, self.k.finish, m * c.dim, stream, |l| {
398            l.arg_ptr(self.acc)
399                .arg_ptr(self.out)
400                .arg_u32((m * c.dim) as u32)
401        })
402    }
403}
spark_model/layers/moe_v41/forward.rs

spark_model/layers/moe_v41/
forward.rs
