Module moe 

DeepSeek-V4.1 MoE (Gate, Expert, MoE in the reference), CPU reference against the golden: sqrt-softplus routing with a correction bias that picks experts but does not scale them, top-k renormalisation, route_scale, the SwiGLU clamp, and one shared expert every token goes through.

Precision follows the reference line by line. The gate runs in f32 from a bf16 x. Each expert Linear is a bf16 GEMM (f32 accumulate, bf16 result), the SwiGLU is f32, the routing weight multiplies in f32, the product is cast back to bf16 before w2, and the per-expert outputs accumulate into an f32 y that is cast to bf16 once at the end. Accumulation order inside a GEMM is the only thing that differs from torch, so outputs are compared at bf16 tolerance while the routing weights and indices are compared exactly.

Structs

MoeCfg

MoeWeights

The routed experts’ parameters, experts[i] = (w1, w2, w3).

Functions

expert

One SwiGLU expert. w1, w3: [inter, dim]; w2: [dim, inter]; all bf16 values. x: [rows, dim] bf16 values; weights: per-row routing weight (f32) or None for the shared expert. Returns [rows, dim] bf16 values.

gate

Gate.forward for score_func = "sqrtsoftplus". x: [tokens, dim] bf16 values; w: [n_routed, dim] bf16 values; bias: [n_routed] f32. Returns (weights[tokens, topk] f32, indices[tokens, topk]), indices in torch’s topk order (descending by score + bias).

linear_bf16

F.linear on bf16 operands: y[i][o] = bf16(sum_d x[i][d] * w[o][d]), f32 accumulate. w is [out, in] row-major, as nn.Linear stores it.

moe

MoE.forward (single rank). Returns (y[tokens, dim] bf16 values, routing weights, routing indices).