CompilerCache::find is O(n) over per-constant entries — degrades decode linearly when compiled fns receive changing scalar args

[aidiffuser]

Summary

mx.compile keys its cache on non-array (scalar) arguments. Each distinct
scalar value appends a new entry to the compiled function's entry vector, and
detail::CompilerCache::find locates entries by a linear scan with memcmp
over the constants blob. When user code passes a per-step-changing scalar
(the classic case: a KV-cache offset during LLM decode), per-call lookup cost
grows linearly with the number of distinct values seen — and the entries are
retained forever, so host memory grows without bound too.

Real-world impact

mlx-lm's deepseek_v4 (PR ml-explore/mlx-lm#1189) passes the cache offset as
a python int into three @mx.compile'd helpers. Measured on an M-series Mac
(mlx 0.32.0.dev, Metal):

• decode tok/s decays linearly with generated tokens: 28.8 → 17.7 by 6k
  generated (~4.3 µs/token added per generated token; ~43 layers × 3 calls).
• sample profiles across the same run: CompilerCache::find + memcmp
  grow from 12.4% → 28.6% of decode-thread wall time between ~token 2k
  and ~7k. GPU-side leaves shrink proportionally — pure CPU lookup cost.
• likely the mechanism behind mlx-lm#1332 (unbounded residency growth during
  DeepSeek-V4 decode) — ~3 entries retained per generated token, forever.

The user-side fix is passing a 0-d mx.array instead of the int (verified:
byte-identical outputs, flat decode). But the failure mode is easy to hit and
hard to diagnose — the API accepts the int silently and everything is correct,
just quadratically slower over a generation.

Suggested remedies (any one helps)

1. Hash the constants blob into a map (or bucket by hash) instead of the
   linear memcmp scan.
2. Move-to-front on hit — decode's access pattern is monotonic, so even MTF
   makes the common case O(1)-ish.
3. A cache-size cap / LRU eviction to bound the memory growth.
4. (Cheapest) a debug-mode warning when one compiled function accumulates

   N constant-keyed entries — points users straight at the offending arg.

A 20-line repro (one compiled function called in a loop with an incrementing
int) shows the growth cleanly; the full sample profiles and the script are
sitting ready — ask and they're attached.

Related: we previously reported the reduction bug in #3784; same test setup
(2× Mac Studio, Metal, long-context LLM workloads).

Co-authored with Claude Fable 5 (Anthropic)

[Vlor999]

I opened #3992 for the lookup half of this (remedy 1 — hash the constants blob
into buckets instead of the linear memcmp scan). Posting the retention numbers
here separately, because #3992 deliberately does not address them and I did
not want that to be silent.

Retention is unbounded, and independent of the lookup fix

Entries are only ever appended; CompilerCache::erase fires when the compiled
function itself is destroyed, so within the life of one compiled function nothing
is ever reclaimed. Each CacheEntry retains inputs, outputs and tape —
whole traced graphs, not just the key.

Measured with the reproducer below on an Apple M4 Pro (14-core), macOS 26.6,
Metal, 0.32.1.dev20260803+fb5133e. Peak RSS, one process:

distinct constants	peak RSS	growth
0	31.4 MB	—
1,000	38.5 MB	+7.1 MB
2,000	45.2 MB	+13.8 MB
4,000	59.5 MB	+28.1 MB

Linear, ~7 KB per distinct constant, never reclaimed. Extrapolating the decode
pattern this issue describes (~3 compiled helpers per generated token), a 6k-token
generation retains ~18k entries ≈ 126 MB of tape held for values that will never
be seen again. That is consistent with the unbounded residency growth reported in
mlx-lm#1332, though I have not reproduced that specific workload.

import resource, mlx.core as mx

@mx.compile
def step(x, offset):
    return (x * 2) + offset

x = mx.zeros((64,), mx.float32)
step(x, -1); mx.eval(step(x, -1))

rss = lambda: resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024 * 1024)
base, seen = rss(), 0
for n in (0, 1000, 2000, 4000):
    while seen < n:
        seen += 1
        step(x, seen)
    print(f"{n:>6} constants  peak RSS {rss():7.1f} MB  (+{rss() - base:5.1f})")

Bucketing by hash does not change any of this — the same entries are retained,
just indexed differently. I confirmed the numbers above are identical before and
after #3992.

Why I left it out of the PR

Remedy 3 (a cap or LRU eviction) is a semantics change, not an optimisation: it
makes a previously-guaranteed cache hit into a possible re-trace. Picking a bound
that is too small converts this issue into a recompile storm, and the right policy
depends on assumptions about workloads that maintainers own, not me. Happy to
implement whichever direction you prefer:

• a size cap per fun_id with LRU eviction (needs a chosen bound),
• eviction only of entries whose constants differ from the current call (targets
  exactly the changing-scalar pattern, leaves shape-keyed entries alone),
• or remedy 4 only — a debug-mode warning past N entries, which fixes nothing but
  points users at the offending argument, since the documented workaround
  (passing a 0-d mx.array instead of an int) is invisible today.

One note on remedy 2

I implemented and measured move-to-front before settling on the hash: it makes no
measurable difference here. This access pattern never produces a cache hit — the
scalar is new on every call — so the reordering never executes and every lookup
still scans the whole list. Numbers are in #3992 if useful.

[aidiffuser]

Thanks for #3992 and for quantifying the retention half — the two sides line up: what we measured as throughput decay (12→29% of decode wall going to the scan over a 6k-token generation) coincided with exactly the residency growth you describe, and our fix (0-d mx.array offsets) removed both at once. Your read on mlx-lm#1332 matches what we saw in production before the workaround.

Agreed on retiring move-to-front — with a fresh scalar per call there is never a hit to promote.

One data point for the eviction question, from the LLM-decode workload that motivated this issue: of the directions you list, evicting only entries whose constants differ from the current call is the one that exactly matches the failure pattern — decode constants are monotonically fresh, so that policy keeps the cache at ~one entry per shape while leaving shape-keyed reuse untouched. A size cap would also work but inherits the bound-picking problem you note. Whichever direction the maintainers choose, we can validate it against the original workload (batched decode, ~3 compiled helpers per token, 6k+ generations) and report throughput plus RSS before/after.

Also seconding your remedy-4 observation independent of the fix: the 0-d array workaround deserves a line in the mx.compile docs — it is currently discoverable only by profiling.

Co-authored with Claude Fable 5 (Anthropic)

[zcbenz]

If you keep re-compiling functions with different constants, lots of things will blow and the cache lookup time would not even be important, and I doubt you would get much performance gain because compilation is not cheap. You should pass the changing constants as mx.array which is basically all ML compilers do, and you would get much better performance.

[aidiffuser]

Agreed on the guidance — passing changing constants as 0-d arrays is exactly the fix we shipped on our side, and the ask was never to make accidental recompilation fast. The two residual concerns: (a) the int-passing pattern ships in first-party code (mlx_lm's deepseek_v4 rope offsets, fix proposed in mlx-lm#1189), so users inherit this without writing a compile call themselves; and (b) the failure is silent — no warning, just gradual slowdown plus the unbounded retention measured above. A debug-mode warning past N per-function entries, or a line in the mx.compile docs, would fully close this from our side.

Co-authored with Claude Fable 5 (Anthropic)