spark_model/layers/moe/
forward_atomic_c4.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Purpose-built C=4 atomic-add MoE decode experiment.
4
5use super::*;
6
7impl MoeLayer {
8    /// C=4 NVFP4 routed MoE decode with FP32 atomic accumulation.
9    ///
10    /// Gate/top-K remain batched. Gate+up reuses the token-major kernel, then
11    /// routed down projections atomic-add weighted FP32 contributions into a
12    /// tiny `[4,H]` scratch accumulator. Finalization casts routed output to
13    /// BF16 and optionally blends shared expert output.
14    pub fn forward_atomic_c4_decode(
15        &self,
16        input: DevicePtr,
17        num_tokens: usize,
18        ctx: &ForwardContext,
19        stream: u64,
20    ) -> Result<()> {
21        // LongCat zero-experts are wired only on the single-token decode
22        // + prefill paths (v1); this variant would silently mis-route the
23        // 384-wide router. Named refusal, not silent wrongness.
24        anyhow::ensure!(
25            self.router_logits_n as usize == ctx.config.num_experts,
26            "zero-expert MoE routing is not wired on this dispatch variant yet (forward_atomic_c4)"
27        );
28
29        // Feature-1 phase-1: decode does not yet fold the expert delta.
30        self.reject_decode_lora(ctx, "forward_atomic_c4_decode")?;
31        let has_shared = self.weights.shared_expert.gate_proj.weight.0 != 0
32            && self.weights.shared_expert.up_proj.weight.0 != 0
33            && self.weights.shared_expert.down_proj.weight.0 != 0;
34        let nvfp4_supported = num_tokens == 4
35            && self.moe_decode_atomic_c4_silu_down_accum_k.0 != 0
36            && self.moe_decode_atomic_c4_finalize_k.0 != 0
37            && self.bf16_gate_weight_ptrs.is_none()
38            && !self.has_mixed_bf16_shared_expert()
39            && self.fp8_gate_weight_ptrs.is_none()
40            && !self.use_t_layout_for_decode()
41            && self.pre_expert_norm.is_none()
42            && has_shared;
43        if !nvfp4_supported {
44            return self.forward_batched(input, num_tokens, ctx, stream);
45        }
46
47        let h = ctx.config.hidden_size as u32;
48        let inter = ctx.config.moe_intermediate_size as u32;
49        let num_experts = ctx.config.num_experts as u32;
50        let top_k = ctx.config.num_experts_per_tok as u32;
51        let n = num_tokens as u32;
52
53        let router_in = self.router_input(input, n, h, ctx, stream)?;
54        let gate_logits = ctx.buffers.gate_logits();
55        if let Some(ref nvfp4) = self.gate_nvfp4 {
56            ops::w4a16_gemm(
57                ctx.gpu,
58                self.w4a16_gemm,
59                router_in,
60                nvfp4,
61                gate_logits,
62                n,
63                num_experts,
64                h,
65                stream,
66            )?;
67        } else {
68            ops::dense_gemm(
69                ctx.gpu,
70                self.dense_gemm,
71                router_in,
72                &self.weights.gate,
73                gate_logits,
74                n,
75                num_experts,
76                h,
77                stream,
78            )?;
79        }
80
81        let scratch = ctx.buffers.scratch();
82        let topk_bytes = num_tokens * top_k as usize * 4;
83        let indices_dev = scratch;
84        let weights_dev = scratch.offset(topk_bytes);
85        let accum_off = (topk_bytes * 2 + 255) & !255;
86        let accum_bytes = num_tokens * h as usize * 4;
87        anyhow::ensure!(
88            ctx.buffers.scratch_bytes() >= accum_off + accum_bytes,
89            "scratch too small for ATLAS_MOE_ATOMIC_C4_DECODE: need {} bytes, have {}",
90            accum_off + accum_bytes,
91            ctx.buffers.scratch_bytes()
92        );
93        let routed_accum = scratch.offset(accum_off);
94
95        if let Some(bias) = self.correction_bias_dev {
96            ops::moe_topk_sigmoid_batched(
97                ctx.gpu,
98                self.moe_topk_sigmoid_batched_k,
99                gate_logits,
100                bias,
101                indices_dev,
102                weights_dev,
103                num_experts,
104                top_k,
105                ctx.config.norm_topk_prob,
106                ctx.config.routed_scaling_factor as f32,
107                n,
108                stream,
109            )?;
110        } else {
111            ops::moe_topk_softmax_batched(
112                ctx.gpu,
113                self.moe_topk_batched,
114                gate_logits,
115                indices_dev,
116                weights_dev,
117                num_experts,
118                top_k,
119                ctx.config.norm_topk_prob,
120                n,
121                stream,
122            )?;
123        }
124
125        let expert_gate_out = ctx.buffers.expert_gate_out();
126        let expert_up_out = ctx.buffers.expert_up_out();
127        let shared_gate_scratch = ctx.buffers.logits();
128        let shared_up_scratch = ctx.buffers.ssm_qkvz();
129        let shared_down_out = ctx.buffers.attn_output();
130        let output = ctx.buffers.moe_output();
131
132        ops::moe_expert_gate_up_shared_prefill(
133            ctx.gpu,
134            self.moe_expert_gate_up_shared_token_major,
135            input,
136            self.gate_ptrs.packed_ptrs,
137            self.gate_ptrs.scale_ptrs,
138            self.gate_ptrs.scale2_vals,
139            expert_gate_out,
140            self.up_ptrs.packed_ptrs,
141            self.up_ptrs.scale_ptrs,
142            self.up_ptrs.scale2_vals,
143            expert_up_out,
144            indices_dev,
145            &self.weights.shared_expert.gate_proj,
146            shared_gate_scratch,
147            &self.weights.shared_expert.up_proj,
148            shared_up_scratch,
149            inter,
150            h,
151            top_k,
152            n,
153            stream,
154        )?;
155
156        ctx.gpu.memset_async(routed_accum, 0, accum_bytes, stream)?;
157        ops::moe_decode_atomic_c4_silu_down_accum(
158            ctx.gpu,
159            self.moe_decode_atomic_c4_silu_down_accum_k,
160            expert_gate_out,
161            expert_up_out,
162            self.down_ptrs.packed_ptrs,
163            self.down_ptrs.scale_ptrs,
164            self.down_ptrs.scale2_vals,
165            indices_dev,
166            weights_dev,
167            routed_accum,
168            shared_gate_scratch,
169            shared_up_scratch,
170            &self.weights.shared_expert.down_proj,
171            shared_down_out,
172            h,
173            inter,
174            top_k,
175            n,
176            stream,
177        )?;
178
179        let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1;
180        ops::moe_decode_atomic_c4_finalize(
181            ctx.gpu,
182            self.moe_decode_atomic_c4_finalize_k,
183            output,
184            routed_accum,
185            shared_down_out,
186            input,
187            self.weights.shared_expert_gate.weight,
188            h,
189            n,
190            !is_ep,
191            stream,
192        )?;
193
194        if let Some(comm) = ctx.comm
195            && ctx.config.ep_world_size > 1
196        {
197            if ctx.graph_capture {
198                comm.all_reduce(output.0, num_tokens * h as usize * 2)?;
199            } else {
200                comm.all_reduce_async(output.0, num_tokens * h as usize * 2, stream)?;
201            }
202            ops::moe_batched_blend(
203                ctx.gpu,
204                self.moe_batched_blend,
205                output,
206                shared_down_out,
207                input,
208                self.weights.shared_expert_gate.weight,
209                h,
210                n,
211                stream,
212            )?;
213        }
214
215        Ok(())
216    }
217}
spark_model/layers/moe/forward_atomic_c4.rs

spark_model/layers/moe/
forward_atomic_c4.rs
