Bound the caching allocator: +6.04 MiB/token leak -> flat, for 0.15 ms/token

[heydryft]

Stacked on aeonmindai/candle#7, which is where the allocator lives. This branch pins that rev and adds the counters.

The defect

perf/alloc-cache-reachable-on-decode (#177) turned candle's caching allocator on for decode. It is worth 7.80 ms/token and that is real. But the allocator had no capacity and no eviction — the only ways memory went back to the driver were set_alloc_cache_enabled(false) and drain_alloc_cache_and_free(), and this pipeline only calls those on a prefill. One long generation has one prefill, so nothing ever drained.

Measured on an H200, 2 600 tokens, memory.used at 2 Hz

Same binary for all three legs, card verified empty before each. Legs that started on a non-empty card were refused rather than reported.

	tok/s	VRAM over 2 600 tok	2nd-half slope	alloc/step	free/step	held	hit rate
cache off	17.12	+212 MiB	1.52 MiB/s	—	—	—	—
unbounded (before)	19.76	+15 700 MiB	110.58 MiB/s	3.4	0.0	15 672 MiB	0.9996
capped 1 GiB (after)	19.70	+2 004 MiB	2.40 MiB/s	43.3	173.0	937 MiB	0.9961

Confirmation leg on the exact shipping tree (arc 576e599, binary dc35a565…): 19.77 tok/s, free/step 173.0, held 937.1 MiB / cap 1024 MiB, server survived.

VRAM by decile, MiB above the start — unbounded is a straight line, capped fills once and plateaus:

unbounded    0  2772  4148  5492  6964  8340  9780 11220 12724 14132 15700
capped 1GiB  0  1716  1780  1780  1812  1844  1844  1876  1940  1940  2004
cache off    0     0     0    32    32    32    32    32    96    96   160

Bounding costs 0.15 ms/token — 50.76 vs 50.61 ms/token — against the 7.80 ms/token the cache buys over having it off. It keeps 98% of the gain.

Two things the report asked about that did not reproduce

The server does not OOM on the first completion. It survives at --max-seq-len 4096 on an empty card, before and after, in every leg. One load-time CUDA_ERROR_OUT_OF_MEMORY was seen (1 of 4 loads) and it is not the cache: the cache is only enabled from inside a decode forward, long after load. Auto-device-map plans 140 GB of a 143 GB card, so load sits on a knife edge. Separate defect.

The unboundedness is not what blocks CUDA-graph capture. Capture aborts identically with the cap on and off, and the allocator counters are byte-identical between the two legs — alloc/step 44.1, free/step 78.7, hit-rate 0.9954, held 315.4 MiB. The cap never engages at 315 MiB. The blocker is one line earlier:

[alloc-cache] WARNING: MISS during capture, size 393216 bytes (not pre-warmed) -> graph will be unstable
ARC capture: instantiate/launch failed: First cuGraphLaunch failed: 700   (CUDA_ERROR_ILLEGAL_ADDRESS)

An allocation size warmup never covered becomes an unstable graph memory node, the first launch reads an unbacked address, and the glibc abort is downstream of the poisoned context. Warmup coverage for 393 216 bytes is the fix, not the cap.

What is here

ARC_ALLOC_CACHE_MAX_MB sets the cap (unset = candle's 1 GiB; 0 = old unbounded, for A/B). A typo deliberately does not fall back to unbounded.

ARC_ALLOC_CACHE_STATS=N prints allocations per step, frees per step, hit rate, and bytes held against the cap, every N decode steps. Those are the numbers this has to be judged on — a green log is not evidence. An earlier arena here reported "accounting OK" and bit-identical output over 52 steps while silently bypassing itself for every buffer under 128 bytes (KERNEL_RULES.md:977-984); what separates a working bounded cache from that is allocations staying low and frees being non-zero, and both are printed.

14 tests (10 in candle for the bound, 4 here for the cap's parse), all host-runnable.

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     74        25734        18232         4703         2799
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     74         4600         4597            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  145        15139        12482          811         1846
 Shell                    40         9948         6637         2655          656
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1498         1294           54          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                204        45280            0        35143        10137
 |- BASH                  72         1655         1203          331          121
 |- C                      3           17           17            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  18          708          708            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  66         2051         1716           77          258
 |- TOML                   6          207          164            0           43
 |- YAML                   5           41           36            5            0
 (Total)                            51052         4688        35685        10679
─────────────────────────────────────────────────────────────────────────────────
 Rust                    673       330618       284971        16725        28922
 |- Markdown             491        29234          471        25204         3559
 (Total)                           359852       285442        41929        32481
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1324       494711       352182        90209        52320
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

CI fixed (3c4bb1d2d). Green and ready. Not merging solo — see the ordering note.

What I fixed. Only red check was Typos: unparseable -> unparsable at mistralrs-core/src/pipeline/normal.rs:200 (doc comment) and :2823 (test name). A real spelling fix, no behaviour change.

Ordering matters more than the fix. This PR bounds candle's caching allocator (+6.04 MiB/token leak -> flat). #177 turns that same allocator ON by default and its own docs concede the KV-length-tracking buffers grow O(context^2) and are 'not bounded by this policy'. So:

• Correct order: this one first, or both together.
• Landing perf(arckv): make candle's caching allocator reachable, on for decode, and BOUNDED #177 alone — default-ON, unbounded — is the worst of the three orderings, and CI will never catch it because CI does not run long contexts.

I have flagged #177 to invert its polarity to opt-in. If that happens, this PR becomes a straightforward improvement that can land on its own merits. If #177 lands as-is, this becomes urgent rather than nice-to-have.

Blast radius of this PR alone: bounds an allocator that is currently off by default, so it is close to inert today and becomes load-bearing the moment #177 lands. The 0.15 ms/token figure was measured on this branch's base — re-measure before quoting.

[heydryft]

Superseded by #177, which now carries these commits.

The ladder was rebased onto current master and this PR's bounded-allocator work was folded into #177 (the branch it built on), so the stack ships reachable AND bounded as one change rather than as a fast path plus a follow-up nobody merges. #177 is now based on #175.

One defect was fixed while folding it in: Cargo.toml here pinned candle 89ab14ef but Cargo.lock still resolved 88d86a2. cargo metadata --locked failed on this branch. A --locked CI build resolves the LOCK, so this lane would have compiled the OLD, UNBOUNDED allocator while reporting green on the PR whose entire subject is bounding it. The lock is regenerated in #177.

Candle ancestry, verified: 88d86a2 -> b2a4dd80 -> 859c49c8 -> 89ab14ef is a clean linear chain, so the bump is safe for master's pin. arcgraph/prewarm-profile (candle 1fa534d) is a SIBLING off 88d86a2, not an ancestor, and it conflicts with 89ab14ef in four hunks of candle-core/src/cuda_backend/device.rs — the AllocCache struct body, set_alloc_cache_enabled, drain_alloc_cache_and_free and cache_take. The substantive part is that prewarm assumes free's value type is a Vec, which bounded changes to a SizeClass, so three of its sites will not compile even after the markers are resolved. Reconcile, do not clobber — whoever merges second owns that merge.

[heydryft]

Two things from tonight's H200 legs that bear directly on this PR.

1. This PR's content is already on master — it is not waiting to land. 7d91c1127 (the normal.rs half) and 65dc0fb4a (the candle bump to 89ab14ef, this PR's target rev) are both ancestors of master; git merge-base --is-ancestor 7d91c1127 origin/master passes, and 37 of this PR's added symbols are in master's normal.rs. Worth closing as landed-by-another-route rather than rebasing the conflict.

2. ⚠️ The 6.04 MiB/token figure this PR's case rests on is contradicted, and in a direction context length cannot explain.

This PR records +6.04 MiB per decoded token for the unbounded cache (V4, 2 600 tokens, memory.used at 2 Hz), and that number is now quoted in mistralrs-core/src/pipeline/normal.rs. Tonight, V4 on an H200 at the default 1 GiB cap measured ~12.8 MiB/step — twice as bad with a cap as the figure recorded for no cap at all.

A cap that never binds is behaviourally identical to no cap, and this one never bound: the leg died at held 236.9 MiB / cap 1024 MiB with free/step 0.0 on every step after step 3. So the two numbers describe the same allocator behaviour. And context length cannot close the gap: the 6.04 leg had the longer context, which should have made its rate higher, not lower.

The control that is not in dispute, measured the same night on the same box: ARC_CANDLE_ALLOC_CACHE=0 ran to ntok=1000, finish=length at 0.19 MiB/token with held 0.0 MiB / sizes 0, against an OOM at 19 tokens with the cache on. Your central claim — that retention is the mechanism — is confirmed. It is the rate that does not reconcile.

I have annotated the source entry as disputed with both provenances rather than deleting it (#213). If you still have the 2 600-token leg's config, the discriminator is whatever differed between the two runs, because it is not the context length.

3. The bound itself is inert on the configuration Arc ships. 1 GiB is ~4x the 262 MiB V4 leaves free on that box, so evict_if_over_capacity has never executed once in a V4 serve. #213 derives the cap from measured free VRAM instead of a constant. That is a follow-on to this work, not a contradiction of it — the bound was the right idea denominated in the wrong units.

[heydryft]

Follow-up to my previous comment — the 6.04 MiB/token figure reconciles, and I was wrong to call it disputed. Correcting it here so the retraction sits next to the claim.

6.04 MiB/token x 2 600 tokens is 15.3 GiB retained, so your leg had at least that much free at decode start. Tonight's box had 844 MiB (262 MiB on the TCFRAG leg). ARC_QTIP_TCFRAG landed as #203 on 08-20 — after your measurement — and its repack retains ~63 GB against a ~79.5 GB model (ARC_QTIP_TCFRAG=0 frees 64 262 MiB vs 262 MiB; #209 has since made it opt-in).

So it is the same per-token retention in both runs, against two orders of magnitude difference in headroom. That is why a run survived 2 600 tokens and another died at 19, and it needs no second mechanism. Your number stands and your central claim — that retention is the mechanism — is confirmed by tonight's control leg (ARC_CANDLE_ALLOC_CACHE=0: 1 000 tokens, finish=length, 0.19 MiB/token, held 0.0 MiB / sizes 0, against an OOM at 19 tokens with the cache on).

One residual, stated rather than buried: 6.04 against ~12.8 MiB/step is still a 2x gap in the rate, unexplained. It is not load-bearing — what needed explaining was the survival difference — but neither number should be quoted without its configuration. Likeliest candidate is TCFRAG itself, since tonight's leg had it on and yours predates it, and it is directly testable: the base256 arm already runs ARC_QTIP_TCFRAG=0 with the cache on, so its held/step is the discriminator. If you still have your leg's config, that would close it.

The source annotation in #213 now reads as reconciled-with-a-caveat rather than disputed.

My other two points stand: this PR's content is already on master (7d91c1127 + 65dc0fb4a, both ancestors), so it can be closed as landed rather than rebased; and the 1 GiB bound is inert on the shipping configuration, which #213 fixes by deriving the cap from measured free VRAM.

[heydryft]

Closing as superseded by 538a955a7 (the merge of #177), verified rather than assumed.

Both of this PR's substantive commits are on master under that merge, as rebased equivalents with identical subjects:

this PR	on master	under 538a955a7?
7d8760f99 perf(arckv): make candle's caching allocator reachable and on for decode	dd5583c4b	git merge-base --is-ancestor dd5583c4b 538a955a7 → 0
b0cf25a73 perf(arckv): bound the caching allocator, and print the counters that prove it	7d91c1127	git merge-base --is-ancestor 7d91c1127 538a955a7 → 0

and git merge-base --is-ancestor 538a955a7 origin/master → 0.

The bounded-allocator surface this PR introduced is live on master today: alloc_cache_action, plan_alloc_cache_capacity, parse_alloc_cache_capacity, arm_alloc_cache_capacity, alloc_cache_capacity_bytes and report_alloc_cache_step all exist in mistralrs-core/src/pipeline/normal.rs.

The one thing left that was arguably unique here — the Cargo.toml candle bump to 89ab14ef1216331a782539f533b99f6816708604 ("no eviction inside the capture window") — is also already master's pinned rev, so nothing in this branch is unmerged.

Follow-on work has since built on top of it: #213 (e976a7753) replaced the fixed cap with one computed against real free VRAM, because a 1 GiB cap cannot bind a card with 262 MiB free.

No unmerged work is being discarded by this close.