Repository navigation
chore(rocm): remove dead overlay env knobs and range-check the dequant byte cap #2242
Copy link
Copy link
Closed
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:lowLow priorityLow prioritystatus:in-progressCurrently being worked onCurrently being worked ontype:choreMaintenance tasks (build, CI, etc.)Maintenance tasks (build, CI, etc.)
Description
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/main 1b4e3657.
Current Behavior
- Eleven variables read only on dead HIP-graph code.
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.md entry (next free number; stays in mlxcelverse, not proposed to any fork upstream).
- Graph knobs: remove each
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). - Epilogue: delete the unreachable
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. - Byte cap: add an unsigned 64-bit reader to
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. - Docs: delete the inert paragraph (
: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. - 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-rocmRelated: #2152, PR #2193, PR #2222, #2232 (changes the cache default, independent), #1140 (env-doc sync check).
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificLinux (CUDA / packaging) specificpriority:lowLow priorityLow prioritystatus:in-progressCurrently being worked onCurrently being worked ontype:choreMaintenance tasks (build, CI, etc.)Maintenance tasks (build, CI, etc.)
Activity
chore(rocm): remove dead overlay env knobs and range-check the byte c…