PagedWeight: Efficient MoE LLM Serving with Dynamic Quality-Aware Weight Quantization

Overview§

PagedWeight is a runtime weight quantization method for Mixture-of-Experts (MoE) LLMs that dynamically adjusts precision per expert based on its importance during serving. It aims to reduce GPU memory pressure from model weights to accommodate the growing KV cache, thereby improving throughput without sacrificing quality.

Core Methodology§

PagedWeight operates at the granularity of individual experts in an MoE layer. During inference, it selects a target total memory budget $B$ for weights. The key idea is to assign higher precision to more important experts and lower precision to less important ones, all while keeping the total weight memory under $B$.

Quality-Aware Quantization§

The importance of an expert $e$ is measured by its contribution to the model's output quality. PagedWeight uses a lightweight quality proxy: the average routing probability over a calibration dataset. Let $p_e$ be the average routing score for expert $e$. Then the allocated precision $b_e$ (bits) is determined by solving an optimization problem:

$$ \min_{b_e} \sum_e p_e \cdot L(b_e) \quad \text{s.t.} \quad \sum_e b_e \cdot W_e \leq B, $$

where $L(b_e)$ is the expected quantization loss (e.g., mean squared error) for precision $b_e$, and $W_e$ is the number of weight elements in expert $e$. The solution leads to higher $b_e$ for experts with large $p_e$ or high sensitivity.

Dynamic Paging§

To enable runtime adjustments, PagedWeight introduces a paging mechanism similar to virtual memory. The GPU memory is divided into pages, and expert weights are quantized and stored in pages of varying bit-widths (e.g., 8-bit, 4-bit, 3-bit). During a forward pass, the router determines which experts are needed. PagedWeight fetches the corresponding pages, dequantizes them on-the-fly, and keeps only the active experts in full precision temporarily. Inactive expert pages can be swapped out to CPU memory or compressed further.

Implementation Details§

• Quantization Scheme: Supports uniform affine quantization with per-expert scaling factors. For a given bit-width $b$, the quantization operation is:

$$Q(w) = \text{round}\left(\frac{w}{s}\right) \cdot s, \quad s = \frac{\max(|w|)}{2^{b-1} - 1}.$$

• Calibration: Importance scores $p_e$ are computed once from a small calibration set. A calibration set can be reused across serving runs.
• Memory Management: Uses a page table mapping expert IDs to physical page addresses. The scheduler dynamically adjusts page sizes based on current KV cache occupancy.

Code Illustration§

import torch

class PagedWeightExpert(torch.nn.Module):
    def __init__(self, expert_weight, importance_score):
        super().__init__()
        self.importance = importance_score
        # Store original weight for potential fine-tuning
        self.register_buffer('weight_fp16', expert_weight)
        self.page_size = 0  # in bytes, determined by allocator

def compute_budgeted_precision(importances, weights_size, budget_bytes):
    # Simplified: allocate bits proportional to importance
    total_importance = importances.sum()
    allocation = budget_bytes * importances / total_importance
    bits = allocation / weights_size
    # Clamp to available bit-widths
    bit_options = torch.tensor([3, 4, 8])
    idx = torch.bucketize(bits, bit_options)
    return bit_options[idx.clamp(max=2)]

# Usage:
# budgets = compute_budgeted_precision(importances, expert_params, target_memory)
# Then quantize each expert accordingly.

Experimental Results§

• Memory Savings: Up to 72% GPU memory reduction vs. FP16 with equivalent accuracy on WikiText-2 perplexity.
• Throughput: Up to 1.94x improvement over FP16 serving.
• Quality vs. Memory: Outperforms uniform quantization baselines (e.g., GPTQ, SmoothQuant) by up to 39.3% in perplexity at similar memory budgets.

Summary§

PagedWeight dynamically trades off expert precision for KV cache space, enabling larger effective batch sizes and higher throughput without retraining. It is orthogonal to other compression techniques and can be combined with them.