Skip to content

update(rocm): enable fused add-RMSNorm and RoPE-append by default on ROCm - #2189

Merged
inureyes merged 4 commits into
mainfrom
update/issue-2145-rocm-fused-defaults
Oct 7, 2026
Merged

inureyes merged 4 commits into
mainfrom
update/issue-2145-rocm-fused-defaults

Conversation

@inureyes

@inureyes inureyes commented Oct 6, 2026 •

Copy link
Copy Markdown
Member

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_DEFAULT and FUSED_ROPE_APPEND_DEFAULT in src/lib/mlxcel-core/src/layers.rs are now cfg!(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, so verify-kernel-port-dispatch passes. fused_flag_enabled_from keeps its behavior: =0 disables and =1 enables on every backend.
  • The new layers::fused_flag_explicit_value returns the setting a value spells out, or none. The Llama 3.1 rope_scaling bypass notice (fused_rope_launcher_requested in src/models/llama3.rs) now uses it, so the notice fires only for an explicit truthy MLXCEL_FUSED_ROPE_APPEND and default ROCm runs stay quiet.
  • Tests: fused_905_defaults_are_on_for_rocm_builds_only pins both defaults to cfg!(feature = "rocm") and checks the override in each direction. fused_flag_explicit_value_separates_a_setting_from_the_default covers the parser, and the_rope_append_bypass_notice_needs_an_explicit_truthy_value covers the unset case.
  • scripts/rocm_decode_profile.py now 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.
  • Docs: environment-variables.md (Default column now reads "off (Metal, CUDA) / on (ROCm)"), the ROCm row in installation.md, benchmarks.md, the correctness-page note, and a new Decision section with the re-measurement on rocm-fused-norm-rope-gfx1151-2026-10-05.md. Raw CSVs are in benchmarks/rocm_strixhalo-gfx1151_2026-10-07_fused-default-{on,off}.csv.

Re-measurement (gfx1151, pp512/tg128, medians of 3, arms alternated)

Model Default (on) =0 (off) Change
Llama-3.1-8B-Instruct-4bit 38.09 38.15 -0.2% (noise; paired +0.02 / -0.49 / -0.05)
Qwen2.5-7B-Instruct-4bit 47.90 47.51 +0.8% (paired +0.39 / +0.30 / -0.04)

Both changes are within run-to-run noise, so neither is a speedup claim. Every run went through scripts/rocm_gpu_guard.sh with 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_trace w8 for both models, default vs both variables at 0: the files are byte-identical (cmp).
  • rocprofv3 kernel trace with no variables set: Qwen2.5 dispatches fused_add_rms_norm and fused_rope_qk_append 252 times each and no graph RoPE kernel. Llama 3.1 dispatches the norm port 288 times and keeps its graph RoPE.
  • Bypass notice: absent with the default on Llama 3.1, present with 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, including greedy_decode_is_token_identical_to_the_unfused_baseline. That greedy test also passes with both variables at 0.
  • make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt: pass.
  • make verify-rocm at 9874c3c4 (rebased on fc21784f): 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 false as 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

…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
@inureyes inureyes added type:enhancement New features, capabilities, or significant additions priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:linux Linux (CUDA / packaging) specific status:review Under review labels Oct 6, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Oct 7, 2026
@inureyes
inureyes merged commit bbd0509 into main Oct 7, 2026
26 checks passed
@inureyes
inureyes deleted the update/issue-2145-rocm-fused-defaults branch October 7, 2026 00:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:linux Linux (CUDA / packaging) specific priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf(rocm): enable fused add-RMSNorm and RoPE-append by default on ROCm

1 participant