metrale_model_layers/layers/moe/forward/
zero_expert.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `MoeLayer::fp32_routing_active` and `MoeLayer::apply_zero_expert`, moved whole from
4//! `forward.rs`.
5//!
6//! Owner: model-layers (MoE).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11impl MoeLayer {
12    /// 2026-09-25: Whether the FP32 routing path is on for this layer: the pre-MoE
13    /// norm then also writes an FP32 copy of the router input
14    /// (`residual_add_rms_norm_gatef32` into `moe_router_in_f32`), which
15    /// `batched_gate_logits` reads through `dense_gemm_f32in`. Needs a dense gate
16    /// (no `gate_nvfp4`), no correction bias, the `dense_gemm_f32in` and
17    /// `moe_topk_f32` kernels, and, checked last, the `fp32_routing` lever (on
18    /// only when `METRALE_FP32_ROUTING=1`).
19    pub fn fp32_routing_active(&self, levers: &crate::layers::ops::ModelLevers) -> bool {
20        self.gate_nvfp4.is_none()
21            && self.correction_bias_dev.is_none()
22            && self.dense_gemm_f32in.0 != 0
23            && self.moe_topk_f32.0 != 0
24            && levers.fp32_routing
25    }
26
27    /// 2026-09-25: Add the zero-expert term `out[t,:] += zero_accum[t] * x[t,:]`,
28    /// where `zero_accum` holds the weights of the selected zero experts that
29    /// the softmax+bias router folded out. It adds into `out` in place, so it
30    /// runs after the routed blend has written `out`, for the tokens whose
31    /// routing wrote `zero_accum`. No launch when `router_logits_n` equals
32    /// `num_experts`; an error when the kernel is missing.
33    pub fn apply_zero_expert(
34        &self,
35        out: metrale_gpu_runtime::gpu::DevicePtr,
36        x: metrale_gpu_runtime::gpu::DevicePtr,
37        n: u32,
38        ctx: &ForwardContext,
39        stream: u64,
40    ) -> Result<()> {
41        if self.router_logits_n as usize == ctx.config.num_experts {
42            return Ok(());
43        }
44        anyhow::ensure!(
45            self.moe_zero_expert_add_k.0 != 0,
46            "zero-expert model but moe_zero_expert_add kernel is absent from this build"
47        );
48        ops::moe_zero_expert_add(
49            ctx.gpu,
50            self.moe_zero_expert_add_k,
51            out,
52            x,
53            self.zero_accum_dev,
54            n,
55            ctx.config.hidden_size as u32,
56            stream,
57        )
58    }
59}
metrale_model_layers/layers/moe/forward/zero_expert.rs

metrale_model_layers/layers/moe/forward/
zero_expert.rs
