MoE Expert Kernels: Grouped Prefill and Mixed-Quant Compact Experts

Six weeks of MoE kernel work, kernel by kernel: the shared router and expert building blocks, the Qwen3.5-MoE grouped/bucketed expert prefill path (PR #376), the Laguna compact mixed-quant expert family (PRs #404, #422), and two throughput-oriented provider enhancements that are not MoE but landed in the same window (PRs #412, #330). Every provider ID below names a real map in version/v8/kernel_maps and a function you can read in src/kernels.

1. The shared building blocks: routers and the SwiGLU expert shape

Every MoE family in the registry is the same two pieces with different contracts: a router that picks experts and weights per token, and an expert SwiGLU that computes down(silu(gate(x)) * up(x)) per routed expert and accumulates routing-weighted results in FP32. The two router providers in src/kernels/topk_kernels.c are deliberately different numerical contracts, not tuning variants:

group_limited_topk_router_sigmoid_f32 (topk_kernels.c:728) choice_scores = sigmoid(logits) + correction_bias # bias steers selection only group_scores = sum of top-2 choice_scores within each group selected_groups = topk(group_scores, topk_group) selected = topk(choice_scores masked to selected_groups, k) weights = gather(raw_scores, selected) # NOT the biased scores moe_softmax_topk_router_llama_f32 (topk_kernels.c:269, candidate) p = softmax(logits) over ALL experts # full softmax first selected = topk(p, k) weights = p[selected] / max(sum(p[selected]), 6.103515625e-5) weights *= routed_scaling_factor
Why two routers can never merge. The correction bias changes which experts fire while leaving the weights gathered from raw scores; the llama router changes the weights themselves (full-softmax then renormalize). Different token routings and different weights are different numerical contracts, so they live in different equivalence groups and the resolver fails closed rather than substituting one for the other.
The expert shape everything below shares. Routed experts read the hidden state once per token, run gate/up projections, apply the ggml-compatible SwiGLU, run the down projection, and add the result into the output row scaled by the routing weight. What varies per provider is the quant contract of each leg: gate/up storage, down storage, activation quantization, and reduction order.

Registry-level router context: kernel architecture — MoE routers and experts. The router selection math above is verified against topk_kernels.c.

2. Qwen3.5-MoE: grouped and bucketed expert prefill (PR #376)

Qwen3.5-35B-A3B is a 40-layer hybrid MoE with 256 routed experts and top-8 routing, Q4_K expert gate/up, Q5_K expert down, and a Q8_0 gated shared expert (real GGUF audit: 733 tensors). Naive per-token dispatch pays T×k scattered expert calls and re-quantizes the hidden state for every (token, expert) pair. The bucketed path turns that into one quantization pass plus contiguous per-expert GEMM segments on the persistent thread pool.

moe_swiglu_expert_forward_q4k_q5k production priority 200
Routed experts: Q4_K gate/up, Q5_K down, Q8_K activations, routing-weighted FP32 accumulate. Serial workspace at src/kernels/axpy_kernels.c:578, thread-pool _parallel_workspace at axpy_kernels.c:1105. Equivalence group q4_k_q8_k_gate_up_ggml_swiglu_q5_k_q8_k_down_weighted_fp32, phases [prefill, decode].
moe_swiglu_expert_forward_q4k_q5k_bucketed candidate priority 210
Same contract, grouped execution (axpy_kernels.c:1951 + version/v8/src/ck_parallel_prefill_v8.c): count-sort tokens by expert per top-k slot, quantize the hidden state to Q8_K once, then run contiguous per-expert GEMM segments. Priority 210 sits above the production non-bucketed provider at 200 — inside the same equivalence group, so it can only ever replace that provider, never a different contract.
moe_swiglu_shared_forward_q8_0_gated candidate priority 200
Always-on shared expert in Q8_0 with a scalar sigmoid gate applied before the routed add (axpy_kernels.c:2174, parallel at :2285). Equivalence group q8_0_q8_0_shared_swiglu_fp32_sigmoid_gate_then_routed_add.
VNNI weight preparation leaf.
gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4 (src/kernels/gemm_kernels_q4k_q8k_vnni.c:1123) plus the map-owned prepare hook ck_moe_prepare_q4k_gate_up_vnni_x8: AVX-VNNI ×8-packed Q4_K gate/up weight layout, built once at load time, so the bucketed segments run packed integer dot products instead of per-call unpacking.
Bucketed expert prefill (moe_swiglu_expert_forward_q4k_q5k_bucketed) tokens + top-k routing T tokens x k slots (Qwen3.5-MoE: top-8 of 256) router weights w[t,j] from sigmoid router count-sort by expert per top-k slot: histogram then scatter planner-owned scratch - no malloc per-expert contiguous segments expert e owns rows [off_e, off_e + cnt_e) dense row batching replaces scatter quantize hidden ONCE hidden fp32 -> Q8_K: one pass per token not one pass per (token, expert) pair packed VNNI GEMM segments Q4_K gate/up x8-packed at load time thread pool runs expert row ranges weighted FP32 accumulate out[t] += w[t,j] * expert(x_t) selected-slot reduction order preserved CONTRAST - naive per-token dispatch T x k scattered expert calls, one (token, expert) pair at a time hidden re-quantized per pair: T x k Q8_K passes instead of 1 production 200: ..._q4k_q5k (serial/parallel) candidate 210: ..._q4k_q5k_bucketed same equivalence_group - priority ranks within it Evidence (PR #376): 10-row routed-expert path 13.5 ms -> 2.5 ms per layer, byte-exact; routed leaf 1.36x at 512 rows, 1.47x at 4096 rows; 4096-token model traffic 1.12 TB -> 504 GB (62.5% fewer DRAM reads). Evidence (PR #376): 512-token P3 prefill at 35.35 tok/s with the same top-1 token as pinned llama.cpp; Q8_0 gated shared expert 10.572 s -> 1.079 s at 4096 tokens on Ryzen 9950X3D. Caveats: bucketed + shared-gated providers are candidate (opt-in routes only); CKE is ~0.73x matched llama.cpp on this model - attention and recurrent projections still dominate the gap.
click / tap the diagram to expand

The preparation hook is declared in the bucketed map itself (version/v8/kernel_maps/moe_swiglu_expert_forward_q4k_q5k_bucketed.json), so weight layout is a map-owned decision, not a model-name branch in codegen.

3. Laguna: compact-MoE mixed-quant experts (PRs #404, #422)

Laguna-XS 2.1 keeps routed and shared experts in mixed compact storage: Q4_K gate/up with either Q4_K or Q6_K down. Activations are quantized fp32→Q8_K at the kernel boundary, the SwiGLU is the ggml-compatible formulation, and the accumulate is routing-weighted FP32. PR #422 added the *_parallel_workspace variants that dispatch rows (prefill) and experts (decode) over the persistent thread pool — measured bit-exact against the serial providers.

ProviderSourceContractSelection
moe_swiglu_expert_forward_q4k_q4k axpy_kernels.c:743 / parallel :1461 Q4_K gate/up × Q8_K acts → ggml SwiGLU → Q4_K down → weighted FP32 production 205, [prefill, decode], group ..._q4_k_q8_k_down_weighted_fp32
moe_swiglu_expert_forward_q4k_q6k axpy_kernels.c:666 / parallel :1438 same, Q6_K down production 205, group ..._q6_k_q8_k_down_weighted_fp32
moe_swiglu_shared_forward_q4k_q4k axpy_kernels.c:877 / parallel :1603 always-on shared expert, Q4_K down, no routing weight production 205, group q4_k_q8_k_shared_gate_up_ggml_swiglu_q4_k_q8_k_down_fp32
moe_swiglu_shared_forward_q4k_q6k axpy_kernels.c:820 / parallel :1584 same, Q6_K down production 205, group ..._shared_gate_up_ggml_swiglu_q6_k_q8_k_down_fp32
attn_gate_softplus_mul_forward src/kernels/hybrid_attention_kernels.c:132 per-head softplus attention-output gate, head-major; linear branch above 20 declared by the circuit as attention_gate_projection + gate op

The circuit that binds them, version/v8/circuits/laguna.json: sigmoid router with correction bias → top-8 routed SwiGLU plus shared SwiGLU; hybrid sliding/global attention with attention_gate_projection and the softplus gate; a YaRN rope cache (yarn_rope_cache_contiguous_positions_f32) on global layers; prefill chunk 4096.

One dataflow, two quant contracts Q4_K gate/up x Q8_K activations acts quantized once (fp32 -> Q8_K) ggml SwiGLU silu(gate) * up fp32 intermediate Q4_K or Q6_K down x Q8_K activations down quant = the contract split FP32 weighted accumulate out[t] += w[t,j] * expert_out shared expert: always-on, unweighted equivalence_group A (Q4_K down) q4_k_q8_k_gate_up_ggml_swiglu_q4_k_q8_k_down_weighted_fp32 moe_swiglu_expert_forward_q4k_q4k (:743) moe_swiglu_expert_forward_q4k_q4k_parallel_workspace (:1461) moe_swiglu_shared_forward_q4k_q4k[_parallel_workspace] (:877/:1603) production - priority 205 - phases [prefill, decode] parallel == serial, bit-exact (PR #422) equivalence_group B (Q6_K down) q4_k_q8_k_gate_up_ggml_swiglu_q6_k_q8_k_down_weighted_fp32 moe_swiglu_expert_forward_q4k_q6k (:666) moe_swiglu_expert_forward_q4k_q6k_parallel_workspace (:1438) moe_swiglu_shared_forward_q4k_q6k[_parallel_workspace] (:820/:1584) production - priority 205 - phases [prefill, decode] parallel == serial, bit-exact (PR #422) priority ranks providers INSIDE one equivalence group - a Q6_K-down provider can never be selected for a Q4_K-down contract (fail-closed) Evidence (PRs #404, #422): parallel providers bit-exact vs serial; Laguna 512-token prefill 23.29 s -> 2.86 s (22.0 -> 179.2 tok/s), occupancy 1.59 -> 13.78 of 16 Ryzen 9950X3D cores; 8192-token prefill 138.3 tok/s. Caveats: router replay parity is within one ULP, not long-trajectory (near-tied route flips vs llama.cpp); the softplus gate formula is observed from the model, not oracle-validated.
click / tap the diagram to expand

The structural view (layer kinds, split RoPE, gated GQA) is in Architecture Variants — Laguna; this page covers the provider layer only.

4. Throughput kernels beyond MoE

Two provider enhancements from the same window are not MoE, but they are the same engineering pattern: find the serialization point, partition the independent dimension across the persistent thread pool, and keep the numerical contract identical so selection stays honest.

Nemotron-H: parallel recurrent prefill (PR #412).
Mamba2 prefill alternated parallel projections with serial recurrent valleys. mamba2_conv1d_f32_parallel_dispatch partitions the conv1d by independent channel ranges (sequential fallback below 8 rows / 16 conv channels); mamba2_selective_scan_f32_parallel_dispatch chunks by independent head ranges (fallback below seq_len 8 / 2 heads). Kill switch: CK_DISABLE_MAMBA2_PARALLEL_PREFILL. The maps declare this through a new impl.variants[].parallel_head_ranges metadata pattern (shape constraints T_min: 8, H_min: 2, named serial fallback) instead of hardcoded thresholds in the resolver.
evidence Ryzen 9950X3D 8K prefill: 252.0 s → 135.0 s (32.5 → 60.7 tok/s), 7.51 → 13.83 active cores, exact first-logit and token agreement. Caveat: llama.cpp remains 2.12× faster on this host.
Gemma3: Q5_1 unpack reuse (PR #330).
gemm_nt_q5_1_q8_1_m4 (src/kernels/gemm_kernels_q5_1_q8_1.c:365) unpacks each Q5_1 weight block once and dots it against four pre-quantized Q8_1 activation rows with AVX2 VNNI dpbusd, instead of re-unpacking per row. The map (gemm_nt_q5_1) is production priority 200 with prefill-only phases — decode keeps the per-row path.
evidence production-shape leaf 45.93 ms → 26.58 ms (1.73×); end-to-end Gemma prefill improved 6.4% / 5.1% / 2.3% at 32 / 128 / 512 tokens with unchanged tolerances. Caveat: the Q5_1 oracle is an observed 1e-4 contract, not bit-exact to the llama.cpp graph.
gemm_nt_q5_1_q8_1_m4 - unpack once, dot four rows Q5_1 weight block 32 values, one block scale K multiple of 32 unpack ONCE qs + qh -> int8 registers amortized over 4 rows r0: Q8_1 row - dpbusd dot r1: Q8_1 row - dpbusd dot r2: Q8_1 row - dpbusd dot r3: Q8_1 row - dpbusd dot 4 x FP32 accumulators ascending block reduction order same contract as serial Q5_1 selection: production, priority 200, phases [prefill] only - decode keeps the per-row provider equivalence_group q5_1_weight_q8_1_internal_fp32_output Evidence (PR #330): production-shape leaf 45.93 ms -> 26.58 ms (1.73x); Gemma prefill +6.4% / +5.1% / +2.3% at 32 / 128 / 512 tokens. Caveats: Q5_1 is an observed 1e-4 oracle contract, not bit-exact; Gemma remains slower than llama.cpp end to end.
click / tap the diagram to expand

5. Why this matters

Everything on this page is the same expert-GEMM pattern specialized per quant contract: Q4_K/Q5_K for Qwen3.5-MoE, Q4_K/Q4_K and Q4_K/Q6_K for Laguna, Q8_0 for the gated shared expert. The performance work — bucketing, weight preparation, parallel row/expert dispatch, unpack reuse — never changes the arithmetic contract; it changes the schedule. Selection metadata is what keeps that honest: each quant combination lives in its own equivalence_group, priority only ranks providers that already share a contract, and anything unproven stays candidate behind opt-in routes. A faster provider can therefore never silently change what the model computes — when in doubt, the resolver fails closed.

Related: v8 Kernel Architecture — MoE routers and experts, Kernel Maps and Provider Selection, Architecture Variants, Cohere Kernel Story (North composes the Q4_K/Q5_K experts documented here).

Image
100% | |
Scroll to zoom | Drag to pan | W/H to fit | 0 to reset | ESC to close