metrale_model_layers/layers/moe/
forward_prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `MoeLayer::forward_prefill`, the MoE block for N prefill rows.
4//!
5//! Owner: model-layers (MoE).
6//! Invariants: none beyond the types.
7
8use super::*;
9
10impl MoeLayer {
11    /// 2026-09-25: MoE for `num_tokens` rows of the normed [num_tokens, H] BF16
12    /// `input`, with grouped GEMMs over the rows sorted by expert.
13    ///
14    /// BF16 and FP8 experts go to `forward_prefill_bf16` / `forward_prefill_fp8`
15    /// above 64 rows when that grouped kernel resolved, and to `forward_batched`
16    /// otherwise. The NVFP4 body here:
17    /// shared expert, gate GEMM, top-K, sort by expert, grouped gate/up GEMM,
18    /// activation, grouped down GEMM, unpermute + weighted reduce, shared blend.
19    #[allow(unused_assignments)]
20    pub fn forward_prefill(
21        &self,
22        input: DevicePtr,
23        num_tokens: usize,
24        ctx: &ForwardContext,
25        stream: u64,
26    ) -> Result<()> {
27        // 2026-09-25: Native HIP (`metrale_hip`, the strix-hip target) has no BF16
28        // grouped MoE GEMM (kernels/strix-hip/common ships the FP8 and W4A16 ones
29        // only), so BF16 experts always take `forward_batched` there. FP8 experts
30        // keep the grouped path.
31        let hip_force_batched = cfg!(metrale_hip);
32        let hip_force_batched_fp8 = false;
33
34        // 2026-09-25: MoE LoRA needs no refusal here. Each grouped body (this one,
35        // `forward_prefill_bf16`, `forward_prefill_fp8`) folds the router delta
36        // before top-k and the routed-expert gate/up and down deltas on the sorted
37        // buffers; `forward_batched` folds them per row.
38
39        // 2026-09-25: BF16 and FP8 experts: the grouped GEMM above 64 rows when its
40        // kernel resolved, else `forward_batched`. 2026-10-01: FP8 experts take the grouped
41        // GEMM at every row count under `prefill_row_invariant`, so a prompt's rows get the
42        // same bits in a cold pass and in a warm request's short replay of its tail.
43        if self.bf16_gate_weight_ptrs.is_some() {
44            if self.moe_bf16_grouped_gemm_k.0 != 0 && num_tokens > 64 && !hip_force_batched {
45                return self.forward_prefill_bf16(input, num_tokens, ctx, stream);
46            }
47            return self.forward_batched(input, num_tokens, ctx, stream);
48        }
49
50        if self.fp8_gate_weight_ptrs.is_some() {
51            let grouped_rows =
52                num_tokens > 64 || (!ctx.decode_step && crate::layers::prefill_row_invariant());
53            if self.moe_fp8_grouped_gemm_k.0 != 0 && grouped_rows && !hip_force_batched_fp8 {
54                return self.forward_prefill_fp8(input, num_tokens, ctx, stream);
55            }
56            return self.forward_batched(input, num_tokens, ctx, stream);
57        }
58
59        // 2026-09-25: Transpose this layer's down_proj into the shared scratch on
60        // the compute stream; a no-op unless the lazy scratch is wired.
61        let _t_xpose = if ctx.profile && self.down_t_scratch_packed.is_some() {
62            ctx.gpu.synchronize(stream)?;
63            Some(std::time::Instant::now())
64        } else {
65            None
66        };
67        self.transpose_down_into_scratch(ctx, stream)?;
68        if let Some(t0) = _t_xpose {
69            ctx.gpu.synchronize(stream)?;
70            tracing::info!(
71                "  MoE prefill [lazy_transpose_down] N={}: {}µs",
72                num_tokens,
73                t0.elapsed().as_micros(),
74            );
75        }
76
77        let h = ctx.config.hidden_size as u32;
78        let inter = ctx.config.moe_intermediate_size as u32;
79        let shared_inter = ctx.config.shared_expert_intermediate_size as u32;
80        let num_experts = ctx.config.num_experts as u32;
81        let top_k = ctx.config.num_experts_per_tok as u32;
82        let n = num_tokens as u32;
83        let total_expanded = n * top_k;
84
85        #[allow(unused_macros)]
86        macro_rules! prof {
87            ($label:expr) => {
88                if ctx.profile {
89                    ctx.gpu.synchronize(stream)?;
90                    let _t = std::time::Instant::now();
91                    tracing::info!("  MoE prefill [{}] N={}", $label, num_tokens);
92                }
93            };
94        }
95        #[allow(unused_assignments)]
96        let mut t0 = if ctx.profile {
97            ctx.gpu.synchronize(stream)?;
98            Some(std::time::Instant::now())
99        } else {
100            None
101        };
102        macro_rules! prof_step {
103            ($label:expr) => {
104                if let Some(t) = t0.take() {
105                    ctx.gpu.synchronize(stream)?;
106                    let elapsed = t.elapsed().as_micros();
107                    tracing::info!("  MoE prefill [{}] N={}: {}µs", $label, num_tokens, elapsed);
108                    t0 = Some(std::time::Instant::now());
109                }
110            };
111        }
112
113        // 2026-09-25: The shared expert runs first. It reads `input` and writes
114        // only ssm_deinterleaved, ssm_qkvz and attn_output. `use_overlap` is
115        // false, so it runs on `stream`. Skipped when shared_inter == 0.
116        let has_shared = shared_inter > 0;
117        let use_overlap = false;
118        let aux = if use_overlap {
119            self.prefill_stream
120        } else {
121            stream
122        };
123
124        if has_shared {
125            self.run_shared_expert_prefill(
126                input,
127                n,
128                h,
129                shared_inter,
130                aux,
131                stream,
132                use_overlap,
133                ctx,
134            )?;
135        }
136        prof_step!("shared_expert");
137
138        let router_in = self.router_input(input, n, h, ctx, stream)?;
139        super::dump::dump_gate_input(ctx.gpu, stream, router_in, n, h)?;
140        let gate_logits = ctx.buffers.gate_logits();
141        if let Some(fp8) = self.gate_fp8 {
142            ops::fp8_gemm_n128(
143                ctx.gpu,
144                self.fp8_gemm_k,
145                router_in,
146                fp8,
147                gate_logits,
148                n,
149                self.router_logits_n,
150                h,
151                stream,
152            )?;
153        } else if let Some(ref nvfp4) = self.gate_nvfp4 {
154            ops::w4a16_gemm(
155                ctx.gpu,
156                self.w4a16_gemm,
157                router_in,
158                nvfp4,
159                gate_logits,
160                n,
161                self.router_logits_n,
162                h,
163                stream,
164            )?;
165        } else {
166            // 2026-09-25: The router GEMM keeps the scalar kernel's accumulation
167            // order (`router_gate_gemm_dense`); it has no cuBLAS arm.
168            self.router_gate_gemm_dense(
169                router_in,
170                gate_logits,
171                n,
172                self.router_logits_n,
173                h,
174                ctx,
175                stream,
176            )?;
177        }
178        super::dump::dump_gate_logits(ctx.gpu, stream, gate_logits, n, num_experts)?;
179        prof_step!("gate_gemm");
180
181        // 2026-09-25: Fold the router LoRA delta onto the logits before top-k. A
182        // no-op without an installed router delta (loaded only with
183        // METRALE_LORA_EXPERTS=1) or when the request is not routed to it.
184        self.apply_router_lora_prefill(router_in, gate_logits, n, ctx, stream)?;
185
186        // 2026-09-25: Routing: the static hash table when `tid2eid_dev` is set,
187        // else a correction bias with sqrtsoftplus, softmax or sigmoid scoring by
188        // `scoring_func`, else plain softmax top-K.
189        let scratch = ctx.buffers.scratch();
190        let indices_dev = scratch;
191        let weights_dev = scratch.offset(total_expanded as usize * 4);
192        if let Some(tid2eid) = self.tid2eid_dev {
193            let token_ids = ctx.token_ids.ok_or_else(|| {
194                anyhow::anyhow!(
195                    "DeepSeek-V4 hash-MoE layer requires ForwardContext.token_ids (prefill grouped)"
196                )
197            })?;
198            ops::moe_hash_route_batched(
199                ctx.gpu,
200                self.moe_hash_route_batched_k,
201                gate_logits,
202                tid2eid,
203                token_ids,
204                indices_dev,
205                weights_dev,
206                num_experts,
207                top_k,
208                ctx.config.norm_topk_prob,
209                ctx.config.routed_scaling_factor as f32,
210                n,
211                stream,
212            )?;
213        } else if let Some(bias) = self.correction_bias_dev {
214            if ctx.config.scoring_func == "sqrtsoftplus" {
215                ops::moe_topk_sqrtsoftplus_batched(
216                    ctx.gpu,
217                    self.moe_topk_sqrtsoftplus_batched_k,
218                    gate_logits,
219                    bias,
220                    indices_dev,
221                    weights_dev,
222                    num_experts,
223                    top_k,
224                    ctx.config.norm_topk_prob,
225                    ctx.config.routed_scaling_factor as f32,
226                    n,
227                    stream,
228                )?;
229            } else if ctx.config.scoring_func == "softmax" {
230                self.router_softmax_bias_batched(
231                    gate_logits,
232                    bias,
233                    indices_dev,
234                    weights_dev,
235                    num_experts,
236                    top_k,
237                    n,
238                    ctx,
239                    stream,
240                )?;
241            } else {
242                ops::moe_topk_sigmoid_batched(
243                    ctx.gpu,
244                    self.moe_topk_sigmoid_batched_k,
245                    gate_logits,
246                    bias,
247                    indices_dev,
248                    weights_dev,
249                    num_experts,
250                    top_k,
251                    ctx.config.norm_topk_prob,
252                    ctx.config.routed_scaling_factor as f32,
253                    n,
254                    stream,
255                )?;
256            }
257        } else {
258            ops::moe_topk_softmax_batched(
259                ctx.gpu,
260                self.moe_topk_batched,
261                gate_logits,
262                indices_dev,
263                weights_dev,
264                num_experts,
265                top_k,
266                ctx.config.norm_topk_prob,
267                n,
268                stream,
269            )?;
270        }
271        super::dump::dump_expert_ids(ctx.gpu, stream, indices_dev, weights_dev, n, top_k)?;
272        prof_step!("topk");
273
274        // 2026-09-25: Sort the (row, slot) pairs by expert. The sort outputs reuse
275        // `gate_logits`, which top-K has consumed: sorted_token_ids [te],
276        // sorted_expert_ids [te], expert_offsets [E + 1], token_to_perm [te], u32.
277        let te = total_expanded as usize;
278        let ne = num_experts as usize;
279        let sorted_token_ids = gate_logits;
280        let sorted_expert_ids = gate_logits.offset(te * 4);
281        let expert_offsets = gate_logits.offset(te * 4 * 2);
282        let token_to_perm = gate_logits.offset(te * 4 * 2 + (ne + 1) * 4);
283        ops::moe_sort_by_expert(
284            ctx.gpu,
285            self.moe_sort_by_expert,
286            indices_dev,
287            sorted_token_ids,
288            sorted_expert_ids,
289            expert_offsets,
290            token_to_perm,
291            total_expanded,
292            num_experts,
293            top_k,
294            stream,
295        )?;
296        prof_step!("sort");
297
298        // 2026-09-25: With `pre_expert_norm`, the router has read the un-normed
299        // `input` and the experts read a normed copy in ssm_deinterleaved. `input`
300        // itself stays intact: the shared-expert blend reads it.
301        let expert_input = if let Some(ref norm_w) = self.pre_expert_norm {
302            let normed_buf = ctx.buffers.ssm_deinterleaved();
303            let n_tokens = num_tokens as u32;
304            let eps = ctx.config.rms_norm_eps as f32;
305            ops::rms_norm(
306                ctx.gpu,
307                self.pre_expert_norm_k,
308                input,
309                norm_w,
310                normed_buf,
311                n_tokens,
312                h,
313                eps,
314                stream,
315            )?;
316            normed_buf
317        } else {
318            input
319        };
320        prof_step!("pre_expert_norm");
321
322        // 2026-09-25: Grouped gate+up GEMM, activation and grouped down GEMM; the
323        // sorted result lands in `expert_down_out`.
324        self.run_routed_grouped_gemm(
325            expert_input,
326            expert_offsets,
327            sorted_token_ids,
328            n,
329            h,
330            inter,
331            num_experts,
332            top_k,
333            num_tokens,
334            ne,
335            &mut t0,
336            ctx,
337            stream,
338        )?;
339        let expert_down_out = ctx.buffers.expert_down_out();
340
341        // 2026-09-25: Fold the routed-expert down LoRA delta onto the sorted
342        // `expert_down_out` before the weighted reduce, so the routing weight
343        // scales base + delta. A no-op without installed down deltas.
344        self.apply_expert_lora_prefill_down(
345            ctx.buffers.expert_gate_out(),
346            expert_down_out,
347            expert_offsets,
348            sorted_token_ids,
349            total_expanded,
350            ctx,
351            stream,
352        )?;
353
354        let output = ctx.buffers.moe_output();
355        ops::moe_unpermute_reduce_indexed(
356            ctx.gpu,
357            self.moe_unpermute_reduce,
358            expert_down_out,
359            output,
360            token_to_perm,
361            weights_dev,
362            h,
363            n,
364            top_k,
365            stream,
366        )?;
367
368        // 2026-09-25: Shared-expert blend: output += sigmoid(input . gate) * shared,
369        // weight 1 when the layer has no shared_expert_gate. With EP it runs
370        // once after the all-reduce instead, so the reduce does not sum it per rank.
371        let is_ep_prefill = ctx.comm.is_some() && ctx.config.ep_world_size > 1;
372        if has_shared && !is_ep_prefill {
373            let shared_down_out = ctx.buffers.attn_output();
374            if use_overlap {
375                ctx.gpu.stream_wait_event(stream, self.event_b)?;
376            }
377            super::dump::dump_routed_only(ctx.gpu, stream, output, n, h)?;
378            super::dump::dump_shared_out(ctx.gpu, stream, shared_down_out, n, h)?;
379            super::dump::dump_shared_gate(
380                ctx.gpu,
381                stream,
382                input,
383                self.weights.shared_expert_gate.weight,
384                n,
385                h,
386            )?;
387            ops::moe_batched_blend(
388                ctx.gpu,
389                self.moe_batched_blend,
390                output,
391                shared_down_out,
392                input,
393                self.weights.shared_expert_gate.weight,
394                h,
395                n,
396                stream,
397            )?;
398        }
399        super::dump::dump_moe_out(ctx.gpu, stream, output, n, h)?;
400        prof_step!("unpermute_blend");
401
402        if let Some(comm) = ctx.comm
403            && ctx.config.ep_world_size > 1
404        {
405            let _t0 = if ctx.profile {
406                ctx.gpu.synchronize(stream)?;
407                Some(std::time::Instant::now())
408            } else {
409                None
410            };
411            if ctx.graph_capture {
412                comm.all_reduce(output.0, num_tokens * h as usize * 2)?;
413            } else {
414                comm.all_reduce_async(output.0, num_tokens * h as usize * 2, stream)?;
415            }
416            if let Some(t0) = _t0 {
417                ctx.gpu.synchronize(stream)?;
418                tracing::info!(
419                    "  EP allreduce (moe out) N={}: {}µs",
420                    num_tokens,
421                    t0.elapsed().as_micros(),
422                );
423            }
424            if has_shared {
425                let shared_down_out = ctx.buffers.attn_output();
426                if use_overlap {
427                    ctx.gpu.stream_wait_event(stream, self.event_b)?;
428                }
429                ops::moe_batched_blend(
430                    ctx.gpu,
431                    self.moe_batched_blend,
432                    output,
433                    shared_down_out,
434                    input,
435                    self.weights.shared_expert_gate.weight,
436                    h,
437                    n,
438                    stream,
439                )?;
440            }
441        }
442
443        Ok(())
444    }
445}
metrale_model_layers/layers/moe/forward_prefill.rs

metrale_model_layers/layers/moe/
forward_prefill.rs
