metrale_model_layers/layers/moe/
init.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `MoeLayer::new` / `new_with_hash`: validate the routing config, build
4//! the expert pointer tables and resolve the kernels.
5//!
6//! Owner: model-layers (MoE).
7//! Invariants:
8//! - A constructed layer has `1 <= num_experts_per_tok <= num_experts`,
9//!   `num_experts_per_tok <= MOE_TOPK_SIGMOID_MAX_TOP_K` and
10//!   `num_experts <= MOE_TOPK_SIGMOID_MAX_EXPERTS`; otherwise construction errors.
11
12use super::*;
13use crate::layers::try_target_kernel;
14
15impl MoeLayer {
16    pub fn new(
17        weights: MoeWeights,
18        num_experts: usize,
19        gate_nvfp4: Option<QuantizedWeight>,
20        gpu: &dyn GpuBackend,
21        config: &metrale_config::ModelConfig,
22    ) -> Result<Self> {
23        Self::new_with_hash(weights, num_experts, gate_nvfp4, None, gpu, config)
24    }
25
26    /// 2026-09-25: Like [`MoeLayer::new`], with an optional hash-routing table
27    /// `tid2eid` (`[vocab_size, top_k]` i64 on the device). `Some` makes this a
28    /// hash-routed layer.
29    #[allow(clippy::too_many_arguments)]
30    pub fn new_with_hash(
31        weights: MoeWeights,
32        num_experts: usize,
33        gate_nvfp4: Option<QuantizedWeight>,
34        tid2eid_dev: Option<DevicePtr>,
35        gpu: &dyn GpuBackend,
36        config: &metrale_config::ModelConfig,
37    ) -> Result<Self> {
38        anyhow::ensure!(
39            config.num_experts_per_tok <= num_experts && num_experts > 0,
40            "MoE config invalid: num_experts_per_tok={} must be in 1..={}",
41            config.num_experts_per_tok,
42            num_experts,
43        );
44        // 2026-09-25: The sigmoid routing kernels hold at most MAX_TOP_K
45        // selections and MAX_EXPERTS experts in shared memory and ignore the
46        // rest, so such configs are refused here.
47        anyhow::ensure!(
48            config.num_experts_per_tok <= crate::layers::ops::MOE_TOPK_SIGMOID_MAX_TOP_K
49                && num_experts <= crate::layers::ops::MOE_TOPK_SIGMOID_MAX_EXPERTS,
50            "MoE config exceeds the routing kernels' fixed shared-memory bounds: \
51             num_experts_per_tok={} (max {}), num_experts={} (max {}). Raise \
52             MAX_TOP_K / MAX_EXPERTS in kernels/gb10/common/moe_topk_sigmoid.cu \
53             and their mirrors in layers::ops together.",
54            config.num_experts_per_tok,
55            crate::layers::ops::MOE_TOPK_SIGMOID_MAX_TOP_K,
56            num_experts,
57            crate::layers::ops::MOE_TOPK_SIGMOID_MAX_EXPERTS,
58        );
59        let gate_ptrs = build_ptr_table(&weights.experts, |e| &e.gate_proj, gpu)?;
60        let up_ptrs = build_ptr_table(&weights.experts, |e| &e.up_proj, gpu)?;
61        let down_ptrs = build_ptr_table(&weights.experts, |e| &e.down_proj, gpu)?;
62
63        // 2026-09-25: Read before the struct literal moves `weights`.
64        let weights_correction_bias: Option<DevicePtr> =
65            weights.correction_bias.map(|dw| dw.weight);
66
67        let _ = num_experts;
68        let rms_norm_k = gpu.kernel("norm", "rms_norm")?;
69        let grouped = super::forward_fp8_grouped_decode::GroupedKernels::resolve(gpu);
70        let nvfp4_grouped = super::forward_nvfp4_grouped_decode::Nvfp4GroupedKernels::resolve(gpu);
71        let moe_e4m3 = super::forward_prefill_fp8::E4m3Kernels::resolve(gpu)?;
72        Ok(Self {
73            weights,
74            // 2026-09-25: NVFP4 until a loader says otherwise; the DeepSeek-V4
75            // loader sets both kinds after construction (deepseek_v4/assemble/layer.rs).
76            experts_scale_kind: crate::weight_map::WeightQuantFormat::Nvfp4,
77            shared_experts_scale_kind: crate::weight_map::WeightQuantFormat::Nvfp4,
78            gate_nvfp4,
79            pre_expert_norm: None,
80            pre_expert_norm_k: rms_norm_k,
81            dense_gemv: gpu.kernel("gemv", "dense_gemv_bf16")?,
82            w4a16_gemv: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
83            w4a16_gemv_sw: super::super::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_sw"),
84            w4a16_gemm: gpu.kernel("w4a16", "w4a16_gemm")?,
85            dense_gemm: gpu.kernel("gemm", "dense_gemm_bf16")?,
86            dense_gemm_router: super::super::try_kernel(gpu, "gemm", "dense_gemm_bf16_router"),
87            moe_router_rt_k: if std::env::var_os("METRALE_NO_MOE_ROUTER_RT").is_some() {
88                KernelHandle(0)
89            } else {
90                try_target_kernel(gpu, "moe_router_gemm_prefill", "moe_router_gemm_rt")
91            },
92            moe_router_gemm_k: super::super::try_kernel(
93                gpu,
94                "moe_router_gemm",
95                "moe_router_gemm_bf16",
96            ),
97            dense_gemm_pipelined: super::super::try_kernel(
98                gpu,
99                "gemm",
100                "dense_gemm_bf16_pipelined",
101            ),
102            // 2026-09-25: FP32 router kernels (METRALE_FP32_GATE, METRALE_FP32_ROUTING);
103            // 0 where the target lacks them.
104            dense_gemm_f32out: super::super::try_kernel(gpu, "gemm", "dense_gemm_bf16_f32out"),
105            dense_gemm_f32in: super::super::try_kernel(gpu, "gemm", "dense_gemm_f32in_f32out"),
106            moe_topk_f32: super::super::try_kernel(gpu, "moe_topk", "moe_topk_softmax_f32"),
107            moe_expert_gate_up_shared: gpu
108                .kernel("moe_shared_expert_fused", "moe_expert_gate_up_shared")?,
109            moe_expert_silu_down_shared: gpu
110                .kernel("moe_shared_expert_fused", "moe_expert_silu_down_shared")?,
111            moe_topk: gpu.kernel("moe_topk", "moe_topk_softmax")?,
112            moe_weighted_sum_blend: gpu.kernel("moe_expert_gemv", "moe_weighted_sum_blend")?,
113            residual_add: gpu.kernel("residual_add", "bf16_residual_add")?,
114            moe_topk_batched: gpu.kernel("moe_topk", "moe_topk_softmax_batched")?,
115            moe_expert_gate_up_shared_batch2: gpu
116                .kernel("moe_fused_batch2", "moe_expert_gate_up_shared_batch2")?,
117            moe_expert_silu_down_shared_batch2: gpu
118                .kernel("moe_fused_batch2", "moe_expert_silu_down_shared_batch2")?,
119            moe_weighted_sum_blend_batch2: gpu
120                .kernel("moe_fused_batch2", "moe_weighted_sum_blend_batch2")?,
121            w4a16_gemv_batch2: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch2")?,
122            moe_expert_gate_up_shared_batch3: gpu
123                .kernel("moe_fused_batch3", "moe_expert_gate_up_shared_batch3")?,
124            moe_expert_silu_down_shared_batch3: gpu
125                .kernel("moe_fused_batch3", "moe_expert_silu_down_shared_batch3")?,
126            moe_weighted_sum_blend_batch3: gpu
127                .kernel("moe_fused_batch3", "moe_weighted_sum_blend_batch3")?,
128            w4a16_gemv_batch3: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch3")?,
129            moe_expert_gate_up_shared_token_major: gpu
130                .kernel("moe_prefill", "moe_expert_gate_up_shared_prefill")?,
131            moe_expert_silu_down_shared_token_major: gpu
132                .kernel("moe_prefill", "moe_expert_silu_down_shared_prefill")?,
133            moe_weighted_sum_blend_token_major: gpu
134                .kernel("moe_prefill", "moe_weighted_sum_blend_prefill")?,
135            moe_decode_atomic_c4_silu_down_accum_k: super::super::try_kernel(
136                gpu,
137                "moe_decode_atomic_c4",
138                "moe_decode_atomic_c4_silu_down_accum",
139            ),
140            moe_decode_atomic_c4_finalize_k: super::super::try_kernel(
141                gpu,
142                "moe_decode_atomic_c4",
143                "moe_decode_atomic_c4_finalize",
144            ),
145            moe_sort_by_expert: gpu.kernel("moe", "moe_sort_by_expert")?,
146            moe_sorted_gate_up: gpu.kernel("moe_sorted", "moe_sorted_gate_up")?,
147            moe_sorted_silu_down: gpu.kernel("moe_sorted", "moe_sorted_silu_down")?,
148            moe_grouped_gemm: gpu.kernel("moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable")?,
149            moe_grouped_gemm_k32: if std::env::var("METRALE_MOE_GROUPED_K32").as_deref() == Ok("1")
150            {
151                super::super::try_kernel(gpu, "moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_k32")
152            } else {
153                KernelHandle(0)
154            },
155            moe_grouped_gemm_m256: if std::env::var("METRALE_MOE_GROUPED_M256").as_deref()
156                == Ok("1")
157            {
158                super::super::try_kernel(gpu, "moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_m256")
159            } else {
160                KernelHandle(0)
161            },
162            moe_grouped_gemm_t: gpu.kernel("moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_t")?,
163            moe_grouped_gemm_t_k64: gpu
164                .kernel("moe_w4a16", "moe_w4a16_grouped_gemm_ptrtable_t_k64")?,
165            moe_fused_gate_up_t: gpu.kernel("moe_w4a16", "moe_w4a16_fused_gate_up_t")?,
166            moe_fused_gate_up_t_k64: gpu.kernel("moe_w4a16", "moe_w4a16_fused_gate_up_t_k64")?,
167            // 2026-09-25: E8M0 prefill kernels; only
168            // kernels/gb10/deepseek-v4-flash/nvfp4/moe_w4a16_grouped_gemm.cu defines them.
169            moe_grouped_gemm_e8m0: super::super::try_kernel(
170                gpu,
171                "moe_w4a16",
172                "moe_w4a16_grouped_gemm_ptrtable_e8m0",
173            ),
174            moe_grouped_gemm_t_e8m0: super::super::try_kernel(
175                gpu,
176                "moe_w4a16",
177                "moe_w4a16_grouped_gemm_ptrtable_t_e8m0",
178            ),
179            moe_grouped_gemm_t_k64_e8m0: super::super::try_kernel(
180                gpu,
181                "moe_w4a16",
182                "moe_w4a16_grouped_gemm_ptrtable_t_k64_e8m0",
183            ),
184            moe_fused_gate_up_t_e8m0: super::super::try_kernel(
185                gpu,
186                "moe_w4a16",
187                "moe_w4a16_fused_gate_up_t_e8m0",
188            ),
189            moe_fused_gate_up_t_k64_e8m0: super::super::try_kernel(
190                gpu,
191                "moe_w4a16",
192                "moe_w4a16_fused_gate_up_t_k64_e8m0",
193            ),
194            // 2026-09-25: Only kernels/gb10/minimax-m2-229b defines the M=128 kernel;
195            // elsewhere the handle is 0 and prefill uses the M=64 kernel.
196            moe_fused_gate_up_t_k64_m128: super::super::try_kernel(
197                gpu,
198                "moe_w4a16",
199                "moe_w4a16_fused_gate_up_t_k64_m128",
200            ),
201            // 2026-09-25: FP4 gate_up kernel (METRALE_HOLO_MOE_GATEUP_FP4); the
202            // dispatch requires a non-zero handle.
203            moe_fused_gate_up_t_k64_fp4: super::super::try_kernel(
204                gpu,
205                "moe_w4a16",
206                "moe_w4a16_fused_gate_up_t_k64_fp4",
207            ),
208            moe_fp8_grouped_gemm_t: gpu.kernel("moe_w4a16", "moe_fp8_grouped_gemm_ptrtable_t")?,
209            // 2026-09-25: The routed FP8 prefill GEMM over the compacted work-list;
210            // 0 where the target lacks it (FP8 prefill then takes forward_batched).
211            moe_fp8_grouped_gemm_k: super::super::try_kernel(
212                gpu,
213                "moe_fp8_grouped_gemm",
214                "moe_fp8_grouped_gemm",
215            ),
216            // 2026-09-25: The work-list builder, in moe_permute.cu (module "moe").
217            moe_build_tile_worklist_k: super::super::try_kernel(
218                gpu,
219                "moe",
220                "moe_build_tile_worklist",
221            ),
222            moe_w8a8_m16_k: try_target_kernel(gpu, "moe_w8a8_m16", "pm4_m16"),
223            moe_bucket_builder_k: try_target_kernel(gpu, "moe_bucket_builder", "bucket_builder"),
224            moe_adaptive_sms: adaptive_fp8::adaptive_sm_count(gpu)?,
225            moe_w8a8_grouped_gemm_k: super::super::try_kernel(
226                gpu,
227                "moe_w8a8_grouped_gemm",
228                "moe_w8a8_grouped_gemm",
229            ),
230            // 2026-09-25: PM4 W8A8 grouped GEMM; with a 0 handle the W8A8 path uses
231            // `moe_w8a8_grouped_gemm` over the dense grid.
232            moe_w8a8_grouped_gemm_pm4_k: super::super::try_kernel(
233                gpu,
234                "moe_w8a8_grouped_gemm",
235                "moe_w8a8_grouped_gemm_pm4",
236            ),
237            moe_e4m3,
238            moe_unpermute_blend_k: if std::env::var_os("METRALE_NO_MOE_UNPERMUTE_BLEND").is_some() {
239                KernelHandle(0)
240            } else {
241                try_target_kernel(gpu, "moe_unpermute_blend", "moe_unpermute_blend")
242            },
243            per_token_group_quant_fp8_k: ops::Fp8ActQuant::resolve(gpu),
244            // 2026-09-25: 0 when the model's moe_silu_mul module lacks this entry;
245            // the unfused pair runs then.
246            silu_mul_quant_fp8_k: super::super::try_kernel(
247                gpu,
248                "moe_silu_mul",
249                "silu_mul_quant_fp8",
250            ),
251            fp8_gemm_t_blockscaled_k: super::super::try_kernel(
252                gpu,
253                "fp8_gemm_t_blockscaled",
254                "fp8_gemm_t_blockscaled",
255            ),
256            moe_bf16_grouped_gemm_k: super::super::try_kernel(
257                gpu,
258                "moe_bf16_grouped_gemm",
259                "moe_bf16_grouped_gemm",
260            ),
261            moe_expert_gate_up_shared_bf16_k: super::super::try_kernel(
262                gpu,
263                "moe_shared_expert_fused_bf16",
264                "moe_expert_gate_up_shared_bf16",
265            ),
266            moe_expert_silu_down_shared_bf16_k: super::super::try_kernel(
267                gpu,
268                "moe_shared_expert_fused_bf16",
269                "moe_expert_silu_down_shared_bf16",
270            ),
271            moe_expert_gate_up_shared_bf16_batch2_k: super::super::try_kernel(
272                gpu,
273                "moe_shared_expert_fused_bf16_batch2",
274                "moe_expert_gate_up_shared_bf16_batch2",
275            ),
276            moe_expert_silu_down_shared_bf16_batch2_k: super::super::try_kernel(
277                gpu,
278                "moe_shared_expert_fused_bf16_batch2",
279                "moe_expert_silu_down_shared_bf16_batch2",
280            ),
281            w8a16_gemm_k: super::super::try_kernel(gpu, "w8a16_gemm", "w8a16_gemm"),
282            w8a16_gemm_pipelined_k: super::super::try_kernel(
283                gpu,
284                "w8a16_gemm_pipelined",
285                "w8a16_gemm_pipelined",
286            ),
287            moe_gate_topk_fused_k: super::super::try_kernel(
288                gpu,
289                "moe_gate_topk",
290                "moe_gate_topk_fused",
291            ),
292            w4a16_gemm_t: gpu.kernel("w4a16", "w4a16_gemm_t")?,
293            bf16_to_fp8_k: gpu.kernel("w4a16", "bf16_to_fp8")?,
294            fp8_gemm_k: gpu.kernel("w4a16", "fp8_gemm_t")?,
295            moe_silu_mul: gpu.kernel("moe_silu_mul", "moe_silu_mul")?,
296            moe_act_mul: gpu.kernel("moe_silu_mul", "moe_silu_mul")?,
297            gelu_activation: false,
298            moe_unpermute_reduce: gpu.kernel("moe", "moe_unpermute_reduce_indexed")?,
299            moe_batched_blend: gpu.kernel("moe", "moe_batched_blend")?,
300            gate_ptrs,
301            up_ptrs,
302            down_ptrs,
303            gate_ptrs_t: None,
304            up_ptrs_t: None,
305            down_ptrs_t: None,
306            cutlass_grouped_host: None,
307            _cutlass_sfb_owned: Vec::new(),
308            down_t_scratch_packed: None,
309            down_t_scratch_scale: None,
310            moe_transpose_u8_batched_k: gpu
311                .kernel("moe_transpose_batched", "moe_transpose_u8_batched")?,
312            moe_expert_gate_up_shared_t_k: gpu
313                .kernel("moe_shared_expert_fused_t", "moe_expert_gate_up_shared_t")?,
314            moe_expert_silu_down_shared_t_k: gpu
315                .kernel("moe_shared_expert_fused_t", "moe_expert_silu_down_shared_t")?,
316            // 2026-09-25: E8M0-routed / NVFP4-shared decode kernels from
317            // moe_shared_expert_fused_t.cu; 0 where the target lacks them.
318            moe_expert_gate_up_shared_t_e8m0_k: super::super::try_kernel(
319                gpu,
320                "moe_shared_expert_fused_t",
321                "moe_expert_gate_up_shared_t_e8m0",
322            ),
323            moe_expert_silu_down_shared_t_e8m0_k: super::super::try_kernel(
324                gpu,
325                "moe_shared_expert_fused_t",
326                "moe_expert_silu_down_shared_t_e8m0",
327            ),
328            // 2026-09-25: Optional routing kernels (sqrtsoftplus here, softmax + bias
329            // and hash routing below): 0 where the target lacks them.
330            moe_topk_sqrtsoftplus_k: super::super::try_kernel(
331                gpu,
332                "moe_topk_sqrt",
333                "moe_topk_sqrtsoftplus",
334            ),
335            moe_topk_sqrtsoftplus_batched_k: super::super::try_kernel(
336                gpu,
337                "moe_topk_sqrt",
338                "moe_topk_sqrtsoftplus_batched",
339            ),
340            router_logits_n: (config.num_experts + config.zero_expert_num) as u32,
341            moe_topk_softmax_bias_k: super::super::try_kernel(
342                gpu,
343                "moe_topk_softmax_bias",
344                "moe_topk_softmax_bias",
345            ),
346            moe_topk_softmax_bias_batched_k: super::super::try_kernel(
347                gpu,
348                "moe_topk_softmax_bias",
349                "moe_topk_softmax_bias_batched",
350            ),
351            moe_zero_expert_add_k: super::super::try_kernel(
352                gpu,
353                "moe_topk_softmax_bias",
354                "moe_zero_expert_add",
355            ),
356            // 2026-09-25: 16384 f32 (64 KiB), one per row, on every layer; the
357            // softmax + bias router writes it.
358            zero_accum_dev: gpu.alloc(16384 * 4)?,
359            moe_hash_route_k: super::super::try_kernel(gpu, "moe_hash_route", "moe_hash_route"),
360            moe_hash_route_batched_k: super::super::try_kernel(
361                gpu,
362                "moe_hash_route",
363                "moe_hash_route_batched",
364            ),
365            tid2eid_dev,
366            moe_expert_gate_up_shared_batch2_t_k: gpu.kernel(
367                "moe_shared_expert_fused_batch2_t",
368                "moe_expert_gate_up_shared_batch2_t",
369            )?,
370            moe_expert_silu_down_shared_batch2_t_k: gpu.kernel(
371                "moe_shared_expert_fused_batch2_t",
372                "moe_expert_silu_down_shared_batch2_t",
373            )?,
374            moe_expert_gate_up_shared_batch3_t_k: gpu.kernel(
375                "moe_shared_expert_fused_batch3_t",
376                "moe_expert_gate_up_shared_batch3_t",
377            )?,
378            moe_expert_silu_down_shared_batch3_t_k: gpu.kernel(
379                "moe_shared_expert_fused_batch3_t",
380                "moe_expert_silu_down_shared_batch3_t",
381            )?,
382            moe_expert_gate_up_shared_fp8_t_k: gpu.kernel(
383                "moe_shared_expert_fused_fp8_t",
384                "moe_expert_gate_up_shared_fp8_t",
385            )?,
386            moe_expert_silu_down_shared_fp8_t_k: gpu.kernel(
387                "moe_shared_expert_fused_fp8_t",
388                "moe_expert_silu_down_shared_fp8_t",
389            )?,
390            moe_expert_gate_up_shared_fp8_batch2_t_k: gpu.kernel(
391                "moe_shared_expert_fused_fp8_batch2_t",
392                "moe_expert_gate_up_shared_fp8_batch2_t",
393            )?,
394            moe_expert_silu_down_shared_fp8_batch2_t_k: gpu.kernel(
395                "moe_shared_expert_fused_fp8_batch2_t",
396                "moe_expert_silu_down_shared_fp8_batch2_t",
397            )?,
398            moe_expert_gate_up_shared_fp8_batch3_t_k: gpu.kernel(
399                "moe_shared_expert_fused_fp8_batch3_t",
400                "moe_expert_gate_up_shared_fp8_batch3_t",
401            )?,
402            moe_expert_silu_down_shared_fp8_batch3_t_k: gpu.kernel(
403                "moe_shared_expert_fused_fp8_batch3_t",
404                "moe_expert_silu_down_shared_fp8_batch3_t",
405            )?,
406            unified_layout: std::env::var("METRALE_UNIFIED_MOE_LAYOUT")
407                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
408                .unwrap_or(false),
409            hybrid_layout: std::env::var("METRALE_HYBRID_MOE_LAYOUT")
410                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
411                .unwrap_or(false),
412            nvfp4_gate_up_m128: std::env::var("METRALE_NVFP4_GATE_UP_M128")
413                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
414                .unwrap_or(false),
415            gateup_fp4: std::env::var("METRALE_HOLO_MOE_GATEUP_FP4")
416                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
417                .unwrap_or(false),
418            down_fp4: std::env::var("METRALE_HOLO_MOE_DOWN_FP4")
419                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
420                .unwrap_or(false),
421            shared_gate_t: None,
422            shared_up_t: None,
423            shared_down_t: None,
424            gate_fp8: None,
425            shared_gate_fp8: None,
426            shared_up_fp8: None,
427            shared_down_fp8: None,
428            prefill_stream: gpu.create_stream()?,
429            event_a: gpu.create_event()?,
430            event_b: gpu.create_event()?,
431            moe_expert_gate_up_shared_fp8: gpu.kernel(
432                "moe_shared_expert_fused_fp8",
433                "moe_expert_gate_up_shared_fp8",
434            )?,
435            moe_expert_silu_down_shared_fp8: gpu.kernel(
436                "moe_shared_expert_fused_fp8",
437                "moe_expert_silu_down_shared_fp8",
438            )?,
439            // 2026-09-25: FP8 batch2/3 kernels for `forward_k2` / `forward_k3`.
440            moe_expert_gate_up_shared_fp8_batch2: gpu.kernel(
441                "moe_shared_expert_fused_fp8_batch2",
442                "moe_expert_gate_up_shared_fp8_batch2",
443            )?,
444            moe_expert_silu_down_shared_fp8_batch2: gpu.kernel(
445                "moe_shared_expert_fused_fp8_batch2",
446                "moe_expert_silu_down_shared_fp8_batch2",
447            )?,
448            moe_weighted_sum_blend_fp8_batch2: gpu.kernel(
449                "moe_shared_expert_fused_fp8_batch2",
450                "moe_weighted_sum_blend_fp8_batch2",
451            )?,
452            moe_expert_gate_up_shared_fp8_batch3: gpu.kernel(
453                "moe_shared_expert_fused_fp8_batch3",
454                "moe_expert_gate_up_shared_fp8_batch3",
455            )?,
456            moe_expert_silu_down_shared_fp8_batch3: gpu.kernel(
457                "moe_shared_expert_fused_fp8_batch3",
458                "moe_expert_silu_down_shared_fp8_batch3",
459            )?,
460            moe_weighted_sum_blend_fp8_batch3: gpu.kernel(
461                "moe_shared_expert_fused_fp8_batch3",
462                "moe_weighted_sum_blend_fp8_batch3",
463            )?,
464            moe_expert_gate_up_act_fp8_grouped_k: grouped.gate_up,
465            moe_expert_down_act_fp8_grouped_k: grouped.silu_down,
466            moe_weighted_sum_blend_fp8_grouped_k: grouped.blend,
467            moe_fp8_grouped_sort_k: grouped.sort,
468            moe_topk_softmax_rows_k: grouped.topk_rows,
469            router_gemv_batchm_k: grouped.router_gemv,
470            nvfp4_grouped,
471            fp8_grouped_tc: super::fp8_grouped_tc::Fp8GroupedTcKernels::resolve(gpu),
472            fp8_gate_weight_ptrs: None,
473            fp8_up_weight_ptrs: None,
474            fp8_down_weight_ptrs: None,
475            bf16_gate_weight_ptrs: None,
476            bf16_up_weight_ptrs: None,
477            bf16_down_weight_ptrs: None,
478            bf16_shared_expert: None,
479            fp8_shared_expert: None,
480            moe_down_t_k64_fp4: super::super::try_kernel(
481                gpu,
482                "moe_w4a16",
483                "moe_w4a16_down_t_k64_fp4",
484            ),
485            moe_permute_tokens_k: super::super::try_kernel(gpu, "moe", "moe_permute_tokens"),
486            // 2026-09-25: Set after construction by the qwen35 loader.
487            is_dflash_capture_layer: false,
488            lora: None,
489            correction_bias_dev: weights_correction_bias,
490            // 2026-09-25: Optional, so a target without the `moe_topk_sig` module
491            // still constructs; only a layer with a correction bias dispatches it.
492            moe_topk_sigmoid_k: super::super::try_kernel(gpu, "moe_topk_sig", "moe_topk_sigmoid"),
493            moe_topk_sigmoid_batched_k: super::super::try_kernel(
494                gpu,
495                "moe_topk_sig",
496                "moe_topk_sigmoid_batched",
497            ),
498        })
499    }
500}
metrale_model_layers/layers/moe/init.rs

metrale_model_layers/layers/moe/
init.rs
