metrale_model_layers/layers/moe/forward_batched.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `MoeLayer::forward_batched`: one gate GEMM for N tokens, then
4//! per-token routing and expert dispatch.
5//!
6//! Owner: model-layers (MoE).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11impl MoeLayer {
12 /// 2026-09-25: The gate projection is one GEMM with M = N; top-k, the
13 /// expert kernels and the blend then run once per token.
14 pub fn forward_batched(
15 &self,
16 input: DevicePtr,
17 num_tokens: usize,
18 ctx: &ForwardContext,
19 stream: u64,
20 ) -> Result<()> {
21 // 2026-09-25: A zero-expert router is served only through the
22 // softmax+bias arm (`router_bias_one`). `forward_prefill` sends BF16
23 // and FP8 expert prefills here unless it takes their grouped GEMM
24 // (more than 64 rows and the kernel resolved).
25 anyhow::ensure!(
26 self.router_logits_n as usize == ctx.config.num_experts
27 || (self.correction_bias_dev.is_some() && ctx.config.scoring_func == "softmax"),
28 "zero-expert MoE routing is not wired on this dispatch variant yet (forward_batched)"
29 );
30 // 2026-09-26: FP8 experts with a plain BF16 softmax router: the per-token
31 // loop below, grouped by expert, with this function's router arithmetic
32 // (`GroupedRouting::PerToken`), so the output bytes are the loop's. Each
33 // routed and shared expert then streams its weights once for all rows.
34 // `METRALE_NO_FP8_MOE_GROUPED_DECODE` (presence) turns it off.
35 if self.fp8_grouped_routing_ok(num_tokens, super::GroupedRouting::PerToken, ctx) {
36 return self.forward_fp8_grouped_decode_routed(
37 input,
38 num_tokens,
39 super::GroupedRouting::PerToken,
40 ctx,
41 stream,
42 );
43 }
44
45 // 2026-09-25: The router LoRA delta folds onto the whole batch's
46 // `gate_logits` before top-k (`batched_gate_logits`), and the expert
47 // gate/up and down deltas fold per token below. `row_adapter_base` is
48 // the per-step `[padded_n]` i32 adapter map (`moe_row_adapter`), with
49 // which base rows are skipped on the device. It is null when no LoRA
50 // weights are loaded or `ctx.attn_metadata` carries none, and the
51 // fold hooks then use `moe_route_gate`, which refuses a `Refuse` route.
52 // In batched decode a batch routed to a non-active adapter (`Refuse`) is
53 // refused host-side by `ensure_decode_route_servable` before this pass.
54 let row_adapter_base = ctx
55 .attn_metadata
56 .as_ref()
57 .map_or(DevicePtr::NULL, |m| m.moe_row_adapter);
58 let h = ctx.config.hidden_size as u32;
59 let inter = ctx.config.moe_intermediate_size as u32;
60 let shared_inter = ctx.config.shared_expert_intermediate_size as u32;
61 let num_experts = ctx.config.num_experts as u32;
62 let top_k = ctx.config.num_experts_per_tok as u32;
63 let n = num_tokens as u32;
64 let bf16 = 2usize;
65
66 // 2026-09-25: The router width (`num_experts + zero_expert_num`), not the
67 // expert count. It is also the row stride of `gate_t` below; a wrong
68 // stride would read each token's logits from the wrong offset without
69 // failing.
70 let router_n = self.router_logits_n;
71 let (gate_logits, fp32_gate, gate_elem) =
72 self.batched_gate_logits(input, n, h, router_n, row_adapter_base, ctx, stream)?;
73
74 let h_usize = h as usize;
75 let expert_gate_out = ctx.buffers.expert_gate_out();
76 let expert_up_out = ctx.buffers.expert_up_out();
77 let expert_down_out = ctx.buffers.expert_down_out();
78 // 2026-09-25: The start of `logits()` is the shared-expert gate scratch,
79 // as in `forward.rs`.
80 let shared_gate_scratch = ctx.buffers.logits();
81 let shared_up_scratch = ctx.buffers.ssm_qkvz();
82
83 for t in 0..num_tokens {
84 let input_t = input.offset(t * h_usize * bf16);
85 let gate_t = gate_logits.offset(t * router_n as usize * gate_elem);
86 let output_t = ctx.buffers.moe_output().offset(t * h_usize * bf16);
87
88 let scratch = ctx.buffers.scratch();
89 let indices_dev = scratch;
90 let weights_dev = scratch.offset(top_k as usize * 4);
91
92 // 2026-09-25: The fold sees `n_slots == top_k` rows, all of token `t`,
93 // so its map is token `t`'s i32 entry at `row_adapter_base + 4t`, a
94 // fixed address for a captured graph. A null map stays null.
95 let ra_t = if row_adapter_base.0 != 0 {
96 row_adapter_base.offset(t * 4)
97 } else {
98 DevicePtr::NULL
99 };
100
101 if let Some(tid2eid) = self.tid2eid_dev {
102 // 2026-09-25: Hash routing (`moe_hash_route.cu`), reading token
103 // `t`'s u32 id at `token_ids + 4t`.
104 let token_ids = ctx.token_ids.ok_or_else(|| {
105 anyhow::anyhow!(
106 "DeepSeek-V4 hash-MoE layer requires ForwardContext.token_ids (prefill)"
107 )
108 })?;
109 ops::moe_hash_route(
110 ctx.gpu,
111 self.moe_hash_route_k,
112 gate_t,
113 tid2eid,
114 token_ids.offset(t * 4),
115 indices_dev,
116 weights_dev,
117 num_experts,
118 top_k,
119 ctx.config.norm_topk_prob,
120 ctx.config.routed_scaling_factor as f32,
121 stream,
122 )?;
123 } else if let Some(bias) = self.correction_bias_dev {
124 self.router_bias_one(
125 gate_t,
126 bias,
127 indices_dev,
128 weights_dev,
129 num_experts,
130 top_k,
131 t,
132 ctx,
133 stream,
134 )?;
135 } else {
136 ops::moe_topk_softmax(
137 ctx.gpu,
138 if fp32_gate {
139 self.moe_topk_f32
140 } else {
141 self.moe_topk
142 },
143 gate_t,
144 indices_dev,
145 weights_dev,
146 num_experts,
147 top_k,
148 ctx.config.norm_topk_prob,
149 stream,
150 )?;
151 }
152 // 2026-09-25: Dump the last token's routing (no-op unless
153 // `METRALE_DUMP_EXPERT_IDS=1`).
154 if t == num_tokens - 1 {
155 super::dump::dump_expert_ids(ctx.gpu, stream, indices_dev, weights_dev, 1, top_k)?;
156 }
157
158 let shared_out = ctx.buffers.attn_output();
159 if let (Some(gp), Some(up), Some(dp), Some(shared)) = (
160 self.bf16_gate_weight_ptrs,
161 self.bf16_up_weight_ptrs,
162 self.bf16_down_weight_ptrs,
163 self.bf16_shared_expert,
164 ) {
165 // 2026-09-25: BF16 experts, with the same kernels as `forward`.
166 ops::moe_expert_gate_up_shared_bf16(
167 ctx.gpu,
168 self.moe_expert_gate_up_shared_bf16_k,
169 input_t,
170 gp,
171 expert_gate_out,
172 up,
173 expert_up_out,
174 indices_dev,
175 shared.gate_proj.weight,
176 shared_gate_scratch,
177 shared.up_proj.weight,
178 shared_up_scratch,
179 inter,
180 h,
181 top_k,
182 stream,
183 )?;
184 self.apply_expert_lora_decode_gateup(
185 expert_gate_out,
186 expert_up_out,
187 input_t,
188 indices_dev,
189 top_k,
190 top_k,
191 ra_t,
192 ctx,
193 stream,
194 )?;
195 ops::moe_expert_silu_down_shared_bf16(
196 ctx.gpu,
197 self.moe_expert_silu_down_shared_bf16_k,
198 expert_gate_out,
199 expert_up_out,
200 dp,
201 expert_down_out,
202 indices_dev,
203 shared_gate_scratch,
204 shared_up_scratch,
205 shared.down_proj.weight,
206 shared_out,
207 h,
208 inter,
209 top_k,
210 stream,
211 )?;
212 } else if let (Some(gp), Some(up), Some(dp), Some(sh)) = (
213 &self.fp8_gate_weight_ptrs,
214 &self.fp8_up_weight_ptrs,
215 &self.fp8_down_weight_ptrs,
216 &self.fp8_shared_expert,
217 ) {
218 ops::moe_expert_gate_up_shared_fp8(
219 ctx.gpu,
220 self.moe_expert_gate_up_shared_fp8,
221 input_t,
222 gp.weight_ptrs,
223 gp.scale_ptrs,
224 expert_gate_out,
225 up.weight_ptrs,
226 up.scale_ptrs,
227 expert_up_out,
228 indices_dev,
229 &sh.gate_proj,
230 shared_gate_scratch,
231 &sh.up_proj,
232 shared_up_scratch,
233 inter,
234 h,
235 top_k,
236 stream,
237 )?;
238 self.apply_expert_lora_decode_gateup(
239 expert_gate_out,
240 expert_up_out,
241 input_t,
242 indices_dev,
243 top_k,
244 top_k,
245 ra_t,
246 ctx,
247 stream,
248 )?;
249 ops::moe_expert_silu_down_shared_fp8(
250 ctx.gpu,
251 self.moe_expert_silu_down_shared_fp8,
252 expert_gate_out,
253 expert_up_out,
254 dp.weight_ptrs,
255 dp.scale_ptrs,
256 expert_down_out,
257 indices_dev,
258 shared_gate_scratch,
259 shared_up_scratch,
260 &sh.down_proj,
261 shared_out,
262 h,
263 inter,
264 top_k,
265 stream,
266 )?;
267 } else if self.use_t_layout_for_prefill() {
268 // 2026-09-25: Transposed-layout kernels; `use_t_layout_for_prefill`
269 // holds in both the unified and the hybrid layout.
270 let gate_t = self
271 .gate_ptrs_t
272 .as_ref()
273 .expect("gate_ptrs_t under unified_t");
274 let up_t = self.up_ptrs_t.as_ref().expect("up_ptrs_t under unified_t");
275 let down_t = self
276 .down_ptrs_t
277 .as_ref()
278 .expect("down_ptrs_t under unified_t");
279 let null_qw = QuantizedWeight::null();
280 let sh_gate_t = self.shared_gate_t.as_ref().unwrap_or(&null_qw);
281 let sh_up_t = self.shared_up_t.as_ref().unwrap_or(&null_qw);
282 let sh_down_t = self.shared_down_t.as_ref().unwrap_or(&null_qw);
283 // 2026-09-25: The `_e8m0` kernels are `<32, true, GROUP_SIZE, false>`
284 // (`moe_shared_expert_fused_t.cu`): routed experts E8M0, shared
285 // expert NVFP4. Check the shared expert's format before using them.
286 if self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Mxfp4E8m0 {
287 self.shared_experts_scale_kind.expect(
288 crate::weight_map::WeightQuantFormat::Nvfp4,
289 "decode fused _e8m0 kernel assumes an NVFP4 shared expert",
290 );
291 }
292 ops::moe_expert_gate_up_shared_t(
293 ctx.gpu,
294 self.e8m0_or(
295 self.moe_expert_gate_up_shared_t_k,
296 self.moe_expert_gate_up_shared_t_e8m0_k,
297 "decode gate_up_shared_t",
298 ),
299 input_t,
300 gate_t.packed_ptrs,
301 gate_t.scale_ptrs,
302 gate_t.scale2_vals,
303 expert_gate_out,
304 up_t.packed_ptrs,
305 up_t.scale_ptrs,
306 up_t.scale2_vals,
307 expert_up_out,
308 indices_dev,
309 sh_gate_t,
310 shared_gate_scratch,
311 sh_up_t,
312 shared_up_scratch,
313 inter,
314 h,
315 top_k,
316 stream,
317 )?;
318 self.apply_expert_lora_decode_gateup(
319 expert_gate_out,
320 expert_up_out,
321 input_t,
322 indices_dev,
323 top_k,
324 top_k,
325 ra_t,
326 ctx,
327 stream,
328 )?;
329 ops::moe_expert_silu_down_shared_t(
330 ctx.gpu,
331 self.e8m0_or(
332 self.moe_expert_silu_down_shared_t_k,
333 self.moe_expert_silu_down_shared_t_e8m0_k,
334 "decode silu_down_shared_t",
335 ),
336 expert_gate_out,
337 expert_up_out,
338 down_t.packed_ptrs,
339 down_t.scale_ptrs,
340 down_t.scale2_vals,
341 expert_down_out,
342 indices_dev,
343 shared_gate_scratch,
344 shared_up_scratch,
345 sh_down_t,
346 shared_out,
347 h,
348 inter,
349 top_k,
350 stream,
351 )?;
352 } else {
353 self.refuse_lean_layout("MoeLayer::forward_batched (row-major experts)")?;
354 ops::moe_expert_gate_up_shared(
355 ctx.gpu,
356 self.moe_expert_gate_up_shared,
357 input_t,
358 self.gate_ptrs.packed_ptrs,
359 self.gate_ptrs.scale_ptrs,
360 self.gate_ptrs.scale2_vals,
361 expert_gate_out,
362 self.up_ptrs.packed_ptrs,
363 self.up_ptrs.scale_ptrs,
364 self.up_ptrs.scale2_vals,
365 expert_up_out,
366 indices_dev,
367 &self.weights.shared_expert.gate_proj,
368 shared_gate_scratch,
369 &self.weights.shared_expert.up_proj,
370 shared_up_scratch,
371 inter,
372 h,
373 top_k,
374 stream,
375 )?;
376 self.apply_expert_lora_decode_gateup(
377 expert_gate_out,
378 expert_up_out,
379 input_t,
380 indices_dev,
381 top_k,
382 top_k,
383 ra_t,
384 ctx,
385 stream,
386 )?;
387 ops::moe_expert_silu_down_shared(
388 ctx.gpu,
389 self.moe_expert_silu_down_shared,
390 expert_gate_out,
391 expert_up_out,
392 self.down_ptrs.packed_ptrs,
393 self.down_ptrs.scale_ptrs,
394 self.down_ptrs.scale2_vals,
395 expert_down_out,
396 indices_dev,
397 shared_gate_scratch,
398 shared_up_scratch,
399 &self.weights.shared_expert.down_proj,
400 shared_out,
401 h,
402 inter,
403 top_k,
404 stream,
405 )?;
406 }
407
408 // 2026-09-25: Fold the expert down_proj delta into this token's
409 // `expert_down_out` (`[top_k, hidden]`) in place, before the blend, so
410 // the routing weight scales base plus delta.
411 self.apply_expert_lora_decode_down(
412 expert_gate_out,
413 expert_up_out,
414 expert_down_out,
415 indices_dev,
416 top_k,
417 top_k,
418 ra_t,
419 ctx,
420 stream,
421 )?;
422
423 if self.has_mixed_bf16_shared_expert() {
424 self.run_bf16_shared_expert(
425 input_t,
426 1,
427 h,
428 shared_inter,
429 shared_gate_scratch,
430 shared_up_scratch,
431 shared_out,
432 ctx,
433 stream,
434 )?;
435 }
436
437 ops::moe_weighted_sum_blend(
438 ctx.gpu,
439 self.moe_weighted_sum_blend,
440 output_t,
441 expert_down_out,
442 weights_dev,
443 shared_out,
444 input_t,
445 self.weights.shared_expert_gate.weight,
446 h,
447 top_k,
448 h,
449 stream,
450 )?;
451
452 if let Some(comm) = ctx.comm
453 && ctx.config.ep_world_size > 1
454 {
455 if ctx.graph_capture {
456 comm.all_reduce(output_t.0, h as usize * 2)?;
457 } else {
458 comm.all_reduce_async(output_t.0, h as usize * 2, stream)?;
459 }
460 }
461 }
462
463 Ok(())
464 }
465}