perf(rocm): port Metal-only kernels (xielu, add3 norm, mamba1, relu2) #2069

Part of #1814 (epic #1801)

Lowest priority of the #1814 ports. The moe_fc1_relu2 part depends on #2065 (HIP moe_down and the fused_relu2 table guard); the other three parts have no dependency.

Problem / Background

Four kernels have a Metal port only (.cuda = nullptr, .rocm = nullptr), so a HIP port has to be written from Metal source, and each sits behind a support check that does not read its table. They run only on narrow model families (Apertus, Cohere2, Mamba1 and Jamba, and an opt-in Nemotron-H path), which is why they come last.

Current Behavior

All in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:

Proposed Solution

Each kernel is independent and may land as its own PR. For each: a *_HIP_SOURCE written from the Metal source (translating simdgroup operations to wave32 shuffles), a HIP holder, a filled .rocm, and the gate change below. .cuda stays nullptr, so CUDA keeps its current fallback.

1. fused_xielu. Replace the early return's condition with !mlxcel::has_kernel_port(xielu_ports()), and update the comment at apertus.rs:261-263. Parity: new test in src/models/apertus_tests.rs comparing mlxcel_core::fused_xielu with apertus_xielu (apertus.rs:136) on random f32, f16 and bf16 inputs, within the fused_norm_parity_tests.rs:57 budget (f32 1e-6 / 1e-5, f16 2e-3 / 1.2e-2, bf16 1.6e-2 / 7e-2 normalized RMS / max).
2. fused_add3_layer_norm. New bridge predicate fused_add3_layer_norm_available() returning has_kernel_port(add3_layer_norm_ports()); layers.rs:1125 reads it and the expect message at :1157 names it. The existing test residual_add3_layer_norm_matches_the_unfused_pair (layers.rs:9630) asserts byte-identical output against the unfused pair on the same backend, so take the kernel's structure from the Metal source and its reduction order from the overlay's src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/layer_norm.hip, and cover dims up to FUSED_ADD3_LAYER_NORM_MAX_DIM (6656). If byte identity cannot be reached, stop and report it on this issue; do not loosen the test.
3. mamba1_selective_scan. mamba1_scan_kernel_available() returns false when MLXCEL_MAMBA1_SCAN_KERNEL=0 on any platform, otherwise has_kernel_port(mamba1_scan_ports()); remove the #ifdef __APPLE__ split. Parity: the existing src/lib/mlxcel-core/src/mamba1_scan_parity_tests.rs (f32 against a scalar reference, bf16 no less accurate than the graph scan), which starts running on ROCm.
4. moe_fc1_relu2. HIP port only; the default stays off. Parity: Nemotron-3-Nano-30B-A3B-4bit greedy 128-token output with MLXCEL_FUSED_MOE_RELU2=1 matches the default path on gfx1151, as it does on Metal.

Port requirements (the same for every port split from #1814)

Filling a .rocm slot follows bitlinear_matmul (#1862, PR #1870), the one kernel with all three ports, in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp: a HIP source string (:159), a holder that calls mlx::core::fast::hip_kernel(...) under #ifdef MLXCEL_BRIDGE_ROCM_BACKEND and throws otherwise (:214-237), and a .rocm getter in the KernelPorts table (:240-255). fast::hip_kernel compiles through hipRTC at runtime, like cuda_kernel, so there is no ahead-of-time build cost.

• Wave32 guard. gfx1151 is wave32; CDNA parts (gfx90a, gfx942) are wave64, where a fold that starts at lane 16 drops half the lanes and still returns a finite, plausible, wrong result. Every ported kernel that reduces or shuffles across lanes carries both #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32 plus #error and the same for the non-underscored __AMDGCN_WAVEFRONT_SIZE, as at mlx_cxx_kernels.cpp:170-175. static_assert(warpSize == 32) does not compile in HIP. Shuffles pass an explicit width of 32.
• Dtype keys. Every input dtype is in template_args, so the hipRTC cache key ("custom_kernel_" + name + template_arguments_hash(template_args), src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp:231) differs per dtype. make verify-kernel-dtype-keys passes with EXPECTED_IN_SCOPE (scripts/ci/check_kernel_dtype_keys.py:111) and the 9 in scope control count asserted in scripts/ci/check_kernel_dtype_keys_test.sh:113 both unchanged, because the HIP launch stays in the file that already holds the CUDA launch. If a launch has to live in a new file, pin that file and update the count in the same change.
• Predicate in the same change. The kernel's support predicate returns mlxcel::has_kernel_port(<its table>), and every gate in front of the kernel reads that predicate, so the port is reachable and the predicate and the dispatch cannot disagree. The dispatch keeps mlxcel::select_kernel_port, and make verify-kernel-port-dispatch passes.
• Parity. On gfx1151 the port matches the graph fallback it replaces, on the same inputs, within the tolerance named in this issue, not only against the CUDA arm.
• Measured. Decode (or prefill, where stated) tok/s before and after with scripts/bench_decode.sh on an otherwise idle GPU (90 s with /sys/class/kfd/kfd/proc empty and no compiler running, as in docs/benchmark_results/rocm-baseline-gfx1151-2026-09-30.md), compared with scripts/compare_bench_csv.py --before <old.csv> --after <new.csv>, with the numbers in the PR.
• Metal and CUDA unchanged. The .metal and .cuda entries and their sources are untouched, and the predicate still answers true on both.

Scope

In scope: the four HIP ports and the three gate changes above, the fused_xielu parity test.

Out of scope: CUDA ports of these kernels; making MLXCEL_FUSED_MOE_RELU2 a default.

Acceptance Criteria

• All four tables have a .rocm getter with the wave32 guard, and grep -rn '\.rocm = nullptr' --include=*.cpp src/ reports four fewer lines.
• grep -n 'metal::is_available' src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp no longer matches in fused_xielu or mamba1_scan_kernel_available, and residual_add3_layer_norm no longer calls metal_is_available.
• On gfx1151: the new xielu test, residual_add3_layer_norm_matches_the_unfused_pair (taking the kernel path) and mamba1_scan_parity_tests pass, and the relu2 greedy check matches.
• Decode tok/s before and after on gfx1151, GPU otherwise idle, for each family with a checkpoint on the host; any family without one is named in the PR.
• make verify-kernel-dtype-keys and make verify-kernel-port-dispatch pass with the pin and its count unchanged; Metal and CUDA behavior unchanged.

Verification

cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 mamba1_scan_parity_tests residual_add3_layer_norm
cargo test --release --features rocm -p mlxcel --lib -- --test-threads=1 apertus
make verify-kernel-dtype-keys verify-kernel-port-dispatch
make verify-rocm