metrale_model_arch/moe_v41/
forward.rs

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

metrale_model_arch/moe_v41/
forward.rs
