fix(server): raise the CUDA graph-cache default to survive long-lived shape-diverse decode (#818)#820
Merged
Merged
Conversation
On the CUDA backend, MLX keeps an LRU cache of captured CUDA graphs keyed by graph shape (mlx/backend/cuda/device.cpp, default capacity 400), and mlx/backend/cuda/lru_cache.h also counts a lifetime graph-cache miss counter that never resets, throwing a fatal std::runtime_error("Cache thrashing is happening ...") once lifetime misses pass 2 * capacity (800 at the default). Any long-lived, shape-diverse CUDA server crosses that threshold over its lifetime, and speculative or batched decode reaches it fastest because the draft and verify phases across varying batch sizes and sequence-length buckets multiply the number of distinct graph shapes. The throw is a whole-process abort, not a request-level error, so it drops every in-flight request. MLX_CUDA_GRAPH_CACHE_SIZE=2000 was verified to eliminate the crash (13/13 requests across bursts on GB10).
Mirror the existing apply_metal_ops_per_buffer_default startup pattern with a CUDA analogue in mlxcel-core: cuda_graph_cache_default() returns Some(2000) on a cuda-feature build and None otherwise, and apply_cuda_graph_cache_default() sets MLX_CUDA_GRAPH_CACHE_SIZE to that value before any MLX op unless the operator already set it (an explicit value always wins), logging one info line. The three entry points that already apply the Metal default (src/main.rs, src/bin/mlx_server.rs, src/bin/bench_decode.rs) now apply the CUDA default immediately after. 2000 is an LRU capacity cap, not a preallocation, so it only costs memory as distinct graph shapes accumulate, and the variable is read only by MLX's CUDA backend so it is a harmless no-op on Metal/CPU.
The upstream lifetime-miss fix and the FFI-boundary exception catch are left as separate follow-ups (proposals 2 and 3 in the issue); this startup default is the shippable fix. Documented in docs/environment-variables.md and the mlxcel --help env-var listing.
Verified on mlxcel-core with the cuda feature: cargo check passes, 26 hardware unit tests pass (including cuda_graph_cache_default_matches_build_feature), and clippy is clean under -D warnings.
Refs #818
inureyes
added
type:bug
…oo (#818) The #818 startup default (MLX_CUDA_GRAPH_CACHE_SIZE=2000 on CUDA builds) was wired into mlxcel, mlxcel-server, and the decode benchmark, but not into speculative_bench, which runs the exact workload issue #818 names as the fastest trigger of MLX's fatal "Cache thrashing" abort: a long-lived process that sweeps many distinct draft/verify graph shapes across pairings and K values. Without the default, speculative_bench on a CUDA build can still cross MLX's 800-lifetime-miss threshold at capacity 400 and abort mid-sweep. Call apply_cuda_graph_cache_default() at the top of main, right after arg parsing and before any model load or MLX op, mirroring the placement already used in bench_decode. Scoped to the CUDA default only (the metal MLX_MAX_OPS_PER_BUFFER knob is #353, out of scope here). No-op off CUDA; an explicit MLX_CUDA_GRAPH_CACHE_SIZE always wins. Refs #818
The operator notice requested by issue #818 was dead: apply_cuda_graph_cache_default emitted a tracing::info! from the top of every main(), before the server installs its tracing subscriber (initialize_server_logging in start_server), so the line was discarded on every path while the env var itself was still set correctly. Move the notice into start_server, right after the "Runtime device" log where a subscriber is guaranteed installed and runtime.device is in scope. It is gated on the cuda feature plus runtime.device == Gpu so it never fires on Metal or CPU, reads the effective MLX_CUDA_GRAPH_CACHE_SIZE, and reports that mlxcel raises MLX's default of 400 to 2000 unless an operator override is set. The env-setting helper keeps its early-main() behavior on all four binaries; only the dead log leaves. Observability lives on the server path only, matching the silent sibling apply_metal_ops_per_buffer_default. Verified with cargo check -p mlxcel-core --lib --features cuda, cargo clippy -p mlxcel-core --lib --features cuda -- -D warnings, and cargo check --bin mlxcel-server --features cuda. Refs #818
Complete the polish flagged by review on PR #820. The SAFETY comment in `apply_cuda_graph_cache_default` (src/lib/mlxcel-core/src/hardware.rs) listed only three in-tree callers, missing the fourth added later, `src/bin/speculative_bench.rs`; the comment now lists all four so the single-threaded-access rationale it documents is accurate. `docs/CONTINUOUS_BATCHING.md`'s speculative-decoding operator-controls section already cross-references environment-variables.md for the MTP/QMV knobs, so it gets a concise pointer to the new `MLX_CUDA_GRAPH_CACHE_SIZE` default, since speculative decode's draft/verify shape diversity is exactly what drives MLX's captured-graph-cache LRU hardest on CUDA. Validation: - cargo fmt --all -- --check (clean) - cargo clippy -p mlxcel-core --lib --features cuda -- -D warnings (clean) Refs #818
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On the CUDA backend, a long-lived shape-diverse
mlxcel-server(batched decode, and especially speculative decoding) aborts the whole process with an uncaught C++std::runtime_error("Cache thrashing is happening, please set the environment variable MLX_CUDA_GRAPH_CACHE_SIZE to a larger value than 400 ..."), dropping every in-flight request. Root cause in vendored MLX 0.32.1:mlx/backend/cuda/lru_cache.hthrows once a graph-cache miss counter that never resets passes2 * capacity, andmlx/backend/cuda/device.cppsets that capacity to 400, so the throw threshold is 800 lifetime misses. Speculative and batched decode multiply distinct graph shapes (draft/verify phases times batch sizes times sequence-length buckets), so the abort is inevitable at the default.MLX_CUDA_GRAPH_CACHE_SIZE=2000was verified to eliminate the crash (13/13 requests across bursts on GB10).This ships the short-term startup default (proposal 1 in the issue), mirroring the existing
apply_metal_ops_per_buffer_defaultpattern. The upstream lifetime-miss fix (proposal 3) and the FFI-boundary exception catch (proposal 2) remain separate follow-ups; the process is unrecoverable after the throw, so the startup default is the shippable fix.What changed
src/lib/mlxcel-core/src/hardware.rs: addedcuda_graph_cache_default() -> Option<u32>, a pure, cfg-gated function returningSome(2000)on acuda-feature build andNoneotherwise, andapply_cuda_graph_cache_default(), which setsMLX_CUDA_GRAPH_CACHE_SIZEto that value before any MLX op unless the operator already set it (an explicit value always wins), emitting onetracing::info!line. Both carry the same doc-comment rationale and the same SAFETY comment shape as the Metal analogue.src/main.rs,src/bin/mlx_server.rs,src/bin/bench_decode.rs: callapply_cuda_graph_cache_default()immediately after the existingapply_metal_ops_per_buffer_default()at the top ofmain, before any model load or MLX op.src/main.rs--helpenv-var listing: documentedMLX_CUDA_GRAPH_CACHE_SIZEnext toMLX_MAX_OPS_PER_BUFFER.docs/environment-variables.md: added anMLX_CUDA_GRAPH_CACHE_SIZErow (CUDA-only, raised default 2000, why MLX's own 400 aborts long-lived servers, explicit value overrides).src/lib/mlxcel-core/src/hardware.rstests: addedcuda_graph_cache_default_matches_build_feature, mirroring the Metal gating test and cfg-gated to both feature states (avoids a flaky global-env test, matching the existing pattern that does not test theapply_*mutator).2000 is an LRU capacity cap, not a preallocation, so it only costs memory as distinct graph shapes accumulate, and the variable is read only by MLX's CUDA backend so it is a harmless no-op on Metal/CPU.
Test plan
cargo check -p mlxcel-core --lib --features cuda(exit 0)cargo test -p mlxcel-core --lib --features cuda hardware(26 passed, 0 failed)cargo clippy -p mlxcel-core --lib --tests --features cuda -- -D warnings(exit 0)Closes #818