🤖 machine learning

MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

MoE-nD introduces a per-layer mixture-of-experts routing framework that dynamically assigns unique token eviction ratios and quantization bit-widths to each transformer layer under a global memory budget, achieving up to 14x KV cache compression while maintaining accuracy comparable to uncompressed baselines on long-context reasoning tasks.

Original authors: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

Published 2026-04-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Overstuffed Backpack"

Imagine a Large Language Model (LLM) as a brilliant student trying to solve a very long story or a complex math problem. To do this, the student needs to keep a running list of everything they've read so far in their "working memory" (called the KV Cache).

For short stories, this memory fits easily in a backpack. But for long contexts (like a whole book or a 16,000-word document), the backpack gets so heavy and big that:

It slows the student down (latency).
It rips the backpack open (Out-of-Memory errors) on standard computers.

The Old Solution: "One Size Fits All"

To fix this, researchers tried to shrink the backpack. They used two main tools:

Eviction (Throwing things out): Deleting old notes to make space.
Quantization (Writing in shorthand): Writing notes in smaller, less precise handwriting (e.g., using 4-bit numbers instead of 16-bit).

The Flaw: Previous methods treated the student's brain like a cookie cutter. They applied the same rule to every part of the brain.

Example: "Throw away 50% of the notes from the entire story, and write all remaining notes in tiny shorthand."

Why this fails: The student's brain isn't uniform.

Layer 1 (The beginning): Might be crucial for understanding the plot. If you throw away 50% of these notes, the student forgets the story entirely.
Layer 20 (The middle): Might be full of fluff. Throwing away 50% here does almost no harm.
Layer 28 (The end): Might need high precision to solve the final math equation. Writing in "tiny shorthand" here causes calculation errors.

Treating every layer the same is like wearing a winter coat in the summer and a t-shirt in the winter because you decided to wear the "average" outfit for the whole year.

The New Solution: MoE-nD (The "Smart Wardrobe")

The authors propose MoE-nD, which stands for Mixture-of-Experts for n-Dimensional routing.

Think of this as a Smart Wardrobe Manager for the student's backpack. Instead of one rule for everyone, the manager looks at every single layer of the brain individually and asks: "What does this specific part of the brain need right now?"

It has three "Experts" (options) to choose from for each layer:

The Trash Can (Eviction): How much of this layer's memory can we delete?
The Shorthand Pen (K-Quantization): How small can we write the "Key" notes?
The Shorthand Pen (V-Quantization): How small can we write the "Value" notes?

How it works:

The Map (Offline Calibration): Before the student starts the test, the system runs a quick, cheap test to see how sensitive each layer is. It creates a map: "Layer 5 hates being deleted. Layer 10 hates being written in shorthand. Layer 20 doesn't care about either."
The Budget (Global Limit): The student has a strict backpack size limit (e.g., 136 MB).
The Greedy Solver: A smart algorithm looks at the map and the budget. It plays a game of "best value."
- It says: "Okay, Layer 5 is super sensitive to deletion, so let's keep it full size. But Layer 10 is fine with tiny shorthand, so let's shrink that. Layer 20 can lose half its notes without crying."
The Result: The backpack is compressed to 1/14th of its original size, but the student's performance is exactly the same as if they had the full, uncompressed backpack.

The Key Discovery: "The Trash Can is the Hero"

The paper did a cool experiment to see which part of the strategy mattered most. They tested three versions:

Uniform: Same rules for everyone. (Terrible results).
Smart Shorthand: Different shorthand sizes for different layers, but same trash rules. (Still bad).
Smart Trash: Same shorthand sizes, but different trash rules for different layers. (This was the winner!)

The Analogy:
Imagine you are packing for a trip.

Uniform: You throw away 50% of your clothes and shrink the rest. You freeze and can't dress properly.
Smart Shorthand: You keep all your clothes but shrink the fabric. You still have too many clothes; the bag is too heavy.
MoE-nD (Smart Trash): You realize your "socks" (some layers) are useless and throw 90% of them away, but you keep your "jacket" (critical layers) perfectly intact. The bag becomes light, but you are still warm and ready.

The Results: Magic on Long Stories

When they tested this on hard benchmarks (like reading a 16,000-word story or solving complex math):

Old methods: When they tried to shrink the memory, the model got confused and failed (accuracy dropped to near zero).
MoE-nD: It shrank the memory by 14 times (from 1.9 GB down to 136 MB) and the model performed just as well as the giant, uncompressed version.

When Does It Not Work?

The paper is honest about its limits. It doesn't help with short tasks (like a simple 500-word math problem).

Why? If the story is short, the backpack is already small. The "Smart Manager" looks at the layers and says, "Hey, we have plenty of room! Let's keep everything."
In this case, the system just defaults to "keep everything," and the fancy routing doesn't add any value. It's like using a GPS to navigate a room you can see clearly; you don't need the complex directions.

Summary

MoE-nD is a technique that stops treating AI brains like a single block of clay. Instead, it treats every layer as a unique individual, giving some layers a "delete" button, others a "shrink" button, and others a "keep safe" button. By customizing the compression for every single layer, it allows AI to remember long stories without needing a supercomputer's memory.

1. Problem Statement

Large Language Models (LLMs) processing long contexts face a critical bottleneck: the KV (Key-Value) cache consumes tens of gigabytes of memory per sequence, leading to latency issues and out-of-memory (OOM) failures on commodity hardware.

Existing compression strategies typically operate on a single axis (e.g., token eviction, quantization, or low-rank projection) and apply a uniform policy across all transformer layers. The authors argue this "homogeneity" is suboptimal because:

Different layers respond differently to compression operations (e.g., some layers tolerate aggressive token eviction, while others are sensitive to quantization).
The optimal trade-off between eviction (reducing sequence length) and quantization (reducing precision) varies significantly from layer to layer.
Current methods fail to jointly optimize these axes on a per-layer basis under a global memory budget.

2. Methodology: MoE-nD

The authors propose MoE-nD, a framework that treats different compression configurations as "experts" and routes each layer to its optimal configuration.

Core Components

Multi-Axis Design Space: For each layer $\ell$ $ℓ$ , the system selects a tuple $(keep\_ratio, k\_bits, v\_bits)$ $(k ee p_r a t i o, k_bi t s, v_bi t s)$ :
- Eviction: A keep ratio (fraction of tokens retained) ranging from 10% to 100%.
- Quantization: Bit-widths for Key (K) and Value (V) tensors (e.g., 16, 8, or 4 bits).
Offline Sensitivity Calibration:
- The system measures the relative L2 error between the full-precision attention output and the compressed output for every layer under every candidate configuration.
- This creates a sensitivity table $S_{\ell, c}$ where $\ell$ is the layer and $c$ is the configuration.
- Key Finding: Aggressive eviction (75–90%) causes error variance of 500–700× across layers, whereas V-quantization variance is minimal (~1.2×). This confirms that per-layer routing is essential for eviction but less critical for V-quantization.
Greedy Budget Solver:
- Finding the globally optimal configuration is an intractable discrete optimization problem ( $\approx 10^{42}$ candidates for a 28-layer model).
- MoE-nD uses a greedy solver that iteratively upgrades the layer offering the highest marginal reduction in sensitivity per unit of memory saved ( $\Delta S / \Delta m$ ) until the global memory budget $M$ is exhausted.
- This process takes <50ms on a CPU.
Heterogeneous Attention Patch:
- To handle layers with different cache lengths and bit-widths during inference, the authors implement a custom attention kernel.
- It uses per-layer position arrays to track token indices and applies RoPE (Rotary Positional Embedding) inversion based on the specific layer's retained tokens, ensuring correct attention calculation despite heterogeneous cache states.

3. Key Contributions

Joint Multi-Axis Routing: First work to jointly route eviction ratios and K/V quantization bits on a per-layer basis under a single global memory budget.
Empirical Discovery of Layer Heterogeneity: Demonstrated that the trade-off between eviction and quantization is not global but layer-specific. For example, Layer 0 might prefer eviction (low error cost), while Layer 6 prefers quantization (low error cost) for the same memory savings.
Ablation of Routing Axes: Through an ablation chain ( $2d_{uniform} \to 2d \to 2d_{hetero}$ ), the authors prove that per-layer eviction routing is the primary driver of performance gains, while per-layer quantization routing alone offers negligible improvement.
Defined Scope of Applicability: Identified that the method yields significant gains only when input length and budget tightness force non-trivial routing decisions (i.e., long-context or tight-budget scenarios). In short-context scenarios, the solver naturally defaults to uniform behavior (keep=1.0), rendering the heterogeneity moot.

4. Experimental Results

The method was evaluated on DeepSeek-R1-Distill-Qwen-7B across LongBench-v1 (16k tokens), AIME (math reasoning), and MATH-500.

LongBench-v1 (Long Context):
- MoE-nD ( $2d_{hetero}$ ) achieved 14× compression (136 MB vs. 1.9 GB full cache) with no detectable accuracy loss (12.0 vs. 11.5 baseline).
- Competing methods at similar memory budgets (single-axis eviction or uniform quantization) suffered massive accuracy drops (staying under 8/100).
- At matched memory (136 MB), MoE-nD outperformed the best non-hetero baseline ( $2d$ ) by 6.1 points.
AIME Reasoning Benchmarks:
- MoE-nD consistently outperformed all baselines across all budgets.
- At tight budgets (e.g., $b=64$ ), MoE-nD achieved 16.7–30% accuracy, while other methods collapsed to near 0%.
- Notably, at $b=256$ on AIME-25, MoE-nD surpassed the full-cache baseline (36.7 vs. 30.0), likely due to the removal of self-distracting intermediate tokens via aggressive eviction.
Null Results (Scope Validation):
- On MATH-500 and LongBench TREC (shorter inputs), MoE-nD showed no advantage over uniform baselines.
- Reason: The greedy solver determined that keeping 100% of tokens (keep=1.0) was optimal for most layers given the loose memory budget, causing the heterogeneous method to degenerate into a uniform one. This validates the method's theoretical boundaries.

5. Significance and Implications

Efficiency: Enables running long-context LLMs on consumer hardware (e.g., 136 MB KV cache vs. 1.9 GB) without sacrificing reasoning quality.
Paradigm Shift: Moves KV cache compression from "one-size-fits-all" to adaptive, layer-specific optimization.
Practical Insight: The paper clarifies that per-layer quantization routing is less critical than per-layer eviction routing, guiding future research toward optimizing token selection strategies.
Limitations: Currently requires a custom attention kernel (not yet FlashAttention compatible) and needs re-calibration for different model families (though the authors expect the heterogeneity pattern to hold across architectures).

In conclusion, MoE-nD demonstrates that intelligent, per-layer allocation of memory resources across multiple compression axes is the key to unlocking high-performance, long-context inference on resource-constrained devices.