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:
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 200Routed 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 210Same 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 200Always-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.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.
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.
| Provider | Source | Contract | Selection |
|---|---|---|---|
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.
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.
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.
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.
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).