Module forward

Module forward 

Source
Expand description

2026-09-25: The GLM-5.3 MLP forward over rows rows: the dense SwiGLU FFN (also used for the shared expert) and the routed MoE site.

Owner: model-arch (GLM-5.3). Invariants:

  • forward_dense and forward_moe leave a partial sum in out when the site is split over TP or EP ranks; the caller all-reduces.
  • forward_moe zeroes expert_out before any expert writes it, so an expert another EP rank owns adds zero.

Structs§

Glm5NextMlpWorkspace
2026-09-25: MLP scratch for up to max_rows rows, allocated once at load. The buffer sizes are listed in mlp_ws_bytes.

Constants§

MOE_ROW_BATCH_MAX_ROWS
2026-09-25: Widest compiled w4a16_gemv_sw_moe_batchm_mR tier: the METRALE_MOE_BATCHM_ENTRY(2..=8) instances in kernels/gb10/common/w4a16_gemv.cu, held at index R - 2 of Glm5NextMlpKernels::w4a16_gemv_sw_moe_batchm. forward_moe splits a wider row group into sub-groups no wider than this.
MOE_ROW_UNION_MAX_IDS
2026-09-25: Most ids glm5next_moe_row_union resolves: it runs as one block of rows * top_k threads, and forward_moe does not take the row-batched path above this.

Functions§

forward_dense
2026-09-25: A BF16 SwiGLU MLP of width inter, down(clamped_swiglu(gate(x), up(x))), for a dense layer or the shared expert. Errors when inter or m does not fit the workspace.
forward_moe
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.
mlp_ws_bytes
2026-09-25: Byte size of each Glm5NextMlpWorkspace buffer, in the order new allocates them: a_gate, a_up, a_act, logits, ids, wts, expert_out, shared_out, u_eid, u_slot, sorted_token_ids, sorted_expert_ids, expert_offsets, token_to_perm. The loader logs their sum. new computes the same sizes on its own, so the two must change together.
mlp_ws_shared
2026-09-25: Whether the text layers share one MLP workspace: yes unless METRALE_GLM_MLP_WS_SHARED=0. Read once.
mlp_ws_total_bytes
2026-09-25: Total device bytes of one MLP workspace of max_rows rows.
row_batch_max
2026-09-25: Widest sub-group of the row-batched routed path: METRALE_GLM_MOE_ROW_BATCH_MAX, default MOE_ROW_BATCH_MAX_ROWS, clamped to 1..=MOE_ROW_BATCH_MAX_ROWS. Read once.