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

TL;DR

MoE-nD introduces a per-layer mixture-of-experts routing framework for KV cache compression in long-context LLM inference, significantly improving compression efficiency while maintaining accuracy.

Contribution

It proposes a novel per-layer heterogeneous eviction and quantization method, optimizing compression by routing each layer to its best configuration under a global memory budget.

Findings

Achieves 14x compression with no accuracy loss on LongBench-v1 tasks.

Outperforms other compression baselines by a large margin in memory efficiency.

Improves reasoning benchmark scores by 6 to 27 points over uniform quantization methods.

Abstract

KV cache memory is the dominant bottleneck for long-context LLM inference. Existing compression methods each act on a single axis of the four-dimensional KV tensor -- token eviction (sequence), quantization (precision), low-rank projection (head dimension), or cross-layer sharing -- but apply the same recipe to every layer. We show that this homogeneity leaves accuracy on the table: different layers respond very differently to each compression operation, and the optimal per-layer mix of eviction and quantization is far from uniform. We propose MoE-nD, a mixture-of-experts framework that routes each layer to its own (eviction-ratio, K-bits, V-bits) tuple under a global memory budget. An offline-calibrated greedy solver chooses the routing that minimizes predicted quality loss; at inference time, per-layer heterogeneous eviction and quantization are applied jointly through a single attention patch. On a 4-task subset of LongBench-v1 (16k inputs, n=50 per task, adapted reasoning-model protocol; see section Experiments), MoE-nD's hetero variant matches our uncompressed 1.9~GB baseline at 14x compression (136~MB) while every other compressed baseline we tested (1d, 2d_uniform, 2d) at comparable or smaller memory stays under 8/100. The gains hold on AIME reasoning benchmarks (+6 to +27 pts over the strongest per-layer-quantization baseline across eight configurations). Two null results -- MATH-500 and LongBench's TREC -- share a principled cause (short inputs, solver picks keep=1.0 on most layers), cleanly characterizing when per-layer eviction routing has headroom to help.