forward_moe

Function forward_moe 

Source
pub fn forward_moe(
    gpu: &dyn GpuBackend,
    k: &Glm5NextMlpKernels,
    cfg: &Glm5NextMlpConfig,
    w: &Glm5NextMoeWeights,
    x: DevicePtr,
    out: DevicePtr,
    rows: usize,
    ws: &Glm5NextMlpWorkspace,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: One routed MoE site over rows rows of x into out: router, top-k, routed experts, shared expert, and the weighted combine.

The routed experts take one of three paths:

  • the grouped GEMM (forward_prefill_gemm), for rows > MOE_ROW_BATCH_MAX_ROWS and rows >= prefill_gemm_min_rows();
  • the row-batched union GEMV over the sub-groups of moe_row_groups, for rows >= 2;
  • per row otherwise, device-dispatched (w4a16_gemv_sw_moe) or host-dispatched.

Route tracing (METRALE_GLM_ROUTE_TRACE=1) forces the per-row path, the one that reads ids back. Errors when rows does not fit the workspace or the bound expert count is not this rank’s share.