Skip to content

chore(rocm): remove dead overlay env knobs and range-check the dequant byte cap #2242

Description

@inureyes

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

  1. 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_DECODE device.cpp:1419; MLX_GRAPH_PREFILL_REPLAY device.cpp:734, :1091, device.h:637; MLX_GRAPH_REPLAY_SLOTS device.cpp:1100; MLX_GRAPH_FREE_LAG device.cpp:1170; MLX_GRAPH_DEFER_MAX_MB device.cpp:1189; MLX_GRAPH_NODEFER allocator.cpp:739 (deferral only when graph_active()); MLX_GRAPH_POISON_FREE allocator.cpp:674 (drains the deferred list, never filled); MLX_GRAPH_SPLIT_LOG device.h:677; MLX_PURE_DEBUG device.cpp:944, :955 (inside the uncalled decode capture); MLX_NO_CONCAT_SPLIT eval.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.
  2. 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.
  3. MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES is parsed by parse_non_negative_size_t_env (quantized/qmm.hip:438-452, its only caller is dequant_cache_max_bytes() at :985-989) with bare strtoull: -1 wraps to SIZE_MAX (cache cap effectively unlimited), ERANGE is not checked, and junk is ignored without a message. Every other integer knob goes through env_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.
  4. 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

Related: #2152, PR #2193, PR #2222, #2232 (changes the cache default, independent), #1140 (env-doc sync check).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificpriority:lowLow prioritystatus:in-progressCurrently being worked ontype:choreMaintenance tasks (build, CI, etc.)

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions