update(rocm): port paged attention (v1, v2 partial, merge) to HIP

Makefile

@@ -795,9 +795,9 @@ verify-rocm-smoke: ## ROCm smoke: build, link and generate on the GPU, asserting
 # runs BEFORE `verify-test-rocm` because make stops at the first failing
 # prerequisite, and the ROCm test gate can still carry known failures while
 # the backend is experimental. (The mlxcel-core tests of the fused
-# paged-attention kernels, which have no ROCm port yet, no longer fail: they
-# skip through `test_support::kernel_ports` and print a `skipping ... #1814`
-# line, #1809.) With the test gate first the
+# paged-attention kernels used to skip here through
+# `test_support::kernel_ports`, printing a `skipping ... #1814` line, #1809;
+# since the HIP ports, #2068, they run.) With the test gate first the
 # smoke was never reached at all, which was measured rather than reasoned: the
 # first full run of this target ended without the smoke having executed once.
 # The smoke is also the fastest high-signal part, so a developer gets the

docs/benchmark_results/rocm-paged-attention-gfx1151-2026-10-05.md

@@ -0,0 +1,40 @@
+# HIP port of the paged-attention kernels on gfx1151 (2026-10-05)
+
+lablup/mlxcel#2068, part of #1814. Before this change the three paged-attention port tables (v1 decode, v2 partial, merge) had no `.rocm` entry, so on ROCm the server's batched paged decode, MLA split-KV and the sparse paged decode all ran gather-then-SDPA. The ports (`src/lib/mlx-cpp/turbo/paged_attention_hip.h`) are the CUDA bodies compiled through hipRTC, with `__shfl_xor(v, o, 32)` for the lane fold. The decode profile ([rocm-decode-profile-gfx1151-2026-09-30.md](rocm-decode-profile-gfx1151-2026-09-30.md)) put paged attention at 0% of decode because `bench_decode.sh` decodes into a dense KV cache; the path these kernels serve is the server's paged decode, so that is what is measured here.
+
+## Environment
+
+AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0. mlxcel `f4ca4b9a` (the PR branch on `origin/main` `c05d5438`), release build with `--features rocm`. The PR's later commits change only host-side shape checks, comments and the autotune runner label, not the kernel bodies or the dispatch.
+
+## Method
+
+The matrix of `scripts/benchmark_paged_decode_production.sh` (issue #899) with `Meta-Llama-3.1-8B-Instruct-4bit`: `mlxcel-server --parallel 4 --ctx-size 131072`, `scripts/bench_serving_concurrency.py` with 128 decode tokens per request, and the same binary for both arms:
+
+- before: `MLXCEL_PAGED_ATTENTION_NATIVE=0`, the gather-then-SDPA path, which is also what every ROCm build ran before this change;
+- after: unset, the fused v2 kernel.
+
+Each case ran with a freshly started server per arm. The before and after arm of one case ran back to back inside one `scripts/rocm_gpu_guard.sh` window (55 s with `/sys/class/kfd/kfd/proc` empty and no compiler; the host was shared with other port units, and an attempt that saw a foreign GPU process or a compiler was rejected and rerun). The server log of every after arm announces `paged decode v2: fused v2 launch`, and no before arm does. The batch-4 cases ran three times; the two single-sequence cases once each, because a guarded window long enough for them was rare on the shared host.
+
+## Results
+
+`decode tok/s` is the mean per-request decode rate (completion tokens after the first, over the time after the first token); `aggregate tok/s` is all completion tokens over the level's wall-clock span, which includes prefill.
+
+| Case | Decode tok/s, before | Decode tok/s, after | Decode, after / before (medians) | Aggregate tok/s, before | Aggregate tok/s, after |
+|---|---|---|---|---|---|
+| batch 4, ~1K prompt | 5.7 / 5.5 / 5.6 | 6.1 / 6.3 / 6.0 | 1.09x | 19.4 / 19.2 / 19.4 | 20.7 / 20.6 / 20.7 |
+| batch 4, ~4K prompt | 10.1 / 10.6 / 10.2 | 10.4 / 10.6 / 10.5 | 1.03x (within run spread) | 10.8 / 11.2 / 10.9 | 11.0 / 11.3 / 11.2 |
+| batch 4, ~16K prompt | 10.9 / 10.9 / 10.9 | 13.9 / 13.9 / 13.9 | 1.28x | 1.9 / 2.0 / 2.0 | 1.7 / 1.7 / 1.7 |
+| batch 1, ~16K prompt | 17.2 | 24.1 | 1.40x (one run) | 2.0 | 1.7 |
+| batch 1, ~32K prompt | 11.2 | 19.8 | 1.77x (one run) | 0.3 | 0.3 |
+
+Time to first token is unchanged within a few percent in every case (prefill does not use these kernels): for example 155.1 s before and 151.1 s after (medians) at batch 4 and ~16K.
+
+What the numbers say:
+
+- The fused path is faster at decode wherever the context is long, and the gain grows with context: 1.28x at batch 4 and ~16K tokens per request, 1.40x and 1.77x for a single sequence at ~16K and ~32K. Run-to-run spread at batch 4 is at most 0.5 tok/s, so the ~16K and ~1K gains are outside it; the ~4K gain (1.03x) is not and is reported as no measurable change.
+- The aggregate column does not follow decode at long context because it is dominated by prefill (the mean time to first token alone is about 150 s at batch 4 and ~16K) and depends on how many tokens each request generated before it stopped, which differs between the two numeric paths. It is not a measure of the kernel.
+- Batch-4 decode on this server is far below the single-stream `bench_decode.sh` rate for the same model (35 tok/s): the per-request rate includes time spent waiting while the other three requests prefill. That is the server's scheduling, not this change, and is the same in both arms.
+
+## Correctness
+
+The 36 mlxcel-core tests that skipped on ROCm for lack of these ports (`skipping ... lablup/mlxcel#1814`) run and pass on gfx1151 with no test edits: the v1 decode against the gather path over 200 steps and a GQA and batch matrix, the v2 partial and merge against host references and the gather path (f32 and f16 pools, empty requests, trimmed windows, GQA head mapping, two pool dtypes at one geometry), cascade, sparse and MLA split-KV. `make verify-rocm` reports 11869 passed, 0 failed. With the HIP lane fold started at 8 instead of 16 and the HIP merge computed in base e instead of base 2, 29 of those tests fail.

docs/environment-variables.md

@@ -520,7 +520,7 @@ recommended as normal deployment settings.
 | `MLXCEL_DISABLE_SOFTCAP_GQA_DECODE_GROUPED` | `1` disables, `0` enables | unset | Legacy rollback knob for grouped softcap-GQA decode, checked before the enable variable above. Set to `1` to restore the repeat-based path. |
 | `MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS` | truthy disables | maskless path on | Disables the single-query maskless attention path. |
 | `MLXCEL_EXPERIMENTAL_BOOL_CAUSAL_MASK` | truthy enables | off | Enables an experimental boolean causal-mask path. |
-| `MLXCEL_PAGED_ATTENTION_NATIVE` | `1`/`true`/`on`/`yes` force the native kernel; `0`/`false`/`off`/`no` force gather (case-insensitive); unset or any other value defers to the adaptive selector | selector-governed (server: the issue #899 token floors; library entry point on Metal: native only for batch>=4 + ctx<=4096 + single-slab; on CUDA: native for any single-slab layer since #634) | Overrides the fused paged-attention decode kernel (Metal since epic #116 Phase 6/#123, and since #634 also CUDA via `mx.fast.cuda_kernel`). **Two consumers**, per `resolve_dispatch_decision` and `resolve_paged_v2_dispatch` in `src/lib/mlxcel-core/src/layers.rs`: (1) the library-only `paged_decode_attention_pooled` entry point, where since #331 an unset value no longer means "always gather" (`select_pooled_paged_dispatch` picks the kernel only inside the regime ADR 0001 measured it winning) and this variable force-pins either arm for A/B testing; and (2) **the server's pool-backed batched paged decode**, where issue #899 made the fused v2 kernel the production path and named this variable's force-off values its kill switch: setting one restores the pre-#899 gather-then-SDPA behaviour end to end, and a force-on value pins v2 for every servable shape, bypassing the measured token floors (`MLXCEL_PAGED_V2_MIN_KV_TOKENS`, `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST`). The override is checked before the selector in both consumers, and a forced dispatch still goes through the kernel's structural declines (single-slab layer, servable geometry, non-empty batch). History: #710 retired the *library* entry point to a library-only API, which is where the "not a server knob" reading came from; #899 gave the variable its second, server-side consumer. The separate `DecodeBatchContext::use_native_paged_kernel` request still governs the dense-compat block-table decode, a different path. See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md) and [Continuous batching](CONTINUOUS_BATCHING.md#seeing-which-path-ran). |
+| `MLXCEL_PAGED_ATTENTION_NATIVE` | `1`/`true`/`on`/`yes` force the native kernel; `0`/`false`/`off`/`no` force gather (case-insensitive); unset or any other value defers to the adaptive selector | selector-governed (server: the issue #899 token floors; library entry point on Metal: native only for batch>=4 + ctx<=4096 + single-slab; on CUDA: native for any single-slab layer since #634; on ROCm: the CUDA rule since #2068) | Overrides the fused paged-attention decode kernel (Metal since epic #116 Phase 6/#123, since #634 also CUDA via `mx.fast.cuda_kernel`, and since #2068 also ROCm via `fast::hip_kernel`). **Two consumers**, per `resolve_dispatch_decision` and `resolve_paged_v2_dispatch` in `src/lib/mlxcel-core/src/layers.rs`: (1) the library-only `paged_decode_attention_pooled` entry point, where since #331 an unset value no longer means "always gather" (`select_pooled_paged_dispatch` picks the kernel only inside the regime ADR 0001 measured it winning) and this variable force-pins either arm for A/B testing; and (2) **the server's pool-backed batched paged decode**, where issue #899 made the fused v2 kernel the production path and named this variable's force-off values its kill switch: setting one restores the pre-#899 gather-then-SDPA behaviour end to end, and a force-on value pins v2 for every servable shape, bypassing the measured token floors (`MLXCEL_PAGED_V2_MIN_KV_TOKENS`, `MLXCEL_PAGED_V2_MIN_KV_TOKENS_PER_REQUEST`). The override is checked before the selector in both consumers, and a forced dispatch still goes through the kernel's structural declines (single-slab layer, servable geometry, non-empty batch). History: #710 retired the *library* entry point to a library-only API, which is where the "not a server knob" reading came from; #899 gave the variable its second, server-side consumer. The separate `DecodeBatchContext::use_native_paged_kernel` request still governs the dense-compat block-table decode, a different path. See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md) and [Continuous batching](CONTINUOUS_BATCHING.md#seeing-which-path-ran). |
 | `MLXCEL_PAGED_ATTENTION_V2` | `1`/`true`/`on`/`yes` enable (case-insensitive, trimmed); unset or anything else disables | off | **Diagnostic.** Issue #898's comparison gate: makes the **library-only** v1 entry point (`PagedBlockPool::paged_decode_fused`, reached from `paged_decode_attention_pooled`) try the fused v2 kernel first. It does **not** gate the server's production v2 decode, which is issue #899, is default-on, and is governed by the token floors and `MLXCEL_PAGED_ATTENTION_NATIVE` instead. Read once per process, and a typo degrades to v1 rather than enabling a path nobody asked for. See [ADR 0001](adr/0001-paged-attention-gather-vs-fused-kernel.md). |
 | `MLXCEL_PAGED_DECODE_V2_CHUNK` | positive integer pages | unset (planned or autotuned chunk size) | **Diagnostic.** Pins `pages_per_chunk` for every paged decode v2 launch, overriding both the plan's occupancy heuristic and any autotuned value (issue #906), matching the `MLXCEL_PAGED_DECODE_SPLITS` / `MLXCEL_QMM_TILE_*` escape-hatch convention. The plan clamps the value into the feasible range, so an over-large value is safe; a non-positive or unparseable value warns and is ignored in favour of the planned size. Read once per process. |
 | `MLXCEL_PAGED_DECODE_V2_TARGET_CTAS` | positive integer | derived (Apple: `gpu_core_count * 8`, floored at 64; every other host including CUDA: `512`) | **Diagnostic.** Moves the occupancy target the v2 plan binary-searches its chunk size against, as opposed to `MLXCEL_PAGED_DECODE_V2_CHUNK`, which pins the chunk size itself. This is the knob to turn when the derived device parallelism is wrong. The Apple figure is the `gpu_core_count` device-scale proxy, not a calibrated core count, and the fixed non-Apple target is **unvalidated on CUDA** (a CUDA-specific target should come from the SM count), so both are starting points. A non-positive or unparseable value warns and uses the derived target. Read once per process. |

docs/installation.md

@@ -553,7 +553,8 @@ so `mlxcel inspect` and `--estimate-memory` give the same answers as before.
 | NVFP4 checkpoints | No native kernel. ModelOpt NVFP4 checkpoints (for example the Gemma 4 NVFP4 exports) are converted to affine 4-bit at load with no environment variables, and the load log names the route and the reason. A layer that cannot be converted fails the load with its name and the reason. MLX-native NVFP4 exports (`"mode": "nvfp4"` in `config.json`, such as `mlx-community/*-nvfp4`) have no load-time conversion and are refused at load; use an affine export instead (lablup/mlxcel#1806). |
 | Affine MoE models (for example Qwen3-30B-A3B) | Run natively. Single-token decode takes the fused two-kernel MoE path (gate-up, then down) through its HIP port (lablup/mlxcel#2065), as on Metal and CUDA: Qwen3-30B-A3B-4bit decode on `gfx1151` went from 61.3 to 62.5 tok/s, with logits matching Metal's fused kernel at every decided position ([results](benchmark_results/rocm-fused-moe-gfx1151-2026-10-05.md)); `MLXCEL_FUSED_MOE=0` selects the `gather_qmm` graph path instead. Experts wider than `MLXCEL_FUSED_MOE_MAX_DFF` (8192 by default, as on CUDA; Mixtral-8x7B's 14336) stay on `gather_qmm`. |
 | Model-specific fused kernels (Apertus xIELU, Cohere2 add3 + LayerNorm, Mamba1 selective scan, Nemotron-H's opt-in `MLXCEL_FUSED_MOE_RELU2`) | Run as HIP ports (lablup/mlxcel#2069). xIELU and the add3 LayerNorm are byte-identical to ROCm's own unfused graph; the Mamba1 scan keeps its state in f32 as on Metal; the relu2 path is performance-neutral on Nemotron-H and stays opt-in. Apertus, Cohere2 and the Mamba1 families have no checkpoint on the tested host, so they are covered by kernel tests only ([results](benchmark_results/rocm-metal-only-ports-gfx1151-2026-10-05.md)). |
-| mlxcel's other fused kernels (fused norm, RoPE + KV append, paged attention) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. |
+| mlxcel's other fused kernels (fused norm, RoPE + KV append) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. |
+| Paged attention (the server's batched paged decode, MLA split-KV, sparse paged decode) | The v1 decode, v2 partial and merge kernels run as HIP ports (lablup/mlxcel#2068), so these paths take the fused kernels instead of gather-then-SDPA, as on Metal and CUDA; the kernels' tests match the gather and host references on `gfx1151`. Server decode on `gfx1151` with Meta-Llama-3.1-8B-Instruct-4bit, against the gather path on the same binary: 1.28x per-request decode at batch 4 and ~16K tokens, 1.40x and 1.77x for one sequence at ~16K and ~32K, no measurable change at batch 4 and ~4K ([results](benchmark_results/rocm-paged-attention-gfx1151-2026-10-05.md)). `MLXCEL_PAGED_ATTENTION_NATIVE=0` forces the gather path. |
 | Sampled decode (`temperature > 0`) | The fused Gumbel-max sampler (no filter) and the dual-pivot rejection sampler (top-p, and the other filter combinations its routing admits) run as HIP kernels, as on Metal and CUDA (lablup/mlxcel#2064). On `gfx1151` with Meta-Llama-3.1-8B-Instruct-4bit, decode at `--temperature 0.8 --top-p 0.95` went from 33.98 to 37.56 tok/s (medians of three; the before runs drifted from 36.26 down to 32.81, the after runs held 37.06 to 37.82); at `--temperature 0.8` alone the change is inside the noise ([results](benchmark_results/rocm-samplers-gfx1151-2026-10-05.md)). `MLXCEL_SAMPLING_GUMBEL=0` and `MLXCEL_SAMPLING_REJECTION=0` force the graph. |
 | Hybrid SSM decode (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) | The single-token Mamba2 SSM update runs the fused HIP kernel instead of the ~55-op SSD graph (lablup/mlxcel#2067): decode on `gfx1151` went from 60.4 to 88.4 tok/s on granite-4.0-h-tiny-4bit and from 51.4 to 74.4 tok/s on Nemotron-3-Nano-30B-A3B-4bit, with model logits matching the graph and Metal's kernel at every decided position ([results](benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md)). Measured on granite and Nemotron-H; falcon-h1 and plamo-2 take the same gate but have not been run on ROCm. `MLXCEL_SSM_KERNEL=0` forces the graph. |
 | GPU faults | Reported as errors (lablup/mlxcel#1804). A launch HIP rejects (an oversized block, no code object for the device) fails the evaluation that issued it and the device stays usable. An asynchronous fault (an out-of-bounds access) fails the evaluation waiting on it within about a second instead of hanging or returning NaN; the server fails that request and answers later ones with the same error, because after a queue fault the HIP runtime rejects every call for the rest of the process, so the process has to restart, and its shutdown may need a SIGKILL (HIP's teardown waits on callbacks the faulted queue never runs). A kernel that never finishes is not detected; `MLX_ROCM_GPU_WATCHDOG_SECS` (default off) fails any single host wait that outlives it. |

examples/paged_attention_kernel_bench.rs

@@ -62,13 +62,20 @@ const WARMUP: usize = 20;
 const ITERS: usize = 50;
 
 /// Backend the fused kernel actually dispatches on: Metal on Apple, the #634
-/// CUDA port on NVIDIA, else CPU (no native kernel). Drives the selector label
-/// so the printed `select=` column matches production dispatch on this host.
+/// CUDA port on NVIDIA, the #2068 HIP port on ROCm, else CPU (no native
+/// kernel). Drives the selector label so the printed `select=` column matches
+/// production dispatch on this host.
 fn bench_backend() -> PagedDecodeBackend {
-    if mlxcel_core::metal_is_available() {
+    if !mlxcel_core::paged_attention_decode_available() {
+        PagedDecodeBackend::Other
+    } else if mlxcel_core::metal_is_available() {
         PagedDecodeBackend::Metal
     } else if mlxcel_core::cuda_is_available() {
         PagedDecodeBackend::Cuda
+    } else if mlxcel_core::hardware::gpu_backend_kind()
+        == mlxcel_core::hardware::GpuBackendKind::Rocm
+    {
+        PagedDecodeBackend::Rocm
     } else {
         PagedDecodeBackend::Other
     }

src/lib/mlx-cpp/turbo/CMakeLists.txt

@@ -40,6 +40,7 @@ set(MLXCEL_TURBO_KERNEL_HEADERS
     ${CMAKE_CURRENT_SOURCE_DIR}/turbo4_delegated_sdpa.h
     ${CMAKE_CURRENT_SOURCE_DIR}/paged_attention.h
     ${CMAKE_CURRENT_SOURCE_DIR}/paged_attention_v2.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/paged_attention_hip.h
     ${CMAKE_CURRENT_SOURCE_DIR}/sampling.h
     ${CMAKE_CURRENT_SOURCE_DIR}/sampling_rejection.h
     ${CMAKE_CURRENT_SOURCE_DIR}/fused_norm.h