spark_model/layers/moe/forward_batched.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! MoeLayer::forward_batched.
4
5use super::*;
6
7impl MoeLayer {
8 /// Batched forward: GEMM gate for N tokens, per-token expert dispatch.
9 ///
10 /// Gate projection reads weights once for N tokens (GEMM M=N).
11 /// Expert dispatch remains per-token (data-dependent routing).
12 pub fn forward_batched(
13 &self,
14 input: DevicePtr,
15 num_tokens: usize,
16 ctx: &ForwardContext,
17 stream: u64,
18 ) -> Result<()> {
19 // LongCat zero-experts ARE wired here (softmax+bias arm below); the
20 // other scoring functions on this variant are not, and would silently
21 // mis-route the 384-wide router. Named refusal, not silent wrongness.
22 //
23 // This path is load-bearing for LongCat: `forward_prefill` only takes
24 // the grouped GEMM above 64 tokens, so every FP8/BF16-expert prefill
25 // shorter than that — i.e. most conversational turns — lands here.
26 anyhow::ensure!(
27 self.router_logits_n as usize == ctx.config.num_experts
28 || (self.correction_bias_dev.is_some() && ctx.config.scoring_func == "softmax"),
29 "zero-expert MoE routing is not wired on this dispatch variant yet (forward_batched)"
30 );
31
32 // SOLID Incr-4: batched decode folds the routed-expert gate/up + down
33 // LoRA delta per token (below) AND the router (mlp.gate) delta on the
34 // whole-batch gate_logits before top-k (`apply_router_lora_batched`,
35 // after the gate GEMM), all gated by the device per-row `row_adapter` map
36 // so base rows no-op and the launches are route-agnostic under capture.
37 // The `reject_decode_lora` / `reject_batched_router_lora` bails are lifted
38 // here; a batch containing a NON-active adapter (`Refuse`) is bailed
39 // host-side in `decode_batch_compute_main` before graph lookup, so it
40 // never reaches these folds. `row_adapter_base` is the fixed-address
41 // `[num_seqs]` i32 map uploaded per step (MoE `-1 = base` semantics);
42 // `DevicePtr(0)` when no adapter is resident, in which case the fold hooks
43 // fall back to the request-granularity `moe_route_gate` (a homogeneous
44 // batch folds all rows; base skips) — see
45 // `apply_expert_lora_decode_{gateup,down}` and `apply_router_lora_batched`.
46 let row_adapter_base = ctx
47 .attn_metadata
48 .as_ref()
49 .map_or(DevicePtr::NULL, |m| m.moe_row_adapter);
50 let h = ctx.config.hidden_size as u32;
51 let inter = ctx.config.moe_intermediate_size as u32;
52 let shared_inter = ctx.config.shared_expert_intermediate_size as u32;
53 let num_experts = ctx.config.num_experts as u32;
54 let top_k = ctx.config.num_experts_per_tok as u32;
55 let n = num_tokens as u32;
56 let bf16 = 2usize;
57
58 // Router width, not expert count — they differ only on LongCat, where
59 // the router also scores the zero-computation experts. This value also
60 // strides `gate_t` below; getting it wrong reads each token's logits
61 // from the wrong offset rather than failing.
62 let router_n = self.router_logits_n;
63 let (gate_logits, fp32_gate, gate_elem) =
64 self.batched_gate_logits(input, n, h, router_n, row_adapter_base, ctx, stream)?;
65
66 // Per-token: topK routing + expert dispatch + weighted sum
67 let h_usize = h as usize;
68 let expert_gate_out = ctx.buffers.expert_gate_out();
69 let expert_up_out = ctx.buffers.expert_up_out();
70 let expert_down_out = ctx.buffers.expert_down_out();
71 // ⚠ logits buffer aliased — see warning in moe/forward.rs:208-219
72 // and project_batch_decode_corruption.md (bug 2). Concurrent
73 // callers using `buffers.logits()` during the forward loop MUST
74 // offset past `shared_expert_intermediate_size * 2` bytes.
75 let shared_gate_scratch = ctx.buffers.logits();
76 let shared_up_scratch = ctx.buffers.ssm_qkvz();
77
78 for t in 0..num_tokens {
79 let input_t = input.offset(t * h_usize * bf16);
80 let gate_t = gate_logits.offset(t * router_n as usize * gate_elem);
81 let output_t = ctx.buffers.moe_output().offset(t * h_usize * bf16);
82
83 let scratch = ctx.buffers.scratch();
84 let indices_dev = scratch;
85 let weights_dev = scratch.offset(top_k as usize * 4);
86
87 // Per-row LoRA map slice for THIS token: the fold's `n_slots == top_k`
88 // (`row/top_k == 0`), so token `t`'s entry is `row_adapter_base + t*4`
89 // (i32). The offset is a structural loop constant → a fixed address
90 // baked correctly per captured graph. NULL base stays NULL (no
91 // per-row map; the hooks then use the request gate).
92 let ra_t = if row_adapter_base.0 != 0 {
93 row_adapter_base.offset(t * 4)
94 } else {
95 DevicePtr::NULL
96 };
97
98 if let Some(tid2eid) = self.tid2eid_dev {
99 // DeepSeek-V4 hash routing: expert selection is static
100 // `tid2eid[token_id]`; the learned gate weights the selection.
101 // token IDs are uploaded [num_tokens] u32 in the SAME order as
102 // this loop, so token t lives at offset t.
103 let token_ids = ctx.token_ids.ok_or_else(|| {
104 anyhow::anyhow!(
105 "DeepSeek-V4 hash-MoE layer requires ForwardContext.token_ids (prefill)"
106 )
107 })?;
108 ops::moe_hash_route(
109 ctx.gpu,
110 self.moe_hash_route_k,
111 gate_t,
112 tid2eid,
113 token_ids.offset(t * 4),
114 indices_dev,
115 weights_dev,
116 num_experts,
117 top_k,
118 ctx.config.norm_topk_prob,
119 ctx.config.routed_scaling_factor as f32,
120 stream,
121 )?;
122 } else if let Some(bias) = self.correction_bias_dev {
123 self.router_bias_one(
124 gate_t,
125 bias,
126 indices_dev,
127 weights_dev,
128 num_experts,
129 top_k,
130 t,
131 ctx,
132 stream,
133 )?;
134 } else {
135 ops::moe_topk_softmax(
136 ctx.gpu,
137 if fp32_gate {
138 self.moe_topk_f32
139 } else {
140 self.moe_topk
141 },
142 gate_t,
143 indices_dev,
144 weights_dev,
145 num_experts,
146 top_k,
147 ctx.config.norm_topk_prob,
148 stream,
149 )?;
150 }
151 // Last-token routing dump (no-op unless ATLAS_DUMP_EXPERT_IDS=1):
152 // the token whose top-K determines the next prediction.
153 if t == num_tokens - 1 {
154 super::dump::dump_expert_ids(ctx.gpu, stream, indices_dev, weights_dev, 1, top_k)?;
155 }
156
157 let shared_out = ctx.buffers.attn_output();
158 if let (Some(gp), Some(up), Some(dp), Some(shared)) = (
159 self.bf16_gate_weight_ptrs,
160 self.bf16_up_weight_ptrs,
161 self.bf16_down_weight_ptrs,
162 self.bf16_shared_expert,
163 ) {
164 // BF16 path (FP8-dequant-on-load): same fused kernels as decode.
165 ops::moe_expert_gate_up_shared_bf16(
166 ctx.gpu,
167 self.moe_expert_gate_up_shared_bf16_k,
168 input_t,
169 gp,
170 expert_gate_out,
171 up,
172 expert_up_out,
173 indices_dev,
174 shared.gate_proj.weight,
175 shared_gate_scratch,
176 shared.up_proj.weight,
177 shared_up_scratch,
178 inter,
179 h,
180 top_k,
181 stream,
182 )?;
183 // SOLID Incr-4: fold gate/up delta BEFORE the fused silu+down.
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 // FP8 path for batched decode
219 ops::moe_expert_gate_up_shared_fp8(
220 ctx.gpu,
221 self.moe_expert_gate_up_shared_fp8,
222 input_t,
223 gp.weight_ptrs,
224 gp.scale_ptrs,
225 expert_gate_out,
226 up.weight_ptrs,
227 up.scale_ptrs,
228 expert_up_out,
229 indices_dev,
230 &sh.gate_proj,
231 shared_gate_scratch,
232 &sh.up_proj,
233 shared_up_scratch,
234 inter,
235 h,
236 top_k,
237 stream,
238 )?;
239 // SOLID Incr-4: fold gate/up delta BEFORE the fused silu+down.
240 self.apply_expert_lora_decode_gateup(
241 expert_gate_out,
242 expert_up_out,
243 input_t,
244 indices_dev,
245 top_k,
246 top_k,
247 ra_t,
248 ctx,
249 stream,
250 )?;
251 ops::moe_expert_silu_down_shared_fp8(
252 ctx.gpu,
253 self.moe_expert_silu_down_shared_fp8,
254 expert_gate_out,
255 expert_up_out,
256 dp.weight_ptrs,
257 dp.scale_ptrs,
258 expert_down_out,
259 indices_dev,
260 shared_gate_scratch,
261 shared_up_scratch,
262 &sh.down_proj,
263 shared_out,
264 h,
265 inter,
266 top_k,
267 stream,
268 )?;
269 } else if self.use_t_layout_for_prefill() {
270 // Phase 8a unified-layout NVFP4 batched prefill — transposed
271 // kernels coalesce well at large N. Hybrid mode lands here too.
272 let gate_t = self
273 .gate_ptrs_t
274 .as_ref()
275 .expect("gate_ptrs_t under unified_t");
276 let up_t = self.up_ptrs_t.as_ref().expect("up_ptrs_t under unified_t");
277 let down_t = self
278 .down_ptrs_t
279 .as_ref()
280 .expect("down_ptrs_t under unified_t");
281 let null_qw = QuantizedWeight::null();
282 let sh_gate_t = self.shared_gate_t.as_ref().unwrap_or(&null_qw);
283 let sh_up_t = self.shared_up_t.as_ref().unwrap_or(&null_qw);
284 let sh_down_t = self.shared_down_t.as_ref().unwrap_or(&null_qw);
285 // ARM-2 Phase-K RIDER A1: the _e8m0 fused decode kernel is
286 // <32,true,GROUP_SIZE,false> — routed E8M0, shared NVFP4. Assert
287 // the shared expert really is NVFP4 before trusting that.
288 if self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Mxfp4E8m0 {
289 self.shared_experts_scale_kind.expect(
290 crate::weight_map::WeightQuantFormat::Nvfp4,
291 "decode fused _e8m0 kernel assumes an NVFP4 shared expert",
292 );
293 }
294 ops::moe_expert_gate_up_shared_t(
295 ctx.gpu,
296 self.e8m0_or(
297 self.moe_expert_gate_up_shared_t_k,
298 self.moe_expert_gate_up_shared_t_e8m0_k,
299 "decode gate_up_shared_t",
300 ),
301 input_t,
302 gate_t.packed_ptrs,
303 gate_t.scale_ptrs,
304 gate_t.scale2_vals,
305 expert_gate_out,
306 up_t.packed_ptrs,
307 up_t.scale_ptrs,
308 up_t.scale2_vals,
309 expert_up_out,
310 indices_dev,
311 sh_gate_t,
312 shared_gate_scratch,
313 sh_up_t,
314 shared_up_scratch,
315 inter,
316 h,
317 top_k,
318 stream,
319 )?;
320 // SOLID Incr-4: fold gate/up delta BEFORE the fused silu+down.
321 self.apply_expert_lora_decode_gateup(
322 expert_gate_out,
323 expert_up_out,
324 input_t,
325 indices_dev,
326 top_k,
327 top_k,
328 ra_t,
329 ctx,
330 stream,
331 )?;
332 ops::moe_expert_silu_down_shared_t(
333 ctx.gpu,
334 self.e8m0_or(
335 self.moe_expert_silu_down_shared_t_k,
336 self.moe_expert_silu_down_shared_t_e8m0_k,
337 "decode silu_down_shared_t",
338 ),
339 expert_gate_out,
340 expert_up_out,
341 down_t.packed_ptrs,
342 down_t.scale_ptrs,
343 down_t.scale2_vals,
344 expert_down_out,
345 indices_dev,
346 shared_gate_scratch,
347 shared_up_scratch,
348 sh_down_t,
349 shared_out,
350 h,
351 inter,
352 top_k,
353 stream,
354 )?;
355 } else {
356 // NVFP4 path
357 ops::moe_expert_gate_up_shared(
358 ctx.gpu,
359 self.moe_expert_gate_up_shared,
360 input_t,
361 self.gate_ptrs.packed_ptrs,
362 self.gate_ptrs.scale_ptrs,
363 self.gate_ptrs.scale2_vals,
364 expert_gate_out,
365 self.up_ptrs.packed_ptrs,
366 self.up_ptrs.scale_ptrs,
367 self.up_ptrs.scale2_vals,
368 expert_up_out,
369 indices_dev,
370 &self.weights.shared_expert.gate_proj,
371 shared_gate_scratch,
372 &self.weights.shared_expert.up_proj,
373 shared_up_scratch,
374 inter,
375 h,
376 top_k,
377 stream,
378 )?;
379 // SOLID Incr-4: fold gate/up delta BEFORE the fused silu+down.
380 self.apply_expert_lora_decode_gateup(
381 expert_gate_out,
382 expert_up_out,
383 input_t,
384 indices_dev,
385 top_k,
386 top_k,
387 ra_t,
388 ctx,
389 stream,
390 )?;
391 ops::moe_expert_silu_down_shared(
392 ctx.gpu,
393 self.moe_expert_silu_down_shared,
394 expert_gate_out,
395 expert_up_out,
396 self.down_ptrs.packed_ptrs,
397 self.down_ptrs.scale_ptrs,
398 self.down_ptrs.scale2_vals,
399 expert_down_out,
400 indices_dev,
401 shared_gate_scratch,
402 shared_up_scratch,
403 &self.weights.shared_expert.down_proj,
404 shared_out,
405 h,
406 inter,
407 top_k,
408 stream,
409 )?;
410 }
411
412 // SOLID Incr-4: fold the routed-expert down_proj delta into this
413 // token's `expert_down_out` (slot-major [top_k, hidden]) IN PLACE,
414 // recomputing x = silu(gate)*up from the still-materialized
415 // gate/up out — BEFORE the weighted-sum blend (so the router weight
416 // scales base+delta). Route-agnostic via `ra_t` (base rows no-op).
417 self.apply_expert_lora_decode_down(
418 expert_gate_out,
419 expert_up_out,
420 expert_down_out,
421 indices_dev,
422 top_k,
423 top_k,
424 ra_t,
425 ctx,
426 stream,
427 )?;
428
429 if self.has_mixed_bf16_shared_expert() {
430 self.run_bf16_shared_expert(
431 input_t,
432 1,
433 h,
434 shared_inter,
435 shared_gate_scratch,
436 shared_up_scratch,
437 shared_out,
438 ctx,
439 stream,
440 )?;
441 }
442
443 ops::moe_weighted_sum_blend(
444 ctx.gpu,
445 self.moe_weighted_sum_blend,
446 output_t,
447 expert_down_out,
448 weights_dev,
449 shared_out,
450 input_t,
451 self.weights.shared_expert_gate.weight,
452 h,
453 top_k,
454 h,
455 stream,
456 )?;
457
458 // EP all-reduce per-token partial output
459 if let Some(comm) = ctx.comm
460 && ctx.config.ep_world_size > 1
461 {
462 if ctx.graph_capture {
463 comm.all_reduce(output_t.0, h as usize * 2)?;
464 } else {
465 comm.all_reduce_async(output_t.0, h as usize * 2, stream)?;
466 }
467 }
468 }
469
470 Ok(())
471 }
472}