metrale_model_layers/layers/moe/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The MoE feed-forward block `MoeLayer`: router, routed experts, optional
4//! shared expert, and the kernel handles and pointer tables its forward paths use.
5//!
6//! Owner: model-layers (MoE).
7//! Invariants:
8//! - Construction (`new_with_hash`) fails unless `num_experts > 0`,
9//!   `num_experts_per_tok <= num_experts`, and both are within the sigmoid
10//!   routing kernels' fixed bounds.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14
15use crate::layer::ForwardContext;
16use crate::layers::ops;
17use crate::weight_map::{DenseWeight, Fp8ExpertWeight, MoeWeights, QuantizedWeight};
18
19/// 2026-09-25: The MoE feed-forward block. It is a component that layers call for
20/// their FFN step; it does not implement `TransformerLayer`.
21#[allow(dead_code)]
22pub struct MoeLayer {
23    pub weights: MoeWeights,
24    /// 2026-09-25: Quant format of the routed experts. `new_with_hash` sets `Nvfp4`;
25    /// the DeepSeek-V4 loader sets the detected format, and `Mxfp4E8m0` makes
26    /// `e8m0_or` and the prefill dispatch pick the `_e8m0` kernels.
27    #[allow(dead_code)]
28    pub experts_scale_kind: crate::weight_map::WeightQuantFormat,
29    /// 2026-09-25: Quant format of the shared expert, detected separately from the
30    /// routed experts. `new_with_hash` sets `Nvfp4`. With E8M0 routed experts,
31    /// the dual-format decode paths check it with `WeightQuantFormat::expect`.
32    #[allow(dead_code)]
33    pub shared_experts_scale_kind: crate::weight_map::WeightQuantFormat,
34    // 2026-09-25: NVFP4 copy of the router gate. When set, the router GEMM runs on
35    // it instead of the BF16 `weights.gate`.
36    gate_nvfp4: Option<QuantizedWeight>,
37    /// 2026-09-25: RMS norm applied to the expert input after routing; the router reads
38    /// its own input (`router_input`). Only the Gemma-4 loader sets it.
39    pub pre_expert_norm: Option<crate::weight_map::DenseWeight>,
40    pre_expert_norm_k: metrale_gpu_runtime::gpu::KernelHandle,
41    dense_gemv: KernelHandle,
42    w4a16_gemv: KernelHandle,
43    w4a16_gemv_sw: KernelHandle,
44    w4a16_gemm: KernelHandle,
45    dense_gemm: KernelHandle,
46    /// 2026-09-25: `dense_gemm_bf16_router`, read only by `router_gate_gemm_dense`,
47    /// which uses `dense_gemm` instead when this handle is 0.
48    dense_gemm_router: KernelHandle,
49    /// 2026-09-28: `moe_router_gemm_rt`, the same bits as `dense_gemm_router` on a register
50    /// tile, taken from `MOE_ROUTER_RT_MIN_ROWS` rows; 0 when the target lacks it or
51    /// `METRALE_NO_MOE_ROUTER_RT` is set.
52    moe_router_rt_k: KernelHandle,
53    // 2026-09-26: `moe_router_gemm_bf16` (`router_gemm_bf16`); 0 on targets that
54    // do not ship it.
55    moe_router_gemm_k: KernelHandle,
56    dense_gemm_pipelined: KernelHandle,
57    /// 2026-09-25: FP32-output router GEMM for the `fp32_gate` lever
58    /// (`METRALE_FP32_GATE`). 0 when the target lacks it; the gate then stays BF16.
59    dense_gemm_f32out: KernelHandle,
60    /// 2026-09-25: FP32-in/FP32-out router GEMM for the `fp32_routing` lever
61    /// (`METRALE_FP32_ROUTING`); `fp32_routing_active` is false while it is 0.
62    dense_gemm_f32in: KernelHandle,
63    moe_topk_f32: KernelHandle,
64    moe_expert_gate_up_shared: KernelHandle,
65    moe_expert_silu_down_shared: KernelHandle,
66    moe_topk: KernelHandle,
67    moe_weighted_sum_blend: KernelHandle,
68    residual_add: KernelHandle,
69    moe_topk_batched: KernelHandle,
70    moe_expert_gate_up_shared_batch2: KernelHandle,
71    moe_expert_silu_down_shared_batch2: KernelHandle,
72    moe_weighted_sum_blend_batch2: KernelHandle,
73    w4a16_gemv_batch2: KernelHandle,
74    moe_expert_gate_up_shared_batch3: KernelHandle,
75    moe_expert_silu_down_shared_batch3: KernelHandle,
76    moe_weighted_sum_blend_batch3: KernelHandle,
77    w4a16_gemv_batch3: KernelHandle,
78    // 2026-09-25: Token-major kernels for `forward_token_major_decode`. The
79    // `atomic_c4` pair serves `forward_atomic_c4_decode`
80    // (`METRALE_MOE_ATOMIC_C4_DECODE=1`).
81    moe_expert_gate_up_shared_token_major: KernelHandle,
82    moe_expert_silu_down_shared_token_major: KernelHandle,
83    moe_weighted_sum_blend_token_major: KernelHandle,
84    moe_decode_atomic_c4_silu_down_accum_k: KernelHandle,
85    moe_decode_atomic_c4_finalize_k: KernelHandle,
86    moe_sort_by_expert: KernelHandle,
87    moe_sorted_gate_up: KernelHandle,
88    moe_sorted_silu_down: KernelHandle,
89    moe_grouped_gemm: KernelHandle,
90    /// 2026-09-25: Wider-K twin of `moe_grouped_gemm`. Loaded only when
91    /// `METRALE_MOE_GROUPED_K32=1` and the target ships it; `grouped_gemm_kernel`
92    /// prefers it whenever it is non-zero.
93    moe_grouped_gemm_k32: KernelHandle,
94    /// 2026-09-25: M_TILE=256 twin of `moe_grouped_gemm`. Loaded only when
95    /// `METRALE_MOE_GROUPED_M256=1` and the target ships it; used whenever it is
96    /// non-zero.
97    moe_grouped_gemm_m256: KernelHandle,
98    moe_silu_mul: KernelHandle,
99    /// 2026-09-25: Activation kernel for the sorted/unfused paths: `moe_silu_mul`, or
100    /// `gelu_mul` after `set_gelu_activation`.
101    moe_act_mul: KernelHandle,
102    /// 2026-09-25: Set by `set_gelu_activation`. The fused SiLU+FP8-quant kernel is used
103    /// only while it is false (`fused_silu_quant_ok`).
104    gelu_activation: bool,
105    moe_unpermute_reduce: KernelHandle,
106    moe_batched_blend: KernelHandle,
107    gate_ptrs: ExpertPtrTable,
108    up_ptrs: ExpertPtrTable,
109    down_ptrs: ExpertPtrTable,
110    /// 2026-09-25: Transposed-layout pointer tables; `None` until a transpose pass or
111    /// `set_down_transpose_scratch` fills them.
112    gate_ptrs_t: Option<ExpertPtrTable>,
113    up_ptrs_t: Option<ExpertPtrTable>,
114    down_ptrs_t: Option<ExpertPtrTable>,
115    /// 2026-09-25: Host-side per-expert packed/SFB pointers and scale2 for the CUTLASS
116    /// grouped NVFP4 path, built once by `build_cutlass_grouped_sfb`. The grouped
117    /// entry reads them on the host; they live on the layer so they are dropped
118    /// with the model. `None` means that path is unavailable.
119    cutlass_grouped_host: Option<ops::MoeCutlassHostTables>,
120    /// 2026-09-25: Owns the per-expert SFB buffers that `cutlass_grouped_host` points at.
121    _cutlass_sfb_owned: Vec<DevicePtr>,
122    /// 2026-09-25: Down-projection transpose scratch shared by all MoE layers, wired by
123    /// `set_down_transpose_scratch`. `transpose_down_into_scratch` refills it at
124    /// the start of each layer's prefill, and `down_ptrs_t` then points into it.
125    /// `None` unless that scratch is wired.
126    down_t_scratch_packed: Option<DevicePtr>,
127    down_t_scratch_scale: Option<DevicePtr>,
128    moe_transpose_u8_batched_k: KernelHandle,
129    moe_expert_gate_up_shared_t_k: KernelHandle,
130    moe_expert_silu_down_shared_t_k: KernelHandle,
131    // 2026-09-25: Decode variants for E8M0-scaled routed experts with an NVFP4 shared
132    // expert. 0 on targets that do not ship them.
133    moe_expert_gate_up_shared_t_e8m0_k: KernelHandle,
134    moe_expert_silu_down_shared_t_e8m0_k: KernelHandle,
135    moe_topk_sqrtsoftplus_k: KernelHandle,
136    moe_topk_sqrtsoftplus_batched_k: KernelHandle,
137    moe_hash_route_k: KernelHandle,
138    moe_hash_route_batched_k: KernelHandle,
139    /// 2026-09-25: Router logit width, `num_experts + zero_expert_num`. It equals
140    /// `num_experts` unless the model has zero-computation experts.
141    pub(crate) router_logits_n: u32,
142    moe_topk_softmax_bias_k: KernelHandle,
143    moe_topk_softmax_bias_batched_k: KernelHandle,
144    moe_zero_expert_add_k: KernelHandle,
145    /// 2026-09-25: Per-token folded zero-expert weight (f32), written by the
146    /// softmax+bias router kernels. Allocated once, for 16384 tokens.
147    zero_accum_dev: DevicePtr,
148    /// 2026-09-25: Static hash-routing table `[vocab_size, top_k]` i64. `Some` marks a
149    /// hash-routed layer: expert selection reads this table instead of the
150    /// gate's top-k.
151    tid2eid_dev: Option<DevicePtr>,
152    moe_expert_gate_up_shared_batch2_t_k: KernelHandle,
153    moe_expert_silu_down_shared_batch2_t_k: KernelHandle,
154    moe_expert_gate_up_shared_batch3_t_k: KernelHandle,
155    moe_expert_silu_down_shared_batch3_t_k: KernelHandle,
156    moe_expert_gate_up_shared_fp8_t_k: KernelHandle,
157    moe_expert_silu_down_shared_fp8_t_k: KernelHandle,
158    moe_expert_gate_up_shared_fp8_batch2_t_k: KernelHandle,
159    moe_expert_silu_down_shared_fp8_batch2_t_k: KernelHandle,
160    moe_expert_gate_up_shared_fp8_batch3_t_k: KernelHandle,
161    moe_expert_silu_down_shared_fp8_batch3_t_k: KernelHandle,
162    /// 2026-09-25: `METRALE_UNIFIED_MOE_LAYOUT` (`1` or `true`), read at construction.
163    /// With all three transposed tables built and no down scratch, decode runs
164    /// the `moe_expert_*_shared_t` kernels unless `hybrid_layout` is also set.
165    unified_layout: bool,
166    /// 2026-09-25: `METRALE_NVFP4_GATE_UP_M128` (`1` or `true`). Prefill's fused
167    /// gate+up uses the M=128 kernel when this is set and
168    /// `moe_fused_gate_up_t_k64_m128` is non-zero, else the M=64 kernel.
169    nvfp4_gate_up_m128: bool,
170    /// 2026-09-25: `METRALE_HOLO_MOE_GATEUP_FP4` (`1` or `true`). Prefill's fused
171    /// gate+up runs `moe_fused_gate_up_t_k64_fp4` over the existing
172    /// `gate_ptrs_t`/`up_ptrs_t` tables when this is set and the handle is non-zero.
173    gateup_fp4: bool,
174    /// 2026-09-25: `METRALE_HOLO_MOE_DOWN_FP4` (`1` or `true`): the same choice for the
175    /// prefill down projection, with `moe_down_t_k64_fp4` over `down_ptrs_t`.
176    down_fp4: bool,
177    /// 2026-09-25: `METRALE_HYBRID_MOE_LAYOUT` (`1` or `true`), read at construction.
178    /// Prefill uses the transposed tables (`use_t_layout_for_prefill`) while
179    /// decode keeps the original `[N, K/2]` tables, even when `unified_layout`
180    /// is also set.
181    hybrid_layout: bool,
182    shared_gate_t: Option<QuantizedWeight>,
183    shared_up_t: Option<QuantizedWeight>,
184    shared_down_t: Option<QuantizedWeight>,
185    moe_grouped_gemm_t: KernelHandle,
186    moe_grouped_gemm_t_k64: KernelHandle,
187    moe_fused_gate_up_t: KernelHandle,
188    moe_fused_gate_up_t_k64: KernelHandle,
189    // 2026-09-25: E8M0-scale prefill variants of the routed-expert W4A16 GEMMs. 0 on
190    // targets that do not ship them.
191    moe_grouped_gemm_e8m0: KernelHandle,
192    moe_grouped_gemm_t_e8m0: KernelHandle,
193    moe_grouped_gemm_t_k64_e8m0: KernelHandle,
194    moe_fused_gate_up_t_e8m0: KernelHandle,
195    moe_fused_gate_up_t_k64_e8m0: KernelHandle,
196    /// 2026-09-25: M=128 variant of the K64 fused gate+up kernel; 0 on targets that do
197    /// not ship it.
198    moe_fused_gate_up_t_k64_m128: KernelHandle,
199    /// 2026-09-25: Block-scaled FP4 variant of the K64 fused gate+up kernel, launched
200    /// through the same wrapper as `moe_fused_gate_up_t_k64`; 0 on targets that
201    /// do not ship it.
202    moe_fused_gate_up_t_k64_fp4: KernelHandle,
203    moe_fp8_grouped_gemm_t: KernelHandle,
204    w4a16_gemm_t: KernelHandle,
205    bf16_to_fp8_k: KernelHandle,
206    /// 2026-09-25: FP8 copies of the NVFP4 router gate and shared expert, made by
207    /// `predequant_for_prefill`. The shared-expert copies are skipped when a BF16
208    /// shared expert is installed.
209    gate_fp8: Option<DevicePtr>,
210    shared_gate_fp8: Option<DevicePtr>,
211    shared_up_fp8: Option<DevicePtr>,
212    shared_down_fp8: Option<DevicePtr>,
213    fp8_gemm_k: KernelHandle,
214    /// 2026-09-25: Auxiliary stream. Nothing launches on it today: the prefill
215    /// shared-expert overlap is switched off (`use_overlap = false` in
216    /// `forward_prefill`), and `kick_off_lazy_transpose`, the down transpose that
217    /// would run on it, has no caller.
218    prefill_stream: u64,
219    /// 2026-09-25: `event_a` marks the input ready for the auxiliary stream; `event_b`
220    /// marks the auxiliary work (the overlapped shared expert, or
221    /// `kick_off_lazy_transpose`'s transpose) done. Neither is recorded while
222    /// both of those are off.
223    event_a: u64,
224    event_b: u64,
225    /// 2026-09-25: `[num_experts]` routing correction bias, from
226    /// `MoeWeights.correction_bias`. `Some` selects bias-corrected routing
227    /// (sigmoid, or sqrtsoftplus when the config says so); `None` selects softmax.
228    correction_bias_dev: Option<DevicePtr>,
229    moe_topk_sigmoid_k: KernelHandle,
230    moe_topk_sigmoid_batched_k: KernelHandle,
231    moe_expert_gate_up_shared_fp8: KernelHandle,
232    moe_expert_silu_down_shared_fp8: KernelHandle,
233    moe_expert_gate_up_shared_fp8_batch2: KernelHandle,
234    moe_expert_silu_down_shared_fp8_batch2: KernelHandle,
235    moe_weighted_sum_blend_fp8_batch2: KernelHandle,
236    moe_expert_gate_up_shared_fp8_batch3: KernelHandle,
237    moe_expert_silu_down_shared_fp8_batch3: KernelHandle,
238    moe_weighted_sum_blend_fp8_batch3: KernelHandle,
239    // 2026-09-25: Grouped FP8 decode kernels from `GroupedKernels::resolve`; 0 on
240    // targets that do not ship them.
241    moe_expert_gate_up_act_fp8_grouped_k: KernelHandle,
242    moe_expert_down_act_fp8_grouped_k: KernelHandle,
243    moe_weighted_sum_blend_fp8_grouped_k: KernelHandle,
244    // 2026-09-27: `moe_fp8_grouped_sort`: the slot sort and the active-expert list in
245    // one launch.
246    moe_fp8_grouped_sort_k: KernelHandle,
247    // 2026-09-26: The exact-router kernels of the grouped decode
248    // (`GroupedRouting::PerRow` / `PerToken`): `moe_topk_softmax_rows` and
249    // `dense_gemv_bf16_batchm`; 0 on targets that do not ship them.
250    moe_topk_softmax_rows_k: KernelHandle,
251    router_gemv_batchm_k: KernelHandle,
252    // 2026-09-27: The grouped NVFP4 decode's expert kernels (NVFP4 `--expert-quantization`
253    // tiers); 0 on
254    // targets that do not ship them.
255    nvfp4_grouped: forward_nvfp4_grouped_decode::Nvfp4GroupedKernels,
256    // 2026-09-28: The grouped FP8 decode's tensor-core expert kernels (`fp8_grouped_tc.rs`).
257    fp8_grouped_tc: fp8_grouped_tc::Fp8GroupedTcKernels,
258    // 2026-09-25: Routed-expert FP8 grouped GEMM over the tile work-list that
259    // `moe_build_tile_worklist` builds; 0 on targets that do not ship it.
260    moe_fp8_grouped_gemm_k: KernelHandle,
261    moe_build_tile_worklist_k: KernelHandle,
262    // 2026-09-25: The Hopper short-prefill M16 expert GEMM and bucket builder
263    // (`adaptive_fp8.rs`), and the SM count their grids are sized from; 0 when the
264    // target lacks either module.
265    moe_w8a8_m16_k: KernelHandle,
266    moe_bucket_builder_k: KernelHandle,
267    moe_adaptive_sms: u32,
268    // 2026-09-25: W8A8 grouped GEMM with an FP32 epilogue over per-token-quantized
269    // activations. The FP8 prefill uses it when `fp8_blockscaled_prefill` is on
270    // (the default; `METRALE_FP8_SINGLE_SCALE=1` turns it off) and this handle
271    // and the per-token quant kernel are present.
272    moe_w8a8_grouped_gemm_k: KernelHandle,
273    // 2026-09-25: Work-list variant of the W8A8 grouped GEMM. Preferred when it and
274    // `moe_build_tile_worklist_k` are non-zero; otherwise the W8A8 path runs
275    // `moe_w8a8_grouped_gemm_k`.
276    moe_w8a8_grouped_gemm_pm4_k: KernelHandle,
277    // 2026-09-27: The PM4 computation on the native e4m3 MMA (`forward_prefill_fp8/e4m3.rs`),
278    // preferred over PM4 for non-decode prefill.
279    moe_e4m3: forward_prefill_fp8::E4m3Kernels,
280    // 2026-09-28: `moe_unpermute_blend` (unpermute + shared blend in one launch, same bits);
281    // 0 when absent or under `METRALE_NO_MOE_UNPERMUTE_BLEND`.
282    moe_unpermute_blend_k: KernelHandle,
283    per_token_group_quant_fp8_k: ops::Fp8ActQuant,
284    /// 2026-09-25: Fused SiLU·mul + per-token-group FP8 quant for the W8A8 prefill
285    /// down inputs. When `fused_silu_quant_ok` is false (handle 0, GeGLU, or an
286    /// unsupported K) the path runs `silu_mul` then `per_token_group_quant_fp8`.
287    silu_mul_quant_fp8_k: KernelHandle,
288    // 2026-09-25: Dense W8A8 GEMM for the shared expert on the W8A8 FP8 prefill path.
289    fp8_gemm_t_blockscaled_k: KernelHandle,
290    // 2026-09-25: BF16 grouped GEMM for experts installed with `set_bf16_experts`.
291    // Prefill of more than 64 tokens uses it when it is non-zero, except on HIP
292    // builds; otherwise BF16 experts run `forward_batched`.
293    moe_bf16_grouped_gemm_k: KernelHandle,
294    moe_expert_gate_up_shared_bf16_k: KernelHandle,
295    moe_expert_silu_down_shared_bf16_k: KernelHandle,
296    // 2026-09-25: Two-row BF16 decode kernels. `forward_k2` uses them only when both
297    // are non-zero, BF16 experts are installed and expert parallelism is off.
298    moe_expert_gate_up_shared_bf16_batch2_k: KernelHandle,
299    moe_expert_silu_down_shared_bf16_batch2_k: KernelHandle,
300    w8a16_gemm_k: KernelHandle,
301    w8a16_gemm_pipelined_k: KernelHandle,
302    moe_gate_topk_fused_k: KernelHandle,
303    // 2026-09-25: FP8 expert pointer tables; `None` until `set_fp8_experts`.
304    fp8_gate_weight_ptrs: Option<Fp8ExpertPtrTable>,
305    fp8_up_weight_ptrs: Option<Fp8ExpertPtrTable>,
306    fp8_down_weight_ptrs: Option<Fp8ExpertPtrTable>,
307    // 2026-09-25: BF16 expert pointer tables; `None` until `set_bf16_experts`. When
308    // set, prefill takes the BF16 expert path before any FP8 or NVFP4 one.
309    bf16_gate_weight_ptrs: Option<DevicePtr>,
310    bf16_up_weight_ptrs: Option<DevicePtr>,
311    bf16_down_weight_ptrs: Option<DevicePtr>,
312    // 2026-09-25: BF16 shared expert, installed independently of the routed experts'
313    // precision (`set_bf16_shared_expert`).
314    bf16_shared_expert: Option<Bf16SharedExpert>,
315    // 2026-09-25: FP8 shared expert; `None` until `set_fp8_experts`.
316    fp8_shared_expert: Option<Fp8ExpertWeight>,
317    /// 2026-09-25: `moe_w4a16_down_t_k64_fp4`; 0 on targets that do not ship it.
318    pub(crate) moe_down_t_k64_fp4: KernelHandle,
319    /// 2026-09-25: `moe_permute_tokens`. Loaded, but no forward path reads it.
320    #[allow(dead_code)]
321    pub(crate) moe_permute_tokens_k: KernelHandle,
322    // 2026-09-25: True when this layer's index is in `config.dflash_capture_layers`
323    // (set by the qwen35 loader). With the `frankenstein_decode_via_prefill`
324    // lever (`METRALE_FRANKENSTEIN_DECODE_VIA_PREFILL`), `forward` runs this
325    // layer's single-token decode through `forward_prefill` with one row.
326    pub is_dflash_capture_layer: bool,
327    /// 2026-09-25: This layer's installed MoE LoRA (router and routed-expert deltas plus
328    /// apply scratch), set by [`MoeLayer::set_lora_weights`]; `None` means no
329    /// adapter. Decode paths that call `reject_decode_lora` refuse a batch
330    /// routed to the adapter.
331    pub(crate) lora: Option<MoeLoraWeights>,
332}
333
334impl MoeLayer {
335    /// 2026-09-25: Returns `e8m0` when the routed experts are `Mxfp4E8m0`, else `nvfp4`.
336    /// Panics when `e8m0` is selected but its handle is 0, so E8M0 weights never
337    /// run through an NVFP4 kernel.
338    #[inline]
339    fn e8m0_or(
340        &self,
341        nvfp4: metrale_gpu_runtime::gpu::KernelHandle,
342        e8m0: metrale_gpu_runtime::gpu::KernelHandle,
343        site: &str,
344    ) -> metrale_gpu_runtime::gpu::KernelHandle {
345        if self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Mxfp4E8m0 {
346            assert!(
347                e8m0.0 != 0,
348                "ARM-2 Phase-K: routed experts tagged Mxfp4E8m0 at {site}, but the \
349                 _e8m0 kernel handle is unresolved (not compiled into this target)."
350            );
351            e8m0
352        } else {
353            nvfp4
354        }
355    }
356}
357
358mod adaptive_fp8;
359mod tables;
360pub(crate) use tables::{Bf16SharedExpert, ExpertPtrTable, Fp8ExpertPtrTable};
361
362mod dump;
363mod forward;
364mod lora;
365mod lora_gateup;
366mod lora_router;
367pub(crate) use lora::MoeLoraWeights;
368mod forward_atomic_c4;
369mod forward_batched;
370mod forward_batched_gate;
371mod forward_ep;
372mod forward_fp8_grouped_decode;
373mod fp8_grouped_tc;
374mod fp8_grouped_tc_w8a8;
375pub use forward_fp8_grouped_decode::fp8_grouped_decode_shape_ok;
376pub use fp8_grouped_tc_w8a8::{moe_expert_fp8_act, set_moe_expert_fp8_act};
377mod forward_fp8_grouped_router;
378pub use forward_fp8_grouped_router::GroupedRouting;
379mod forward_bf16_grouped_decode;
380mod forward_nvfp4_grouped_decode;
381mod nvfp4_lean;
382pub use forward_nvfp4_grouped_decode::{
383    NVFP4_GROUPED_DECODE_MAX_ROWS, NVFP4_GROUPED_DECODE_TC_MAX_ROWS, nvfp4_grouped_decode_shape_ok,
384};
385mod forward_k2;
386mod forward_k3;
387mod forward_phase;
388mod forward_prefill;
389mod forward_prefill_bf16;
390mod forward_prefill_fp8;
391mod forward_prefill_phase;
392mod forward_prefill_routed;
393mod forward_prefill_router;
394mod forward_token_major;
395mod helpers_a;
396mod helpers_b;
397mod helpers_c;
398mod init;
399#[cfg(test)]
400mod mod_tests;
401mod ptr_table_build;
402mod union_stats;
403pub(crate) use ptr_table_build::*;
metrale_model_layers/layers/moe/mod.rs

metrale_model_layers/layers/moe/
mod.rs
