Skip to content

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

Description

@inureyes

Part of #1801

Problem / Background

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:855 FUSED_ADD_RMSNORM_DEFAULT: bool = false and :871 FUSED_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):

pub(crate) const FUSED_ADD_RMSNORM_DEFAULT: bool = cfg!(feature = "rocm");
pub(crate) const FUSED_ROPE_APPEND_DEFAULT: bool = cfg!(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.

Activity

  1. added
    type:enhancementNew features, capabilities, or significant additions
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    platform:linuxLinux (CUDA / packaging) specific
    on Oct 6, 2026
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:mediumMedium prioritystatus:doneCompletedtype:enhancementNew features, capabilities, or significant additions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions