JustFit: Just-in-Time State Management for Local LLM Serving
Abstract
Local agents need memory for model execution and working history. We present JustFit, an MLX runtime that coordinates their overlapping allocations: KVExec executes and checkpoints four-bit KV with bounded workspace, PhaseSwap loads phase-dependent components, and StateTrans preserves history across execution modes. With Qwen3.8-27B MXFP4 on a 24 GiB M4 Pro MacBook, JustFit reaches 327,680 retained positions across two requests, 10.67 times the evaluated baseline's 30,720-position single-request record. One request completes the full 262,144-position native window at median 5.986 tokens/s. Each shape completes 16,384 outputs per request in three fresh processes: B1 uses one cold build and two prefix extensions; B2 uses three ordered prefix extensions (Section 4). Image encoding can proceed while preserving a live 196,608-input text request. A controlled, repetitive 32K+6K workload reaches median 18.284 tokens/s at 15,626 MiB; a separate AIME 2026 evaluation scores 29/30. Coordinating execution and state lifetimes makes longer histories feasible on personal hardware.
Figures & tables
| Workload | Complete | TG (tok/s) | Peak (MiB) |
|---|---|---|---|
| B1: 256K | 3/3 | 5.986 | 20,357–20,450 |
| B2: 288K | 1/1 | 10.341 | 20,858 |
| B2: 320K | 3/3 | 9.735 | 20,432–20,953 |
| Text input | Image E2E (s) | Peak (MiB) | Guard (MiB) |
|---|---|---|---|
| 98,304 | 16.419 | 20,346 | 20,700 |
| 131,072 | 18.094 | 20,335 | 20,700 |
| 163,840 | 19.873 | 20,635 | 21,000 |
| 196,608 | 21.636 | 20,852 | 21,000 |
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
| Parameter | Value or procedure |
|---|---|
| Quantizer | MSE scalar quantization after normalization and randomized Hadamard rotation; K/V use four bits each. |
| Per-vector storage | 32 U32 words (eight indices per word) and one FP16 norm computed from an FP32 norm. |
| Codebook | 16 FP32 centroids for density proportional to . |
| Construction | 32,768-point grid on ; quantile initialization; at most 100 centroid updates; stop when maximum change is below . |
| Random signs | NumPy default_rng(seed + d*7919) , sampling ; K/V default seeds are 0/1. |
| Quantization | Midpoint thresholds between ordered centroids, with strict comparisons. |
| EVO | Added mechanism | Capacity evidence | Peak (MiB) | PP | TG |
|---|---|---|---|---|---|
| (tok/s) | |||||
| 0 | mlx-vlm baseline | 30K (B1) | — | — | |
| 1 | TQ4 + MTP | 38K | 18,225 | 110.82 | 5.975 |
| 2 | Fused dequant | 80K | 18,220 | 115.86 | 6.541 |
| 3 | Segmented KV | 128K | 17,806 | 115.89 | 13.291 |
| 4 | Paged KV + attention | 128K (carried) | 17,044 | 110.57 | 11.446 |
| EVO | Peak (MiB) | PP (tok/s) | TG (tok/s) |
|---|---|---|---|
| 1 | 18,221–18,227 | 109.912–110.840 | 5.973–5.980 |
| 2 | 18,099–18,220 | 115.858–115.862 | 6.539–6.542 |
| 3 | 17,795–17,806 | 115.878–115.889 | 13.288–13.294 |
| 4 | 17,040–17,051 | 110.559–110.578 | 11.441–11.446 |
| 5 | 16,545–16,559 | 110.549–110.560 | 11.445–11.445 |
| 6 | 16,365–16,403 | 108.850–108.877 | 11.539–11.545 |
| Run | PP (tok/s) | TG (tok/s) | Wall (s) | Wall TPS | Peak (MiB) | Margin (MiB) |
|---|---|---|---|---|---|---|
| 1 | 68.321 | 7.0149 | 5,213.943 | 3.1423 | 20,862 | 138 |
| 2 | 68.264 | 7.0158 | 5,216.103 | 3.1410 | 20,857 | 143 |
| 3 | 68.287 | 7.0177 | 5,214.499 | 3.1420 | 20,959 | 41 |
| Workload / arrival | Positions | PP | TG | Peak (MiB) |
|---|---|---|---|---|
| B2 (96K+12K), simultaneous | 221,184 | 90.51 | 12.72 | 20,195 |
| B2 (96K+12K), staggered | 221,184 | 90.52 | 12.81 | 20,156 |
| B2 (96K+16K), simultaneous | 229,376 | 90.44 | 12.55 | 20,310 |
| B2 (96K+16K), staggered | 229,376 | 90.44 | 12.51 | 20,247 |
| B4 96K+12K; (8K+4K) | 147,456 | 91.25 | 23.41 | 19,178 |
| B4 128K+12K; (8K+4K) | 180,224 | 83.35 | 21.54 | 20,236 |
| Quantity | Observation |
|---|---|
| Total uncached input / completed output tokens | 155,648 / 49,152 |
| Prefill work / PP | 1,821.472 s / 85.452 tok/s |
| Common-active outputs / interval | 42,945 / 1,902.766 s |
| Common-active aggregate TG | 22.570 tok/s |
| First-request start to last completion / wall rate | 4,046.521 s / 12.147 tok/s |
| Sampled peak / guard margin | 18,915 / 2,085 MiB |
| Input | Runtime | Peak (MiB) | PP (tok/s) | TG (tok/s) |
|---|---|---|---|---|
| 8K | mlx-vlm | 17,180 | 129.65 | 15.63 |
| 8K | JustFit | 15,746 | 127.27 | 24.31 |
| 16K | mlx-vlm | 18,449 | 125.23 | 14.97 |
| 24K | mlx-vlm | 19,898 | 121.12 | 14.97 |
| 32K | mlx-vlm | — | — | |
| 32K | JustFit | 15,899 | 114.76 | 19.11 |
| Workload | PP (tok/s) | TG (tok/s) | Peak (MiB) |
|---|---|---|---|
| B1 24K+512 | |||
| B4 (2K+512) | |||
| B2 2K each, 512/128 output |
| PP (tok/s) | TG (tok/s) | Peak (MiB) | ||||
|---|---|---|---|---|---|---|
| Workload | M3 | M4 | M3 | M4 | M3 | M4 |
| B1 32K+6K ( ) | 87.109 | 114.373 | 12.603 | 18.284 | 15,650 | 15,626 |
| B1 cold 240K+16K | 44.403 | 60.546 | 4.944 | 5.936 | 20,905 | 20,357 |
| B2 cold (128K+16K) | 57.262 | 81.399 | 7.101 | 10.341 | 20,749 | 20,858 |
| B4 staggered, four lanes | 62.888 | 85.452 | 16.800 | 22.570 | 19,285 | 18,915 |
| Profile | AR TG | MTP TG | Change (%) | AR peak (MiB) | MTP peak (MiB) |
|---|---|---|---|---|---|
| B2, 2K+512 | 30.455 | 29.926 | 14,946 | 15,543 | |
| B3, 2K+512 | 38.170 | 31.156 | 15,276 | 16,212 | |
| B4, 2K+512 | 39.980 | 31.822 | 16,041 | 17,269 | |
| B2, 32K+2K | 21.322 | 20.916 | 16,505 | 17,101 | |
| B3, 32K+2K | 25.405 | 22.315 | 17,284 | 18,355 | |
| B4, 32K+2K | 26.138 | 22.984 | 18,017 | 19,538 |
| Run | Preparation | Cached | New input | PP | TG | Peak (MiB) | Margin |
|---|---|---|---|---|---|---|---|
| (tok/s) | (MiB) | ||||||
| 1 | Cold | 0 | 245,760 | 60.546 | 5.936 | 20,357 | 643 |
| 2 | Checkpoint | 161,792 | 83,968 | 44.331 | 5.987 | 20,449 | 551 |
| 3 | Checkpoint | 161,792 | 83,968 | 44.315 | 5.986 | 20,450 | 550 |
| Run | Cached A / B | New A / B | TG (tok/s) | Peak (MiB) | Margin (MiB) |
|---|---|---|---|---|---|
| 1 | 153,600 / 131,071 | 10,240 / 1 | 9.732 | 20,953 | 47 |
| 2 | 163,839 / 131,071 | 1 / 1 | 9.738 | 20,456 | 544 |
| 3 | 163,839 / 131,071 | 1 / 1 | 9.735 | 20,432 | 568 |
| History | All-head (ms) | Two-head (ms) | All-head peak (MiB) | Two-head peak (MiB) |
|---|---|---|---|---|
| 32K | 44.752 | 42.180 | 526.1 | 268.6 |
| 96K | 132.024 | 136.474 | 1,566.4 | 796.9 |
| 192K | 271.220 | 280.897 | 3,123.8 | 1,589.3 |
| Prefix | Old save (s) | Single-pass (s) | Old application bytes | Single-pass bytes |
|---|---|---|---|---|
| 16K | 0.0872 | 0.0479 | 545,270,638 | 272,635,224 |
| 32K | 0.1550 | 0.0869 | 1,090,530,364 | 545,265,080 |
| 224K | 2.9759 | 1.5073 | 7,633,644,836 | 3,816,822,312 |
| Text input | Cached | New input | Gate | Guard | Peak | Margin | Image E2E |
|---|---|---|---|---|---|---|---|
| (positions) | (MiB) | (MiB) | (MiB) | (s) | |||
| 98,304 | 98,303 | 1 | 143,360 | 20,700 | 20,346 | 354 | 16.419 |
| 131,072 | 131,071 | 1 | 143,360 | 20,700 | 20,335 | 365 | 18.094 |
| 163,840 | 163,839 | 1 | 204,800 | 21,000 | 20,635 | 365 | 19.873 |
| 196,608 | 161,792 | 34,816 | 204,800 | 21,000 | 20,852 | 148 | 21.636 |
| 200,704 | 194,560 | 6,144 | 204,800 | 21,000 | 20,980 | 20 | 12.994 |
| Text input | TTFT (s) | Maximum stream gap (s) |
|---|---|---|
| 98,304 | 2.385 | 1.735 |
| 131,072 | 2.313 | 1.109 |
| 163,840 | 2.685 | 1.130 |
| 196,608 | 723.711 | 1.177 |
| 200,704 | 140.741 | 1.187 |