Repository navigation
perf(cuda): allocator policy: wire set_cache_limit and re-evaluate periodic clear_memory_cache #627
Copy link
Copy link
Closed
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layerspriority:mediumMedium priorityMedium prioritystatus:doneCompletedCompletedtype:performancePerformance improvementsPerformance improvements
Description
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:
set_cache_limitis 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 resolvesMLXCEL_WIRED_LIMITandMLXCEL_MEMORY_LIMIT; cache limit is the missing third.- 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
- Wire
MLXCEL_CACHE_LIMITinto runtime init. - 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, alongsideresolve_wired_limit()(line ~145) andresolve_memory_limit()(line ~167), addresolve_cache_limit()readingMLXCEL_CACHE_LIMIT(same size-suffix parsing as the existing two) and callingmlxcel_core->set_cache_limit.initialize_runtime()is invoked from every entry path (serversrc/server/startup.rs:1612, CLIsrc/commands/generate.rs,chat.rs, bench binaries), so one change point covers all. - Document in
docs/environment-variables.mdnext toMLXCEL_WIRED_LIMIT/MLXCEL_MEMORY_LIMIT.
2. Periodic clear policy on CUDA
- Benchmark matrix on GB10 (post pin-upgrade), using
mlxcel-bench-decodewith 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_limitfor bounding. Implement as a backend-conditional constant with an env override (MLXCEL_CACHE_CLEAR_INTERVAL, 0 = disabled), applied at the existing call sitegenerate.rs:1290-1292. Audit the scheduler's equivalent periodic calls (src/server/batch/scheduler.rs, searchclear_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_LIMITfunctional 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 cudapasses.
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; grepclear_memory_cachefor the full list. - Upstream context: [CUDA] Buffer cache and CUDA graph ml-explore/mlx#2358 (buffer cache vs CUDA graph reuse).
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layerspriority:mediumMedium priorityMedium prioritystatus:doneCompletedCompletedtype:performancePerformance improvementsPerformance improvements
Activity
perf(cuda): add MLXCEL_CACHE_LIMIT and disable periodic decode cache-…