DPA-MoE: A Runtime Precision Control Framework for MoE Inference

AI summary

Overview

Research area: Efficient inference for Mixture-of-Experts (MoE) large language models, specifically runtime weight quantization and precision allocation.

Technical level: Intermediate. Readers should be comfortable with quantization basics (bitplanes, scales, post-training quantization) and with how MoE routing works, but the paper's core idea is explained largely through its objective function and results rather than heavy derivations.

Scope (one sentence): The paper proposes DPA-MoE, a framework that encodes MoE expert weights so that a single checkpoint can be run at many effective weight-memory budgets, and that decides at inference time how many precision "chunks" each router-selected expert receives.

What This Paper Is About

Standard MoE models save compute by routing each token to only a few experts, but every selected expert is still read at one fixed precision. That means a deployed checkpoint exposes only a single point on the memory-versus-quality curve, and adapting to a new constraint requires fetching or re-encoding a separate multi-gigabyte quantized file. DPA-MoE's goal is to make expert precision a runtime decision: the pretrained router still decides which experts run, while a new allocator decides how much of each selected expert is used, using live router gate scores plus offline measurements of each expert's quantization error.

Key Contributions

1. Formulation of any-precision MoE. The authors separate sparse routing (which experts fire) from a precision controller (how many weight bits are read from each selected expert), treating the chunk count k rather than a numeric bit-width as the unit of precision.

2. SharedScale AnyBCQ encoding. A prefix-decodable expert encoding built on AnyBCQ that fixes one shared scale/offset table across all precision tiers, so a single checkpoint supports the full 2–8 chunk range without retraining or re-encoding. Per-expert storage overhead drops from 35.4% to 11.1% versus standard AnyBCQ.

3. A two-stage allocation framework. Offline profiling builds a residual lookup table R[L, e, k] measuring each expert's reconstruction error at each chunk count; at inference, these are combined with live gate scores g_e(x)^2 and solved as a multi-choice knapsack problem to distribute a fixed chunk budget across active experts.

4. Empirical evidence across four MoE models that gate-weighted allocation improves the weight-size–quality frontier over uniform chunk allocation, with modest steady-state throughput cost in a lightweight prototype.

Main Findings

• Better fidelity at matched memory: On Qwen3-30B-A3B and DeepSeekMoE-16B-Chat, the DPA curve lies below Uniform-k across the full compression range, reducing mean KL divergence by up to 4.5× at matched size (DeepSeekMoE-16B-Chat, k̄ = 4) and recovering equivalent fidelity at 2–4 GB less memory.

• Higher downstream accuracy at matched budgets: Macro-average accuracy across ARC-Challenge, ARC-Easy, WinoGrande, and PIQA improves at every k̄. On Qwen3-30B-A3B, accuracy rises from 63.38% to 66.12% at k̄ = 3 and from 66.92% to 67.90% at k̄ = 4; on DeepSeekMoE-16B-Chat, from 63.86% to 65.79% at k̄ = 3 and 64.57% to 66.79% at k̄ = 4.

• Throughput cost at fixed budget: At matched k̄, DPA costs 4–6% generation throughput on Qwen3-30B-A3B and 6–8% on DeepSeekMoE-16B-Chat relative to Uniform-k, with TTFT increasing by roughly 30–50%, attributed to the per-dispatch knapsack solve and non-uniform chunk materialization.

• Gains look better at matched quality: On Qwen3-30B-A3B, DPA at k̄ = 4 reaches 67.90% accuracy at 48.72 TPS, close to Uniform k = 5 (67.86%) but about 5% higher throughput than that setting's 46.29 TPS. On DeepSeekMoE-16B-Chat, DPA at k̄ = 4 gives 66.79% at 97.06 TPS versus Uniform k = 5 at 67.20% and 95.17 TPS, trailing by 0.41 accuracy points.

• Generality across architectures: On OLMoE-1B-7B-0125 (8 of 64 experts, no shared experts) and Qwen1.5-MoE-A2.7B (4 of 60 routed plus 4 shared), DPA reduces KL divergence by 26–69% wherever the allocation is not forced. Macro-average accuracy rises from 60.60% to 62.40% at k̄ = 3 on OLMoE, and from 62.04% to 63.45% at k̄ = 4 on Qwen1.5-MoE. The methods coincide at k̄ = 2 and k̄ = 8, where only one allocation is feasible.

• Which signal matters: An ablation on OLMoE-1B-7B-0125 shows gate-only allocation captures most of the gain over Uniform-k (26.4–44.3%), residual-only contributes little (0.2–7.9%), but combining both still beats gate-only by 4.6–5.4% at every budget tested (k̄ = 3 through 6).

• Cheap allocation step: The dynamic program costs O(n² k̄ k_range). For Qwen3-30B-A3B decoding (n = 8, B ≤ 64, k_range = 7) that is at most 3,584 candidate transitions per dispatch; a worst-case batched prefill with n ≤ 128 evaluates at most 917,504 transitions per MoE layer dispatch, amortized across the microbatch.

• Diagonal approximation is safe: Dropping cross-expert error terms preserves allocation rankings, with Spearman correlation of at least 0.999 across all precision levels, though cross terms account for 59.8% of the exact error at k = 2.

Methodology in Plain English

The authors start by changing how expert weights are stored. Instead of one quantized blob per expert, each expert's feed-forward weights become an ordered stack of chunks, where each chunk is a binary bitplane plus a per-group scale (group size 128) and a shared offset. Because chunks are a prefix code, using the first k chunks always yields a valid, more accurate reconstruction, and going from k to k+1 needs only the new bitplane and its scale. The "SharedScale" part means one scale table serves all tiers, rather than storing a separate table per tier as standard AnyBCQ does. The encoding objective combines a layer-level reconstruction term at full precision with a term that protects quality at every intermediate precision.

With precision now a runtime dial, the question becomes how to spend a fixed total chunk budget across the experts the router picked. The authors split this into two strictly separate pieces. Offline, they measure how much each expert's output deviates from its full-precision version at each k, averaged over a calibration set — a gate-free measure of intrinsic sensitivity. Online, they take the router's own gate scores, square them, and multiply by the offline residuals. Summing these per-expert costs and minimizing under a total-chunk constraint is a multi-choice knapsack problem, which they solve exactly with dynamic programming and backtracking. For batched prefill, gate scores are aggregated over the microbatch before the same optimization runs. Because the true objective includes cross-expert interaction terms, they verify empirically that dropping those terms leaves the ranking of allocations essentially unchanged.

Evaluation covers four open-weight MoE models of varying scale, routing width, and shared-expert design. Quality is measured as mean KL divergence between quantized and full-precision output distributions over 10 sequences of 4,096 tokens from WikiText-2, plus accuracy on four standard benchmarks; speed is measured as generation throughput and time-to-first-token over 100 sequences each from WikiText-2, C4-en, and the Korean split of mC4. Baselines are fixed GGUF quantizations at roughly 2-, 5-, and 8-bit (Q2_K, Q5_K_M, Q8_0) and a Uniform-k variant using the same encoding but without gate-aware allocation. Speed runs keep all chunks resident in GPU memory to isolate allocation cost.

Why This Matters

The paper reframes expert precision as a runtime variable rather than a property baked into a checkpoint. If one encoded artifact can serve many memory budgets, deployment no longer requires maintaining, shipping, or swapping a family of separately quantized multi-gigabyte files — which matters most exactly where the constraint shifts while the model stays loaded, such as local and on-device inference. The result also fills a conceptual gap: MoE already decides which experts run, so it is natural that a second controller should decide how much of each one to read, especially since expert FFN weights are about 95% of parameters in Qwen3-30B-A3B.

Real-world applications:

• Local and desktop assistants where available memory fluctuates as other applications come and go, and the user may want speed now and accuracy later.
• Memory-constrained edge or embedded serving of MoE models, where a single checkpoint must cover several device tiers.
• Interactive and coding assistants where tight time-to-first-token budgets can be traded against answer quality per request.
• Private on-premise or defense-sector deployments that cannot fetch additional model variants and need one artifact to cover multiple operating points.

Industry relevance: The method is deliberately positioned as a precision-allocation primitive complementary to existing serving systems such as DeepSpeed-MoE, HOBBIT, and D2MoE, rather than a replacement. It requires no retraining, no extra adapter, and no change to the pretrained router, which lowers the barrier to adoption — although it does require access to router gate scores, which is a real constraint for closed-weight deployments.

Future Directions

• Tighter runtime integration: The current prototype uses unfused allocation and chunk materialization on top of llama.cpp, so the measured TTFT penalty and throughput loss are prototype costs; fusing allocation with expert dispatch kernels is the obvious next step.
• Batch-size-aware fallback policies: The authors note that allocator cost grows with batch size while its benefit shrinks, and suggest deployments should switch back to Uniform-k past a device-dependent threshold — but that threshold is left to be characterized per target hardware.
• Alternative encoding substrates: SharedScale AnyBCQ is presented as one prefix-stable option, not the only one; the allocation method should transfer to other progressive encodings whose precision units can be treated as allocatable chunks with measurable residuals.
• Higher-level policies on top of the precision interface: Phase-aware decoding, adaptive quality control, and speculative decoding are named as candidates that could drive the exposed precision knob, and the surrogate objective (gate scores plus finite calibration residuals) is acknowledged as an imperfect proxy for true downstream loss.

Target Audience

Researchers and engineers working on LLM inference efficiency, MoE serving systems, and post-training quantization will get the most from this paper, particularly those interested in any-precision or progressive encodings and in input-adaptive resource allocation. It is also useful for practitioners deploying MoE models under shifting memory constraints on client or edge hardware, and for systems researchers looking for a precision-control primitive to compose with their own schedulers. Readers seeking a turnkey serving system or a head-to-head comparison against published adaptive-precision methods such as D2MoE will not find one here — the authors state that D2MoE has no public implementation and was not reproduced.

Authors’ abstract

Mixture-of-Experts (MoE) models reduce computation by routing each token to a small set of experts, but each selected expert is still used at a fixed precision. We present DPA-MoE, a runtime precision control framework that makes routed experts any-precision. DPA-MoE encodes expert feed-forward weights as prefix-decodable chunks using SharedScale AnyBCQ, profiles per-expert reconstruction residuals offline, and allocates chunks at inference time using live router gate scores. This post-routing allocation lets a single encoded MoE checkpoint operate at multiple effective weight budgets without retraining or re-encoding. Experiments on Qwen3-30B-A3B and DeepSeekMoE-16B show that DPA-MoE improves KL divergence and downstream accuracy over uniform chunk allocation at matched budgets, with modest steady-state throughput overhead in our prototype.