Skip to content

perf(cuda): allocator policy: wire set_cache_limit and re-evaluate periodic clear_memory_cache #627

Closed

Description

@inureyes

Part of #623. Depends on #625 (measure on the upgraded pin; upstream memory-pool behavior changed on Tegra).

Context

Two allocator findings from the code audit:

  1. set_cache_limit is exposed in the bridge (src/lib/mlxcel-core/src/lib.rs:1765) but never called anywhere. MLX's buffer-cache size is therefore unbounded/default on both backends, with no operator knob. Runtime init (src/execution/runtime.rs, initialize_runtime() at line ~73) already resolves MLXCEL_WIRED_LIMIT and MLXCEL_MEMORY_LIMIT; cache limit is the missing third.
  2. The decode loop calls ffi::clear_memory_cache() every 256 generated tokens (src/lib/mlxcel-core/src/generate.rs:1290-1292), and the same helper is called liberally across the scheduler and CLI commands. On Metal this trims the buffer cache cheaply. On CUDA, dropping cached buffers forces reallocation through the CUDA memory pool on subsequent steps and defeats MLX's CUDA-graph executable cache (graph reuse depends on stable buffer addresses; see upstream issue "Buffer cache and CUDA graph", [CUDA] Buffer cache and CUDA graph ml-explore/mlx#2358). The periodic clear may be a net loss on CUDA.

Scope

  1. Wire MLXCEL_CACHE_LIMIT into runtime init.
  2. Measure and set a CUDA-appropriate policy for the periodic decode-loop cache clear.

Out of scope: changing MLX's allocator itself.

Implementation plan

1. Cache-limit knob

  • In src/execution/runtime.rs, alongside resolve_wired_limit() (line ~145) and resolve_memory_limit() (line ~167), add resolve_cache_limit() reading MLXCEL_CACHE_LIMIT (same size-suffix parsing as the existing two) and calling mlxcel_core -> set_cache_limit. initialize_runtime() is invoked from every entry path (server src/server/startup.rs:1612, CLI src/commands/generate.rs, chat.rs, bench binaries), so one change point covers all.
  • Document in docs/environment-variables.md next to MLXCEL_WIRED_LIMIT / MLXCEL_MEMORY_LIMIT.

2. Periodic clear policy on CUDA

  • Benchmark matrix on GB10 (post pin-upgrade), using mlxcel-bench-decode with 2048-token generations on llama-3.1-8b-4bit and qwen2.5-0.5b-bf16 (one bandwidth-bound, one overhead-bound):
    • cadence 256 (current), 4096, and disabled;
    • record decode tok/s, peak memory (get_peak_memory), and RSS.
  • Decide per-backend: likely keep cadence on Metal, disable or greatly lengthen on CUDA and rely on set_cache_limit for bounding. Implement as a backend-conditional constant with an env override (MLXCEL_CACHE_CLEAR_INTERVAL, 0 = disabled), applied at the existing call site generate.rs:1290-1292. Audit the scheduler's equivalent periodic calls (src/server/batch/scheduler.rs, search clear_memory_cache) and apply the same policy.
  • Verify no memory blow-up over a long run: 16k-token generation plus a 30-minute server soak with the concurrency script from perf(bench): long-prompt prefill benchmark coverage and serving TTFT/decode-rate telemetry #624, watching the new peak-memory numbers.

Acceptance criteria

  • MLXCEL_CACHE_LIMIT functional from all entry points; documented.
  • Cache-clear cadence policy is backend-aware with env override; chosen values justified by the committed benchmark table (extend docs/benchmark_results/ with a dated note).
  • No RSS/peak-memory regression on the soak test; decode tok/s on the two reference models is at or above the pre-change baseline on CUDA.
  • cargo test --features cuda passes.

Validation

cargo build --release --features cuda
MLXCEL_CACHE_CLEAR_INTERVAL=0 ./target/release/mlxcel-bench-decode --model ./models/llama-3.1-8b-4bit --max-tokens 2048
MLXCEL_CACHE_LIMIT=8G ./target/release/mlxcel --model ./models/llama-3.1-8b-4bit generate "..."

References

  • Bridge symbol: src/lib/mlxcel-core/src/lib.rs:1765 (set_cache_limit), memory counters at lines 1749-1771.
  • Runtime init: src/execution/runtime.rs:73-178.
  • Clear sites: src/lib/mlxcel-core/src/generate.rs:1290-1292; grep clear_memory_cache for the full list.
  • Upstream context: [CUDA] Buffer cache and CUDA graph ml-explore/mlx#2358 (buffer cache vs CUDA graph reuse).

Activity

  1. added
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    and removed on Jul 2, 2026
  2. added a commit that references this issue on Jul 30, 2026
    98d4c5d
  3. self-assigned this
    on Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions