perf(rocm): port paged attention (v1 decode, v2 partial, merge) to HIP #2068

Part of #1814 (epic #1801)

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

Problem / Background

The paged-attention kernels (v1 decode, v2 partial, merge) have Metal and CUDA ports and no HIP port, so on ROCm the paged decode paths take the gather fallback, and 36 test skips in make verify-test-rocm (skipping ... lablup/mlxcel#1814) wait on this port. PR #2059 already added the support predicates and the skip macros, so the tests need no edits once the tables are filled.

Current Behavior

All three predicates read has_kernel_port. Rust bridges: paged_attention_kernels_available() (all three), paged_attention_decode_available(), paged_attention_v2_available() (partial and merge), paged_attention_merge_available(). Production gates read them: src/lib/mlxcel-core/src/cache/paged.rs:2006, src/lib/mlxcel-core/src/mla/mod.rs:147, src/lib/mlxcel-core/src/paged_v2/sparse.rs:453, and paged_decode_backend() (src/lib/mlxcel-core/src/layers.rs:5758-5777). The skips live in src/lib/mlxcel-core/src/test_support/kernel_ports.rs (17 require_paged_*_port!() sites in cache/paged_batch_decode_tests.rs, ffi_tests.rs, mla/split_kv_tests.rs, paged_v2/launch_tests.rs).

One gate does not follow the predicate: paged_decode_backend() returns Other unless metal_is_available() or cuda_is_available() is true, so on ROCm a filled v1 table still maps to Other and select_pooled_paged_dispatch (layers.rs:5605-5630) never picks native.

Proposed Solution

1. HIP ports of the three CUDA sources, same kernel names, inputs, outputs, grid and template_args. PAGED_V2_SIMD_WIDTH is 32 (paged_attention_v2.cpp:459), which matches wave32; the wave32 guard is what keeps a wave64 build from silently breaking that assumption. Replace the bool use_cuda in PagedV2PartialHolder and PagedMergeHolder with mlxcel::GpuKernelBackend (src/lib/mlx-cpp/turbo/gpu_backend.h:35) and add a ROCm holder, keeping std::call_once (the server reaches first use concurrently). Fill the three .rocm getters.
2. PagedDecodeBackend::Rocm. Add the variant (layers.rs:5550); paged_decode_backend() returns it when the predicate is true and crate::hardware::gpu_backend_kind() == GpuBackendKind::Rocm. In select_pooled_paged_dispatch it follows the CUDA rule (slab_count <= NATIVE_MAX_SLABS, since the HIP kernel is the CUDA kernel's port), its memo key is 3u64 (:5835-5837), and the gridDim.z <= 65535 guard at cache/paged.rs:2232 applies to Rocm as it does to Cuda. Add select_pooled_paged_dispatch unit tests for Rocm next to the existing ones (layers.rs:9290-9386).
3. Update the "no HIP port yet (perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814)" comments on the tables, paged_decode_backend() and kernel_ports.rs.

Tolerance: the bounds the 34 skipped tests already assert, unchanged.

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 three HIP ports, the two holder refactors, PagedDecodeBackend::Rocm and its tests.

Out of scope: the SDPA tables (turbo4_*, sparse_v_ports()), which are unreachable off macOS through cache::turbo::sparse_v::kernel_enabled(); removing the skip macros themselves (they stay for future ports).

Acceptance Criteria

• All three tables have a .rocm getter with the wave32 guard. (v1 and v2 partial carry it; the merge kernel has no lane-level operation and carries none, per the port requirement's scope. PR update(rocm): port paged attention (v1, v2 partial, merge) to HIP #2103.)
• make verify-test-rocm prints no skipping ... lablup/mlxcel#1814 line, the 34 previously skipped tests pass on gfx1151, and none of them was edited.
• paged_decode_backend() returns Rocm on gfx1151, with unit tests for the new variant.
• Server paged decode throughput on gfx1151, GPU otherwise idle, from scripts/benchmark_paged_decode_production.sh with Llama-3.1-8B-4bit: the fused arm against the MLXCEL_PAGED_ATTENTION_NATIVE=0 arm.
• make verify-kernel-dtype-keys and make verify-kernel-port-dispatch pass with the pin and its count unchanged; Metal and CUDA behavior unchanged. (Metal and CUDA not run on this host; their bodies, template arguments and table entries are unchanged.)

Verification

cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 paged_v2 mla:: cache::paged ffi_tests::test_fused_paged layers::
make verify-kernel-dtype-keys verify-kernel-port-dispatch
make verify-test-rocm 2>&1 | grep -c 'lablup/mlxcel#1814'   # expect 0
make verify-rocm