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