Skip to content

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

Closed

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).

Activity

  1. added
    type:choreMaintenance tasks (build, CI, etc.)
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    platform:linuxLinux (CUDA / packaging) specific
    on Oct 8, 2026
  2. added and removed on Oct 8, 2026
  3. added a commit that references this issue on Oct 12, 2026
    88934d9
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