perf(rocm): port ssm_update_kernel to HIP and fix ssm_kernel_available #2067

Part of #1814 (epic #1801)

Depends on #2061 (the decode profile sets the port order).

Problem / Background

On ROCm the Mamba2 layers of the hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) run the ~55-op SSD graph for every single-token decode step, where Metal and CUDA run one fused ssm_update_kernel launch. The CUDA port measured 2.6-4.5x decode on the granite/falcon family (docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md). The strided_scan fix in PR #2059 (LOCAL_FIXES.md item 22) made granite and Nemotron-H run on gfx1151, so they are now testable.

Current Behavior

• ssm_ports() (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:639) has .metal (SSM_METAL_SOURCE :499) and .cuda (SSM_CUDA_SOURCE :570, holder get_ssm_kernel_cuda() :631), .rocm = nullptr; dispatch at :1625.
• ssm_kernel_available() (:1576-1590) does not read the table: on Apple it returns metal::is_available(), elsewhere it honors MLXCEL_SSM_CUDA_KERNEL=0 and returns mlx::core::cu::is_available(). A filled .rocm would stay unreachable.
• Rust gates read it: src/models/granitemoehybrid.rs:440, src/models/nemotron_h.rs:537 and :635, and the falcon-h1 and plamo-2 mixers (grep -rn ssm_kernel_available src/models). Each then calls ssm_step_kernel and expects the launch, so predicate and dispatch must agree.
• No test compares the kernel with the graph step; parity so far was end-to-end greedy output on CUDA.

Proposed Solution

1. Predicate. ssm_kernel_available() returns false when MLXCEL_SSM_KERNEL=0 or MLXCEL_SSM_CUDA_KERNEL=0 (the old name kept as an alias so existing scripts still work), and otherwise mlxcel::has_kernel_port(ssm_ports()) on every platform. Remove the #ifdef __APPLE__ split. Update the MLXCEL_SSM_CUDA_KERNEL row in docs/environment-variables.md:527 and add MLXCEL_SSM_KERNEL.
2. HIP port. SSM_HIP_SOURCE and get_ssm_kernel_hip() ported from SSM_CUDA_SOURCE, same kernel name, inputs, outputs, grid and template_args; fill .rocm.
3. Parity test. New src/lib/mlxcel-core/src/ssm_update_parity_tests.rs (registered in lib.rs like fused_norm_parity_tests.rs) that builds random single-token inputs at the granite-4.0-h-tiny and Nemotron-H head shapes, runs ssm_update_kernel and a Rust reference of the same recurrence composed from MLX ops (the math in granitemoehybrid.rs ssm_step, :474), and compares output and new state. It returns early only when ssm_kernel_available() is false. Tolerance, normalized RMS / max against the reference: f32 1e-5 / 1e-4, bf16 1.6e-2 / 7e-2 (the bf16 budget of fused_norm_parity_tests.rs:57). The same test then covers Metal and CUDA.

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: ssm_kernel_available(), the HIP port, the parity test, the environment-variable doc.

Out of scope: mamba1_selective_scan (Mamba1 and Jamba, a sibling issue under #1814); metal_gated_delta_forward_impl (unreachable off Apple).

Acceptance Criteria

• ssm_kernel_available() reads has_kernel_port(ssm_ports()); grep -n 'cu::is_available' src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp no longer matches inside it.
• ssm_ports() has a .rocm getter with the wave32 guard; ssm_update_parity_tests passes on gfx1151 for f32 and bf16.
• granite-4.0-h-tiny-4bit and Nemotron-3-Nano-30B-A3B-4bit decode tok/s on gfx1151 before (MLXCEL_SSM_KERNEL=0) and after, GPU otherwise idle.
• A new ROCm granite w1 trace shows 0 decided-position disagreements at --decided 2.0 against benchmarks/logit_traces/rocm_gfx1151_c5fe9a16/ (the graph-path trace). Nemotron-H's w1 trace has no decided position (the comparison reports inconclusive), so for it report the top-1 agreement from the same comparison and the greedy 128-token output with and without MLXCEL_SSM_KERNEL=0.
• make verify-kernel-dtype-keys and make verify-kernel-port-dispatch pass with the pin and its count unchanged; Metal and CUDA behavior unchanged with the environment unset.

Verification

cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 ssm_update_parity_tests
cargo build --release --features rocm --example logit_trace
./target/release/examples/logit_trace models/mlx/granite-4.0-h-tiny-4bit tests/fixtures/wikitext2_excerpt.txt 1 128 8 0 > rocm_granite_w1.tsv
python3 scripts/compare_logit_traces.py benchmarks/logit_traces/rocm_gfx1151_c5fe9a16/rocm_gfx1151_c5fe9a16_granite-4.0-h-tiny_default_w1.tsv rocm_granite_w1.tsv --decided 2.0
MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/granite-4.0-h-tiny-4bit
make verify-kernel-dtype-keys verify-kernel-port-dispatch
make verify-rocm

Progress

Implemented in #2099. Metal and CUDA could not be run on the development host (gfx1151 only); the last criterion is met by construction there and stated as not verified in the PR.