spark_model/layers/moe/forward.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! MoeLayer::forward (decode).
4
5use super::*;
6
7impl MoeLayer {
8 /// True when the ATLAS_FP32_ROUTING path is active: the SSM-side MoE-input
9 /// norm should emit an FP32 `router_in` (residual_add_rms_norm_gatef32) which
10 /// the gate GEMM then consumes at full precision. Requires the f32 kernels to
11 /// be present and the softmax-routed dense-gate config (NVFP4 gate / sigmoid+bias
12 /// stay BF16). Default off → BF16 routing unchanged.
13 /// The lever is the LAST term on purpose: the four preconditions are
14 /// properties of this layer's weights and kernels, and only the final
15 /// one is configuration. `levers` is passed rather than read because
16 /// this is called once per layer per DECODE TOKEN from six sites.
17 pub fn fp32_routing_active(&self, levers: &crate::layers::ops::ModelLevers) -> bool {
18 self.gate_nvfp4.is_none()
19 && self.correction_bias_dev.is_none()
20 && self.dense_gemm_f32in.0 != 0
21 && self.moe_topk_f32.0 != 0
22 && levers.fp32_routing
23 }
24
25 /// Forward pass: gate → top-K routing → batched expert FFN → blend.
26 ///
27 /// All expert dispatch stays on device — zero D2H synchronization.
28 /// 9 kernel launches per MoE layer (down from 58).
29 ///
30 /// When `gelu_activation` is true, falls back to the sorted prefill path
31 /// (which uses separate activation kernel) to avoid fused SiLU decode kernels.
32 /// LongCat zero-computation experts: `out[t,:] += zero_accum[t] * x[t,:]`
33 /// where `zero_accum` was written by the softmax+bias router kernels
34 /// (the folded weights of selected identity experts). MUST run after the
35 /// routed blend for the SAME tokens whose routing wrote `zero_accum`.
36 /// No-op (no launch) when the model has no zero-experts.
37 pub fn apply_zero_expert(
38 &self,
39 out: spark_runtime::gpu::DevicePtr,
40 x: spark_runtime::gpu::DevicePtr,
41 n: u32,
42 ctx: &ForwardContext,
43 stream: u64,
44 ) -> Result<()> {
45 if self.router_logits_n as usize == ctx.config.num_experts {
46 return Ok(());
47 }
48 anyhow::ensure!(
49 self.moe_zero_expert_add_k.0 != 0,
50 "zero-expert model but moe_zero_expert_add kernel is absent from this build"
51 );
52 ops::moe_zero_expert_add(
53 ctx.gpu,
54 self.moe_zero_expert_add_k,
55 out,
56 x,
57 self.zero_accum_dev,
58 n,
59 ctx.config.hidden_size as u32,
60 stream,
61 )
62 }
63
64 pub fn forward(
65 &self,
66 input: DevicePtr,
67 ctx: &ForwardContext,
68 stream: u64,
69 ) -> Result<DevicePtr> {
70 // SOLID Incr-4: a genuine single-token decode (num_seqs == 1) folds the
71 // routed expert down_proj LoRA delta below (before the wsum blend). The
72 // multi-seq per-token reuse of this fn (num_seqs > 1 — decode_batch's
73 // per-token MoE loop, or the attention layers' per-token FFN) shares a
74 // `padded_n` CUDA graph across mixed-batch transitions, so host-gating
75 // the fold there is capture-unsafe (a Fold-captured graph could replay
76 // onto a later base-containing batch): keep the loud refusal until the
77 // device per-row `row_adapter` map is plumbed.
78 let single_seq_decode = ctx.attn_metadata.as_ref().map_or(1, |m| m.num_seqs) <= 1;
79 if !single_seq_decode {
80 // Multi-seq per-token reuse shares a `padded_n` CUDA graph across mixed
81 // base/adapter batches, so host-gating a fold there is capture-unsafe:
82 // keep the loud refusal until the device per-row `row_adapter` map lands.
83 self.reject_decode_lora(ctx, "forward")?;
84 }
85 // Single-seq decode: the router delta folds onto `gate_logits` before top-k
86 // (below), and the routed-expert gate/up/down deltas fold onto their
87 // intermediates — no bail. A `Refuse`/mixed batch still bails inside each
88 // fold's `moe_route_gate`, preserving per-row adapter-identity protection.
89 // ── Phase 2.7 Tier C: Frankenstein decode-via-prefill dispatch ──
90 // For DFlash capture layers only, when `ATLAS_FRANKENSTEIN_DECODE_VIA_PREFILL=1`
91 // is set, route this layer's single-token MoE through `forward_prefill(M=1)`,
92 // which uses the tensor-core grouped GEMM kernel (E2M1→E4M3 MMA) instead of
93 // the scalar FP32 FMA decode path. Tests whether the numerical recipe of the
94 // MoE kernel is the dominant cause of low DFlash drafter acceptance.
95 //
96 // Other (non-capture) layers fall through to the normal scalar decode path,
97 // preserving Atlas's TPS on the bulk of the network. The 5 capture layers
98 // pay ~250 µs each (microbench), totalling ≈1.25 ms per token (negligible
99 // at Atlas's ~58 ms/token decode latency).
100 if self.is_dflash_capture_layer && ctx.levers.frankenstein_decode_via_prefill {
101 // One-time per-process log so we can verify the env-gated route is hit.
102 if ctx.stats.once("log:moe_route") {
103 tracing::info!(
104 "FRANKENSTEIN: routing DFlash capture-layer MoE decode through forward_prefill(M=1) (one-time log)"
105 );
106 }
107 self.forward_prefill(input, 1, ctx, stream)?;
108 return Ok(ctx.buffers.moe_output());
109 }
110
111 // GeGLU models: fused kernels now have GELU activation (model-specific override).
112 // No longer need to redirect through sorted prefill path.
113 // But we still need pre_expert_norm between routing and dispatch.
114 // For the fused decode path, apply pre_expert_norm to the input before experts.
115 // The gate GEMV already completed on the raw input in the fused path below.
116
117 let h = ctx.config.hidden_size as u32;
118 let inter = ctx.config.moe_intermediate_size as u32;
119 let shared_inter = ctx.config.shared_expert_intermediate_size as u32;
120 let num_experts = ctx.config.num_experts as u32;
121 let top_k = ctx.config.num_experts_per_tok as u32;
122 let profile = ctx.profile;
123
124 macro_rules! prof {
125 ($label:expr, $body:expr) => {{
126 if profile {
127 let t = std::time::Instant::now();
128 let r = $body;
129 ctx.gpu.synchronize(stream)?;
130 tracing::info!(" MoE {}: {:.0}μs", $label, t.elapsed().as_micros());
131 r
132 } else {
133 $body
134 }
135 }};
136 }
137
138 let scratch = ctx.buffers.scratch();
139 let indices_dev = scratch;
140 let weights_dev = scratch.offset(top_k as usize * 4);
141
142 // Note: moe_gate_topk_fused exists but uses single-CTA design,
143 // too slow for 256 experts (serializes computation). Separate path is faster.
144 {
145 // Gemma-4 router pre-norm (no-op for other models).
146 let router_in = self.router_input(input, 1, h, ctx, stream)?;
147 let gate_logits = ctx.buffers.gate_logits();
148 prof!("gate", {
149 if let Some(ref nvfp4) = self.gate_nvfp4 {
150 ops::w4a16_decode_gemv(
151 ctx.gpu,
152 self.w4a16_gemv,
153 self.w4a16_gemv_sw,
154 ctx.levers.gemv_sw,
155 router_in,
156 nvfp4,
157 gate_logits,
158 // = num_experts everywhere except LongCat, whose
159 // router also scores the zero-expert logits.
160 self.router_logits_n,
161 h,
162 stream,
163 )
164 } else {
165 ops::dense_gemv(
166 ctx.gpu,
167 self.dense_gemv,
168 router_in,
169 &self.weights.gate,
170 gate_logits,
171 self.router_logits_n,
172 h,
173 stream,
174 )
175 }
176 })?;
177
178 // Feature-1: fold the router `mlp.gate` LoRA delta onto `gate_logits`
179 // BEFORE top-k — the exact decode mirror of the prefill router fold
180 // (`apply_router_lora_prefill` is n-generic; here n=1). Device-clean
181 // (no D2H) so it captures cleanly. No-op unless a router delta is
182 // installed; `Refuse` bails inside the hook. Works for both the
183 // NVFP4-gate and dense-gate branches (same `gate_logits` output).
184 if single_seq_decode {
185 self.apply_router_lora_prefill(router_in, gate_logits, 1, ctx, stream)?;
186 }
187
188 prof!("topk", {
189 if let Some(tid2eid) = self.tid2eid_dev {
190 // DeepSeek-V4 hash routing (hash_moe layer): expert SELECTION
191 // is the static `tid2eid[token_id]` table; the learned gate
192 // still supplies the sqrtsoftplus scores that weight them.
193 let token_ids = ctx.token_ids.ok_or_else(|| {
194 anyhow::anyhow!(
195 "DeepSeek-V4 hash-MoE layer requires ForwardContext.token_ids (decode)"
196 )
197 })?;
198 ops::moe_hash_route(
199 ctx.gpu,
200 self.moe_hash_route_k,
201 gate_logits,
202 tid2eid,
203 token_ids, // decode: single token at offset 0
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 stream,
211 )
212 } else if let Some(bias) = self.correction_bias_dev {
213 if ctx.config.scoring_func == "sqrtsoftplus" {
214 // DeepSeek-V4 sqrtsoftplus + correction bias:
215 // scores = sqrtsoftplus(gate_logits)
216 // indices = topk(scores + bias)
217 // weights = scores[indices] / sum(scores[indices])
218 ops::moe_topk_sqrtsoftplus(
219 ctx.gpu,
220 self.moe_topk_sqrtsoftplus_k,
221 gate_logits,
222 bias,
223 indices_dev,
224 weights_dev,
225 num_experts,
226 top_k,
227 ctx.config.norm_topk_prob,
228 ctx.config.routed_scaling_factor as f32,
229 stream,
230 )
231 } else if ctx.config.scoring_func == "softmax" {
232 // LongCat-Flash: softmax scores + correction bias for
233 // SELECTION, unbiased softmax * scaling for weights,
234 // zero-expert fold into zero_accum (identity experts
235 // are applied by the caller via apply_zero_expert).
236 ops::moe_topk_softmax_bias(
237 ctx.gpu,
238 self.moe_topk_softmax_bias_k,
239 gate_logits,
240 bias,
241 indices_dev,
242 weights_dev,
243 self.zero_accum_dev,
244 self.router_logits_n,
245 num_experts,
246 top_k,
247 ctx.config.norm_topk_prob,
248 ctx.config.routed_scaling_factor as f32,
249 stream,
250 )
251 } else {
252 // DeepSeek-V3 / MiniMax-M2 sigmoid + correction bias:
253 // scores = sigmoid(gate_logits)
254 // indices = topk(scores + bias)
255 // weights = scores[indices] / sum(scores[indices])
256 // Kernel does all three steps; norm_topk_prob toggles
257 // the final divide. scaling_factor comes from the model
258 // config (e.g., Step 3.7 = 3.0, MiniMax M2 = 1.0).
259 ops::moe_topk_sigmoid(
260 ctx.gpu,
261 self.moe_topk_sigmoid_k,
262 gate_logits,
263 bias,
264 indices_dev,
265 weights_dev,
266 num_experts,
267 top_k,
268 ctx.config.norm_topk_prob,
269 ctx.config.routed_scaling_factor as f32,
270 stream,
271 )
272 }
273 } else {
274 ops::moe_topk_softmax(
275 ctx.gpu,
276 self.moe_topk,
277 gate_logits,
278 indices_dev,
279 weights_dev,
280 num_experts,
281 top_k,
282 ctx.config.norm_topk_prob,
283 stream,
284 )
285 }
286 })?;
287 }
288
289 if tracing::enabled!(tracing::Level::DEBUG) && !ctx.graph_capture {
290 ctx.gpu.synchronize(stream)?;
291 // Read expert indices (u32[top_k]) and weights (f32[top_k])
292 let k = top_k as usize;
293 let mut idx_buf = vec![0u8; k * 4];
294 let mut wt_buf = vec![0u8; k * 4];
295 ctx.gpu.copy_d2h(indices_dev, &mut idx_buf)?;
296 ctx.gpu.copy_d2h(weights_dev, &mut wt_buf)?;
297 let indices: Vec<u32> = (0..k)
298 .map(|i| {
299 u32::from_le_bytes([
300 idx_buf[i * 4],
301 idx_buf[i * 4 + 1],
302 idx_buf[i * 4 + 2],
303 idx_buf[i * 4 + 3],
304 ])
305 })
306 .collect();
307 let weights: Vec<f32> = (0..k)
308 .map(|i| {
309 f32::from_le_bytes([
310 wt_buf[i * 4],
311 wt_buf[i * 4 + 1],
312 wt_buf[i * 4 + 2],
313 wt_buf[i * 4 + 3],
314 ])
315 })
316 .collect();
317 tracing::info!(" MoE experts: {:?}, weights: {:.4?}", indices, weights);
318 }
319
320 // Apply pre-expert norm AFTER routing, BEFORE expert dispatch (Gemma-4 26B).
321 // Write to scratch buffer to preserve original `input` (= residual in caller).
322 let expert_input = if let Some(ref norm_w) = self.pre_expert_norm {
323 let normed = ctx.buffers.ssm_deinterleaved();
324 let eps = ctx.config.rms_norm_eps as f32;
325 prof!("pre_expert_norm", {
326 ops::rms_norm(
327 ctx.gpu,
328 self.pre_expert_norm_k,
329 input,
330 norm_w,
331 normed,
332 1,
333 h,
334 eps,
335 stream,
336 )
337 })?;
338 normed
339 } else {
340 input
341 };
342
343 // ── Batched expert FFN: 3 GEMV + 1 activation + 1 weighted sum ──
344 let expert_gate_out = ctx.buffers.expert_gate_out();
345 let expert_up_out = ctx.buffers.expert_up_out();
346 let expert_down_out = ctx.buffers.expert_down_out();
347 // ⚠ `logits` aliased as shared-gate scratch — concurrent users
348 // MUST offset past `shared_expert_intermediate_size * 2`
349 // (decode_b.rs:197 uses .offset(65536)). See bug 2 in memory
350 // `project_batch_decode_corruption.md` (2026-05-10).
351 let shared_gate_scratch = ctx.buffers.logits();
352 let shared_up_scratch = ctx.buffers.ssm_qkvz();
353 let shared_out = ctx.buffers.attn_output();
354
355 if let (Some(gp), Some(up), Some(dp), Some(shared)) = (
356 self.bf16_gate_weight_ptrs,
357 self.bf16_up_weight_ptrs,
358 self.bf16_down_weight_ptrs,
359 self.bf16_shared_expert,
360 ) {
361 // BF16 path: FP8-dequant-on-load. Eliminates the per-layer 0.989
362 // FP8 cosine ceiling by serving experts as BF16 end-to-end.
363 prof!("exp_gate_up_bf16", {
364 ops::moe_expert_gate_up_shared_bf16(
365 ctx.gpu,
366 self.moe_expert_gate_up_shared_bf16_k,
367 expert_input,
368 gp,
369 expert_gate_out,
370 up,
371 expert_up_out,
372 indices_dev,
373 shared.gate_proj.weight,
374 shared_gate_scratch,
375 shared.up_proj.weight,
376 shared_up_scratch,
377 inter,
378 h,
379 top_k,
380 stream,
381 )
382 })?;
383 if single_seq_decode {
384 self.apply_expert_lora_decode_gateup(
385 expert_gate_out,
386 expert_up_out,
387 expert_input,
388 indices_dev,
389 top_k,
390 top_k,
391 DevicePtr::NULL,
392 ctx,
393 stream,
394 )?;
395 }
396 prof!("exp_silu_down_bf16", {
397 ops::moe_expert_silu_down_shared_bf16(
398 ctx.gpu,
399 self.moe_expert_silu_down_shared_bf16_k,
400 expert_gate_out,
401 expert_up_out,
402 dp,
403 expert_down_out,
404 indices_dev,
405 shared_gate_scratch,
406 shared_up_scratch,
407 shared.down_proj.weight,
408 shared_out,
409 h,
410 inter,
411 top_k,
412 stream,
413 )
414 })?;
415 } else if let (Some(gp), Some(up), Some(dp), Some(sh)) = (
416 &self.fp8_gate_weight_ptrs,
417 &self.fp8_up_weight_ptrs,
418 &self.fp8_down_weight_ptrs,
419 &self.fp8_shared_expert,
420 ) {
421 // FP8 path: fused expert gate+up with FP8 weight/scale pointer tables
422 prof!("exp_gate_up_fp8", {
423 ops::moe_expert_gate_up_shared_fp8(
424 ctx.gpu,
425 self.moe_expert_gate_up_shared_fp8,
426 expert_input,
427 gp.weight_ptrs,
428 gp.scale_ptrs,
429 expert_gate_out,
430 up.weight_ptrs,
431 up.scale_ptrs,
432 expert_up_out,
433 indices_dev,
434 &sh.gate_proj,
435 shared_gate_scratch,
436 &sh.up_proj,
437 shared_up_scratch,
438 inter,
439 h,
440 top_k,
441 stream,
442 )
443 })?;
444
445 if single_seq_decode {
446 self.apply_expert_lora_decode_gateup(
447 expert_gate_out,
448 expert_up_out,
449 expert_input,
450 indices_dev,
451 top_k,
452 top_k,
453 DevicePtr::NULL,
454 ctx,
455 stream,
456 )?;
457 }
458 // FP8 path: fused silu+down
459 prof!("exp_silu_down_fp8", {
460 ops::moe_expert_silu_down_shared_fp8(
461 ctx.gpu,
462 self.moe_expert_silu_down_shared_fp8,
463 expert_gate_out,
464 expert_up_out,
465 dp.weight_ptrs,
466 dp.scale_ptrs,
467 expert_down_out,
468 indices_dev,
469 shared_gate_scratch,
470 shared_up_scratch,
471 &sh.down_proj,
472 shared_out,
473 h,
474 inter,
475 top_k,
476 stream,
477 )
478 })?;
479 } else if self.use_t_layout_for_decode() {
480 prof!("exp_unified_t", {
481 self.dispatch_unified_t_decode(
482 ctx,
483 expert_input,
484 expert_gate_out,
485 expert_up_out,
486 expert_down_out,
487 shared_gate_scratch,
488 shared_up_scratch,
489 shared_out,
490 indices_dev,
491 h,
492 inter,
493 top_k,
494 single_seq_decode,
495 stream,
496 )
497 })?;
498 } else {
499 // NVFP4 path: fused routed+shared gate+up
500 prof!("exp_gate_up", {
501 ops::moe_expert_gate_up_shared(
502 ctx.gpu,
503 self.moe_expert_gate_up_shared,
504 expert_input,
505 self.gate_ptrs.packed_ptrs,
506 self.gate_ptrs.scale_ptrs,
507 self.gate_ptrs.scale2_vals,
508 expert_gate_out,
509 self.up_ptrs.packed_ptrs,
510 self.up_ptrs.scale_ptrs,
511 self.up_ptrs.scale2_vals,
512 expert_up_out,
513 indices_dev,
514 &self.weights.shared_expert.gate_proj,
515 shared_gate_scratch,
516 &self.weights.shared_expert.up_proj,
517 shared_up_scratch,
518 inter,
519 h,
520 top_k,
521 stream,
522 )
523 })?;
524
525 if single_seq_decode {
526 self.apply_expert_lora_decode_gateup(
527 expert_gate_out,
528 expert_up_out,
529 expert_input,
530 indices_dev,
531 top_k,
532 top_k,
533 DevicePtr::NULL,
534 ctx,
535 stream,
536 )?;
537 }
538
539 if tracing::enabled!(tracing::Level::DEBUG) && !ctx.graph_capture {
540 ctx.gpu.synchronize(stream)?;
541 // Dump gate/up outputs for expert slot 0
542 let mut gate_buf = vec![0u8; 16];
543 ctx.gpu.copy_d2h(expert_gate_out, &mut gate_buf)?;
544 let gate_vals: Vec<f32> = (0..8)
545 .map(|i| {
546 let bits = u16::from_le_bytes([gate_buf[i * 2], gate_buf[i * 2 + 1]]);
547 f32::from_bits((bits as u32) << 16)
548 })
549 .collect();
550 tracing::info!(" MoE gate_out[slot0,0..8]: {:?}", gate_vals);
551 let mut up_buf = vec![0u8; 16];
552 ctx.gpu.copy_d2h(expert_up_out, &mut up_buf)?;
553 let up_vals: Vec<f32> = (0..8)
554 .map(|i| {
555 let bits = u16::from_le_bytes([up_buf[i * 2], up_buf[i * 2 + 1]]);
556 f32::from_bits((bits as u32) << 16)
557 })
558 .collect();
559 tracing::info!(" MoE up_out[slot0,0..8]: {:?}", up_vals);
560 // Shared expert gate/up scratch outputs
561 let mut sg_buf = vec![0u8; 16];
562 ctx.gpu.copy_d2h(shared_gate_scratch, &mut sg_buf)?;
563 let sg_vals: Vec<f32> = (0..8)
564 .map(|i| {
565 let bits = u16::from_le_bytes([sg_buf[i * 2], sg_buf[i * 2 + 1]]);
566 f32::from_bits((bits as u32) << 16)
567 })
568 .collect();
569 tracing::info!(" MoE shared_gate_scratch[0..8]: {:?}", sg_vals);
570 let mut su_buf = vec![0u8; 16];
571 ctx.gpu.copy_d2h(shared_up_scratch, &mut su_buf)?;
572 let su_vals: Vec<f32> = (0..8)
573 .map(|i| {
574 let bits = u16::from_le_bytes([su_buf[i * 2], su_buf[i * 2 + 1]]);
575 f32::from_bits((bits as u32) << 16)
576 })
577 .collect();
578 tracing::info!(" MoE shared_up_scratch[0..8]: {:?}", su_vals);
579 }
580
581 // NVFP4 path: fused routed+shared silu+down
582 prof!("exp_silu_down", {
583 ops::moe_expert_silu_down_shared(
584 ctx.gpu,
585 self.moe_expert_silu_down_shared,
586 expert_gate_out,
587 expert_up_out,
588 self.down_ptrs.packed_ptrs,
589 self.down_ptrs.scale_ptrs,
590 self.down_ptrs.scale2_vals,
591 expert_down_out,
592 indices_dev,
593 shared_gate_scratch,
594 shared_up_scratch,
595 &self.weights.shared_expert.down_proj,
596 shared_out,
597 h,
598 inter,
599 top_k,
600 stream,
601 )
602 })?;
603 }
604
605 // SOLID Incr-4 decode expert down-fold: land the routed-expert down_proj
606 // LoRA delta into `expert_down_out` (slot-major [top_k, hidden]) IN PLACE,
607 // recomputing `x = silu(gate)*up` from the still-materialized
608 // `expert_gate_out`/`expert_up_out`. Must run BEFORE `moe_weighted_sum_blend`
609 // (so the router weight scales base+delta) AND before the EP zero-temp
610 // memset below (which reuses `expert_gate_out` as scratch). NULL
611 // row_adapter: a genuine single-token decode is one homogeneous request —
612 // `moe_route_gate` (Fold/Skip/Refuse) is the per-request opt-out. No-op
613 // when no MoE LoRA / no expert adapter is installed (base byte-identical).
614 if single_seq_decode {
615 self.apply_expert_lora_decode_down(
616 expert_gate_out,
617 expert_up_out,
618 expert_down_out,
619 indices_dev,
620 top_k,
621 top_k,
622 DevicePtr::NULL,
623 ctx,
624 stream,
625 )?;
626 }
627
628 if self.has_mixed_bf16_shared_expert() {
629 self.run_bf16_shared_expert(
630 input,
631 1,
632 h,
633 shared_inter,
634 shared_gate_scratch,
635 shared_up_scratch,
636 shared_out,
637 ctx,
638 stream,
639 )?;
640 }
641
642 if tracing::enabled!(tracing::Level::DEBUG) && !ctx.graph_capture {
643 ctx.gpu.synchronize(stream)?;
644 // Dump down outputs for expert slot 0
645 let mut down_buf = vec![0u8; 16];
646 ctx.gpu.copy_d2h(expert_down_out, &mut down_buf)?;
647 let down_vals: Vec<f32> = (0..8)
648 .map(|i| {
649 let bits = u16::from_le_bytes([down_buf[i * 2], down_buf[i * 2 + 1]]);
650 f32::from_bits((bits as u32) << 16)
651 })
652 .collect();
653 tracing::info!(" MoE down_out[slot0,0..8]: {:?}", down_vals);
654 // Shared out
655 let mut sh_buf = vec![0u8; 16];
656 ctx.gpu.copy_d2h(shared_out, &mut sh_buf)?;
657 let sh_vals: Vec<f32> = (0..8)
658 .map(|i| {
659 let bits = u16::from_le_bytes([sh_buf[i * 2], sh_buf[i * 2 + 1]]);
660 f32::from_bits((bits as u32) << 16)
661 })
662 .collect();
663 tracing::info!(" MoE shared_out[0..8]: {:?}", sh_vals);
664 }
665
666 // Fused wsum+blend+gate: routed expert weighted sum + sigmoid(gate)*shared
667 // Gate scalar GEMV is computed inline by each block (redundant but negligible).
668 //
669 // EP fix: for EP>1, the shared expert is computed identically on all ranks.
670 // If we include it in the output before all-reduce, it gets summed world_size
671 // times. Solution: pass NULL shared_out for EP, all-reduce the routed sum,
672 // then add shared_out once after all-reduce.
673 let output = ctx.buffers.moe_output();
674 let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1;
675 let shared_for_blend = if is_ep && !shared_out.is_null() {
676 // EP: exclude shared expert from blend (will add after all-reduce).
677 // Zero a temp buffer to pass as shared_out (kernel reads it even with NULL gate).
678 let zero_buf = ctx.buffers.expert_gate_out(); // temp buffer, will be zeroed
679 ctx.gpu.memset_async(zero_buf, 0, h as usize * 2, stream)?;
680 zero_buf
681 } else {
682 shared_out
683 };
684 prof!("wsum_blend", {
685 ops::moe_weighted_sum_blend(
686 ctx.gpu,
687 self.moe_weighted_sum_blend,
688 output,
689 expert_down_out,
690 weights_dev,
691 shared_for_blend,
692 input,
693 self.weights.shared_expert_gate.weight,
694 h,
695 top_k,
696 h,
697 stream,
698 )
699 })?;
700
701 // EP all-reduce: sum partial expert outputs across ranks.
702 // Each rank only computed its local experts (remote → zero), so
703 // SUM gives the correct global result.
704 if let Some(comm) = ctx.comm
705 && ctx.config.ep_world_size > 1
706 {
707 if ctx.graph_capture {
708 comm.all_reduce(output.0, h as usize * 2)?;
709 } else {
710 comm.all_reduce_async(output.0, h as usize * 2, stream)?;
711 }
712 // Now add shared expert contribution ONCE (after all-reduce).
713 // Must apply the sigmoid gate: output += sigmoid(dot(input, gate_w)) * shared_out.
714 // Using moe_batched_blend with num_tokens=1 computes the gate and blends correctly.
715 // BUG #41 fix: previous code used residual_add (ignoring the gate), producing
716 // wrong output that compounded across 48 layers into gibberish.
717 if !shared_out.is_null() {
718 if self.weights.shared_expert_gate.weight.0 == 0 {
719 // No gate weight (e.g., Mistral): shared expert always at full strength.
720 ops::residual_add(ctx.gpu, self.residual_add, output, shared_out, h, stream)?;
721 } else {
722 // Gated shared expert (e.g., Qwen3.5): apply sigmoid gate.
723 ops::moe_batched_blend(
724 ctx.gpu,
725 self.moe_batched_blend,
726 output,
727 shared_out,
728 input,
729 self.weights.shared_expert_gate.weight,
730 h,
731 1,
732 stream,
733 )?;
734 }
735 }
736 }
737
738 if tracing::enabled!(tracing::Level::DEBUG) && !ctx.graph_capture {
739 ctx.gpu.synchronize(stream)?;
740 let mut buf = vec![0u8; 8];
741 ctx.gpu.copy_d2h(output, &mut buf)?;
742 let vals: Vec<f32> = (0..4)
743 .map(|i| {
744 let lo = buf[i * 2];
745 let hi = buf[i * 2 + 1];
746 f32::from_bits(((lo as u32) | ((hi as u32) << 8)) << 16)
747 })
748 .collect();
749 tracing::info!(" MoE output: {:?}", vals);
750 }
751
752 Ok(output)
753 }
754}