fix(rocm): cover the in-flight working set in the pre-load memory estimate #2155

Part of #1801

Problem / Background

The pre-load memory estimate (mlxcel inspect, generate/serve --estimate-memory, --recommend-quant) is below the peak ROCm actually reaches. After PR #2084 bounded in-flight command batches (MLX_ROCM_MAX_INFLIGHT_MB, default 1024), docs/benchmark_results/rocm-memory-gfx1151-2026-09-30.md:124 records that mlxcel inspect --max-tokens 640 estimates 5.58 for Meta-Llama-3.1-8B-Instruct-4bit while the measured pp512/tg128 peak on gfx1151 is 6.14 GB. The page left "recalibrating the 1.20 factor for ROCm" out of scope. A preflight that under-reserves can pass a load that then exceeds the budget. Qwen3-30B-A3B-4bit is already over-estimated (20.72 against an 18.58 GB peak).

Units differ between the two figures and must be compared in bytes: bench_decode prints the peak as bytes / 1e9 (src/bin/bench_decode.rs:621), and format_bytes prints the estimate in GiB (src/execution/memory_estimate.rs:1436-1447). Read as GiB, 5.58 is 5.99e9 bytes, which still leaves the estimate about 0.15 GB under the peak.

Current Behavior

estimate_total_memory (src/execution/memory_estimate.rs:485-555) sums weights, architecture-aware KV, (DEFAULT_HEADROOM_FACTOR - 1) * (weights + kv) with DEFAULT_HEADROOM_FACTOR = 1.20 (:101, calibrated on Apple Silicon only, :65-100), and a workload activation term. Nothing models the ROCm backend's in-flight working set: up to MLX_ROCM_MAX_INFLIGHT_MB of transient allocations held by committed but unfinished batches (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device.cpp:68-104, LOCAL_FIXES.md item 28), mostly the f16 weight copies of the dequantize-and-GEMM qmm path during prefill.

Proposed Solution

Add an explicit ROCm-only additive term instead of retuning the shared factor.

• In memory_estimate.rs, add pub fn rocm_inflight_reserve_bytes(raw: Option<&str>) -> u64, a pure parser that mirrors max_inflight_bytes() exactly: unset or empty means 1024 MiB; a whole-string non-negative decimal integer means that many MiB; anything else (sign, suffix, overflow when shifted by 20) means the default; 0 means 0.
• Gate it the way the ROCm cache-limit default is gated (src/execution/runtime.rs:56-59): #[cfg(feature = "rocm")] reads MLX_ROCM_MAX_INFLIGHT_MB and calls the parser; #[cfg(not(feature = "rocm"))] yields 0. No runtime backend-kind branch, and the estimator stays side-effect-free (no runtime bring-up).
• Add a backend_inflight_bytes: u64 field to MemoryEstimate, include it in total_bytes and runtime_headroom_bytes, add it to the inspect JSON report struct (InspectReport, :359-376), and print a line Backend in-flight: <bytes> (MLX_ROCM_MAX_INFLIGHT_MB) in format_estimate only when it is non-zero.
• MLX_ROCM_MAX_INFLIGHT_MB=0 (bound off) adds 0, and docs/environment-variables.md says the estimate does not model an unbounded in-flight set (peaks of 20.60 GB were measured that way).

Rejected: raising the factor on ROCm. It scales with weights and over-reserves for large models (the 30B MoE already sits above its peak), while the measured excess tracks the in-flight budget, not model size.

Scope

• In scope: src/execution/memory_estimate.rs, its tests, docs/environment-variables.md (MLXCEL_HEADROOM_FACTOR row and the MLX_ROCM_MAX_INFLIGHT_MB paragraph), and a new "Estimate against peak" table in docs/benchmark_results/rocm-memory-gfx1151-2026-09-30.md.
• Out of scope: Metal and CUDA calibration (their totals must not change), the in-flight bound itself.

Implementation Notes

• Reuse: resolve_headroom_factor style for env reading; the existing #[cfg(feature = "rocm")] constant pattern in runtime.rs.
• Downstream: src/server/prompt_cache/snapshot_sizing.rs:157-166 uses slack_bytes(), so ROCm snapshot capacity shrinks by the reserve; that is intended and must be stated in the PR.
• Edge cases: invalid env values fall back to 1024 MiB like the overlay does (the estimator must agree with what the backend will enforce); huge values saturate rather than overflow.

Acceptance Criteria

• Unit tests for rocm_inflight_reserve_bytes: unset, empty, 0, 256, -1, 8abc, 1G, and a value that overflows when shifted, each with the expected bytes.
• A #[cfg(feature = "rocm")] test: estimate_total_memory(...).total_bytes equals the sum of weights, KV, factor headroom, activation and backend_inflight_bytes, with backend_inflight_bytes == 1 << 30 by default. A #[cfg(not(feature = "rocm"))] test: backend_inflight_bytes == 0, so Metal and CUDA totals are unchanged.
• On gfx1151, for Llama-3.1-8B-Instruct-4bit, Qwen3-30B-A3B-4bit and one more dense model, at --max-tokens 640 and 4096, with the default budget and with MLX_ROCM_MAX_INFLIGHT_MB=256 and 4096: the estimate in bytes is at least the measured bench_decode peak in bytes. The table goes on the results page.
• mlxcel inspect on a ROCm build prints the new line, and the JSON report carries the field.

Verification

cargo test --profile test-fast --features metal,accelerate -- memory_estimate --test-threads=1
cargo test --profile test-fast --features rocm -- memory_estimate --test-threads=1
./target/release/mlxcel inspect models/mlx/Meta-Llama-3.1-8B-Instruct-4bit --max-tokens 640
make verify-rocm   # gfx1151, inside scripts/rocm_gpu_guard.sh for the peak runs

Related: #2062, #2084, #1805.