Repository navigation
update(rocm): enable fused add-RMSNorm and RoPE-append by default on ROCm - #2189
Merged
Merged
Conversation
…ROCm FUSED_ADD_RMSNORM_DEFAULT and FUSED_ROPE_APPEND_DEFAULT become cfg!(feature = "rocm"): on in a ROCm build, off on Metal and CUDA. The HIP ports from #2063 are byte-identical to the graph and were never slower on gfx1151. The env vars still override in both directions. The Llama 3.1 rope_scaling bypass notice now fires only for an explicit truthy MLXCEL_FUSED_ROPE_APPEND, so default ROCm runs stay silent. New layers::fused_flag_explicit_value separates a setting from the default. Tests pin the per-build defaults and the notice decision; docs updated. Closes #2145
This was referenced Oct 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Turns the fused add-RMSNorm and RoPE + KV-append kernels on by default in a ROCm build, as the maintainer decided after #2107. Metal and CUDA defaults are unchanged.
FUSED_ADD_RMSNORM_DEFAULTandFUSED_ROPE_APPEND_DEFAULTinsrc/lib/mlxcel-core/src/layers.rsare nowcfg!(feature = "rocm"), following the build-flag pattern PR update(rocm): port the fused MoE decode kernels to HIP #2098 used for the ROCm SGY default. There is no runtime backend comparison, soverify-kernel-port-dispatchpasses.fused_flag_enabled_fromkeeps its behavior:=0disables and=1enables on every backend.layers::fused_flag_explicit_valuereturns the setting a value spells out, or none. The Llama 3.1rope_scalingbypass notice (fused_rope_launcher_requestedinsrc/models/llama3.rs) now uses it, so the notice fires only for an explicit truthyMLXCEL_FUSED_ROPE_APPENDand default ROCm runs stay quiet.fused_905_defaults_are_on_for_rocm_builds_onlypins both defaults tocfg!(feature = "rocm")and checks the override in each direction.fused_flag_explicit_value_separates_a_setting_from_the_defaultcovers the parser, andthe_rope_append_bypass_notice_needs_an_explicit_truthy_valuecovers the unset case.scripts/rocm_decode_profile.pynow reports the perf(rocm): port fused_add_rms_norm and fused_rope_qk_append to HIP #2063 roles as reached with the shipped defaults on ROCm, and its test is updated to match. The README's ROCm bullet no longer calls these fusions opt-in.environment-variables.md(Default column now reads "off (Metal, CUDA) / on (ROCm)"), the ROCm row ininstallation.md,benchmarks.md, the correctness-page note, and a new Decision section with the re-measurement onrocm-fused-norm-rope-gfx1151-2026-10-05.md. Raw CSVs are inbenchmarks/rocm_strixhalo-gfx1151_2026-10-07_fused-default-{on,off}.csv.Re-measurement (gfx1151, pp512/tg128, medians of 3, arms alternated)
=0(off)Both changes are within run-to-run noise, so neither is a speedup claim. Every run went through
scripts/rocm_gpu_guard.shwith the host-wide lock from #2146 (main's copy after #2161 merged). All 12 runs were clean on the first attempt.Verification
All on gfx1151:
logit_tracew8for both models, default vs both variables at0: the files are byte-identical (cmp).rocprofv3kernel trace with no variables set: Qwen2.5 dispatchesfused_add_rms_normandfused_rope_qk_append252 times each and no graph RoPE kernel. Llama 3.1 dispatches the norm port 288 times and keeps its graph RoPE.MLXCEL_FUSED_ROPE_APPEND=1, absent with=0.cargo test --profile test-fast --features rocm -p mlxcel-core --lib -- --test-threads=1 layers::tests::fused fused_norm_parity_tests fused_rope_parity_tests: 32 passed.cargo test --profile test-fast --features rocm --lib -- --test-threads=1 models::llama3::: 27 passed, includinggreedy_decode_is_token_identical_to_the_unfused_baseline. That greedy test also passes with both variables at0.make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt: pass.make verify-rocmat9874c3c4(rebased onfc21784f): OK. Clippy and the ROCm smoke test passed, and the tests had 11996 passed, 0 failed and 383 ignored across 152 test binaries.python3 -m unittest tests.test_rocm_decode_profile: 26 passed.Not verified: Metal and CUDA (not available on this host). On those builds both constants evaluate to
falseas before, and the notice change only narrows when it prints. Gemma and IQuest Loop Coder now take the norm port by default on ROCm, but no checkpoint for either is available here.Closes #2145