Kurate.org
Rankings
Back to Rankings

MOSAIC: Multi-Objective Search for Allocating Inference Cache Across LLM Decoder Layers

Sravanth Kodavanti, Manjunath Arveti

Oct 18, 2026arXiv:2610.25431v1
cs.LGcs.AI
Share
Scorecard· 16/16
5.0/10 impact

Thorough and honestly reported extension of EvolKV-style layer budget search, but incremental, with test-set selection inflating headline gains and per-task search costs limiting practical adoption

Abstract

The Key-Value (KV) cache is the main memory bottleneck of long-context LLM inference. Token eviction methods such as SnapKV, H2O, AdaKV and L2Norm shrink it, but they give every transformer layer the same budget, although layers differ in which tokens they need. We present MOSAIC, a multi-objective neural architecture search that learns per-layer KV-cache budgets for any eviction method. One search maps the trade-off between average budget and task score; its best allocation transfers to every target budget without further search, and an optional budget-constrained search refines it. Across four eviction methods on Llama-3-8B-Instruct and Mistral-7B-Instruct-v0.2, MOSAIC beats uniform allocation in 146 of 154 (task, method, budget) settings, by up to +48 points on RULER and +11.9 on LongBench (LCC). On both models, it reaches the score of uniform allocation with up to 4x less KV cache (75% smaller). Against EvolKV, MOSAIC explores a ~10^40 x larger search space with 4x fewer evaluations and 18% fewer GPU-hours per budget, and scores higher in 5 of 6 settings. The learned allocations consistently favour a small set of middle layers.

AI Impact Assessments

(1 model)

Core Contribution

MOSAIC treats per-layer KV-cache budgets as a black-box search problem. It runs a multi-objective Bayesian optimization (classifier-guided, with an MLP surrogate and differential evolution) over all 32 layer budgets jointly. The search trades average budget against task score. A winner is selected from the Pareto front on full data and rescaled to any target budget. An optional per-budget constrained search then refines it. The method is eviction-agnostic and was tested on SnapKV, H2O, AdaKV and L2Norm. The framing is an offline calibration, "static quantization for KV allocation". The closest prior work is EvolKV, which uses CMA-ES on groups of 8 layers at a single budget. MOSAIC's novelty is the joint search, the budget–score Pareto exploration, and the reusable anchor. It is an incremental but sensible step beyond EvolKV, PyramidKV and CAKE. The layer-allocation idea itself is not new.

Methodological Rigor

The paper is unusually candid, which helps credibility, but there are serious design issues.

  • Selection on test data. The authors state that "LongBench and RULER have no validation split, so configurations are selected and scored on the same examples: our results show gains on the evaluated examples, not generalization." Stage 2 winner selection and Stage 4 candidate selection use the full evaluation data. The headline "146 of 154 wins" and "up to +48 points" are therefore optimistically biased. Appendix B quantifies this. Calibration-to-full gaps reach +5 points, and calibration picks the full-data best in only a minority of LongBench cells.
  • Biased "final" configuration. MOSAIC's reported number is the better of Winner and NAS-refined, chosen on the test set. Uniform gets no equivalent selection. The Winner-only numbers (75/87 wins on Llama, 51/75 on Mistral) are a fairer measure and noticeably weaker.
  • Weak transfer evidence. Shapes do not transfer across tasks (they fall below uniform), so each task needs its own search on task-specific data. This undercuts the "one search serves every budget" claim as a practical deployment story, because it still needs labeled data per task.
  • Inconsistent protocol. Floors (16 vs 64), calibration fractions, grids (5 vs 7 levels) and stopping rules (manual, or "when improvement stalls") vary across runs. Some Mistral Summarization cells are missing.
  • EvolKV comparison. The EvolKV comparison covers only SnapKV on two tasks, with a reproduced baseline. Gains on Single-Doc QA are within noise (+0.01 to +0.06, with a loss in one cell). The "10^40× larger search space" claim is somewhat rhetorical, since search-space size is not a measure of efficiency. The compute comparison (43.8 vs 53.7 GPU-hours) excludes the one-time Stages 1–2, and the full pipeline costs more than EvolKV at three budgets.
  • No seeds or error bars. There are no variance estimates and no random-search baseline for the full pipeline. The 57 LHS points in Stage 4 act only as a control.
  • Gain magnitudes. Gains on LongBench are mostly small (0.1–2 points) apart from L2Norm. The big RULER gains (+20–48 points) come mostly from weak baselines (H2O and L2Norm uniform). They concentrate on needle tasks, where allocating more cache to a few layers rescues retrieval. This is plausible, but it is selected and scored on the same 4K RULER examples.

    • Potential Impact

    Practitioners deploying KV eviction could use offline-calibrated per-layer budgets, which add no inference overhead. The finding that a few middle layers consistently need more cache is useful and consistent with prior layer-importance work, but it is not deeply analyzed. There is no mechanistic explanation and no end-to-end latency or memory measurement; only score vs. nominal budget is reported. Evaluation covers two 7–8B models and 4K RULER contexts. Modern long-context settings (128K, GQA-heavy newer models, quantization combined with eviction) are not tested. The search cost of 9–40 GPU-hours per task and method is non-trivial, and the per-task requirement limits adoption. The artifact (a search recipe) is more likely to be used as a baseline or tuning tool than as a widely adopted primitive.

    Timeliness & Relevance

    KV-cache compression is a very active area, and per-layer or head allocation is a current topic (Ada-KV, CAKE, PyramidKV, EvolKV). The paper is relevant, but the space is crowded and moves fast. The NAS-style framing is a modest differentiator.

    Strengths

  • Broad evaluation grid: 4 methods, 2 models, 2 benchmarks, many budgets, with full per-cell tables in the appendices.
  • Honest reporting of losses, the test-set selection caveat, calibration noise analysis, and floor ablation.
  • Method-agnostic design.
  • Interesting observations on shape transfer across eviction methods but not across tasks.
  • Transparent search-cost accounting and a head-to-head against the closest prior work.

    • Limitations

  • Selection on test data inflates the headline numbers, and there is no held-out generalization.
  • Small gains on most LongBench cells.
  • Task-specific searches are required.
  • No code release is mentioned, and manual stopping hurts reproducibility.
  • Only two older 7–8B models, with short contexts and no systems-level measurements.
  • Incremental over EvolKV, with a modest, non-uniform advantage.
  • Inconsistent experimental settings across cells.

    • Overall

    A competent, thorough, carefully reported engineering-style contribution to a crowded subfield. Its impact is limited by the incremental novelty, the same-data selection, and the per-task search requirement. Expect moderate citation as a layer-allocation baseline.

    Rating:5/ 10
    Significance 4.5Rigor 4.5Novelty 4.5Clarity 6.5

    Generated Oct 20, 2026

    Comparison History (0)

    No comparisons yet.