Module moe_v41 

DeepSeek-V4.1 Flash MoE on streamed, still-quantized experts.

The routed experts never leave their Q2_K / Q3_K blocks: per token the router picks topk of 384, the ExpertLru gathers those experts’ raw slices into its device-visible slots (misses read from the SSD), and each expert runs as three K-quant GEMVs on the raw blocks (kquant_mmvq_q2_k_w for gate and up, kquant_mmvq_q3_k_w for down) with the activation quantised to q8_1 in between. The shared expert every token goes through is resident bf16 and runs on the dense GEMM.

Routing follows the CPU reference (deepseek_v41_ref::moe::gate) exactly: the logits are an f32-accumulated GEMM of the bf16 input against the bf16 gate weight, then sqrt(softplus(logit / temp)), and the top-k is chosen by score + correction_bias while the weights are the unbiased scores, renormalised and scaled by route_scale. The selection runs on the CPU from the downloaded logits ([tokens, 384] f32), which keeps the tie-breaking the reference’s and costs nothing at decode.

Oracle: moe_v41_tests.rs, synthetic Q2_K / Q3_K experts on a pinned arena against the CPU decoders + q8_1 emulation + the reference’s expert math.

Structs

MoeV41

MoeV41Cfg

MoeV41LayerWeights

One layer’s resident MoE weights: the router and the shared expert.

MoeV41Stage

What stage_m1 leaves for compute_m1: the distinct expert count and the routing (for callers and diagnostics), with the host-span timing.

MoeV41Timing

Where one call’s time went (wall clock, host side).

RouterWeights

The router of one layer alone: what predicting a layer’s selection from an earlier layer’s input needs (see MoeV41::predict_launch).

Functions

bf16_bytes

A [tokens, dim] bf16 device buffer’s bytes, for callers staging inputs.

prefetch_k

ATLAS_DS41_PREFETCH_K (default 0 = off), read once: how many of the next layer’s predicted experts to start reading while this layer computes (needs ATLAS_DS41_READER_POOL=1). Off by default: on the 09-19 standard the predictor’s router launch (2.7 ms a token) and the wasted reads on the shared disk cost more than the caught misses saved (K=4: 12.50/16.67, K=6: 13.63/16.39, against 14.20/17.73 without).

route_from_logits

The reference’s Gate.forward on f32 logits: returns (weights[tokens, topk], indices[tokens, topk]) in torch’s top-k order.

shared_side

ATLAS_DS41_SHARED_SIDE (default on; 0 = off), read once: at decode the shared expert runs on a side stream while the main stream does the router and the host picks the experts (the GPU is otherwise idle for that span); its output is added into acc after the routed experts, as before, so the accumulation order and the bits are unchanged.