Module moe 

2026-09-25: DeepSeek-V4.1 MoE (the reference’s Gate, Expert, MoE), 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: the gate runs in f32 from a bf16 x. Each expert projection is linear_bf16 (f32 accumulate, bf16 result), the SwiGLU is f32, the routing weight multiplies in f32, the product is rounded to bf16 before w2, and the per-expert outputs accumulate into an f32 y rounded to bf16 once at the end.

Owner: model-arch, DeepSeek-V4.1 reference. Invariants: none beyond the types.

Structs

MoeCfg

MoeWeights

2026-09-25: The MoE parameters; experts[i] and shared are (w1, w2, w3).

Functions

expert

2026-09-25: 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

2026-09-25: The gate 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 descending by score + bias (a stable sort).

linear_bf16

2026-09-25: Linear on bf16 operands: y[i][o] = bf16(sum_d x[i][d] * w[o][d]), f32 accumulate. w is [out, in] row-major.

moe

2026-09-25: The MoE forward on one rank. Returns (y[tokens, dim] bf16 values, routing weights, routing indices).