[inureyes]

Summary

On ROCm the pre-load memory estimate (mlxcel inspect, --estimate-memory, the paged KV auto budget, the prompt-cache capacity defaults) was below the measured peak for dense models, because committed but unfinished command batches may hold up to MLX_ROCM_MAX_INFLIGHT_MB of transients and nothing modeled them. ROCm builds now add that budget as its own term; Metal and CUDA totals are unchanged.

• rocm_inflight_reserve_bytes parses the variable exactly as the overlay's max_inflight_bytes() does (strtoull semantics: leading whitespace and + accepted, a leading -, suffix or overflow falls back to 1024 MiB, 0 adds nothing). A test reads device.cpp and fails if the default or the variable name drifts.
• The term is selected with #[cfg(feature = "rocm")], like the cache-limit default in runtime.rs, so the estimator stays side-effect-free. It is in total_bytes and runtime_headroom_bytes, printed as Backend in-flight (only when non-zero), and carried as backend_inflight_bytes in inspect --json.
• The paged KV auto budget subtracts the reserve, so the budget it picks still fits the estimate.
• --recommend-quant is unaffected: it reads only the estimate's weight and KV figures (the env-var page said otherwise in the first revision; corrected during review).

Units: the issue's correction does not hold. The page's 5.58 is total_bytes / 1e9 (inspect prints 5.20 GiB, 5,579,686,809 bytes), so the original 0.56 GB gap was real. All comparisons below are in bytes.

Verification (gfx1151)

• Guarded peaks (scripts/rocm_gpu_guard.sh, one clean run per row) for Llama-3.1-8B, Qwen2.5-7B and Qwen3-30B-A3B at 640 and 4096 tokens with budgets 256, 1024 and 4096 MiB: before the change the estimate missed 9 of the 18 peaks; now it covers all 18, the smallest margin being 221,877,772 bytes (Llama, 4096 tokens, default). The table and raw lines are in docs/benchmark_results/rocm-memory-gfx1151-2026-09-30.md, "Estimate against peak".
• rocm_estimate_covers_every_measured_gfx1151_peak rebuilds all 18 estimates (equal to inspect byte for byte) and asserts estimate >= peak bound. With the term forced to 0, it and two other tests fail.
• Prompt-cache sizing side effect: on ROCm the ceiling (a quarter of the slack) drops by 268,435,456 bytes at the default budget. On this host the chosen capacities for all three models do not change, since six entries (2.82 to 6.44 GB) sit well under the ceilings (14.95 to 18.91 GB). In general a capacity drops by at most the larger of budget/4 and one entry, and never below the compiled-in default.
• cargo test --features rocm --lib -- execution::memory_estimate server::prompt_cache::snapshot_sizing (70 passed), clippy on lib and tests with -D warnings, dead_doc_pointers, and the fast make verify-* script gates pass. Full make verify-rocm at 9374ae72 (on main 87538835): OK, 12004 passed, 0 failed, 384 ignored over 153 test binaries.

Not verified: Metal and CUDA (not available on this host). On those builds the code path is the #[cfg(not(feature = "rocm"))] stub that returns 0, plus the new zero-valued field and JSON key; non_rocm_estimate_reserves_nothing_for_inflight was not compiled or run here.

Closes #2155

[inureyes]

Implementation Review Summary

Intent

On ROCm builds, add the MLX_ROCM_MAX_INFLIGHT_MB budget to the pre-load memory estimate as its own term, parsed exactly as the backend parses it, with Metal and CUDA totals unchanged.

Findings Addressed

None needed fixing at CRITICAL or HIGH. I checked parse_inflight_mb_like_strtoull against max_inflight_bytes() in device.cpp (leading-byte - rejection, C whitespace, +, a - after whitespace wrapping past the shift limit with -0 giving 0, ERANGE, trailing bytes, the shift limit) and found no divergence. Every MemoryEstimate consumer handles the term the same way: inspect text and JSON, the generate and serve preflights through total_bytes and fits, the generate delta log, auto_kv_budget_bytes, and snapshot_sizing through slack_bytes().

Remaining Items

• The docs/environment-variables.md ROCm paragraph says --recommend-quant adds the budget. It does not: advise_quantization passes only weight and KV bytes to recommend_quantization, which never reads total_bytes. The same is true of the older MLXCEL_HEADROOM_FACTOR row. Fix: drop --recommend-quant from the new sentence. (MEDIUM, left for finalization)
• The module doc at the top of memory_estimate.rs still lists four components and leaves out the backend in-flight term. (LOW)

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass (execution::memory_estimate on ROCm: 56 passed; clippy on lib and tests with -D warnings is clean; the non-ROCm path is not compiled on this host)

[inureyes]

Security and performance review of the in-flight reserve change: no CRITICAL or HIGH findings, no commits made.

• Arithmetic: mb << 20 cannot drop bits because the parser only accepts counts up to u64::MAX >> 20 and the default is 1024. The new term is folded in with saturating_add, and auto_kv_budget_bytes and format_estimate subtract it with saturating_sub, so a maximal budget saturates the total to u64::MAX and reports not fitting (pinned by rocm_largest_accepted_budget_saturates_the_total). The test-only total_bytes - backend_inflight_bytes cannot underflow because the total always includes the term.
• Env parsing: the parser mirrors max_inflight_bytes() in device.cpp case for case (first-byte -, whitespace then sign, trailing bytes, ERANGE, shift bound), and the slice after the sign is safe because that byte is ASCII. Non-UTF-8 values fall back to the default, which the C side also does. The value comes from the local operator's environment, not from a request, so it is not attacker-controlled. It is read once per estimate call, and estimates run only before load, so there is no hot-path cost.
• Test env mutation: every setter holds crate::test_support::env_lock and declares it before the EnvGuards, so the guards restore the variables before the lock is released. Existing unlocked estimate tests read the variable once per call and only check self-consistent sums, so a concurrent setter cannot make them flaky.
• LOW, not fixed: EnvGuard saves the previous value with std::env::var(..).ok(), so a non-UTF-8 pre-existing value would be removed instead of restored. This only affects tests; var_os would close it.