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
chore(rocm): remove dead overlay env knobs and range-check the dequant byte cap #2242
PR #2222 documented the runtime MLX_ROCM_* variables and found four groups of environment knobs in the ROCm overlay that are dead, unsafe to parse, or undocumented. They are documented as inert today, which is a record of the problem, not a fix. All references below are against origin/main 1b4e3657.
Current Behavior
Eleven variables read only on dead HIP-graph code.use_hip_graphs() returns a constant false (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device.cpp:46-57), and nothing in mlxcel (src/, src/lib/mlxcel-core/) calls decode_capture_begin/_end_record/_replay/_destroy (device.cpp:1563-1577) or set_graph_decode_mode (device.cpp:1424). The reads (paths relative to patches-rocm/mlx/backend/rocm/): MLX_GRAPH_DECODEdevice.cpp:1419; MLX_GRAPH_PREFILL_REPLAYdevice.cpp:734, :1091, device.h:637; MLX_GRAPH_REPLAY_SLOTSdevice.cpp:1100; MLX_GRAPH_FREE_LAGdevice.cpp:1170; MLX_GRAPH_DEFER_MAX_MBdevice.cpp:1189; MLX_GRAPH_NODEFERallocator.cpp:739 (deferral only when graph_active()); MLX_GRAPH_POISON_FREEallocator.cpp:674 (drains the deferred list, never filled); MLX_GRAPH_SPLIT_LOGdevice.h:677; MLX_PURE_DEBUGdevice.cpp:944, :955 (inside the uncalled decode capture); MLX_NO_CONCAT_SPLITeval.cpp:50 (consulted only when use_hip_graphs(), eval.cpp:62); MLX_MAX_MB_PER_BUFFER through get_graph_limits()device.cpp:116-120, stored in max_mb_per_graph_ and read only by the graph branch of needs_commit() (device.cpp:923-924). docs/environment-variables.md:408 lists them as inert.
MLX_ROCM_NO_HIPBLASLT_EPILOGUE (gemms/hipblaslt_gemm.cpp:471-474) only changes calls that carry a bias or a non-default epilogue. The only such caller is the bias_vec branch of AddMM::eval_gpu (matmul.cpp:1077-1119, requires c.ndim() == 1), which is unreachable: addmm in patches-rocm/mlx/ops.cpp broadcasts c to the output shape before building AddMM (ops.cpp:5951 and :5988-5989). docs/environment-variables.md:328 says so.
MLXCEL_ACTIVATION_MULT (src/execution/memory_estimate.rs:128, default ACTIVATION_BUFFER_MULT = 2.0 at :124, parsed by resolve_activation_mult() at :760-777) has no row in docs/environment-variables.md; it is only named inside the MLXCEL_HEADROOM_FACTOR row (:122).
Proposed Solution
One overlay change plus docs, recorded as one new src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md entry (next free number; stays in mlxcelverse, not proposed to any fork upstream).
Graph knobs: remove each getenv above and fold in the value the code takes when the variable is unset (MLX_GRAPH_DECODE on, MLX_GRAPH_PREFILL_REPLAY off, MLX_GRAPH_REPLAY_SLOTS 4, MLX_GRAPH_DEFER_MAX_MB 2048, debug prints removed, get_graph_limits() returning the 50/200 defaults without reading env). Branches that only an opt-in value enabled (prefill replay capture, poison, split log) are deleted. Keep use_hip_graphs(), the graph build path and the decode-capture API: removing the HIP-graph machinery is a separate decision (see the TrainArena note at device.cpp:52-55). Do not touch MLX_MAX_OPS_PER_BUFFER, which the eager needs_commit() branch uses (device.cpp:911-915).
Epilogue: delete the unreachable bias_vec branch in AddMM::eval_gpu and the MLX_ROCM_NO_HIPBLASLT_EPILOGUE read; keep hipblaslt_gemm_epilogue and its bias_ptr plumbing, and update the comment at gemms/hipblaslt_gemm.h:44 to say it currently has no caller.
Byte cap: add an unsigned 64-bit reader to env_int.h with the same rule and message as env_int_or_default (whole string base 10, reject a leading -, check ERANGE, [ROCm] ignoring invalid NAME="value" (expected ...); using the default N), read once from a function-local static; replace parse_non_negative_size_t_env and delete it.
Docs: delete the inert paragraph (:408) and the MLX_ROCM_NO_HIPBLASLT_EPILOGUE row (:328), fix the MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES row (:315) to describe the warning, and add an MLXCEL_ACTIVATION_MULT row after MLXCEL_HEADROOM_FACTOR (positive finite f64, default 2.0, invalid or non-positive warns and uses the default, affects the same estimator callers).
Acceptance Criteria
git grep on the overlay finds none of the eleven graph variables nor MLX_ROCM_NO_HIPBLASLT_EPILOGUE outside LOCAL_FIXES.md.
tests/rocm_qmm_env.rs gains cases for MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES = -1, 12abc and 18446744073709551616: each prints exactly one warning line and the cache keeps the 256 MiB default; 0 still disables it.
docs/environment-variables.md has an MLXCEL_ACTIVATION_MULT row and no inert-variable paragraph.
Decode throughput on gfx1151 unchanged within run-to-run spread (scripts/bench_decode.sh, one dense and one MoE checkpoint).
Verification
make verify-rocm-overlay
cargo test --profile test-fast --features rocm --test rocm_qmm_env -- --test-threads=1
make verify-rocm
Part of #1801
Problem / Background
PR #2222 documented the runtime
MLX_ROCM_*variables and found four groups of environment knobs in the ROCm overlay that are dead, unsafe to parse, or undocumented. They are documented as inert today, which is a record of the problem, not a fix. All references below are against origin/main1b4e3657.Current Behavior
use_hip_graphs()returns a constantfalse(src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device.cpp:46-57), and nothing in mlxcel (src/,src/lib/mlxcel-core/) callsdecode_capture_begin/_end_record/_replay/_destroy(device.cpp:1563-1577) orset_graph_decode_mode(device.cpp:1424). The reads (paths relative topatches-rocm/mlx/backend/rocm/):MLX_GRAPH_DECODEdevice.cpp:1419;MLX_GRAPH_PREFILL_REPLAYdevice.cpp:734,:1091,device.h:637;MLX_GRAPH_REPLAY_SLOTSdevice.cpp:1100;MLX_GRAPH_FREE_LAGdevice.cpp:1170;MLX_GRAPH_DEFER_MAX_MBdevice.cpp:1189;MLX_GRAPH_NODEFERallocator.cpp:739(deferral only whengraph_active());MLX_GRAPH_POISON_FREEallocator.cpp:674(drains the deferred list, never filled);MLX_GRAPH_SPLIT_LOGdevice.h:677;MLX_PURE_DEBUGdevice.cpp:944,:955(inside the uncalled decode capture);MLX_NO_CONCAT_SPLITeval.cpp:50(consulted only whenuse_hip_graphs(),eval.cpp:62);MLX_MAX_MB_PER_BUFFERthroughget_graph_limits()device.cpp:116-120, stored inmax_mb_per_graph_and read only by the graph branch ofneeds_commit()(device.cpp:923-924).docs/environment-variables.md:408lists them as inert.MLX_ROCM_NO_HIPBLASLT_EPILOGUE(gemms/hipblaslt_gemm.cpp:471-474) only changes calls that carry a bias or a non-default epilogue. The only such caller is thebias_vecbranch ofAddMM::eval_gpu(matmul.cpp:1077-1119, requiresc.ndim() == 1), which is unreachable:addmminpatches-rocm/mlx/ops.cppbroadcastscto the output shape before buildingAddMM(ops.cpp:5951and:5988-5989).docs/environment-variables.md:328says so.MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTESis parsed byparse_non_negative_size_t_env(quantized/qmm.hip:438-452, its only caller isdequant_cache_max_bytes()at:985-989) with barestrtoull:-1wraps toSIZE_MAX(cache cap effectively unlimited),ERANGEis not checked, and junk is ignored without a message. Every other integer knob goes throughenv_int_or_default(env_int.h, fix(rocm): range-check GEMM env integers and fix a stale HIP graph comment #2152 / PR fix(rocm): range-check GEMM env integers and fix a stale graph comment #2193), which warns once and falls back.MLXCEL_ACTIVATION_MULT(src/execution/memory_estimate.rs:128, defaultACTIVATION_BUFFER_MULT = 2.0at:124, parsed byresolve_activation_mult()at:760-777) has no row indocs/environment-variables.md; it is only named inside theMLXCEL_HEADROOM_FACTORrow (:122).Proposed Solution
One overlay change plus docs, recorded as one new
src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.mdentry (next free number; stays in mlxcelverse, not proposed to any fork upstream).getenvabove and fold in the value the code takes when the variable is unset (MLX_GRAPH_DECODEon,MLX_GRAPH_PREFILL_REPLAYoff,MLX_GRAPH_REPLAY_SLOTS4,MLX_GRAPH_DEFER_MAX_MB2048, debug prints removed,get_graph_limits()returning the 50/200 defaults without reading env). Branches that only an opt-in value enabled (prefill replay capture, poison, split log) are deleted. Keepuse_hip_graphs(), the graph build path and the decode-capture API: removing the HIP-graph machinery is a separate decision (see the TrainArena note atdevice.cpp:52-55). Do not touchMLX_MAX_OPS_PER_BUFFER, which the eagerneeds_commit()branch uses (device.cpp:911-915).bias_vecbranch inAddMM::eval_gpuand theMLX_ROCM_NO_HIPBLASLT_EPILOGUEread; keephipblaslt_gemm_epilogueand itsbias_ptrplumbing, and update the comment atgemms/hipblaslt_gemm.h:44to say it currently has no caller.env_int.hwith the same rule and message asenv_int_or_default(whole string base 10, reject a leading-, checkERANGE,[ROCm] ignoring invalid NAME="value" (expected ...); using the default N), read once from a function-local static; replaceparse_non_negative_size_t_envand delete it.:408) and theMLX_ROCM_NO_HIPBLASLT_EPILOGUErow (:328), fix theMLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTESrow (:315) to describe the warning, and add anMLXCEL_ACTIVATION_MULTrow afterMLXCEL_HEADROOM_FACTOR(positive finitef64, default2.0, invalid or non-positive warns and uses the default, affects the same estimator callers).Acceptance Criteria
git grepon the overlay finds none of the eleven graph variables norMLX_ROCM_NO_HIPBLASLT_EPILOGUEoutsideLOCAL_FIXES.md.tests/rocm_qmm_env.rsgains cases forMLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES=-1,12abcand18446744073709551616: each prints exactly one warning line and the cache keeps the 256 MiB default;0still disables it.docs/environment-variables.mdhas anMLXCEL_ACTIVATION_MULTrow and no inert-variable paragraph.scripts/bench_decode.sh, one dense and one MoE checkpoint).Verification
make verify-rocm-overlay cargo test --profile test-fast --features rocm --test rocm_qmm_env -- --test-threads=1 make verify-rocmRelated: #2152, PR #2193, PR #2222, #2232 (changes the cache default, independent), #1140 (env-doc sync check).