Module glm5next_mlp 

2026-09-25: GLM-5.3-Flash MLP: the dense SwiGLU FFN and the routed NVFP4 MoE, their kernels and per-rank geometry.

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

• Glm5NextMlpConfig::from_config returns a config only when the dense and shared widths divide over tp_world_size, num_experts divides over ep_world_size, and validate passes.
• For such configs, the local_expert_ranges of ranks 0..ep_world_size partition 0..num_experts, and local_slot is None for every id outside this rank’s range.

One all-reduce for both EP and TP

With TP or EP above 1, a routed site’s output is a partial sum on every rank: the shared expert is TP-sharded (local_shared_intermediate) and each rank runs only the experts it owns. forward::forward_moe adds the shared output and the routed slots in one combine kernel, and the layer then reduces that output once. The combine has to come before the reduce: adding a TP-sharded shared expert after the reduce would drop the other ranks’ part of it.

Numerics the kernels rely on

• The SwiGLU clamp is asymmetric: gate is bounded above only, up on both sides. The limit is ModelConfig::swiglu_limit, which the glm5_next parser refuses to default.
• The router’s correction bias steers selection only. The emitted weight is the chosen expert’s unbiased sigmoid score.
• routed_scaling_factor multiplies the top-k weights; the shared expert is added unscaled.
• Every rank holds the whole router and ranks all num_experts, so for the same input every rank selects the same ids. Sharding the router would give each rank different partial logits and a different top-k.
• num_experts is the full routed-expert count; local_experts is this rank’s share. glm5next_router_topk must be given the full count.

Re-exports

pub use weights::Glm5NextDenseMlpWeights;

pub use weights::Glm5NextExpertWeights;

pub use weights::Glm5NextMoeWeights;

Modules

build

2026-09-25: Bind one GLM-5.3 MLP site for a rank: TP-slice the dense FFN or shared expert, and build the routed experts’ pointer tables over GLOBAL expert ids.

forward

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.

forward_prefill_gemm

2026-09-25: Routed-MoE prefill for GLM-5.3 through the shared grouped W4A16 GEMM (moe_w4a16_grouped_gemm_ptrtable and its tile variants): one launch per projection for all experts. tile.rs holds the tile table and the env levers, dispatch.rs the launch.

weights

2026-09-25: Device weight handles for one GLM-5.3 MLP site, already sharded for this rank.

Structs

Glm5NextMlpConfig

2026-09-25: GLM MLP geometry for one rank, built by Self::from_config.

Glm5NextMlpKernels

2026-09-25: The kernels a GLM MLP site launches.

Enums

Glm5NextMlpKind

2026-09-25: Which MLP a layer runs; the same variants as crate::glm5next_skeleton::Mlp.

Constants

FFN_MODULE

2026-09-25: Module of kernels/gb10/common/glm5next_ffn.cu. A .cu file not listed in common/KERNEL.toml’s [modules] resolves under its file stem; the listed ones below do not, so each module name here is checked against that table.

GEMM_MODULE

2026-09-25: [modules]: dense_gemm_bf16 = "gemm".

KERNEL_MAX_TOP_K

2026-09-25: The most experts glm5next_router_topk can select per token: its per-token selection lives in the shared arrays sel_id[16] and sel_w[16].

MOE_GROUPED_MODULE

2026-09-25: [modules]: moe_w4a16_grouped_gemm = "moe_w4a16", the tensor-core grouped W4A16 GEMM.

MOE_MODULE

2026-09-25: [modules]: moe_permute = "moe", the token sort / permute / unpermute kernels.

W4A16_GEMV_MODULE

2026-09-25: Module of w4a16_gemv.cu (not in [modules], so its file stem). Its GEMV kernels run every routed-expert projection except the grouped prefill GEMM.

W4A16_MODULE

2026-09-25: [modules]: w4a16_gemm = "w4a16".