perf(arckv): make candle's caching allocator reachable, on for decode, and BOUNDED

[heydryft]

Errors in the brief this had to correct first

Two premises I was given did not survive checking, and both change what "just turn it on" means:

1. "with the correctness tests it already has" — there are none. The candle fork (aeonmindai/candle rev 88d86a2, which is what Cargo.lock pins) has zero tests for the allocator: cuda_backend/device.rs and cuda_backend/mod.rs are the only files that mention it and neither carries a #[cfg(test)]. Every arc-side exercise is #[cfg(feature = "cuda")] and unrunnable in CI. The only executable one is the capture_probe example, which has no assertions — it prints a max-abs-diff for a human to read.
2. The measured pairs are not in the tree. host calls 25,622 -> 737, layout H2D 2,361.6 -> 6.0, driver frees exactly 0, bit-identical 52/52 — none of the "after" numbers exist on any of the 602 refs. 25,622 appears once, as a pre-arena baseline in COMPETITIVE_TEARDOWN.md:201-215. The nearest 52 is KERNEL_RULES.md:977-984, and it is a cautionary record, not a success: an arena that reported accounting OK, 4,976 cache hits/step and bit-identical output over 52 steps while silently bypassing itself for every buffer under 128 bytes. The only tell was a driver_frees that was supposed to be impossible.

So the arena is not a measured, tested thing waiting for a default flip. Consequently this PR flips what is safe to flip and says exactly what a GPU run must falsify.

The defect that is real

if probe && seq_len == 1 && std::env::var_os("ARC_CANDLE_ALLOC_CACHE").is_some()

The first conjunct tied a general-purpose caching allocator to the V4 capture probe. Recycling a decode step's frees has nothing to do with capture; it is worth having on its own. Three stacked default-off gates meant the ~11k allocations per token were a disabled feature, not a missing one.

What lands

A pure policy, alloc_cache_action(seq_len, enabled, killed):

forward	action
decode (seq_len == 1)	Enable
prefill (seq_len != 1)	DrainAndDisable
already in that state	Leave

ARC_CANDLE_ALLOC_CACHE=0 is the kill switch. Any other value, and unset, leave the policy in force — so the =1 the ops scripts pass (arcgraph_heap_probe.sh:191) still means what it always meant.

The capture probe keeps its own gate for the graph-mode positions. Only the allocator moved out from under it.

Why prefill must drain rather than coast

Not incidental, and it is the part I would have got wrong without reading the fork. The cache is keyed on exact byte count — free: HashMap<usize, Vec<CUdeviceptr>>, looked up with free.get_mut(&bytes). No bucketing, no smallest-fit, no splitting. It also has no capacity bound and no eviction. Leaving it on across a prefill parks that prefill's large one-shot buffers for the process lifetime under a key nothing will ever request again — which is the OOM-on-a-tight-model failure ARC_NO_DEDICATED_DECODE already exists to work around.

Why decode is the right place for it

A decode step allocates the same shapes 43 times over, once per layer, and the shapes that do not depend on KV length — hidden states, MLP and expert intermediates — repeat step after step. Exact-size keying is a good fit for exactly that traffic and a bad fit for everything else.

The number a GPU run should falsify

Predicted, not measured — I have no GPU. Buffers whose size tracks KV length (the causal mask most obviously) change size every decode step, so each step files one more never-reused entry. That is O(context^2) bytes worst case and this policy does not bound it: a [1,1,1,kv] BF16 mask summed over 4k tokens is ~16 MB, and over 128k it is not 16 MB.

The fixed-capacity graph-mode path (deepseek4.rs:4344-4353, which swaps the growing causal mask for graph_mode_length_mask at cfg_full.sliding_window) removes that growth entirely — but it is reached only under ARC_V4_CAPTURE_PROBE. So the allocator's steady-state safety is downstream of the shape-invariance work another agent owns, and the measurement that settles it is the long-context high-water mark. Until then, ARC_CANDLE_ALLOC_CACHE=0.

Allocations removed per token

Unmeasured. I will not put a number here. The mechanism removes a cuMemFreeAsync/cuMemAllocAsync pair for every decode allocation whose exact byte size recurs, against a baseline of 11,436 allocations per token. What fraction of those recur is precisely what has never been measured on a build where the cache was actually on — which is the point of this PR.

Evidence

Six host-runnable tests for the policy — the first tests this subsystem has anywhere. cargo test -p mistralrs-core --lib alloc_cache_policy: 6 passed; 0 failed.

🤖 Generated with Claude Code

https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H

[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        45330            0        35183        10147
 |- 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)                            51102         4688        35725        10689
─────────────────────────────────────────────────────────────────────────────────
 Rust                    673       331228       285444        16819        28965
 |- Markdown             491        29538          471        25464         3603
 (Total)                           360766       285915        42283        32568
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1324       495675       352655        90603        52417
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

MEASURED on an H200 — the win is real, and so is the leak this PR predicted

Box: H200 143,771 MiB, exclusive. Binary 2df0eeb2ae59… = this PR's commit
cherry-picked onto graph @ 0875caa (master + the fused E4M3 kernel), so the
baseline below is the same tree minus this one commit.

Instrument: untraced bench for tok/s; nsys --trace=cuda for the counts, with
the step count pinned from the trace (the 25 s window holds 346 / 462 steps, so
neither leg is the "window caught no decode" void).

The counts — this is what the PR asked to be asserted

per decode step	before	after
cuMemAllocAsync	10,114.68 (7.450 ms host)	131.44 (0.384 ms host)
cuMemFreeAsync	10,114.69 (7.428 ms host)	0.00
both, host time	14.878 ms/step	0.384 ms/step
launches	7,935.9	7,924.1
kernel time	31.266 ms	31.477 ms
GPU-busy (kernel union)	43.3 %	58.2 %
blocking D2H	1.00	1.00

cuMemFreeAsync reaching exactly 0.00 is the load-bearing number: every free
is now recycled rather than returned to the driver. Kernel time is unchanged
(31.27 → 31.48 ms), which is the control — the GPU work is identical, only host
overhead moved.

Throughput

	tok/s	ms/token
baseline (0875caa)	19.3	51.85
this PR	20.8	48.03
this PR + ARC_CANDLE_ALLOC_CACHE=0	19.2	52.06

The kill-switch leg is the same binary with one env var, and it lands back on
baseline. So the +7.4 % is attributable to the allocator and to nothing else
in the rebuild.

25,622 → 737 does not reproduce and I did not try to make it — these are
freshly measured, and the "before" is measured too.

The number this PR asked a GPU run to falsify — it does not survive

The PR's own text: "each step files one more never-reused entry … that is
O(context^2) bytes worst case and this policy does not bound it." Measured,
gen-len 2500, memory.used sampled at 2 Hz, same binary both legs:

cache OFF   t=70s 79,043 MiB ─ 79,043 ─ 79,043 ─ 79,075 ─ 79,107 ─ t=196s 79,139
            FLAT: +96 MiB over 128 s of decode

cache ON    t=84s 81,475 MiB ─ 83,011 ─ 84,579 ─ 86,147 ─ 89,315 ─ t=196s 101,493
            LINEAR, no plateau: +20,018 MiB over 112 s = +178.7 MiB/s

At 20.8 tok/s that is ~8.6 MiB per decoded token, unbounded. From the 101,493
MiB reached at 2,500 tokens there is 42,278 MiB of headroom left, so a single
sequence OOMs at roughly 7,400 tokens on a 143 GB H200.

The PR estimated ~16 MB at 4k tokens from the causal mask alone. The measurement
is ~21 GB at 2.5k, so the mask is not the story — many buffers are KV-length-keyed
and every one of them files a fresh never-reused byte-size key each step.

Note also that DrainAndDisable cannot rescue this in a decode-only workload:
the drain fires on prefill, and a long generation has exactly one of those, at
the start. In a mixed serving loop each prefill would drain — bounding the growth
but also discarding the decode buffers that produce the win.

Verdict

The mechanism is right and the 14.5 ms/step is real. Do not merge with the cache
on by default — as measured it converts a 7.4 % speedup into an OOM at ~7.4k
tokens. Either bound the cache (cap + LRU eviction, or bucket the key instead of
exact-byte), or gate the default on the shape-invariant decode path that stops
KV-length-keyed allocations, and keep the unbounded form behind an opt-in until
then.

Method note, per KERNEL_RULES.md:977-984: the count falling was asserted, and
then the resource was checked anyway. The green count is what this PR wanted; the
memory curve is why "the count fell" is not by itself a merge signal.

[heydryft]

Not merging yet — one blocking change needed. CI is fully green (18/18); this is not a CI objection.

🔴 This is the single largest blast radius in the current queue, and it is default-ON.

mistralrs-core/src/pipeline/normal.rs:183:

std::env::var("ARC_CANDLE_ALLOC_CACHE").is_ok_and(|v| v == "0")

That is a disable predicate. Unset ⇒ killed = false ⇒ want = seq_len == 1 ⇒ AllocCacheAction::Enable ⇒ cd.set_alloc_cache_enabled(true) on every decode step, with no flag set.

The previous condition was probe && seq_len == 1 && env::var_os(...).is_some() — three terms, all default-off. This PR drops probe and inverts the env var from opt-in to opt-out. That is two independent moves toward on-by-default in one change.

Why that matters here specifically: candle's caching allocator is exact-byte-keyed with no eviction and no capacity bound. This PR's own documentation concedes that the KV-length-tracking buffers (the causal mask) grow O(context^2) and are 'not bounded by this policy'. Turning an unbounded cache on by default for every decode is how a long-running server OOMs in production rather than in CI — and CI will never see it, because CI does not run long contexts.

Ask: invert the polarity — ARC_CANDLE_ALLOC_CACHE=1 to opt in, unset means off. Keep the seq_len == 1 narrowing. Then it lands.

Also note #182 (perf/alloc-cache-bounded, 'bound the caching allocator: +6.04 MiB/token leak -> flat') exists and is exactly the missing bound for this. These two should land together, bound first, or this one lands opt-in and #182 relaxes it later. Landing this one alone, default-ON, unbounded, is the worst of the three orderings.

Nothing else here is objectionable: no .cu touched, no format/ABI change, no dependency movement, no PagedAttention or prefix-cache contact.

[heydryft]

Rebased onto current master, re-based on #175, and #182 folded in so this ships reachable and bounded as one change rather than a fast path plus a follow-up nobody merges.

The polarity was already right (ARC_CANDLE_ALLOC_CACHE=0 is a kill switch). The defect was that the cache was unbounded with no eviction: +6.04 MiB/token, forever. It is now bounded, for 0.15 ms/token.

Stale-lock defect fixed. #182 bumped all five candle crates in Cargo.toml to 89ab14ef but left Cargo.lock resolving 88d86a2; cargo metadata --locked failed on that branch. Since a --locked build resolves the LOCK, the lane would have compiled the OLD, UNBOUNDED allocator while reporting green on the PR whose entire subject is bounding it. Regenerated with cargo update -p candle-{core,nn,flash-attn,flash-attn-v3,metal-kernels}; the only non-candle churn is a windows-core 0.61.2/0.62.2 dedup that fell out of re-resolution and that nothing arc builds on Linux or macOS reads.

Candle ancestry verified before bumping: 88d86a2 -> b2a4dd80 -> 859c49c8 -> 89ab14ef is a clean linear chain. 89ab14ef is also the rev that stops eviction inside the capture window — freeing a private-pool pointer after cuMemPoolDestroy corrupts the glibc arena, which is the corrupted size vs. prev_size blocker on #181 — so it is the safer pin, not merely a newer one.

Reconciliation warning for whoever touches candle next. arcgraph/prewarm-profile (candle 1fa534d) is a SIBLING off 88d86a2, not an ancestor. 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 half is that prewarm assumes free's value type is a Vec while bounded changes it to a SizeClass, so three of its sites will not compile even after the markers are resolved, and alloc_cache_free_counts / prewarm_alloc_cache need the same adaptation. Reconcile — do not clobber.

New env vars, all fail-safe (a typo degrades to bounded, never to unbounded): ARC_ALLOC_CACHE_MAX_MB (unset -> candle's 1 GiB default; 0 -> explicitly unbounded), ARC_ALLOC_CACHE_STATS (off unless set; logs every N decode steps). The counter that proves the bound is real is free/step — the unbounded cache's was exactly zero, forever.