Module glm5next_mlp 

GLM-5.3-Flash MLP production surface – dense FFN + routed NVFP4 MoE (TP + EP sharded). GLM-5.3-Flash MLP production surface — dense FFN and routed MoE.

Scoped to LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3.

Same shape as crate::layers::glm5next_dsa: the CUDA kernels already exist and are numerically proven against HF 5.16.1 on real weights — kernels/gb10/common/glm5next_ffn.cu (clamped SwiGLU, sigmoid router top-k, routed/shared combine), gated by examples/glm5next_{ffn,moe}_microtest.rs. What was missing, and what this module adds, is the production surface: config, kernel resolution, a weight contract and a forward a real layer can call, rather than an example wiring pointers by hand.

The stack this layer runs

DENSE (layers 0..first_k_dense_replace)   x -> gate_proj/up_proj -> clamped SwiGLU -> down_proj

ROUTED (every later layer)
  x ─┬─ gate.weight ──── logits(f32) ─ router_topk ─ ids[K], weights[K]
     ├─ experts[id]  ─── NVFP4 gate/up -> clamped SwiGLU -> NVFP4 down  (LOCAL ids only)
     └─ shared_experts ─ BF16 gate/up  -> clamped SwiGLU -> BF16 down
                                      └─> moe_combine -> partial -> all_reduce

🔴 Why one all-reduce covers BOTH EP and TP

The campaign’s topology is world = 2, TP = 2, EP = 2 on the same two ranks. The routed experts are EP-sharded (144/rank, remote ids contribute zero) and the dense/shared FFN is TP-sharded (column-parallel gate/up, row-parallel down). Both leave a partial sum of the same [T, hidden] output, and all_reduce(SUM) is linear, so the two partials are summed by one collective at the end of the site. Reducing them separately would be two collectives for the same answer.

🪤 That is only true because the combine happens BEFORE the reduce. Adding the shared expert after an all-reduce of the routed partial — which is what layers::moe::forward does, for a model whose shared expert is replicated rather than TP-sharded — would add a TP-partial shared output exactly once and lose the other rank’s half.

🪤 Traps this module exists to hold

• The SwiGLU clamp is ASYMMETRIC: gate is upper-bounded only, up is bounded both ways. It is also invisible on well-scaled activations — it fires on the tails. The limit comes from ModelConfig::swiglu_limit, which the glm5_next parser refuses to default.
• The router’s correction bias steers SELECTION ONLY. The emitted weight is the unbiased sigmoid score of the chosen expert. Gathering the biased score still produces a plausible mixture.
• routed_scaling_factor rides on the top-k weights and the shared expert is NOT scaled by it (apply_routed_scale_to_output = false). Scaling the shared output is the same defect with the opposite sign.
• The router is REPLICATED and must stay bit-identical across ranks. Masked-local EP is only equivalent to dispatch when every rank agrees on the same ids. Sharding the gate would give each rank partial logits and a different top-k — no crash, different experts.
• num_experts here is the FULL count (288). The rank’s local range is a separate field. Passing the local count to glm5next_router_topk would rank 144 experts and renormalise over the wrong denominator.

Re-exports

pub use weights::Glm5NextDenseMlpWeights;

pub use weights::Glm5NextExpertWeights;

pub use weights::Glm5NextMoeWeights;

Modules

build

Binding one GLM MLP site for a rank: TP slicing of the dense/shared halves, EP selection of the routed experts.

forward

The GLM MLP decode forward — dense FFN and routed MoE, one token.

forward_prefill_gemm

GLM routed-MoE prefill through the shared tensor-core grouped W4A16 GEMM.

weights

The GLM MLP weight contract, already sharded for this rank.

Structs

Glm5NextMlpConfig

GLM MLP geometry for one rank, read from the checkpoint config — never defaulted.

Glm5NextMlpKernels

Every kernel a GLM MLP site launches.

Enums

Glm5NextMlpKind

Which MLP a layer runs. Mirrors crate::layers::glm5next_skeleton::Mlp; kept separate so the runtime does not depend on the skeleton’s design-artifact types.

Constants

FFN_MODULE

Module name the GLM FFN kernels resolve from. kernels/gb10/common/glm5next_ffn.cu is not listed in common/KERNEL.toml’s [modules], so it takes its file stem.

GEMM_MODULE

[modules]: dense_gemm_bf16 = "gemm".

KERNEL_MAX_TOP_K

float best_w[16] in glm5next_router_topk — the most experts it can select per token.

MOE_GROUPED_MODULE

[modules]: moe_w4a16_grouped_gemm = "moe_w4a16" — the tensor-core grouped W4A16 GEMM.

MOE_MODULE

[modules]: moe_permute = "moe" — the token sort / permute / unpermute family.

W4A16_GEMV_MODULE

🔴 The DECODE weight kernel. w4a16_gemm is a tile GEMM: at M=1 it measured 9.7 GB/s on the routed experts — 26x off the 254 GB/s roofline and 45 % of the whole decode step (2026-08-28 profile). Every routed-expert projection here is M=1.

W4A16_MODULE

[modules]: w4a16_gemm = "w4a16".