Repository navigation
update(rocm): port xIELU, add3 norm, Mamba1 and relu2 kernels to HIP - #2101
Merged
Merged
Conversation
Nemotron-H's opt-in MLXCEL_FUSED_MOE_RELU2 branch needs the fc1 + relu² kernel and the down kernel; ROCm had only the down port (#2065), so the flag declined to gather_qmm there. This adds MOE_FC1_RELU2_HIP_SOURCE, written from the Metal source in the shape of the #2065 gate-up port (one 32-lane wavefront per row, the 16..1 fold with __shfl_down width 32, the wave32 guard, T in the template args), a hip_kernel holder and the table's .rocm getter. CUDA stays without a port. The branch now reads a new bridge predicate, fused_moe_relu2_kernels_available(), built from both kernels' tables, and its rows per block default to 2 on ROCm as run_fused_moe_two_kernel's do, so both kernels of the branch use the geometry #2065 measured. fused_moe_relu2_parity_tests runs fused_moe_forward with the flag unset and set at Nemotron-3-Nano shapes, 4 and 8 bit. On gfx1151 the kernel branch is within 8.1e-6 nrms / 4.2e-4 nmax of a dense f32 reference and within gather_qmm's own jitter of gather_qmm; starting the fold at 8 instead of 16 fails it at 0.76 nrms. Refs #2069
fused_xielu returned its elementwise fallback whenever Metal was absent, so Apertus never reached a kernel on ROCm. The early return now reads fused_xielu_kernel_available(), a new bridge predicate built from xielu_ports(), and the table gains a .rocm getter for XIELU_HIP_SOURCE. The HIP kernel keeps the Metal kernel's goal, byte identity with apertus_xielu, against the ROCm graph's rounding points: each op there is one HIP elementwise kernel that widens to float, computes once and rounds to T, and Expm1 calls the device expm1f. The port rounds every intermediate through T the same way, calls the same expm1f, propagates NaN through minimum as the graph does, and turns off FP contraction so the f32 path cannot fuse x * beta into the final add. New test apertus_tests::fused_xielu_kernel_matches_graph_every_dtype compares the kernel with apertus_xielu in f32, f16 and bf16 within the fused_norm_parity_tests budget and asserts the kernel path is taken on Metal and ROCm. On gfx1151 all three dtypes are byte-identical (0 of 16410 elements differ), and the existing bf16 bit-for-bit test now exercises the HIP kernel. Using alpha_n in the positive branch fails both tests. No Apertus checkpoint is on the host, so decode was not measured. Refs #2069
ROCm had neither Mamba1 selective-scan variant, so Jamba, Mamba and Falcon-Mamba walked the per-step graph scan. This ports the Metal float32-state variant (MAMBA1_SCAN_HIP_SOURCE, holder, mamba1_scan_ports().rocm). The CUDA graph-exact variant is not reachable on ROCm: the graph scan's state @ C has K = N (8 or 16), below gemv's K % 32 == 0, so the overlay sends it to rocBLAS, whose reduction order a custom kernel cannot reproduce. mamba.rs gated the kernel on the Metal backend kind; it now reads a new bridge predicate, mamba1_scan_float_state_kernel_available(), built from the float32-state table, so Mamba and Falcon-Mamba take the port on ROCm too while CUDA keeps the graph scan there. Jamba already went through mamba1_scan_kernel_accepts. The kernel reads the sequence length from X_shape, which faulted: the fork's hip_kernel declared <input>_shape and _strides as pointers while the launch passes them by value. custom_kernel.cpp now declares them as by-value 8-entry arrays, as upstream CUDA does (LOCAL_FIXES item 30, an upstreaming candidate for #1813). On gfx1151 mamba1_scan_parity_tests pass (f32 against the scalar reference within 1e-5; bf16 no less accurate than the graph scan, now run on ROCm too), as do the Mamba and Jamba prefill tests that exercise the kernel. The f32 test gains N = 32, without which a fold starting at 8 still passed; with it that break fails at 0.87 relative error. A new test asserts the predicate is true on Metal and ROCm so these tests cannot skip silently. No Mamba1 or Jamba checkpoint is on the host. Refs #2069
residual_add3_layer_norm required metal_is_available(), so Cohere2 ran compiled_add3 plus fast::layer_norm on ROCm. It now reads a new bridge predicate, fused_add3_layer_norm_available(), built from add3_layer_norm_ports(), which gains a .rocm getter for ADD3_LN_HIP_SOURCE; the expect message names the predicate. The contract is byte identity with the same backend's unfused pair, so the HIP kernel follows ROCm's pair, not the Metal kernel's reduction: the residual is rounded to T after each add as the overlay's compiled Add does, and the norm reproduces the overlay's layer_norm_kernel<T, 256, 4> (256 threads per row, strided groups of 4, 16..1 __shfl_xor folds with width 32, the eight wavefront sums folded again, 1.0f / sqrtf, T(w * norm + b) with the bias read from memory). ROCm builds launch 256 threads per row through MLXCEL_BRIDGE_ROCM_BACKEND, as run_fused_moe_two_kernel sets its ROCm default. residual_add3_layer_norm_matches_the_unfused_pair was compiled only with the metal feature; it now builds everywhere, asserts the kernel path on Metal and ROCm, and adds f32, widths 1025 and 5, and the 6656 limit. On gfx1151 all ten cases are byte-identical; replacing 1.0f / sqrtf with rsqrtf fails it. Cohere2 model tests pass. No Cohere2 checkpoint is on the host. Refs #2069
With MLXCEL_DEVICE=cpu the predicates fused_xielu_kernel_available, fused_add3_layer_norm_available, mamba1_scan_float_state_kernel_available and fused_moe_relu2_kernels_available read only the port table, so these paths built a custom kernel on the CPU stream and threw "Custom kernels only run on GPU" instead of running the graph. Each now also requires the GPU default device, as ssm_kernel_available and mamba1_scan_kernel_accepts do. Metal had the same exposure before (metal::is_available and a backend-kind check say nothing about the device); its GPU behavior is unchanged. tests/cpu_device_custom_kernel_gates.rs, a separate binary so the shared lib test binaries never see the default device move, checks the four predicates and runs fused_xielu and residual_add3_layer_norm on the CPU device; it fails with the device term removed. Mamba's fused-scan gate now also requires mamba1_scan_kernel_accepts, as Jamba's does, so a state wider than 32 takes the graph scan instead of leaving columns unwritten. On ROCm fused_xielu reads the element count from x_shape[0] instead of a template argument, so hipRTC compiles one kernel per dtype rather than one per activation size. Metal's template arguments are unchanged. fused_xielu_kernel_matches_graph_every_dtype now asserts zero differing bits on ROCm, and residual_add3_layer_norm_matches_the_unfused_pair skips visibly when MLXCEL_FUSED_ADD_NORM=0. Comments no longer call the ROCm relu2 branch byte-identical, scope fused_xielu's identity to the dtypes tested per backend, and say add3 identity is pinned by the test rather than by matching compiler flags. Verified on gfx1151 (--release --features rocm, --test-threads=1): xIELU tests (0 differing bits in f32, f16, bf16), residual_add3_layer_norm, mamba1_scan_parity_tests, models::mamba::, models::jamba::, fused_moe_relu2_parity_tests, the new CPU-device test, clippy -D warnings on mlxcel-core and mlxcel, and make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt. Metal and CUDA not verified (not available on this host). Refs #2069
fused_moe_kernels_available() and moe_down_kernel_available() (#2065) read only the port tables, so under MLXCEL_DEVICE=cpu a ROCm build sent SwitchGLU decode into the HIP launch, whose eval_cpu throws after the bridge call has already returned Ok. Both now also require the GPU as the default device, as the four #2069 predicates do; Nemotron-H then takes forward_nonfused on the CPU. Metal had the same exposure through custom_kernels_available() before #2065, and its GPU behavior does not change. cpu_device_custom_kernel_gates asserts both predicates are false on the CPU device. Refs #2069
Add docs/benchmark_results/rocm-metal-only-ports-gfx1151-2026-10-05.md: what each HIP port is held to and its negative check, the alternated Nemotron-H decode runs with MLXCEL_FUSED_MOE_RELU2=1 (75.15 to 75.24 tok/s, neutral), the relu2 greedy divergence with teacher-forced traces showing it lands only on near-ties (0 decided mismatches against ROCm default and Metal), and the CPU-device gating. Raw rows go under benchmarks/, the two traces under benchmarks/logit_traces/rocm_gfx1151_9186b075/. installation.md gains a row for the four ports, and the MLXCEL_FUSED_MOE_RELU2 row states the ROCm numbers. Refs #2069
inureyes
changed the title
inureyes
added
status:review
Bilingual (en, ko) report on the HIP ports of the xIELU, add3 LayerNorm, Mamba1 scan and fc1 relu2 kernels, the CPU-device gating and the measurements. Refs #2069
10 tasks
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
1 participant
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.
Closes #2069
Four kernels had a Metal port only and sat behind gates that did not read their tables. Each now has a HIP source, a filled
.rocmentry and a predicate that reads its table, and its gate reads that predicate. CUDA keeps its fallbacks.Changes
moe_fc1_relu2: HIP port of the fc1 + relu² kernel.fused_moe_relu2_kernels_available()gates theMLXCEL_FUSED_MOE_RELU2branch offused_moe_forward, whose rows per block default to 2 on ROCm as perf(rocm): port the fused MoE decode kernels (moe_gateup, moe_down) #2065 measured. The flag stays opt-in.fused_xielu: the early return readsfused_xielu_kernel_available()instead ofmetal::is_available(). The HIP kernel follows the ROCm graph's rounding (one rounding toTper op, the deviceexpm1f, no FP contraction), so it is byte-identical toapertus_xieluin f32, f16 and bf16. On ROCm the element count comes from the input shape, so hipRTC compiles one kernel per dtype, not one per length.mamba1_selective_scan: HIP port of the Metal float32-state variant. The CUDA graph-exact variant cannot be reproduced on ROCm (the graph'sstate @ Cruns through rocBLAS).mamba.rsnow readsmamba1_scan_float_state_kernel_available()andmamba1_scan_kernel_accepts, so Mamba and Falcon-Mamba take the port and a state wider than 32 falls back.fused_add3_layer_norm:residual_add3_layer_normreadsfused_add3_layer_norm_available()instead ofmetal_is_available(). The HIP kernel reproduces ROCm's own pair (rounded residual adds, the overlay'slayer_norm_kernel<T, 256, 4>), so the byte-identity test holds unchanged; ROCm builds launch 256 threads per row.fused_moe_kernels_available()/moe_down_kernel_available()also require the GPU as the default device, soMLXCEL_DEVICE=cputakes the graph paths instead of throwing "Custom kernels only run on GPU". New test binarytests/cpu_device_custom_kernel_gates.rs.fused_moe_relu2_parity_testsandfused_xielu_kernel_matches_graph_every_dtype; the add3 test module now builds on every backend and adds f32, widths 5 and 1025, and the 6656 limit; the Mamba1 f32 test adds N = 32. No tolerance was loosened.TECHNICAL_REPORTS/2101-*; results pagedocs/benchmark_results/rocm-metal-only-ports-gfx1151-2026-10-05.md, raw rows underbenchmarks/, traces underbenchmarks/logit_traces/rocm_gfx1151_9186b075/, README,docs/installation.mdanddocs/environment-variables.mdrows.grep -rn '\.rocm = nullptr' --include=*.cpp src/goes from 16 to 12; the one left inmlx_cxx_kernels.cppis the graph-exact Mamba1 table. A fork bug that faulted the Mamba1 port (by-value<input>_shape) was fixed independently by #2100 (LOCAL_FIXES item 30); this branch dropped its own copy of that fix on rebase.Measured (gfx1151)
Only Nemotron-H has a checkpoint here; Apertus, Cohere2, Mamba, Falcon-Mamba and Jamba have none and are covered by kernel tests only. Nemotron-3-Nano-30B-A3B-4bit with
MLXCEL_FUSED_MOE_RELU2=1,bench_decode.shpp512/tg128, before (main6668031, the flag declines togather_qmm) and after alternated, every run through the GPU guard: 75.15 / 75.22 / 75.09 (median 75.15) to 75.24 / 75.38 / 75.13 (median 75.24) tok/s, +0.1%, within noise.Deviation from the issue: relu2 greedy text
The issue asked for the relu2 greedy 128-token output to match the default path. It does not: on three prompts the texts agree for about 20 tokens, then part at a near-tie ("user query" against "user request"). The kernel keeps fc1, relu² and fc2 in f32 where
gather_qmmrounds to bf16 (8.1e-6 nrms from the f32 reference againstgather_qmm's 5.1e-3 to 5.4e-3). Teacher-forcedw1ctx512traces show 0 decided-position mismatches against the ROCm default (1/128 top-1, largest gap 0.5) and against Metal M5 (5/128, largest gap 0.25). The default path is byte-identical tomain.Verification (gfx1151)
make verify-rocmon code head800e23dc(rebased on6668031c; later commits are docs only): 147 suites, 11867 passed, 0 failed, 378 ignored.dead_doc_pointersre-run after the docs commits: passed.mamba1_scan_parity_tests5 passed;fused_moe_relu2_parity_testswithin bounds;cpu_device_custom_kernel_gatespasses and fails with the device term removed.models::mamba::andmodels::jamba::22 passed, 2 ignored; Cohere2 model tests 13 passed, 7 ignored.Not verified: Metal and CUDA (not available on this host). Metal's kernel sources and table entries are untouched; Metal now runs the widened add3 and Mamba1 tests and the new relu2 test, and its predicates gained the GPU-device term (GPU behavior unchanged;
MLXCEL_DEVICE=cpunow falls back instead of throwing). CUDA: the fc1_relu2, xIELU, add3 and float-state Mamba1 tables stay null there, and the add3 test compares the unfused pair with itself. Wave64 (CDNA) untested.