Technical Summary: PagedWeight

Problem Statement

Mixture-of-Experts (MoE) Large Language Models (LLMs) offer high efficiency and accuracy by activating only a subset of experts per token. However, in serving scenarios, MoE models face a critical memory tension: the GPU must store both the massive model weights and the growing Key-Value (KV) cache. While techniques like PagedAttention manage KV cache fragmentation, they do not address the significant memory footprint of MoE weights, which can occupy over 60% of GPU memory.

Existing quantization methods are largely static, fixing precision before inference begins. While some runtime precision-adaptation methods exist, they do not utilize precision changes to dynamically reallocate GPU memory between weights and the KV cache. Consequently, there is a lack of systems that can dynamically free memory from less critical expert weights as the KV cache expands, without sacrificing task accuracy or interrupting inference.

Methodology: PagedWeight

The authors propose PagedWeight, a novel management system for MoE LLM serving that dynamically quantizes model weights at runtime to balance expert-weight precision with KV cache sizes. The system treats expert weights similarly to paged KV cache blocks, allowing for dynamic offloading and reloading.

Core Components

1. Paged Weight Representation:

   • PagedWeight operates at the granularity of per-expert linear-blocks (identified by layer, expert, and module type: gate_up or down).
   • It utilizes Any-Precision LLM (APL) technology, representing weights in an overlaid bit-plane format with associated Lookup Tables (LUTs).
   • A Weight Page Table tracks the desired bitwidth ($d_i$) and the committed bitwidth ($q_i$) for each linear-block. Pages can exist in GPU-resident, CPU-resident, or transfer states.

2. Quality-Aware Runtime Planner:
   The planner determines which weights to quantize (reduce bitwidth) based on three factors to minimize accuracy loss:

   • Offline Global Sensitivity: Uses Hessian-weighted scores to estimate the intrinsic sensitivity of each linear-block to quantization.
   • Online Routing Statistics: Tracks the "routing mass" (frequency and weight of selection) of experts. Frequently routed ("hot") experts are protected with higher bitwidth floors and damage multipliers.
   • Prompt Residuals: A prompt-specific correction term predicted via linear regression heads. This adjusts the global sensitivity estimate based on the current input sequence, acknowledging that quantization impact varies by prompt.

   The planner calculates a "damage" score for potential bitwidth reductions. When KV cache pressure triggers a memory target, the planner greedily selects the lowest-damage reductions (highest quality per byte saved) to meet the target.

3. Asynchronous Execution & Kernel Optimization:

   • Asynchronous Page Movement: To hide latency, the system offloads weight pages to CPU RAM and reloads them asynchronously. Bitwidth commitments are only updated at safe boundaries (e.g., after a plan is fully executed), ensuring inference is not interrupted.
   • Fused Mixed-Precision Kernel: A custom CUDA kernel reads Any-Precision bit-planes and LUTs directly, supporting different bitwidths for each linear-block within a single fused operation to reduce overhead.

Key Contributions

1. System Design: The proposal of PagedWeight, a paged-weight system for online Any-Precision MoE weights that manages committed/desired bitwidths, page states, and memory consumption dynamically.
2. Planning Algorithm: A quality-aware runtime planner that synthesizes offline sensitivity, online routing statistics, and prompt residuals to select low-damage weight reduction strategies under memory pressure.
3. Execution Strategy: Implementation of asynchronous page movement to hide offload/reload latency with minimal throughput loss, coupled with a fused mixed-precision MoE kernel.
4. Empirical Validation: Comprehensive evaluation across three MoE models (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) demonstrating superior performance over existing baselines.

Results

The authors evaluated PagedWeight against FP16, Any-Precision LLM (APL), DP-LLM, and MxMoE across language modeling, reasoning, and long-context tasks.

• Quality-Memory Tradeoff: PagedWeight achieves FP16-equivalent accuracy with up to 72.0% GPU memory savings and a 1.94× throughput improvement compared to baselines.
• Quality at Fixed Budget: At similar memory budgets, PagedWeight improves task quality over quantization methods by up to 39.3%, with a throughput loss of at most 4.1%.
• Long-Context Performance: Under tight memory budgets (e.g., 10 GB), PagedWeight maintains long-context quality (Passage Retrieval, NarrativeQA) comparable to FP16, whereas static quantization baselines (like APL) suffer significant degradation.
• Throughput: PagedWeight matches the throughput of uniform quantization baselines, with the largest observed drop being 4.1% at batch size 4.
• Ablation: Removing components like routing statistics, prompt residuals, or page movement consistently degrades performance, confirming the necessity of the full system design.

Significance

The paper claims that PagedWeight effectively navigates the complex tradeoff between model task accuracy, memory consumption, and throughput/latency. By treating expert weights as paged, manageable units rather than static assets, PagedWeight unlocks an unexplored tradeoff space between task accuracy and the allocation of GPU memory between weights and the KV cache. The system demonstrates that dynamic, quality-aware quantization is a viable and superior strategy for MoE serving compared to static quantization or methods that do not adapt to runtime memory pressure.