spark_model/layers/moe/forward_k2.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! MoeLayer::forward_k2 (verify K=2).
4
5use anyhow::Context as _;
6
7use super::*;
8
9mod originals;
10mod unified_t;
11
12impl MoeLayer {
13 /// Fused K=2 forward: process 2 tokens through MoE in 5 kernel launches.
14 ///
15 /// Gate GEMV batch2 → batched topK → fused expert gate+up → fused silu+down → fused wsum+blend.
16 /// Expert buffers sized for 2*top_k slots. Shared expert buffers reuse logits/ssm_qkvz
17 /// (sized for 2 tokens). Output at moe_output() [2, H].
18 pub fn forward_k2(
19 &self,
20 input: DevicePtr, // [2, H] BF16 — normed MoE input for 2 tokens
21 ctx: &ForwardContext,
22 stream: u64,
23 ) -> Result<()> {
24 // LongCat zero-experts are wired only on the single-token decode
25 // + prefill paths (v1); this variant would silently mis-route the
26 // 384-wide router. Named refusal, not silent wrongness.
27 anyhow::ensure!(
28 self.router_logits_n as usize == ctx.config.num_experts,
29 "zero-expert MoE routing is not wired on this dispatch variant yet (forward_k2)"
30 );
31
32 // Feature-1: the fused batch2 fast path has no fold hook. When a MoE
33 // adapter is RESIDENT (install-time-fixed → graph-safe; graphs drain on
34 // rotate/swap), route to the per-row batched fallback which folds
35 // gate/up/down route-agnostically (base rows no-op) — same moe_output[2,H].
36 // forward_batched itself refuses a router-adapted adapter.
37 if self.lora.is_some() {
38 return self.forward_batched(input, 2, ctx, stream);
39 }
40 // BF16 (FP8-dequant-on-load) experts. The FP8/NVFP4 batch2 branches
41 // below read expert weights that were FREED at dequant-load, so they
42 // must NOT run for a dequanted model. When the fused BF16 batch2
43 // kernels are present (and we're not EP), take the dedicated BF16
44 // batch2 path (single-launch 2-token dispatch, same math as the
45 // per-token BF16 decode kernels). Otherwise fall back to the per-token
46 // BF16 batched path (SSOT: reuses the decode BF16 kernels via
47 // forward_batched), which produces the same moe_output()[2,H].
48 let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1;
49 let use_bf16_batch2 = self.bf16_gate_weight_ptrs.is_some()
50 && self.moe_expert_gate_up_shared_bf16_batch2_k.0 != 0
51 && self.moe_expert_silu_down_shared_bf16_batch2_k.0 != 0
52 && !is_ep;
53 if self.bf16_gate_weight_ptrs.is_some() && !use_bf16_batch2 {
54 return self.forward_batched(input, 2, ctx, stream);
55 }
56 // E8M0 (native MXFP4, per-32 E8M0 scale) routed experts MUST NOT reach the
57 // unified-T batch2 kernel `moe_expert_gate_up_shared_batch2_t`: it is an
58 // NVFP4 kernel that hardcodes GROUP_SIZE=16 and would read `inter·h/16`
59 // scale bytes from the correctly-sized `inter·h/32` E8M0 scale buffer — a
60 // 2× over-read → CUDA_ERROR_ILLEGAL_ADDRESS (it also E4M3-decodes E8M0
61 // scale bytes → garbage even in-bounds). No E8M0 batch2 kernel exists, so
62 // route both verify tokens through the per-token unified-T path
63 // (`forward_batched`), whose `use_t_layout_for_prefill` branch selects the
64 // GS32 `_e8m0` kernel via `e8m0_or` — the same correct path ordinary decode
65 // already uses. Mirrors the BF16 fallback above.
66 if k2_e8m0_needs_per_token(self.experts_scale_kind) {
67 return self.forward_batched(input, 2, ctx, stream);
68 }
69 // Mixed NVFP4-routed / BF16-shared (Laguna): the fused batch2 kernels
70 // cannot compute a BF16 shared expert alongside NVFP4 routed weights.
71 // Under the transposed unified layout we still batch the routed half
72 // through the _t kernels and run the shared expert as one batched BF16
73 // GEMM pass afterwards (`mixed_bf16_shared` below). Every other layout
74 // falls back to the per-token loop.
75 let mixed_bf16_shared = self.has_mixed_bf16_shared_expert();
76 // Either fused layout serves the mixed config. The originals-layout
77 // kernels are usable only since 37e818ad NULL-guarded their shared
78 // expert (their `_t` siblings always had that guard); before it, this
79 // faulted with CUDA 700 on the first 2-sequence batch.
80 let mixed_t_ok = self.use_t_layout_for_decode()
81 && self.moe_expert_gate_up_shared_batch2_t_k.0 != 0
82 && self.moe_expert_silu_down_shared_batch2_t_k.0 != 0;
83 let mixed_orig_ok = !self.use_t_layout_for_decode()
84 && self.moe_expert_gate_up_shared_batch2.0 != 0
85 && self.moe_expert_silu_down_shared_batch2.0 != 0
86 && !self.gate_ptrs.packed_ptrs.is_null();
87 if mixed_bf16_shared && !((mixed_t_ok || mixed_orig_ok) && !is_ep) {
88 return self.forward_batched(input, 2, ctx, stream);
89 }
90
91 let h = ctx.config.hidden_size as u32;
92 let inter = ctx.config.moe_intermediate_size as u32;
93 let num_experts = ctx.config.num_experts as u32;
94 let top_k = ctx.config.num_experts_per_tok as u32;
95
96 // DIAG (ATLAS_K2_DIAG=1): synchronize checkpoints to localize the K2-verify
97 // illegal access (the V4 NVFP4 batch2 verify path is exercised for the first
98 // time by MTP). The label of the FIRST failing sync names the bad stage.
99 let k2_diag = std::env::var("ATLAS_K2_DIAG").is_ok_and(|v| v == "1");
100 if k2_diag {
101 ctx.gpu
102 .synchronize(stream)
103 .context("K2 ENTRY: attention+norm BEFORE forward_k2")?;
104 }
105
106 // Gemma-4 router pre-norm (no-op for other models).
107 let router_in = self.router_input(input, 2, h, ctx, stream)?;
108 // 1. Gate GEMV batch2: reads gate weight once for 2 tokens
109 let gate_logits = ctx.buffers.gate_logits(); // [2, 512] BF16
110 if let Some(ref nvfp4) = self.gate_nvfp4 {
111 ops::w4a16_gemv_batch2(
112 ctx.gpu,
113 self.w4a16_gemv_batch2,
114 router_in,
115 nvfp4,
116 gate_logits,
117 num_experts,
118 h,
119 stream,
120 )?;
121 } else {
122 ops::dense_gemm(
123 ctx.gpu,
124 self.dense_gemm,
125 router_in,
126 &self.weights.gate,
127 gate_logits,
128 2,
129 num_experts,
130 h,
131 stream,
132 )?;
133 }
134
135 // 2. Batched topK for 2 tokens: [2, 512] → [2*top_k] indices + [2*top_k] weights.
136 // Sigmoid+bias for MiniMax/DeepSeek-V3, softmax otherwise.
137 let scratch = ctx.buffers.scratch();
138 let indices_dev = scratch; // [2*top_k] u32
139 let weights_dev = scratch.offset(2 * top_k as usize * 4); // [2*top_k] f32
140 if let Some(bias) = self.correction_bias_dev {
141 // DeepSeek-V4 scores experts with sqrt(softplus(.)); sigmoid otherwise
142 // (MiniMax/DeepSeek-V3). Must match the prefill/single-token paths or
143 // decode routing diverges from prefill.
144 if ctx.config.scoring_func == "sqrtsoftplus" {
145 // Use the PROVEN non-batched sqrtsoftplus kernel per token (the
146 // _batched variant is unexercised — the K2 verify is the only
147 // user and it never ran for V4 before). gate_logits is BF16
148 // [2, num_experts] (2-byte stride); indices/weights are
149 // [2, top_k] (u32 / f32, 4-byte stride).
150 for t in 0..2usize {
151 ops::moe_topk_sqrtsoftplus(
152 ctx.gpu,
153 self.moe_topk_sqrtsoftplus_k,
154 gate_logits.offset(t * num_experts as usize * 2),
155 bias,
156 indices_dev.offset(t * top_k as usize * 4),
157 weights_dev.offset(t * top_k as usize * 4),
158 num_experts,
159 top_k,
160 ctx.config.norm_topk_prob,
161 ctx.config.routed_scaling_factor as f32,
162 stream,
163 )?;
164 }
165 } else {
166 ops::moe_topk_sigmoid_batched(
167 ctx.gpu,
168 self.moe_topk_sigmoid_batched_k,
169 gate_logits,
170 bias,
171 indices_dev,
172 weights_dev,
173 num_experts,
174 top_k,
175 ctx.config.norm_topk_prob,
176 ctx.config.routed_scaling_factor as f32,
177 2,
178 stream,
179 )?;
180 }
181 } else {
182 ops::moe_topk_softmax_batched(
183 ctx.gpu,
184 self.moe_topk_batched,
185 gate_logits,
186 indices_dev,
187 weights_dev,
188 num_experts,
189 top_k,
190 ctx.config.norm_topk_prob,
191 2,
192 stream,
193 )?;
194 }
195 super::union_stats::maybe_sample_expert_union(ctx, indices_dev, 2, top_k as usize, stream);
196
197 if k2_diag {
198 ctx.gpu
199 .synchronize(stream)
200 .context("K2: gate-GEMV + topk")?;
201 }
202
203 // 3-5. Fused expert dispatch for 2 tokens
204 let expert_gate_out = ctx.buffers.expert_gate_out();
205 let expert_up_out = ctx.buffers.expert_up_out();
206 let shared_gate_scratch = ctx.buffers.logits();
207 let shared_up_scratch = ctx.buffers.ssm_qkvz();
208 let expert_down_out = ctx.buffers.expert_down_out();
209 let shared_down_out = ctx.buffers.attn_output();
210 let output = ctx.buffers.moe_output();
211
212 if use_bf16_batch2
213 && let (Some(gp), Some(up), Some(dp), Some(shared)) = (
214 self.bf16_gate_weight_ptrs,
215 self.bf16_up_weight_ptrs,
216 self.bf16_down_weight_ptrs,
217 self.bf16_shared_expert,
218 )
219 {
220 // BF16 batch2 path (FP8-dequant-on-load experts, MTP K=2 verify).
221 // Single-launch 2-token dispatch mirroring the FP8 batch2 layout;
222 // identical math to the per-token moe_expert_*_shared_bf16 kernels.
223 // Non-EP only (guaranteed by use_bf16_batch2).
224 ops::moe_expert_gate_up_shared_bf16_batch2(
225 ctx.gpu,
226 self.moe_expert_gate_up_shared_bf16_batch2_k,
227 input,
228 gp,
229 expert_gate_out,
230 up,
231 expert_up_out,
232 indices_dev,
233 shared.gate_proj.weight,
234 shared_gate_scratch,
235 shared.up_proj.weight,
236 shared_up_scratch,
237 inter,
238 h,
239 top_k,
240 stream,
241 )?;
242 ops::moe_expert_silu_down_shared_bf16_batch2(
243 ctx.gpu,
244 self.moe_expert_silu_down_shared_bf16_batch2_k,
245 expert_gate_out,
246 expert_up_out,
247 dp,
248 expert_down_out,
249 indices_dev,
250 shared_gate_scratch,
251 shared_up_scratch,
252 shared.down_proj.weight,
253 shared_down_out,
254 h,
255 inter,
256 top_k,
257 stream,
258 )?;
259 ops::moe_weighted_sum_blend_batch2(
260 ctx.gpu,
261 self.moe_weighted_sum_blend_batch2,
262 output,
263 expert_down_out,
264 weights_dev,
265 shared_down_out,
266 input,
267 self.weights.shared_expert_gate.weight,
268 h,
269 top_k,
270 h,
271 stream,
272 )?;
273 } else if let (Some(gp), Some(up), Some(dp), Some(sh)) = (
274 &self.fp8_gate_weight_ptrs,
275 &self.fp8_up_weight_ptrs,
276 &self.fp8_down_weight_ptrs,
277 &self.fp8_shared_expert,
278 ) {
279 // FP8 batch2 path
280 ops::moe_expert_gate_up_shared_fp8_batch2(
281 ctx.gpu,
282 self.moe_expert_gate_up_shared_fp8_batch2,
283 input,
284 gp.weight_ptrs,
285 gp.scale_ptrs,
286 expert_gate_out,
287 up.weight_ptrs,
288 up.scale_ptrs,
289 expert_up_out,
290 indices_dev,
291 &sh.gate_proj,
292 shared_gate_scratch,
293 &sh.up_proj,
294 shared_up_scratch,
295 inter,
296 h,
297 top_k,
298 stream,
299 )?;
300 ops::moe_expert_silu_down_shared_fp8_batch2(
301 ctx.gpu,
302 self.moe_expert_silu_down_shared_fp8_batch2,
303 expert_gate_out,
304 expert_up_out,
305 dp.weight_ptrs,
306 dp.scale_ptrs,
307 expert_down_out,
308 indices_dev,
309 shared_gate_scratch,
310 shared_up_scratch,
311 &sh.down_proj,
312 shared_down_out,
313 h,
314 inter,
315 top_k,
316 stream,
317 )?;
318 // EP fix: after silu_down, expert_gate_out is free — use as zero buffer
319 // to exclude shared expert from blend (will add after all-reduce).
320 let shared_for_blend = if is_ep && !shared_down_out.is_null() {
321 ctx.gpu
322 .memset_async(expert_gate_out, 0, 2 * h as usize * 2, stream)?;
323 expert_gate_out
324 } else {
325 shared_down_out
326 };
327 ops::moe_weighted_sum_blend_batch2(
328 ctx.gpu,
329 self.moe_weighted_sum_blend_fp8_batch2,
330 output,
331 expert_down_out,
332 weights_dev,
333 shared_for_blend,
334 input,
335 self.weights.shared_expert_gate.weight,
336 h,
337 top_k,
338 h,
339 stream,
340 )?;
341 } else if self.use_t_layout_for_decode() {
342 self.forward_k2_unified_t(
343 input,
344 indices_dev,
345 weights_dev,
346 expert_gate_out,
347 expert_up_out,
348 expert_down_out,
349 shared_gate_scratch,
350 shared_up_scratch,
351 shared_down_out,
352 output,
353 inter,
354 h,
355 top_k,
356 is_ep,
357 mixed_bf16_shared,
358 ctx,
359 stream,
360 )?;
361 } else {
362 self.forward_k2_originals(
363 input,
364 indices_dev,
365 weights_dev,
366 expert_gate_out,
367 expert_up_out,
368 expert_down_out,
369 shared_gate_scratch,
370 shared_up_scratch,
371 shared_down_out,
372 output,
373 inter,
374 h,
375 top_k,
376 is_ep,
377 mixed_bf16_shared,
378 ctx,
379 stream,
380 )?;
381 }
382
383 if k2_diag {
384 ctx.gpu
385 .synchronize(stream)
386 .context("K2: expert dispatch (gate_up/silu_down/blend)")?;
387 }
388
389 // EP all-reduce: sum partial outputs for 2 tokens
390 if let Some(comm) = ctx.comm
391 && ctx.config.ep_world_size > 1
392 {
393 if ctx.graph_capture {
394 comm.all_reduce(output.0, 2 * h as usize * 2)?;
395 } else {
396 comm.all_reduce_async(output.0, 2 * h as usize * 2, stream)?;
397 }
398 // Add shared expert with sigmoid gate (BUG #41 fix)
399 if !shared_down_out.is_null() {
400 if self.weights.shared_expert_gate.weight.0 == 0 {
401 ops::residual_add(
402 ctx.gpu,
403 self.residual_add,
404 output,
405 shared_down_out,
406 2 * h,
407 stream,
408 )?;
409 } else {
410 ops::moe_batched_blend(
411 ctx.gpu,
412 self.moe_batched_blend,
413 output,
414 shared_down_out,
415 input,
416 self.weights.shared_expert_gate.weight,
417 h,
418 2,
419 stream,
420 )?;
421 }
422 }
423 }
424
425 Ok(())
426 }
427}
428
429// Pure dispatch helpers live in a sibling file (500-LoC cap).
430mod forward_k2_helpers;
431pub(crate) use forward_k2_helpers::{batch2_block_width, k2_e8m0_needs_per_token};
432
433// Focused dispatch tests live in a sibling file to keep this file ≤500 LoC.
434#[cfg(test)]
435#[path = "forward_k2_dispatch_tests.rs"]
436mod k2_dispatch_tests;