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
perf(rocm): enable fused add-RMSNorm and RoPE-append by default on ROCm #2145
PR #2107 (closes #2063) ported fused_add_rms_norm and fused_rope_qk_append to HIP. Both ports are byte-identical to the ROCm graph they replace (every on/off logit-trace pair is identical at w1, w8, w1ctx512; 0/230 and 0/274 decided mismatches against Metal). On gfx1151 decode moved within noise: Llama-3.1-8B norm-only 37.85 to 37.95 tok/s (+0.3%), Qwen2.5-7B add / rope / both 46.53 to 47.23 / 46.80 / 47.32 (+1.5 / +0.6 / +1.7%, on beat off in 6 of 7 rounds), Qwen3-30B-A3B control +0.3%. Because the change costs no accuracy and was not slower in any median, the maintainer decided to turn both fusions on by default on ROCm. Metal and CUDA defaults stay off.
Current Behavior
src/lib/mlxcel-core/src/layers.rs:855FUSED_ADD_RMSNORM_DEFAULT: bool = false and :871FUSED_ROPE_APPEND_DEFAULT: bool = false, one value for every backend. The doc comments above them (:819-870) say ROCm keeps the shared default.
fused_add_rmsnorm_enabled() (layers.rs:884) and fused_rope_append_enabled() (:894) read MLXCEL_FUSED_ADD_RMSNORM / MLXCEL_FUSED_ROPE_APPEND once into a OnceLock through fused_flag_enabled_from(value, default) (:908): 0/false/off/no disable, 1/true/on/yes enable, unset or unrecognised takes the default.
Callers: src/models/llama3.rs:1201,1245 (Llama3, Qwen2/2.5 and VLMs on those backbones), src/models/gemma.rs:345, src/models/iquestloopcoder.rs:627 (norm); FusedQKVLinear::forward_fused_rope_append in the Llama3 dense decode path (RoPE).
Proposed Solution
Make the compiled-in default depend on the build, the same way PR #2098 set the ROCm MLXCEL_FUSED_MOE_SGY default to 2. That PR used #ifdef MLXCEL_BRIDGE_ROCM_BACKEND (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:1646-1650 and :3154-3158) because a ROCm build has no Metal or CUDA backend, so the build flag is the backend and no runtime backend-kind comparison is needed. The Rust counterpart is the mlxcel-core Cargo feature rocm (src/lib/mlxcel-core/Cargo.toml:19, enabled by the root feature rocm):
Do not add a gpu_kernel_backend() comparison or a metal_is_available() || cuda_is_available() gate; scripts/ci/check_kernel_port_dispatch.py rejects those, and none is needed. Keep fused_flag_enabled_from unchanged so the env var overrides in both directions. Rewrite both doc comments to state the per-build default and cite the #2107 numbers and the decision.
Implementation Notes
Reuse: fused_flag_enabled_from, the existing OnceLock gates, and the existing tests that read the default rather than a literal (fused_905_flags_follow_the_default_and_respect_explicit_values, layers.rs:9389; fused_norm_parity_tests.rs:404-416).
Edge case, the RoPE bypass notice: Llama 3.1's rope_scaling table routes around the RoPE kernel and report_fused_rope_bypass_once (src/models/llama3.rs:298) prints a one-time stderr notice when a variable asked for the kernel. fused_rope_launcher_requested (llama3.rs:263) asks the gate, not the env, so with the default on the notice would print on every Llama 3.1 run on ROCm. Its doc comment promises "a user who never sets the variables never sees it": make the MLXCEL_FUSED_ROPE_APPEND arm require an explicit truthy value in the environment, and add a test for the unset case.
Edge case, untested families: Gemma and IQuest Loop Coder also take the norm port and have no ROCm checkpoint; the docs must say so.
Constraints: no change on Metal or CUDA builds. greedy_decode_is_token_identical_to_the_unfused_baseline (src/models/llama3_tests.rs:905) must still pass on ROCm, which the ports' byte-identity guarantees.
Acceptance Criteria
On a --features rocm build both fusions run with no env vars set; on Metal and CUDA builds they stay off.
MLXCEL_FUSED_ADD_RMSNORM=0 and MLXCEL_FUSED_ROPE_APPEND=0 turn them off on ROCm; =1 still turns them on elsewhere.
A unit test in layers.rs pins FUSED_ADD_RMSNORM_DEFAULT == cfg!(feature = "rocm") and the same for FUSED_ROPE_APPEND_DEFAULT.
The Llama 3.1 RoPE bypass notice does not appear on a default ROCm run, and still appears with MLXCEL_FUSED_ROPE_APPEND=1.
Docs updated: docs/environment-variables.md:573-574 (Default column becomes "off (Metal, CUDA) / on (ROCm)", ROCm sentences rewritten), the docs/installation.md:556 ROCm table row, the decision section of docs/benchmark_results/rocm-fused-norm-rope-gfx1151-2026-10-05.md, and the note at docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md:280.
Decode re-measured on gfx1151, new default against =0, for Llama-3.1-8B-Instruct-4bit and Qwen2.5-7B-Instruct-4bit: scripts/bench_decode.sh pp512/tg128 through scripts/rocm_gpu_guard.sh, arms alternated, medians of 3, added to the results page.
Verification
cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 layers:: fused_norm_parity_tests fused_rope_parity_tests
cargo test --release --features rocm --lib llama3_tests
make verify-rocm verify-kernel-port-dispatch
On macOS, cargo test --release -p mlxcel-core --lib layers:: must show the default still off.
Part of #1801
Problem / Background
PR #2107 (closes #2063) ported
fused_add_rms_normandfused_rope_qk_appendto HIP. Both ports are byte-identical to the ROCm graph they replace (every on/off logit-trace pair is identical atw1,w8,w1ctx512; 0/230 and 0/274 decided mismatches against Metal). On gfx1151 decode moved within noise: Llama-3.1-8B norm-only 37.85 to 37.95 tok/s (+0.3%), Qwen2.5-7B add / rope / both 46.53 to 47.23 / 46.80 / 47.32 (+1.5 / +0.6 / +1.7%, on beat off in 6 of 7 rounds), Qwen3-30B-A3B control +0.3%. Because the change costs no accuracy and was not slower in any median, the maintainer decided to turn both fusions on by default on ROCm. Metal and CUDA defaults stay off.Current Behavior
src/lib/mlxcel-core/src/layers.rs:855FUSED_ADD_RMSNORM_DEFAULT: bool = falseand:871FUSED_ROPE_APPEND_DEFAULT: bool = false, one value for every backend. The doc comments above them (:819-870) say ROCm keeps the shared default.fused_add_rmsnorm_enabled()(layers.rs:884) andfused_rope_append_enabled()(:894) readMLXCEL_FUSED_ADD_RMSNORM/MLXCEL_FUSED_ROPE_APPENDonce into aOnceLockthroughfused_flag_enabled_from(value, default)(:908):0/false/off/nodisable,1/true/on/yesenable, unset or unrecognised takes the default.src/models/llama3.rs:1201,1245(Llama3, Qwen2/2.5 and VLMs on those backbones),src/models/gemma.rs:345,src/models/iquestloopcoder.rs:627(norm);FusedQKVLinear::forward_fused_rope_appendin the Llama3 dense decode path (RoPE).Proposed Solution
Make the compiled-in default depend on the build, the same way PR #2098 set the ROCm
MLXCEL_FUSED_MOE_SGYdefault to 2. That PR used#ifdef MLXCEL_BRIDGE_ROCM_BACKEND(src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:1646-1650and:3154-3158) because a ROCm build has no Metal or CUDA backend, so the build flag is the backend and no runtime backend-kind comparison is needed. The Rust counterpart is themlxcel-coreCargo featurerocm(src/lib/mlxcel-core/Cargo.toml:19, enabled by the root featurerocm):Do not add a
gpu_kernel_backend()comparison or ametal_is_available() || cuda_is_available()gate;scripts/ci/check_kernel_port_dispatch.pyrejects those, and none is needed. Keepfused_flag_enabled_fromunchanged so the env var overrides in both directions. Rewrite both doc comments to state the per-build default and cite the #2107 numbers and the decision.Implementation Notes
fused_flag_enabled_from, the existingOnceLockgates, and the existing tests that read the default rather than a literal (fused_905_flags_follow_the_default_and_respect_explicit_values,layers.rs:9389;fused_norm_parity_tests.rs:404-416).rope_scalingtable routes around the RoPE kernel andreport_fused_rope_bypass_once(src/models/llama3.rs:298) prints a one-time stderr notice when a variable asked for the kernel.fused_rope_launcher_requested(llama3.rs:263) asks the gate, not the env, so with the default on the notice would print on every Llama 3.1 run on ROCm. Its doc comment promises "a user who never sets the variables never sees it": make theMLXCEL_FUSED_ROPE_APPENDarm require an explicit truthy value in the environment, and add a test for the unset case.greedy_decode_is_token_identical_to_the_unfused_baseline(src/models/llama3_tests.rs:905) must still pass on ROCm, which the ports' byte-identity guarantees.Acceptance Criteria
--features rocmbuild both fusions run with no env vars set; on Metal and CUDA builds they stay off.MLXCEL_FUSED_ADD_RMSNORM=0andMLXCEL_FUSED_ROPE_APPEND=0turn them off on ROCm;=1still turns them on elsewhere.layers.rspinsFUSED_ADD_RMSNORM_DEFAULT == cfg!(feature = "rocm")and the same forFUSED_ROPE_APPEND_DEFAULT.MLXCEL_FUSED_ROPE_APPEND=1.docs/environment-variables.md:573-574(Default column becomes "off (Metal, CUDA) / on (ROCm)", ROCm sentences rewritten), thedocs/installation.md:556ROCm table row, the decision section ofdocs/benchmark_results/rocm-fused-norm-rope-gfx1151-2026-10-05.md, and the note atdocs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md:280.=0, for Llama-3.1-8B-Instruct-4bit and Qwen2.5-7B-Instruct-4bit:scripts/bench_decode.shpp512/tg128 throughscripts/rocm_gpu_guard.sh, arms alternated, medians of 3, added to the results page.Verification
On macOS,
cargo test --release -p mlxcel-core --lib layers::must show the default still off.