Module latent_moe

Module latent_moe 

Source
Expand description

Stable LatentMoE CPU reference.

latent = down_proj(h)                 // hidden → moe_latent (no pre-norm)
scores = sigmoid(router(h)); top-k on scores+bias  // noaux_tc
y = Σ w_i expert_i(latent)            // SiTU-GLU experts
y = RMSNorm(y) if latent_moe_use_norm; y = up_proj(y)
out = y + shared_experts(h)           // shared stay full-width

Structs§

LatentMoeConfig

Functions§

expert_situ
One SiTU-GLU expert: down( situ(w1 x, w3 x) ). Weights are [out, in].
latent_moe_forward
Routed latent path + optional shared expert (identity-scale for tests).
mix_routed_experts
Mix selected SiTU-GLU experts into latent. Empty w1 means packed-only (no host dequant) — CUDA grouped GEMM must have run instead.
moe_from
sigmoid_topk
noaux_tc: sigmoid(logits) for mix weights; scores + bias only ranks.