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.

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.

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".