You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805
Measured 2026-09-14 on the spike host (Radeon 8060S, gfx1151, ROCm 10.0.0), at d67b64d3 on main. The measurements below replace what this issue originally guessed; start from them, not from the guess.
Context
resolve_available_memory (src/execution/memory_estimate.rs:725) walks MLXCEL_MEMORY_LIMIT, then the MLX allocator cap memory_limit(), then HardwareCapabilities::unified_memory_gb, then /proc/meminfo. What that walk actually sees on this host:
Call
Value
mlxcel_core::memory::memory_limit()
82463372083 bytes (76.80 GiB)
gpu_max_memory_size()
103079215104 bytes (96.00 GiB)
get_wired_limit()
0
HardwareCapabilities.unified_memory_gb
0
HardwareCapabilities.silicon_gen
Unknown
The ROCm allocator reports a nonzero memory_limit(), so the walk returns at its second step with 76.80 GiB and never reaches /proc/meminfo. The 28 GiB figure this issue was filed to worry about is not what the estimator sees, and nothing that fits in the carve-out is refused for that reason. gpu_max_memory_size() falls through to total_memory and reports the full 96.00 GiB as predicted (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4630), and the wired limit is a no-op exactly as on CUDA (:4588). Resolving available memory from device_info() ahead of /proc/meminfo is therefore not needed for correctness; anyone who still wants it has to argue it as a policy choice (the allocator's 76.80 GiB default cap against the device's 96.00 GiB) and not as a fix.
Two failures are real and live instead. Both are in the mlxcel-core library test binary and both occur under make verify-test-rocm, but that run aborts on an unrelated nvfp4 case (terminate called ... Unsupported group_size for QuantizedMatmul: 16, issue #1806) before the harness prints its failure list, which is why reading the gate output misses them. Run each test directly to see it:
cuda_arch_tests::the_capability_probe_agrees_with_the_backend (src/lib/mlxcel-core/src/cuda_arch_tests.rs:308) fails with assertion left == right failed: a non-CUDA build has no compute capability to report, left: Some((11, 5)), right: None. The ROCm backend publishes compute_capability_major/minor derived from gfx1151 (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device_info.cpp:80-82), and cuda_compute_capability (src/lib/mlxcel-core/src/cuda_arch.rs:83) reports it as if it were a CUDA compute capability. cuda_arch_startup_summary (:399) consequently prints CUDA compute capability 11.5 (sm_115) for an AMD device.
streams::tests::gpu_backend_available_does_not_track_the_default_device (src/lib/mlxcel-core/src/streams.rs:568) fails with a CPU-only build has no GPU backend to report. The test's premise is that a build is Metal, CUDA or CPU-only.
Both are the same shape: a test encoding "there are two GPU backends and everything else is CPU". That is the same assumption #1803 removed one layer down, where !metal::is_available() was read as "CUDA".
Hardware detection is the real remaining gap. HardwareCapabilities (src/lib/mlxcel-core/src/hardware.rs:125) has no vendor concept: an AMD GPU resolves to silicon_gen: Unknown with every other field at its default (gpu_core_count: 0, unified_memory_gb: 0, memory_bandwidth_gbps: 0.0), so the startup diagnostics, the device listing and the server metadata have nothing to report.
Scope
Give mlxcel a GPU vendor concept, make device reporting correct on a ROCm device, and add the gfx arch-compatibility refusal, including on UMA carve-out hosts. Discrete AMD GPUs (separate VRAM) should also work, but only the UMA host is available for validation. Changing the memory resolution order is out of scope: the measurement shows it already resolves correctly on ROCm.
Implementation plan
Add a vendor to HardwareCapabilities rather than a variant to AppleSiliconGen. AppleSiliconGen::Unknown is currently used as a proxy for "this is not Apple Silicon" at twelve call sites in eleven files, where it gates a bf16 to f16 weight conversion: src/loading/vlm_smolvlm.rs:293, vlm_mllama.rs:114, vlm_internvl.rs:137, vlm_llmjp_vl.rs:249, vlm_idefics2.rs:181, vlm_locateanything.rs:316, vlm_lfm2_vl.rs:94, vlm_jina_vlm.rs:67, vlm_special.rs:542 and vlm_special.rs:1958, src/models/sanitize.rs:2033 (should_convert_bf16_to_f16) and src/lib/mlxcel-core/src/drafter/dflash/drafter.rs:198. Adding an Amd variant to AppleSiliconGen would flip all twelve to true on ROCm and silently enable that conversion on a backend where bf16 is native. Introduce #[non_exhaustive] enum GpuVendor { Apple, Nvidia, Amd, Unknown } as its own field, and give the type an is_apple_silicon() predicate whose body is the existing comparison, so the twelve sites can be converted mechanically with no behavior change. src/server/batch/mtp_policy.rs:1042 uses the same comparison for the MTP compute-bound gate rather than for a dtype decision, and src/models/sanitize_tests.rs:388 mirrors the predicate in a test expectation; convert both to the predicate too, without changing what they mean.
Add a test that no non-Apple vendor makes is_apple_silicon() true, naming the bf16 to f16 conversion the predicate gates. The comment explaining the trap is not enough on its own; the next person to add a vendor needs something that fails.
Populate the AMD arm of the detection: vendor: Amd, silicon_gen: Unknown, the neural-accelerator fields false, carrying the gfx architecture string, the device name and total device memory from MLX device_info() (architecture, device_name, total_memory). Note that the ROCm device_info() publishes no compute-unit count and no max_recommended_working_set_size, so neither can be reported from there.
Stop reporting a CUDA compute capability for a non-CUDA device, and widen the two tests above so their premise is the vendor rather than "Metal, CUDA, or CPU".
Refuse to start when the embedded HIP architectures do not include the detected gfx target, with a named error, bypassable with MLXCEL_DEVICE=cpu, mirroring the CUDA arch-mismatch check (initialize_runtime_checked, src/execution/runtime.rs:117-129, docs/environment-variables.md). Extend MLXCEL_TRACE_ARCH to print the ROCm picture (detected gfx, compiled list).
Leave resolve_available_memory as it is and add unit tests pinning its resolution order (env override, then the MLX allocator cap, then unified_memory_gb, then /proc/meminfo), so the measured behavior is regression-protected on all three backends.
A test pins the is_apple_silicon() invariant: no non-Apple GpuVendor makes it true.
Startup diagnostics show the AMD device name, gfx1151, device memory and the compiled HIP architecture list, and no longer print a CUDA compute capability on an AMD device.
A binary built for a different gfx target refuses to start with an error naming both the detected and the compiled targets; MLXCEL_DEVICE=cpu bypasses it.
Metal and CUDA behavior is unchanged: the memory resolution order is covered by unit tests and unchanged on both, and the twelve converted call sites show an identical predicate in the diff.
Part of #1801. Phase 1. Depends on #1802.
Measured 2026-09-14 on the spike host (Radeon 8060S,
gfx1151, ROCm 10.0.0), atd67b64d3onmain. The measurements below replace what this issue originally guessed; start from them, not from the guess.Context
resolve_available_memory(src/execution/memory_estimate.rs:725) walksMLXCEL_MEMORY_LIMIT, then the MLX allocator capmemory_limit(), thenHardwareCapabilities::unified_memory_gb, then/proc/meminfo. What that walk actually sees on this host:mlxcel_core::memory::memory_limit()gpu_max_memory_size()get_wired_limit()HardwareCapabilities.unified_memory_gbHardwareCapabilities.silicon_genUnknownThe ROCm allocator reports a nonzero
memory_limit(), so the walk returns at its second step with 76.80 GiB and never reaches/proc/meminfo. The 28 GiB figure this issue was filed to worry about is not what the estimator sees, and nothing that fits in the carve-out is refused for that reason.gpu_max_memory_size()falls through tototal_memoryand reports the full 96.00 GiB as predicted (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4630), and the wired limit is a no-op exactly as on CUDA (:4588). Resolving available memory fromdevice_info()ahead of/proc/meminfois therefore not needed for correctness; anyone who still wants it has to argue it as a policy choice (the allocator's 76.80 GiB default cap against the device's 96.00 GiB) and not as a fix.Two failures are real and live instead. Both are in the
mlxcel-corelibrary test binary and both occur undermake verify-test-rocm, but that run aborts on an unrelated nvfp4 case (terminate called ... Unsupported group_size for QuantizedMatmul: 16, issue #1806) before the harness prints its failure list, which is why reading the gate output misses them. Run each test directly to see it:cuda_arch_tests::the_capability_probe_agrees_with_the_backend(src/lib/mlxcel-core/src/cuda_arch_tests.rs:308) fails withassertion left == right failed: a non-CUDA build has no compute capability to report, left: Some((11, 5)), right: None. The ROCm backend publishescompute_capability_major/minorderived fromgfx1151(src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device_info.cpp:80-82), andcuda_compute_capability(src/lib/mlxcel-core/src/cuda_arch.rs:83) reports it as if it were a CUDA compute capability.cuda_arch_startup_summary(:399) consequently printsCUDA compute capability 11.5 (sm_115)for an AMD device.streams::tests::gpu_backend_available_does_not_track_the_default_device(src/lib/mlxcel-core/src/streams.rs:568) fails witha CPU-only build has no GPU backend to report. The test's premise is that a build is Metal, CUDA or CPU-only.Both are the same shape: a test encoding "there are two GPU backends and everything else is CPU". That is the same assumption #1803 removed one layer down, where
!metal::is_available()was read as "CUDA".Hardware detection is the real remaining gap.
HardwareCapabilities(src/lib/mlxcel-core/src/hardware.rs:125) has no vendor concept: an AMD GPU resolves tosilicon_gen: Unknownwith every other field at its default (gpu_core_count: 0,unified_memory_gb: 0,memory_bandwidth_gbps: 0.0), so the startup diagnostics, the device listing and the server metadata have nothing to report.Scope
Give mlxcel a GPU vendor concept, make device reporting correct on a ROCm device, and add the
gfxarch-compatibility refusal, including on UMA carve-out hosts. Discrete AMD GPUs (separate VRAM) should also work, but only the UMA host is available for validation. Changing the memory resolution order is out of scope: the measurement shows it already resolves correctly on ROCm.Implementation plan
HardwareCapabilitiesrather than a variant toAppleSiliconGen.AppleSiliconGen::Unknownis currently used as a proxy for "this is not Apple Silicon" at twelve call sites in eleven files, where it gates a bf16 to f16 weight conversion:src/loading/vlm_smolvlm.rs:293,vlm_mllama.rs:114,vlm_internvl.rs:137,vlm_llmjp_vl.rs:249,vlm_idefics2.rs:181,vlm_locateanything.rs:316,vlm_lfm2_vl.rs:94,vlm_jina_vlm.rs:67,vlm_special.rs:542andvlm_special.rs:1958,src/models/sanitize.rs:2033(should_convert_bf16_to_f16) andsrc/lib/mlxcel-core/src/drafter/dflash/drafter.rs:198. Adding anAmdvariant toAppleSiliconGenwould flip all twelve to true on ROCm and silently enable that conversion on a backend where bf16 is native. Introduce#[non_exhaustive] enum GpuVendor { Apple, Nvidia, Amd, Unknown }as its own field, and give the type anis_apple_silicon()predicate whose body is the existing comparison, so the twelve sites can be converted mechanically with no behavior change.src/server/batch/mtp_policy.rs:1042uses the same comparison for the MTP compute-bound gate rather than for a dtype decision, andsrc/models/sanitize_tests.rs:388mirrors the predicate in a test expectation; convert both to the predicate too, without changing what they mean.is_apple_silicon()true, naming the bf16 to f16 conversion the predicate gates. The comment explaining the trap is not enough on its own; the next person to add a vendor needs something that fails.vendor: Amd,silicon_gen: Unknown, the neural-accelerator fields false, carrying thegfxarchitecture string, the device name and total device memory from MLXdevice_info()(architecture,device_name,total_memory). Note that the ROCmdevice_info()publishes no compute-unit count and nomax_recommended_working_set_size, so neither can be reported from there.src/commands/generate.rs:134-147), themlxceldevice listing and the server metadata endpoints (src/server/startup.rs:1964), alongside the embedded ROCm architecture list from build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #1802.gfxtarget, with a named error, bypassable withMLXCEL_DEVICE=cpu, mirroring the CUDA arch-mismatch check (initialize_runtime_checked,src/execution/runtime.rs:117-129,docs/environment-variables.md). ExtendMLXCEL_TRACE_ARCHto print the ROCm picture (detectedgfx, compiled list).resolve_available_memoryas it is and add unit tests pinning its resolution order (env override, then the MLX allocator cap, thenunified_memory_gb, then/proc/meminfo), so the measured behavior is regression-protected on all three backends.Acceptance criteria
cuda_arch_tests::the_capability_probe_agrees_with_the_backendandstreams::tests::gpu_backend_available_does_not_track_the_default_devicepass on the ROCm build, verified by running them directly (the gate run still aborts on feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806).is_apple_silicon()invariant: no non-AppleGpuVendormakes it true.gfx1151, device memory and the compiled HIP architecture list, and no longer print a CUDA compute capability on an AMD device.gfxtarget refuses to start with an error naming both the detected and the compiled targets;MLXCEL_DEVICE=cpubypasses it.Validation
The two direct binary invocations are required while #1806 aborts the full gate run before its failure list prints.
References
src/lib/mlxcel-core/src/hardware.rs:55(AppleSiliconGen),:125(HardwareCapabilities)src/lib/mlxcel-core/src/cuda_arch.rs:83(cuda_compute_capability),:399(cuda_arch_startup_summary)src/lib/mlxcel-core/src/cuda_arch_tests.rs:308,src/lib/mlxcel-core/src/streams.rs:568(the two failing tests)src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device_info.cpp(which keys the ROCm backend publishes)src/execution/memory_estimate.rs:725(resolve_available_memory),src/execution/runtime.rs:117(initialize_runtime_checked),:266(resolve_wired_limit),:290(resolve_memory_limit),:309(resolve_cache_limit)src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4588(get_wired_limit),:4630(gpu_max_memory_size)