Struct MoeLayer 

pub struct MoeLayer {
    pub weights: MoeWeights,
    pub experts_scale_kind: WeightQuantFormat,
    pub shared_experts_scale_kind: WeightQuantFormat,
    pub pre_expert_norm: Option<DenseWeight>,
    pub is_dflash_capture_layer: bool,
    /* private fields */
}

2026-09-25: The MoE feed-forward block. It is a component that layers call for their FFN step; it does not implement TransformerLayer.

Fields

weights: MoeWeights

experts_scale_kind: WeightQuantFormat

2026-09-25: Quant format of the routed experts. new_with_hash sets Nvfp4; the DeepSeek-V4 loader sets the detected format, and Mxfp4E8m0 makes e8m0_or and the prefill dispatch pick the _e8m0 kernels.

shared_experts_scale_kind: WeightQuantFormat

2026-09-25: Quant format of the shared expert, detected separately from the routed experts. new_with_hash sets Nvfp4. With E8M0 routed experts, the dual-format decode paths check it with WeightQuantFormat::expect.

pre_expert_norm: Option<DenseWeight>

2026-09-25: RMS norm applied to the expert input after routing; the router reads its own input (router_input). Only the Gemma-4 loader sets it.

is_dflash_capture_layer: bool

Implementations

impl MoeLayer

pub fn fp32_routing_active(&self, levers: &ModelLevers) -> bool

2026-09-25: Whether the FP32 routing path is on for this layer: the pre-MoE norm then also writes an FP32 copy of the router input (residual_add_rms_norm_gatef32 into moe_router_in_f32), which batched_gate_logits reads through dense_gemm_f32in. Needs a dense gate (no gate_nvfp4), no correction bias, the dense_gemm_f32in and moe_topk_f32 kernels, and, checked last, the fp32_routing lever (on only when METRALE_FP32_ROUTING=1).

pub fn apply_zero_expert( &self, out: DevicePtr, x: DevicePtr, n: u32, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: Add the zero-expert term out[t,:] += zero_accum[t] * x[t,:], where zero_accum holds the weights of the selected zero experts that the softmax+bias router folded out. It adds into out in place, so it runs after the routed blend has written out, for the tokens whose routing wrote zero_accum. No launch when router_logits_n equals num_experts; an error when the kernel is missing.

impl MoeLayer

pub fn forward( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<DevicePtr>

impl MoeLayer

pub fn forward_atomic_c4_decode( &self, input: DevicePtr, num_tokens: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: 4-row NVFP4 MoE decode, reached from qwen3_ssm/trait_decode_multi_seq.rs when METRALE_MOE_ATOMIC_C4_DECODE=1.

The gate GEMM and top-k run batched. Gate+up uses the token-major kernel, then the routed down projections atomic-add their weighted FP32 results into a [num_tokens, H] accumulator in scratch. The finalize kernel casts it to BF16 and blends the shared expert, except under EP, where the shared expert is blended once after the all-reduce.

Falls back to forward_batched unless num_tokens == 4, both atomic kernels are present, the experts are NVFP4 (no BF16 or FP8 expert tables, no BF16 shared expert) in the non-transposed decode layout, there is no pre-expert norm, and a shared expert exists.

impl MoeLayer

pub fn forward_batched( &self, input: DevicePtr, num_tokens: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: The gate projection is one GEMM with M = N; top-k, the expert kernels and the blend then run once per token.

impl MoeLayer

pub fn forward_ep_dispatch( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<DevicePtr>

2026-09-25: Route one token with the gate projection and plain moe_topk_softmax (unlike forward, no router-LoRA fold, hash routing or correction-bias scoring), copy its top-k to the host after a stream synchronize, build an EpRoutingTable and log its local/remote counts at debug level, then return self.forward(...). The table does not affect the output. Refuses a zero-expert router.

impl MoeLayer

pub fn fp8_grouped_decode_arena_ok( &self, m: usize, cfg: &ModelConfig, b: &BufferArena, ) -> bool

2026-09-25: The weight, kernel, shape and arena terms of Self::fp8_grouped_decode_ok, decidable without a ForwardContext, so the MTP batch cap (mtp_head/batch_caps.rs) can size a batch. Not sufficient alone: the kill switch, the FP32-routing lever and expert parallelism are checked only by the full predicate.

pub fn fp8_grouped_decode_ok(&self, m: usize, ctx: &ForwardContext<'_>) -> bool

2026-09-25: Whether forward_fp8_grouped_decode serves m rows on this layer under ctx: Self::fp8_grouped_decode_arena_ok, the kill switch, FP32 routing off and no expert parallelism.

pub fn forward_fp8_grouped_decode( &self, input: DevicePtr, m: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: The FP8 MoE of m rows: input is [m, H] BF16, the output lands in rows 0..m of moe_output(). Returns an error, launching nothing, when fp8_grouped_decode_ok(m, ctx) is false. Routes with GroupedRouting::Batched.

pub fn forward_fp8_grouped_decode_routed( &self, input: DevicePtr, m: usize, routing: GroupedRouting, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-26: Self::forward_fp8_grouped_decode with the router arithmetic routing names. Returns an error, launching nothing, when fp8_grouped_routing_ok(m, routing, ctx) is false.

impl MoeLayer

pub fn fp8_grouped_routing_ok( &self, m: usize, routing: GroupedRouting, ctx: &ForwardContext<'_>, ) -> bool

2026-09-26: Whether the grouped decode serves m rows with routing. Batched is Self::fp8_grouped_decode_ok. The exact routings also need a BF16 softmax router with no correction bias, pre-router norm or FP32 gate, a layer forward does not reroute (DFlash capture), and their kernels.

impl MoeLayer

pub fn bf16_grouped_decode_arena_ok( &self, m: usize, cfg: &ModelConfig, b: &BufferArena, ) -> bool

2026-10-02: The weight, kernel, shape and arena terms of Self::bf16_grouped_decode_ok, decidable without a ForwardContext (the MTP batch cap sizes a batch with it).

pub fn bf16_grouped_decode_ok(&self, m: usize, ctx: &ForwardContext<'_>) -> bool

2026-10-02: Whether forward_bf16_grouped_decode serves m rows under ctx: the arena terms, the per-row router’s context terms and no expert parallelism.

pub fn forward_bf16_grouped_decode( &self, input: DevicePtr, m: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-10-02: The BF16 MoE of m rows: input is [m, H] BF16, the output lands in rows 0..m of moe_output(). Returns an error, launching nothing, when bf16_grouped_decode_ok(m, ctx) is false.

pub fn any_grouped_decode_ok(&self, m: usize, ctx: &ForwardContext<'_>) -> bool

2026-10-02: The grouped decode of m rows on whichever expert tables this layer holds for it: FP8 (forward_fp8_grouped_decode) or BF16 (forward_bf16_grouped_decode).

pub fn any_grouped_decode_arena_ok( &self, m: usize, cfg: &ModelConfig, b: &BufferArena, ) -> bool

2026-10-02: Self::any_grouped_decode_ok’s arena terms.

pub fn forward_any_grouped_decode( &self, input: DevicePtr, m: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-10-02: Run Self::any_grouped_decode_ok’s path.

impl MoeLayer

pub fn nvfp4_grouped_decode_ok( &self, m: usize, ctx: &ForwardContext<'_>, ) -> bool

2026-09-27: Whether forward_nvfp4_grouped_decode serves m rows on this layer: an NVFP4 --expert-quantization tier is in force or (2026-10-02) the experts are the checkpoint’s declared NVFP4 (set_declared_nvfp4_experts), the routed projections that tier decodes as NVFP4 are present in the row-major decode layout, the shared expert (and, under nvfp4-gate-up, the routed down projections) are there in FP8 or NVFP4 as the tier reads them, the per-row router applies (BF16 softmax gate, no correction bias, pre-router norm, FP32 gate or FP32 routing, no DFlash reroute), the kernels resolved, the shape is admitted, the arena is wide enough, and there is no LoRA, pre-expert norm, hash routing or expert parallelism.

pub fn set_declared_nvfp4_experts(&mut self) -> Result<()>

2026-10-02: Mark this layer’s routed and shared experts as the checkpoint’s declared NVFP4, so nvfp4_grouped_decode_ok admits it under every --expert-quantization tier. Refused for a layer that also holds FP8 experts: the tier decides between those.

pub fn set_draft_nvfp4_experts(&mut self) -> Result<()>

2026-10-02: The same mark for a drafter whose experts were requantized to NVFP4 for drafting (MtpHead::new_nvfp4_draft_moe): the grouped NVFP4 decode serves them at every width.

pub fn forward_nvfp4_grouped_decode( &self, input: DevicePtr, m: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-27: The NVFP4 MoE of m rows: input is [m, H] BF16, the output lands in rows 0..m of moe_output(). Returns an error, launching nothing, when nvfp4_grouped_decode_ok(m, ctx) is false.

impl MoeLayer

pub fn nvfp4_experts_lean(&self) -> bool

2026-10-04: Whether this layer’s NVFP4 expert tables are in the lean layout.

pub fn nvfp4_lean_refusal(&self, config: &ModelConfig) -> Option<&'static str>

2026-10-04: Why this layer keeps the row-major layout, or None when it qualifies for the lean one (module invariants). Call after the load-time prefill passes.

pub fn repack_nvfp4_experts_lean( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, stream: u64, ) -> Result<bool>

2026-10-04: Repacks this layer’s routed and shared NVFP4 experts in place for the lean decode when it qualifies (nvfp4_lean_refusal), on stream, and synchronizes it. Returns whether it did. Fails the load on a negative or NaN scale.

impl MoeLayer

pub fn forward_k2( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: MoE for two rows: input is the normed [2, H] BF16 MoE input.

Gate projection for both rows, top-K, then the fused expert gate+up, silu+down and weighted-sum+blend kernels. The shared expert’s scratch reuses the logits and ssm_qkvz buffers. Output: moe_output() [2, H]. LoRA, BF16 experts that cannot use the BF16 batch2 kernels, E8M0 experts and an unsupported mixed shared expert take forward_batched instead.

impl MoeLayer

pub fn forward_k3( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: MoE for three rows: input is the normed [3, H] BF16 MoE input.

Gate projection for all rows, top-K, then the fused expert gate+up, silu+down and weighted-sum+blend kernels. Output: moe_output() [3, H]. LoRA, BF16 experts, E8M0 experts and an unsupported mixed shared expert take forward_batched instead.

impl MoeLayer

pub fn forward_prefill( &self, input: DevicePtr, num_tokens: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: MoE for num_tokens rows of the normed [num_tokens, H] BF16 input, with grouped GEMMs over the rows sorted by expert.

BF16 and FP8 experts go to forward_prefill_bf16 / forward_prefill_fp8 above 64 rows when that grouped kernel resolved, and to forward_batched otherwise. The NVFP4 body here: shared expert, gate GEMM, top-K, sort by expert, grouped gate/up GEMM, activation, grouped down GEMM, unpermute + weighted reduce, shared blend.

impl MoeLayer

pub fn forward_token_major_decode( &self, input: DevicePtr, num_tokens: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

2026-09-25: MoE for num_tokens decode rows with the token-major kernels of the moe_prefill module, without sorting by expert: batched gate GEMM and top-K, then gate+up, silu+down and weighted-sum+blend over every (row, expert-slot) pair.

Serves untransposed NVFP4 experts with an NVFP4 shared expert. LoRA, BF16 or FP8 experts, a mixed BF16 shared expert, the unified layout, or no shared expert go to forward_batched.

impl MoeLayer

pub fn set_pre_expert_norm(&mut self, norm: DenseWeight)

2026-09-25: Install the pre-expert norm: the router reads the un-normed input and the experts read it normed with this weight (forward, forward_prefill). The Gemma-4 loader passes pre_feedforward_layernorm_2.

pub fn set_gelu_activation(&mut self, gpu: &dyn GpuBackend) -> Result<()>

2026-09-25: Switch the expert activation kernel moe_act_mul to GeGLU (gelu_mul) and set gelu_activation, which only disables the fused SiLU+FP8-quant kernel in forward_prefill_fp8. Fused decode kernels are not steered away by it.

pub fn transpose_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

2026-09-25: Transpose every routed expert’s gate, up and down [N, K/2] weights and scales to [K/2, N] (the *_ptrs_t tables), and the shared expert’s when it has one. The originals are kept, so the expert weights are held twice.

pub fn transpose_gate_up_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

2026-09-25: As transpose_for_prefill, for gate and up only: down_ptrs_t stays None, so prefill runs the fused transposed gate/up kernel and the untransposed grouped down GEMM.

pub fn transpose_for_prefill_unified( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

2026-09-25: Unified layout: transpose gate and up for every expert, free their untransposed copies, then the same for down, so only one half is ever held twice. Afterwards the [N, K/2] kernels cannot run.

The caller sets METRALE_UNIFIED_MOE_LAYOUT=1 (without the hybrid variable) so use_t_layout_for_decode() sends decode to the _t kernels, calls this instead of transpose_for_prefill / transpose_gate_up_for_prefill, and does not build the CUTLASS grouped tables afterwards (they point at the freed weights).

pub fn transpose_for_prefill_hybrid( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

2026-09-25: Hybrid layout: the same build as transpose_for_prefill_unified but the untransposed originals are kept, so decode and verify stay on the [N, K/2] kernels (use_t_layout_for_decode() is false in hybrid mode). The expert weights are held twice; the caller must check that they fit.

pub fn build_cutlass_grouped_sfb( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, stream: u64, ) -> Result<()>

2026-09-25: Build the host tables of the CUTLASS grouped NVFP4 path (METRALE_HOLO_MOE_GROUPED_CUTLASS). Per expert, pack_weight_sfb swizzles the gate/up scales (and down’s, when a down scale table exists) into the CUTLASS SFB layout, from the transposed [K/16, N] scales when gate_ptrs_t/up_ptrs_t exist and from the untransposed [N, K/16] ones otherwise. The snapshot pairs them with the untransposed packed pointers and scale2 values. A no-op when a gate or up scale table is NULL.

impl MoeLayer

pub fn set_down_transpose_scratch( &mut self, scratch_packed: DevicePtr, scratch_scale: DevicePtr, packed_ptrs_t: DevicePtr, scale_ptrs_t: DevicePtr, )

2026-09-25: Wire the shared down_proj scratch and its transposed pointer tables. The factory (m2_setup) calls this for every layer after the gate+up-only transpose; one scratch serves all MoE layers, refilled by each layer’s prefill. scale2_vals is the untransposed down_ptrs table.

impl MoeLayer

pub fn predequant_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, stream: u64, ) -> Result<()>

2026-09-25: Make FP8 copies of the NVFP4 router gate and of the NVFP4 shared expert for the prefill FP8 GEMMs. Routed experts are untouched.

pub fn set_fp8_experts( &mut self, experts: &[Fp8ExpertWeight], shared_expert: Fp8ExpertWeight, gpu: &dyn GpuBackend, ) -> Result<()>

2026-09-25: Install FP8 routed experts: device pointer tables for gate, up and down, indexed by expert id in the FP8 kernels, plus the FP8 shared expert.

pub fn set_bf16_experts( &mut self, gate_experts: &[DenseWeight], up_experts: &[DenseWeight], down_experts: &[DenseWeight], shared_gate: DevicePtr, shared_up: DevicePtr, shared_down: DevicePtr, gpu: &dyn GpuBackend, ) -> Result<()>

2026-09-25: Install BF16 routed experts (FP8 experts dequantized at load; the qwen35 loader does this under METRALE_FP8_DEQUANT_MOE_TO_BF16=1).

shared_* are the shared expert’s BF16 gate/up/down pointers; all three NULL means no shared expert.

pub fn set_bf16_shared_expert( &mut self, gate_proj: DenseWeight, up_proj: DenseWeight, down_proj: DenseWeight, ) -> Result<()>

2026-09-25: Install a BF16 shared expert, whatever the routed experts’ precision.

impl MoeLayer

pub fn new( weights: MoeWeights, num_experts: usize, gate_nvfp4: Option<QuantizedWeight>, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<Self>

pub fn new_with_hash( weights: MoeWeights, num_experts: usize, gate_nvfp4: Option<QuantizedWeight>, tid2eid_dev: Option<DevicePtr>, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<Self>

2026-09-25: Like MoeLayer::new, with an optional hash-routing table tid2eid ([vocab_size, top_k] i64 on the device). Some makes this a hash-routed layer.