spark_model/layers/moe/
init.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! MoeLayer::new constructor.
4
5use super::*;
6
7impl MoeLayer {
8    pub fn new(
9        weights: MoeWeights,
10        num_experts: usize,
11        gate_nvfp4: Option<QuantizedWeight>,
12        gpu: &dyn GpuBackend,
13        config: &atlas_core::config::ModelConfig,
14    ) -> Result<Self> {
15        Self::new_with_hash(weights, num_experts, gate_nvfp4, None, gpu, config)
16    }
17
18    /// Like [`MoeLayer::new`] but with an optional DeepSeek-V4 hash-routing
19    /// `tid2eid` table ([vocab_size, top_k] i64). `Some` marks this as a
20    /// hash-routed layer.
21    #[allow(clippy::too_many_arguments)]
22    pub fn new_with_hash(
23        weights: MoeWeights,
24        num_experts: usize,
25        gate_nvfp4: Option<QuantizedWeight>,
26        tid2eid_dev: Option<DevicePtr>,
27        gpu: &dyn GpuBackend,
28        config: &atlas_core::config::ModelConfig,
29    ) -> Result<Self> {
30        // Sanity-check the routing config: top-k that exceeds the
31        // expert count would index OOB in the topk kernel and produce
32        // silent NaN routing. Catch the misconfiguration at load time.
33        anyhow::ensure!(
34            config.num_experts_per_tok <= num_experts && num_experts > 0,
35            "MoE config invalid: num_experts_per_tok={} must be in 1..={}",
36            config.num_experts_per_tok,
37            num_experts,
38        );
39        // The check above bounds top-k by the expert count, which is the OOB
40        // the topk kernel can READ. It says nothing about the OOB the kernel
41        // can WRITE: the sigmoid routing kernels stage their top-K in a
42        // fixed-size shared array, and `top_k` was passed to them unbounded.
43        anyhow::ensure!(
44            config.num_experts_per_tok <= crate::layers::ops::MOE_TOPK_SIGMOID_MAX_TOP_K
45                && num_experts <= crate::layers::ops::MOE_TOPK_SIGMOID_MAX_EXPERTS,
46            "MoE config exceeds the routing kernels' fixed shared-memory bounds: \
47             num_experts_per_tok={} (max {}), num_experts={} (max {}). Raise \
48             MAX_TOP_K / MAX_EXPERTS in kernels/gb10/common/moe_topk_sigmoid.cu \
49             and their mirrors in layers::ops together.",
50            config.num_experts_per_tok,
51            crate::layers::ops::MOE_TOPK_SIGMOID_MAX_TOP_K,
52            num_experts,
53            crate::layers::ops::MOE_TOPK_SIGMOID_MAX_EXPERTS,
54        );
55        let gate_ptrs = build_ptr_table(&weights.experts, |e| &e.gate_proj, gpu)?;
56        let up_ptrs = build_ptr_table(&weights.experts, |e| &e.up_proj, gpu)?;
57        let down_ptrs = build_ptr_table(&weights.experts, |e| &e.down_proj, gpu)?;
58
59        // Extract the optional correction-bias device pointer before the
60        // struct literal below moves `weights`. `.map(|dw| dw.weight)` turns
61        // an `Option<DenseWeight>` into an `Option<DevicePtr>` for the
62        // `moe_topk_sigmoid` kernel's bias arg.
63        let weights_correction_bias: Option<DevicePtr> =
64            weights.correction_bias.map(|dw| dw.weight);
65
66        let _ = num_experts;
67        let rms_norm_k = gpu.kernel("norm", "rms_norm")?;
68        Ok(Self {
69            weights,
70            // Default: standard NVFP4 (FP8-E4M3 per-16 + f32 global). The
71            // DeepSeek-V4 native-MXFP4 loader overrides this to `Mxfp4E8m0`
72            // after construction (see deepseek_v4/assemble.rs).
73            experts_scale_kind: crate::weight_map::WeightQuantFormat::Nvfp4,
74            shared_experts_scale_kind: crate::weight_map::WeightQuantFormat::Nvfp4,
75            gate_nvfp4,
76            pre_expert_norm: None,
77            pre_expert_norm_k: rms_norm_k,
78            dense_gemv: gpu.kernel("gemv", "dense_gemv_bf16")?,
79            w4a16_gemv: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
80            w4a16_gemv_sw: super::super::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_sw"),
81            w4a16_gemm: gpu.kernel("w4a16", "w4a16_gemm")?,
82            dense_gemm: gpu.kernel("gemm", "dense_gemm_bf16")?,
83            dense_gemm_router: super::super::try_kernel(gpu, "gemm", "dense_gemm_bf16_router"),
84            dense_gemm_pipelined: super::super::try_kernel(
85                gpu,
86                "gemm",
87                "dense_gemm_bf16_pipelined",
88            ),
89            // FP32 gate path (ATLAS_FP32_GATE) — optional; KernelHandle(0) if the
90            // target's kernel set predates these symbols, dispatch then stays BF16.
91            dense_gemm_f32out: super::super::try_kernel(gpu, "gemm", "dense_gemm_bf16_f32out"),
92            dense_gemm_f32in: super::super::try_kernel(gpu, "gemm", "dense_gemm_f32in_f32out"),
93            moe_topk_f32: super::super::try_kernel(gpu, "moe_topk", "moe_topk_softmax_f32"),
94            moe_expert_gate_up_shared: gpu
95                .kernel("moe_shared_expert_fused", "moe_expert_gate_up_shared")?,
96            moe_expert_silu_down_shared: gpu
97                .kernel("moe_shared_expert_fused", "moe_expert_silu_down_shared")?,
98            moe_topk: gpu.kernel("moe_topk", "moe_topk_softmax")?,
99            moe_weighted_sum_blend: gpu.kernel("moe_expert_gemv", "moe_weighted_sum_blend")?,
100            residual_add: gpu.kernel("residual_add", "bf16_residual_add")?,
101            moe_topk_batched: gpu.kernel("moe_topk", "moe_topk_softmax_batched")?,
102            moe_expert_gate_up_shared_batch2: gpu
103                .kernel("moe_fused_batch2", "moe_expert_gate_up_shared_batch2")?,
104            moe_expert_silu_down_shared_batch2: gpu
105                .kernel("moe_fused_batch2", "moe_expert_silu_down_shared_batch2")?,
106            moe_weighted_sum_blend_batch2: gpu
107                .kernel("moe_fused_batch2", "moe_weighted_sum_blend_batch2")?,
108            w4a16_gemv_batch2: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch2")?,
109            moe_expert_gate_up_shared_batch3: gpu
110                .kernel("moe_fused_batch3", "moe_expert_gate_up_shared_batch3")?,
111            moe_expert_silu_down_shared_batch3: gpu
112                .kernel("moe_fused_batch3", "moe_expert_silu_down_shared_batch3")?,
113            moe_weighted_sum_blend_batch3: gpu
114                .kernel("moe_fused_batch3", "moe_weighted_sum_blend_batch3")?,
115            w4a16_gemv_batch3: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch3")?,
116            moe_expert_gate_up_shared_token_major: gpu
117                .kernel("moe_prefill", "moe_expert_gate_up_shared_prefill")?,
118            moe_expert_silu_down_shared_token_major: gpu
119                .kernel("moe_prefill", "moe_expert_silu_down_shared_prefill")?,
120            moe_weighted_sum_blend_token_major: gpu
121                .kernel("moe_prefill", "moe_weighted_sum_blend_prefill")?,
122            moe_decode_atomic_c4_silu_down_accum_k: super::super::try_kernel(
123                gpu,
124                "moe_decode_atomic_c4",
125                "moe_decode_atomic_c4_silu_down_accum",
126            ),
127            moe_decode_atomic_c4_finalize_k: super::super::try_kernel(
128                gpu,
129                "moe_decode_atomic_c4",
130                "moe_decode_atomic_c4_finalize",
131            ),
132            moe_sort_by_expert: gpu.kernel("moe", "moe_sort_by_expert")?,
133            moe_sorted_gate_up: gpu.kernel("moe_sorted", "moe_sorted_gate_up")?,
134            moe_sorted_silu_down: gpu.kernel("moe_sorted", "moe_sorted_silu_down")?,
135            moe_grouped_gemm: gpu.kernel("moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable")?,
136            moe_grouped_gemm_k32: if std::env::var("ATLAS_MOE_GROUPED_K32").as_deref() == Ok("1") {
137                super::super::try_kernel(gpu, "moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_k32")
138            } else {
139                KernelHandle(0)
140            },
141            moe_grouped_gemm_m256: if std::env::var("ATLAS_MOE_GROUPED_M256").as_deref() == Ok("1")
142            {
143                super::super::try_kernel(gpu, "moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_m256")
144            } else {
145                KernelHandle(0)
146            },
147            moe_grouped_gemm_t: gpu.kernel("moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_t")?,
148            moe_grouped_gemm_t_k64: gpu
149                .kernel("moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_t_k64")?,
150            moe_fused_gate_up_t: gpu.kernel("moe_w4a16", "moe_w4a16_fused_gate_up_t")?,
151            moe_fused_gate_up_t_k64: gpu.kernel("moe_w4a16", "moe_w4a16_fused_gate_up_t_k64")?,
152            // ARM-2 Phase-K native-MXFP4 (E8M0) prefill variants — try_kernel:
153            // only the deepseek-v4-flash target's moe_w4a16 module ships them.
154            moe_grouped_gemm_e8m0: super::super::try_kernel(
155                gpu,
156                "moe_w4a16",
157                "moe_w4a16_grouped_gemm_ptrtable_e8m0",
158            ),
159            moe_grouped_gemm_t_e8m0: super::super::try_kernel(
160                gpu,
161                "moe_w4a16",
162                "moe_w4a16_grouped_gemm_ptrtable_t_e8m0",
163            ),
164            moe_grouped_gemm_t_k64_e8m0: super::super::try_kernel(
165                gpu,
166                "moe_w4a16",
167                "moe_w4a16_grouped_gemm_ptrtable_t_k64_e8m0",
168            ),
169            moe_fused_gate_up_t_e8m0: super::super::try_kernel(
170                gpu,
171                "moe_w4a16",
172                "moe_w4a16_fused_gate_up_t_e8m0",
173            ),
174            moe_fused_gate_up_t_k64_e8m0: super::super::try_kernel(
175                gpu,
176                "moe_w4a16",
177                "moe_w4a16_fused_gate_up_t_k64_e8m0",
178            ),
179            // M=128 variant only present in models where Block D #3 has
180            // been ported (currently minimax-m2-229b). Other models keep
181            // KernelHandle(0) and dispatch falls through to M=64.
182            moe_fused_gate_up_t_k64_m128: super::super::try_kernel(
183                gpu,
184                "moe_w4a16",
185                "moe_w4a16_fused_gate_up_t_k64_m128",
186            ),
187            // FUSED FP4 gate_up kernel (ATLAS_HOLO_MOE_GATEUP_FP4). try_kernel:
188            // KernelHandle(0) on images that didn't compile it; the FP4 dispatch
189            // checks this handle != 0 before firing.
190            moe_fused_gate_up_t_k64_fp4: super::super::try_kernel(
191                gpu,
192                "moe_w4a16",
193                "moe_w4a16_fused_gate_up_t_k64_fp4",
194            ),
195            moe_fp8_grouped_gemm_t: gpu.kernel("moe_w4a16", "moe_fp8_grouped_gemm_ptrtable_t")?,
196            // THE routed-expert FP8 prefill kernel: grid-compaction (persistent
197            // 96-CTA grid over a compacted work-list). Handle may be 0 on older
198            // images that don't ship it.
199            moe_fp8_grouped_gemm_k: super::super::try_kernel(
200                gpu,
201                "moe_fp8_grouped_gemm",
202                "moe_fp8_grouped_gemm",
203            ),
204            // Work-list builder (module "moe" = moe_permute.cu). Launched on the
205            // SAME stream as the grouped GEMM (read-after-write of total_tiles).
206            moe_build_tile_worklist_k: super::super::try_kernel(
207                gpu,
208                "moe",
209                "moe_build_tile_worklist",
210            ),
211            moe_w8a8_grouped_gemm_k: super::super::try_kernel(
212                gpu,
213                "moe_w8a8_grouped_gemm",
214                "moe_w8a8_grouped_gemm",
215            ),
216            // PM4-geometry W8A8 grouped GEMM (same module). Handle may be 0 on
217            // targets/images without it; dispatch falls back to the dense grid.
218            moe_w8a8_grouped_gemm_pm4_k: super::super::try_kernel(
219                gpu,
220                "moe_w8a8_grouped_gemm",
221                "moe_w8a8_grouped_gemm_pm4",
222            ),
223            per_token_group_quant_fp8_k: super::super::try_kernel(
224                gpu,
225                "per_token_group_quant_fp8",
226                "per_token_group_quant_fp8",
227            ),
228            // Fused silu_mul + per-token-group quant. Same module as
229            // moe_silu_mul, so a model that shadows moe_silu_mul.cu without
230            // this entry point gets handle 0 → unfused fallback.
231            silu_mul_quant_fp8_k: super::super::try_kernel(
232                gpu,
233                "moe_silu_mul",
234                "silu_mul_quant_fp8",
235            ),
236            fp8_gemm_t_blockscaled_k: super::super::try_kernel(
237                gpu,
238                "fp8_gemm_t_blockscaled",
239                "fp8_gemm_t_blockscaled",
240            ),
241            moe_bf16_grouped_gemm_k: super::super::try_kernel(
242                gpu,
243                "moe_bf16_grouped_gemm",
244                "moe_bf16_grouped_gemm",
245            ),
246            moe_expert_gate_up_shared_bf16_k: super::super::try_kernel(
247                gpu,
248                "moe_shared_expert_fused_bf16",
249                "moe_expert_gate_up_shared_bf16",
250            ),
251            moe_expert_silu_down_shared_bf16_k: super::super::try_kernel(
252                gpu,
253                "moe_shared_expert_fused_bf16",
254                "moe_expert_silu_down_shared_bf16",
255            ),
256            moe_expert_gate_up_shared_bf16_batch2_k: super::super::try_kernel(
257                gpu,
258                "moe_shared_expert_fused_bf16_batch2",
259                "moe_expert_gate_up_shared_bf16_batch2",
260            ),
261            moe_expert_silu_down_shared_bf16_batch2_k: super::super::try_kernel(
262                gpu,
263                "moe_shared_expert_fused_bf16_batch2",
264                "moe_expert_silu_down_shared_bf16_batch2",
265            ),
266            w8a16_gemm_k: super::super::try_kernel(gpu, "w8a16_gemm", "w8a16_gemm"),
267            w8a16_gemm_pipelined_k: super::super::try_kernel(
268                gpu,
269                "w8a16_gemm_pipelined",
270                "w8a16_gemm_pipelined",
271            ),
272            moe_gate_topk_fused_k: super::super::try_kernel(
273                gpu,
274                "moe_gate_topk",
275                "moe_gate_topk_fused",
276            ),
277            w4a16_gemm_t: gpu.kernel("w4a16", "w4a16_gemm_t")?,
278            bf16_to_fp8_k: gpu.kernel("w4a16", "bf16_to_fp8")?,
279            fp8_gemm_k: gpu.kernel("w4a16", "fp8_gemm_t")?,
280            moe_silu_mul: gpu.kernel("moe_silu_mul", "moe_silu_mul")?,
281            moe_act_mul: gpu.kernel("moe_silu_mul", "moe_silu_mul")?, // default: SiLU
282            gelu_activation: false,
283            moe_unpermute_reduce: gpu.kernel("moe", "moe_unpermute_reduce_indexed")?,
284            moe_batched_blend: gpu.kernel("moe", "moe_batched_blend")?,
285            gate_ptrs,
286            up_ptrs,
287            down_ptrs,
288            gate_ptrs_t: None,
289            up_ptrs_t: None,
290            down_ptrs_t: None,
291            cutlass_grouped_host: None,
292            _cutlass_sfb_owned: Vec::new(),
293            down_t_scratch_packed: None,
294            down_t_scratch_scale: None,
295            moe_transpose_u8_batched_k: gpu
296                .kernel("moe_transpose_batched", "moe_transpose_u8_batched")?,
297            // ── Phase 8a transposed-layout decode kernels ──
298            // Module name = file stem (default convention in atlas-kernels).
299            moe_expert_gate_up_shared_t_k: gpu
300                .kernel("moe_shared_expert_fused_t", "moe_expert_gate_up_shared_t")?,
301            moe_expert_silu_down_shared_t_k: gpu
302                .kernel("moe_shared_expert_fused_t", "moe_expert_silu_down_shared_t")?,
303            // ARM-2 Phase-K dual-format decode variants (E8M0 routed / NVFP4
304            // shared). try_kernel — the entries are in the common .cu but load
305            // by name; 0 where a target doesn't compile that module.
306            moe_expert_gate_up_shared_t_e8m0_k: super::super::try_kernel(
307                gpu,
308                "moe_shared_expert_fused_t",
309                "moe_expert_gate_up_shared_t_e8m0",
310            ),
311            moe_expert_silu_down_shared_t_e8m0_k: super::super::try_kernel(
312                gpu,
313                "moe_shared_expert_fused_t",
314                "moe_expert_silu_down_shared_t_e8m0",
315            ),
316            // sqrtsoftplus kernels: lazy-loaded via try_kernel so models that
317            // don't register them (all except DeepSeek-V4) start fine.
318            moe_topk_sqrtsoftplus_k: super::super::try_kernel(
319                gpu,
320                "moe_topk_sqrt",
321                "moe_topk_sqrtsoftplus",
322            ),
323            moe_topk_sqrtsoftplus_batched_k: super::super::try_kernel(
324                gpu,
325                "moe_topk_sqrt",
326                "moe_topk_sqrtsoftplus_batched",
327            ),
328            // Hash routing (DeepSeek-V4 hash_moe layers): lazy-loaded so other
329            // models start fine. `tid2eid_dev` is the per-layer table (Some
330            // only for hash layers).
331            router_logits_n: (config.num_experts + config.zero_expert_num) as u32,
332            moe_topk_softmax_bias_k: super::super::try_kernel(
333                gpu,
334                "moe_topk_softmax_bias",
335                "moe_topk_softmax_bias",
336            ),
337            moe_topk_softmax_bias_batched_k: super::super::try_kernel(
338                gpu,
339                "moe_topk_softmax_bias",
340                "moe_topk_softmax_bias_batched",
341            ),
342            moe_zero_expert_add_k: super::super::try_kernel(
343                gpu,
344                "moe_topk_softmax_bias",
345                "moe_zero_expert_add",
346            ),
347            // 64 KB, unconditional: written by the softmax+bias router even
348            // with zero_expert_num == 0 (always zeros then).
349            zero_accum_dev: gpu.alloc(16384 * 4)?,
350            moe_hash_route_k: super::super::try_kernel(gpu, "moe_hash_route", "moe_hash_route"),
351            moe_hash_route_batched_k: super::super::try_kernel(
352                gpu,
353                "moe_hash_route",
354                "moe_hash_route_batched",
355            ),
356            tid2eid_dev,
357            moe_expert_gate_up_shared_batch2_t_k: gpu.kernel(
358                "moe_shared_expert_fused_batch2_t",
359                "moe_expert_gate_up_shared_batch2_t",
360            )?,
361            moe_expert_silu_down_shared_batch2_t_k: gpu.kernel(
362                "moe_shared_expert_fused_batch2_t",
363                "moe_expert_silu_down_shared_batch2_t",
364            )?,
365            moe_expert_gate_up_shared_batch3_t_k: gpu.kernel(
366                "moe_shared_expert_fused_batch3_t",
367                "moe_expert_gate_up_shared_batch3_t",
368            )?,
369            moe_expert_silu_down_shared_batch3_t_k: gpu.kernel(
370                "moe_shared_expert_fused_batch3_t",
371                "moe_expert_silu_down_shared_batch3_t",
372            )?,
373            moe_expert_gate_up_shared_fp8_t_k: gpu.kernel(
374                "moe_shared_expert_fused_fp8_t",
375                "moe_expert_gate_up_shared_fp8_t",
376            )?,
377            moe_expert_silu_down_shared_fp8_t_k: gpu.kernel(
378                "moe_shared_expert_fused_fp8_t",
379                "moe_expert_silu_down_shared_fp8_t",
380            )?,
381            moe_expert_gate_up_shared_fp8_batch2_t_k: gpu.kernel(
382                "moe_shared_expert_fused_fp8_batch2_t",
383                "moe_expert_gate_up_shared_fp8_batch2_t",
384            )?,
385            moe_expert_silu_down_shared_fp8_batch2_t_k: gpu.kernel(
386                "moe_shared_expert_fused_fp8_batch2_t",
387                "moe_expert_silu_down_shared_fp8_batch2_t",
388            )?,
389            moe_expert_gate_up_shared_fp8_batch3_t_k: gpu.kernel(
390                "moe_shared_expert_fused_fp8_batch3_t",
391                "moe_expert_gate_up_shared_fp8_batch3_t",
392            )?,
393            moe_expert_silu_down_shared_fp8_batch3_t_k: gpu.kernel(
394                "moe_shared_expert_fused_fp8_batch3_t",
395                "moe_expert_silu_down_shared_fp8_batch3_t",
396            )?,
397            unified_layout: std::env::var("ATLAS_UNIFIED_MOE_LAYOUT")
398                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
399                .unwrap_or(false),
400            hybrid_layout: std::env::var("ATLAS_HYBRID_MOE_LAYOUT")
401                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
402                .unwrap_or(false),
403            nvfp4_gate_up_m128: std::env::var("ATLAS_NVFP4_GATE_UP_M128")
404                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
405                .unwrap_or(false),
406            // FP4 prefill MoE over the shared FAST_MOE=full [K/2,N] tables.
407            gateup_fp4: std::env::var("ATLAS_HOLO_MOE_GATEUP_FP4")
408                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
409                .unwrap_or(false),
410            down_fp4: std::env::var("ATLAS_HOLO_MOE_DOWN_FP4")
411                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
412                .unwrap_or(false),
413            shared_gate_t: None,
414            shared_up_t: None,
415            shared_down_t: None,
416            gate_fp8: None,
417            shared_gate_fp8: None,
418            shared_up_fp8: None,
419            shared_down_fp8: None,
420            prefill_stream: gpu.create_stream()?,
421            event_a: gpu.create_event()?,
422            event_b: gpu.create_event()?,
423            moe_expert_gate_up_shared_fp8: gpu.kernel(
424                "moe_shared_expert_fused_fp8",
425                "moe_expert_gate_up_shared_fp8",
426            )?,
427            moe_expert_silu_down_shared_fp8: gpu.kernel(
428                "moe_shared_expert_fused_fp8",
429                "moe_expert_silu_down_shared_fp8",
430            )?,
431            // FP8 batch2/3 kernels for MTP verify
432            moe_expert_gate_up_shared_fp8_batch2: gpu.kernel(
433                "moe_shared_expert_fused_fp8_batch2",
434                "moe_expert_gate_up_shared_fp8_batch2",
435            )?,
436            moe_expert_silu_down_shared_fp8_batch2: gpu.kernel(
437                "moe_shared_expert_fused_fp8_batch2",
438                "moe_expert_silu_down_shared_fp8_batch2",
439            )?,
440            moe_weighted_sum_blend_fp8_batch2: gpu.kernel(
441                "moe_shared_expert_fused_fp8_batch2",
442                "moe_weighted_sum_blend_fp8_batch2",
443            )?,
444            moe_expert_gate_up_shared_fp8_batch3: gpu.kernel(
445                "moe_shared_expert_fused_fp8_batch3",
446                "moe_expert_gate_up_shared_fp8_batch3",
447            )?,
448            moe_expert_silu_down_shared_fp8_batch3: gpu.kernel(
449                "moe_shared_expert_fused_fp8_batch3",
450                "moe_expert_silu_down_shared_fp8_batch3",
451            )?,
452            moe_weighted_sum_blend_fp8_batch3: gpu.kernel(
453                "moe_shared_expert_fused_fp8_batch3",
454                "moe_weighted_sum_blend_fp8_batch3",
455            )?,
456            fp8_gate_weight_ptrs: None,
457            fp8_up_weight_ptrs: None,
458            fp8_down_weight_ptrs: None,
459            bf16_gate_weight_ptrs: None,
460            bf16_up_weight_ptrs: None,
461            bf16_down_weight_ptrs: None,
462            bf16_shared_expert: None,
463            fp8_shared_expert: None,
464            moe_down_t_k64_fp4: super::super::try_kernel(
465                gpu,
466                "moe_w4a16",
467                "moe_w4a16_down_t_k64_fp4",
468            ),
469            moe_permute_tokens_k: super::super::try_kernel(gpu, "moe", "moe_permute_tokens"),
470            // Phase 2.7 Tier C — set by loader after construction (qwen35.rs).
471            is_dflash_capture_layer: false,
472            lora: None,
473            correction_bias_dev: weights_correction_bias,
474            // `moe_topk_sig` is only registered for sigmoid-gated MoE models
475            // (MiniMax-M2, Nemotron-Nano, Nemotron-Super). Softmax-gated MoEs
476            // (Qwen3.5, Qwen3-Next, Gemma-4, Mistral) never hit the sigmoid
477            // dispatch path, so a missing kernel is fine — fail at call time
478            // via the KernelHandle(0) check in ops::moe_topk_sigmoid rather
479            // than at MoeLayer::new(), which would otherwise block all
480            // softmax-MoE model startup (observed on Qwen3.5-35B-A3B-FP8 in
481            // alpha-2.43: "Module 'moe_topk_sig' not loaded" during model
482            // build).
483            moe_topk_sigmoid_k: super::super::try_kernel(gpu, "moe_topk_sig", "moe_topk_sigmoid"),
484            moe_topk_sigmoid_batched_k: super::super::try_kernel(
485                gpu,
486                "moe_topk_sig",
487                "moe_topk_sigmoid_batched",
488            ),
489        })
490    }
491}
spark_model/layers/moe/init.rs

spark_model/layers/moe/
init.rs
