perf(rocm): explain and bound allocator peak memory on UMA hosts #2062

Part of #1814 (epic #1801)

Independent of #2061; can start now.

Problem / Background

On gfx1151 (Radeon 8060S, UMA, 96 GiB carve-out, 31 GiB visible to the host), the ROCm spike measured a peak of about 20 GB for Meta-Llama-3.1-8B-Instruct-4bit, whose weights are about 4.5 GB, and 21.7 GB for Qwen3-30B-A3B-4bit (17 GB of weights). On a UMA host that memory competes with the OS and other GPU tenants. Nobody has explained where the extra 15 GB goes or bounded it.

Current Behavior

• Fork allocator (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.cpp): on an integrated device memory_limit_ is max(0.8 * hipMemGetInfo total, 0.8 * system RAM) (:289-295) and max_pool_size_ = memory_limit_ (:302), so the buffer cache may hold nearly the whole carve-out. malloc deliberately does not reclaim against the soft limit (:504) and free does no eager pool trim (:567, :687-691); only clear_cache() and set_cache_limit() shrink it. Other knobs: the graph-decode arena (decode_arena_begin, :772), the async pool (MLX_ROCM_USE_ASYNC_POOL, MLX_ROCM_NO_ASYNC_POOL, MLX_ROCM_FORCE_ASYNC_POOL, release threshold 0), managed/GTT fallback (MLX_ROCM_ALLOW_MANAGED_FALLBACK, MLX_ROCM_NO_MANAGED_FALLBACK), MLX_ROCM_FINEGRAINED, and MLX_GRAPH_NODEFER.
• mlxcel side: MLXCEL_CACHE_LIMIT and MLXCEL_MEMORY_LIMIT (resolve_memory_limit at src/execution/runtime.rs:330, resolve_cache_limit at :349) call MLX's set_cache_limit / set_memory_limit only when set; unset means no bound. The periodic decode-loop clear_memory_cache runs every 256 tokens on ROCm (DEFAULT_CACHE_CLEAR_INTERVAL, src/lib/mlxcel-core/src/memory.rs:168, the non-CUDA default).

Proposed Solution

1. Explain the peak. For both models, record mlxcel_core::memory::active_memory(), cache_memory() and peak_memory() after load, after prefill and at the end of decode (src/bin/bench_decode.rs already prints the peak at :532), plus the decode arena capacity and high-water mark. State how much of the peak is live tensors, buffer cache, arena, and pool slack.
2. Compare the two control surfaces. For each knob above and for MLX's set_cache_limit / set_memory_limit, measure peak and decode tok/s with scripts/bench_decode.sh at pp512/tg128.
3. Ship a bounded default. When MLXCEL_CACHE_LIMIT is unset on a rocm build, resolve_cache_limit() applies a default cache limit instead of none; an explicit MLXCEL_CACHE_LIMIT (including 0/none to disable) still wins. Pick the value as the smallest bound that keeps decode tok/s within 2% of unbounded on both models; express it as a fraction of weights plus a fixed floor if one constant does not fit both. If step 1 shows the excess is not buffer cache (for example the arena or pool slack), fix it in the overlay instead and add a LOCAL_FIXES.md item, rather than capping a counter that does not hold the memory.

Rejected: changing max_pool_size_ inside the fork only, because it is invisible to operators and not overridable through mlxcel's existing variable.

Scope

In scope: measurement, the ROCm default in src/execution/runtime.rs, overlay allocator changes if step 1 requires them, documentation.

Out of scope: Metal and CUDA defaults; the decode profile (#2061).

Implementation Notes

• Reuse parse_memory_size and the existing resolve_cache_limit() path; do not add a second variable.
• Constraints: Metal and CUDA behavior unchanged (gate the default on the rocm feature). The default must not trigger per-step reallocation that defeats HIP graph reuse; the 2% throughput bound guards this.
• Edge cases: a model whose live working set exceeds the default cache limit still runs (the cache limit bounds only cached free buffers).

Acceptance Criteria

• docs/benchmark_results/rocm-memory-gfx1151-<YYYY-MM-DD>.md breaks down the peak for both models and compares every knob listed above.
• With no environment set, peak memory for Llama-3.1-8B-4bit on gfx1151 is bounded by the documented default, with decode tok/s within 2% of the unbounded run.
• MLXCEL_CACHE_LIMIT still overrides the default, including disabling it; a unit test in runtime.rs covers unset, explicit and disabled on the rocm feature.
• The default and its rationale are documented in docs/installation.md (section "Linux with AMD ROCm (experimental)").

Verification

cargo build --release --features rocm --bin mlxcel-bench-decode
MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/Meta-Llama-3.1-8B-Instruct-4bit
MLXCEL_CACHE_LIMIT=none MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/Meta-Llama-3.1-8B-Instruct-4bit
cargo test --features rocm -p mlxcel --lib execution::runtime
make verify-rocm

Pass: the MLX peak memory line drops to the documented bound in the first run and matches today's figure in the second.