From 16b6b25d58bfa3fddd683b71675317e5d3ee5844 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 03:13:17 +0900 Subject: [PATCH 1/6] update(rocm): port paged attention (v1, v2 partial, merge) to HIP On ROCm the three paged-attention tables had no `.rocm` entry, so every paged decode took the gather-then-SDPA fallback and 36 test skips in the ROCm gate waited on the port. This fills the three slots with hipRTC ports of the CUDA bodies (`paged_attention_hip.h`), launched from the files that already hold the CUDA launches. - The HIP bodies use `__shfl_xor(v, o, 32)` and carry the wave32 `#error` guard where they fold across lanes (v1, v2 partial). The merge body has no lane-level operation and no guard. - The vendored `fast::hip_kernel` faults on `_shape` until #2100, so the HIP bodies read the head counts and block size from new `NumQHeads`/`NumKVHeads`/`PoolBlockSize` template arguments (passed on every backend; Metal and CUDA never read them) and the merge head count from `gridDim.y`. `static_assert`s keep `_shape`/`_strides`/`_ndim` out of the HIP text. - `PagedV2PartialHolder` and `PagedMergeHolder` take a `GpuKernelBackend` instead of `bool use_cuda`. - The launchers refuse shapes the bodies cannot index (rank, head dim) on every backend. - `PagedDecodeBackend::Rocm` follows the CUDA rule in `select_pooled_paged_dispatch`, has memo tag 3, and keeps the `gridDim.z <= 65535` guard. On gfx1151 the previously skipped tests run and pass with no test edits; with the lane fold started at 8 and the merge in base e, 29 of them fail. Refs #2068 --- Makefile | 6 +- docs/environment-variables.md | 2 +- src/lib/mlx-cpp/turbo/CMakeLists.txt | 1 + src/lib/mlx-cpp/turbo/paged_attention.cpp | 67 ++- src/lib/mlx-cpp/turbo/paged_attention.h | 4 +- src/lib/mlx-cpp/turbo/paged_attention_hip.h | 465 ++++++++++++++++++ src/lib/mlx-cpp/turbo/paged_attention_v2.cpp | 112 +++-- src/lib/mlx-cpp/turbo/paged_attention_v2.h | 2 +- .../turbo/paged_attention_v2_merge.cpp | 90 +++- src/lib/mlxcel-core/build.rs | 2 + src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 2 +- src/lib/mlxcel-core/src/cache/paged.rs | 15 +- src/lib/mlxcel-core/src/layers.rs | 124 ++++- src/lib/mlxcel-core/src/lib.rs | 3 +- .../src/test_support/kernel_ports.rs | 20 +- 15 files changed, 821 insertions(+), 94 deletions(-) create mode 100644 src/lib/mlx-cpp/turbo/paged_attention_hip.h diff --git a/Makefile b/Makefile index f554584d2..4147502a9 100644 --- a/Makefile +++ b/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 diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 719e42cc2..bc4812311 100644 --- a/docs/environment-variables.md +++ b/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. | diff --git a/src/lib/mlx-cpp/turbo/CMakeLists.txt b/src/lib/mlx-cpp/turbo/CMakeLists.txt index 97384d0d1..be25199e0 100644 --- a/src/lib/mlx-cpp/turbo/CMakeLists.txt +++ b/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 diff --git a/src/lib/mlx-cpp/turbo/paged_attention.cpp b/src/lib/mlx-cpp/turbo/paged_attention.cpp index 889009e06..a64a443fa 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention.cpp @@ -16,6 +16,7 @@ #include #include "gpu_backend.h" #include "kernel_port.h" +#include "paged_attention_hip.h" #include #include @@ -366,7 +367,8 @@ constexpr const char* PAGED_ATTENTION_DECODE_CUDA_SOURCE = R"( // Apple Silicon SIMD width. Each SIMD group is 32 lanes that partition the head // dimension; `NumSplits` SIMD groups split the token range. On CUDA the same -// constant is the warp width. +// constant is the warp width, and on ROCm the wavefront width of the wave32 +// targets the HIP port is built for (`paged_attention_hip.h`). constexpr int PAGED_ATTENTION_SIMD_WIDTH = 32; // Thread-safe lazy-initialised holder for the JIT-compiled kernel. Mirrors the @@ -420,6 +422,38 @@ inline PagedAttentionKernelHolderCuda& get_paged_attention_kernel_cuda() { return holder; } +// HIP counterpart (issue #2068): the CUDA port compiled through hipRTC, body in +// `paged_attention_hip.h`. On a build without the ROCm backend +// `fast::hip_kernel` is not declared, and `paged_attention_ports()` never +// resolves this entry there (`select_kernel_port` reads the running backend), +// so the throw is unreachable and exists only to keep the table shape uniform. +struct PagedAttentionKernelHolderHip { + std::optional kernel; + std::once_flag init_flag; + + mlx::core::fast::CustomKernelFunction& get() { + std::call_once(init_flag, [this] { +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + kernel = mlx::core::fast::hip_kernel( + "mlxcel_paged_attention_decode", + {"q", "k_pool", "v_pool", "rows", "row_offsets", "logical_starts", + "visible_lens", "scale"}, + {"out"}, + std::string(PAGED_ATTENTION_DECODE_HIP_SOURCE)); +#else + throw std::runtime_error( + "[paged_attention_decode] this build has no ROCm backend"); +#endif + }); + return *kernel; + } +}; + +inline PagedAttentionKernelHolderHip& get_paged_attention_kernel_hip() { + static PagedAttentionKernelHolderHip holder; + return holder; +} + // This kernel's ports, in one place. `has_kernel_port` and // `select_kernel_port` both read it, so a support predicate and the dispatch @@ -432,8 +466,9 @@ const mlxcel::KernelPorts& paged_attention_ports() { .cuda = +[]() -> mlx::core::fast::CustomKernelFunction& { return get_paged_attention_kernel_cuda().get(); }, - // No HIP port yet (#1814). - .rocm = nullptr, + .rocm = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_paged_attention_kernel_hip().get(); + }, }; return ports; } @@ -475,6 +510,21 @@ mlx::core::array paged_attention_decode( const auto& q_shape = q.shape(); // [B, Hq, 1, Dim] const auto& kp_shape = k_pool.shape(); // [num_blocks, block_size, Hkv, Dim] + // Shapes the kernel bodies cannot index are refused here, on every + // backend, instead of being read out of bounds (issue #2068; #2067 did the + // same for the SSM update kernel). The bridge declares this function + // `Result`, so the throw reaches Rust as an `Err`. + // `v_pool` may be narrower than `k_pool` in heads (MiniMax-M3 keeps an + // index key at head `Hkv` of the K allocation; the sparse launch accepts + // that layout at batch 1), so only its rank and head dimension are checked. + if (q.ndim() != 4 || k_pool.ndim() != 4 || v_pool.ndim() != 4 || + q_shape[2] != 1 || kp_shape[3] != q_shape[3] || + v_pool.shape(3) != q_shape[3] || kp_shape[1] < 1 || kp_shape[2] < 1) { + throw std::invalid_argument( + "[paged_attention_decode] expects q [B, Hq, 1, D] and k_pool, v_pool " + "[blocks, block_size, heads, D]"); + } + int batch = q_shape[0]; int hq = q_shape[1]; int dim = q_shape[3]; @@ -520,11 +570,22 @@ mlx::core::array paged_attention_decode( // int-only args, two callers that share this geometry but differ in pool // dtype get one compiled module and the later one reads through the wrong // pointer type. Metal already folds the dtypes into its key. + // + // `NumQHeads`/`NumKVHeads`/`PoolBlockSize` carry the geometry the Metal and + // CUDA bodies read from `q_shape`/`k_pool_shape`, for the HIP body, which + // cannot read `_shape` (see `paged_attention_hip.h`). They are + // passed on every backend, because a launch that branches on the backend + // kind is what `make verify-kernel-port-dispatch` rejects; the Metal and + // CUDA bodies never name them, so there they only extend the module name + // with values that are fixed per model (issue #2068). std::vector> template_args = { {"Dim", dim}, {"NRep", n_rep}, {"DimsPerThread", dims_per_thread}, {"NumSplits", num_splits}, + {"NumQHeads", hq}, + {"NumKVHeads", hkv}, + {"PoolBlockSize", kp_shape[1]}, {"QType", q.dtype()}, {"KVType", k_pool.dtype()}, {"VType", v_pool.dtype()}, diff --git a/src/lib/mlx-cpp/turbo/paged_attention.h b/src/lib/mlx-cpp/turbo/paged_attention.h index a2c365102..3c684f4b3 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention.h +++ b/src/lib/mlx-cpp/turbo/paged_attention.h @@ -90,8 +90,8 @@ int paged_attention_num_splits_cap(int dim); // True when the resolved GPU backend has a port of the v1 paged decode kernel // above. Read from the same `KernelPorts` table `paged_attention_decode` -// dispatches through, so the two cannot disagree. Metal and CUDA today; ROCm -// answers false until lablup/mlxcel#1814 ports it. +// dispatches through, so the two cannot disagree. True on Metal, CUDA and ROCm +// (the HIP port, lablup/mlxcel#2068). bool paged_attention_decode_available(); } // namespace mlxcel::turbo diff --git a/src/lib/mlx-cpp/turbo/paged_attention_hip.h b/src/lib/mlx-cpp/turbo/paged_attention_hip.h new file mode 100644 index 000000000..cba8fb9b5 --- /dev/null +++ b/src/lib/mlx-cpp/turbo/paged_attention_hip.h @@ -0,0 +1,465 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// HIP bodies of the three paged-attention kernels (issue #2068): the v1 split-K +// decode, the v2 partial and the v2 merge. +// +// Kept apart from the launchers only for size. Each launch, and the +// `fast::hip_kernel` call that compiles its string, stays in the `.cpp` that +// already holds the CUDA launch (`paged_attention.cpp`, +// `paged_attention_v2.cpp`, `paged_attention_v2_merge.cpp`), so +// `make verify-kernel-dtype-keys` keeps its pinned scope. This header holds +// data, not a launch. +// +// Each body is a line-for-line port of the CUDA body next to its launcher: +// same thread mapping, inputs, outputs, grid and template arguments. The ROCm +// `CustomKernel::eval_gpu` ceil-divides the Metal-style total-thread grid by +// the threadgroup tuple exactly as CUDA does, so every `blockIdx` / `threadIdx` +// below means what it means in the CUDA body. The differences are spelling: +// +// - The 32-lane all-reduce is `__shfl_xor(v, o, 32)`. `__shfl_xor_sync` exists +// in HIP only as a compatibility shim that ignores its mask, so the native +// form is used with the width stated. +// - Negative infinity is `-__builtin_huge_valf()`. hipRTC compiles these +// bodies with only the headers `fast::hip_kernel` prepends, and `INFINITY` is +// a `` macro that is not guaranteed among them. +// - The KV reads keep the CUDA text's explicit `(float)`: `hip_bfloat16`, the +// type MLX substitutes for bfloat16 on ROCm, converts to float only through +// an `explicit` operator. +// - The geometry the CUDA bodies read from `_shape` comes from template +// arguments instead (`NumQHeads`, `NumKVHeads`, `PoolBlockSize`, which the +// launchers pass on every backend) or, for the merge kernel's head count, +// from `gridDim.y`. The vendored `fast::hip_kernel` declares +// `_shape` as a pointer while the launch passes the shape by value, so +// a HIP body that names one faults the queue (fixed by lablup/mlxcel#2100, +// LOCAL_FIXES item 30). The `static_assert`s at the end of this file keep +// every such name out of these bodies, comments included, because the +// overlay decides whether to declare the argument by searching the source +// text. +// +// Wave32 guard. The v1 and v2-partial bodies fold a dot product across 32 +// lanes with an XOR butterfly that starts at 16. They carry the `#error` guard +// every #1814 port carries, spelled as the bitlinear and SSM ports spell it +// (`static_assert(warpSize == 32)` does not compile in HIP). The guard is inert +// with ROCm 10's AMD clang, which defines neither macro for gfx1151, gfx942 or +// gfx90a (#2067). The fold is written not to need it: the 32 lanes that share a +// `threadIdx.y` are consecutive in the block's linear order, so on a wave64 +// target the explicit width of 32 should keep each butterfly inside its own +// row (not run: this host has no wave64 device). The +// merge body has no lane-level operation (one thread per output element, no +// shuffle and no barrier), so it is correct for any wavefront size and carries +// no guard: an `#error` there would only reject a correct kernel. + +#include + +namespace mlxcel::turbo { + +// v1 split-K flash-decoding (port of `PAGED_ATTENTION_DECODE_CUDA_SOURCE`). +// Grid `(32, NumSplits, B*Hq)` over threadgroup `(32, NumSplits, 1)`: one block +// per (batch, query head), 32 lanes over the head dimension, `NumSplits` warps +// over strided token stripes, merged through shared memory. +inline constexpr const char* PAGED_ATTENTION_DECODE_HIP_SOURCE = R"( + #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32 + #error "mlxcel_paged_attention_decode assumes a 32-lane wavefront" + #endif + #if defined(__AMDGCN_WAVEFRONT_SIZE) && __AMDGCN_WAVEFRONT_SIZE != 32 + #error "mlxcel_paged_attention_decode assumes a 32-lane wavefront" + #endif + const float neg_inf = -__builtin_huge_valf(); + + uint32_t lane = threadIdx.x; // 0 .. 31 (within warp) + uint32_t sg = threadIdx.y; // 0 .. NumSplits-1 + uint32_t bhq = blockIdx.z; // 0 .. B*Hq-1 + + uint32_t hq_count = (uint32_t)NumQHeads; // Hq + uint32_t block_size = (uint32_t)PoolBlockSize; // tokens per block + uint32_t hkv_count = (uint32_t)NumKVHeads; // Hkv + uint32_t dim = (uint32_t)Dim; + uint32_t dpt = (uint32_t)DimsPerThread; // dims this lane owns + uint32_t d0 = lane * dpt; // first dim of this lane + + uint32_t b = bhq / hq_count; // batch index + uint32_t h = bhq % hq_count; // query head + uint32_t kv_head = h / (uint32_t)NRep; // grouped-query KV head + if (kv_head >= hkv_count) { + kv_head = 0; // defensive + } + + int vlen_i = visible_lens[b]; + uint32_t vlen = vlen_i > 0 ? (uint32_t)vlen_i : 0u; + uint32_t logical_start = (uint32_t)logical_starts[b]; + uint32_t row_off = (uint32_t)row_offsets[b]; + + float q_reg[DimsPerThread]; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + q_reg[j] = (d < dim) ? (float)q[bhq * dim + d] : 0.0f; + } + + __shared__ float tg_m[NumSplits]; + __shared__ float tg_l[NumSplits]; + __shared__ float tg_acc[NumSplits * Dim]; + + // Empty window: block-uniform (a function of the batch index only), so the + // whole block returns together and no warp is stranded at the barrier. + if (vlen == 0u) { + if (sg == 0u) { + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + out[bhq * dim + d] = 0.0f; + } + } + } + return; + } + + float scale_v = (float)scale[0]; + + float m = neg_inf; + float l = 0.0f; + float acc[DimsPerThread]; + for (uint32_t j = 0; j < dpt; j++) { + acc[j] = 0.0f; + } + + uint32_t stride_kv = hkv_count * dim; // elements per (block,slot) + for (uint32_t t = sg; t < vlen; t += (uint32_t)NumSplits) { + uint32_t abs_pos = logical_start + t; + uint32_t block_idx = abs_pos / block_size; + uint32_t slot = abs_pos - block_idx * block_size; + uint32_t row = (uint32_t)rows[row_off + block_idx]; + uint32_t base = (row * block_size + slot) * stride_kv + kv_head * dim; + + float partial = 0.0f; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + float kd = (d < dim) ? (float)k_pool[base + d] : 0.0f; + partial += q_reg[j] * kd; + } + // Butterfly all-reduce over the 32 lanes of this row; the trip count + // is warp-uniform, so every lane reaches each shuffle. + #pragma unroll + for (int o = 16; o > 0; o >>= 1) { + partial += __shfl_xor(partial, o, 32); + } + float score = partial * scale_v; + + float m_new = fmaxf(m, score); + float corr = expf(m - m_new); + float p = expf(score - m_new); + l = l * corr + p; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + float vd = (d < dim) ? (float)v_pool[base + d] : 0.0f; + acc[j] = acc[j] * corr + p * vd; + } + m = m_new; + } + + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + tg_acc[sg * dim + d] = acc[j]; + } + } + if (lane == 0u) { + tg_m[sg] = m; + tg_l[sg] = l; + } + __syncthreads(); + + if (sg == 0u) { + float m_g = tg_m[0]; + for (uint32_t s = 1; s < (uint32_t)NumSplits; s++) { + m_g = fmaxf(m_g, tg_m[s]); + } + float l_g = 0.0f; + for (uint32_t s = 0; s < (uint32_t)NumSplits; s++) { + l_g += tg_l[s] * expf(tg_m[s] - m_g); + } + float inv_l = l_g > 0.0f ? (1.0f / l_g) : 0.0f; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + float a = 0.0f; + for (uint32_t s = 0; s < (uint32_t)NumSplits; s++) { + a += tg_acc[s * dim + d] * expf(tg_m[s] - m_g); + } + out[bhq * dim + d] = a * inv_l; + } + } + } +)"; + +// v2 partial (port of `PAGED_ATTENTION_V2_PARTIAL_CUDA_SOURCE`). Grid +// `(32, NumWarps * Hkv * QGroups, num_chunks)` over threadgroup +// `(32, NumWarps, 1)`: one block per (chunk, KV head, query-head group). +// Scores are in base 2, as in the Metal and CUDA bodies, so the emitted LSE is +// in the units the merge kernel consumes. +inline constexpr const char* PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE = R"( + #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32 + #error "mlxcel_paged_attention_v2_partial assumes a 32-lane wavefront" + #endif + #if defined(__AMDGCN_WAVEFRONT_SIZE) && __AMDGCN_WAVEFRONT_SIZE != 32 + #error "mlxcel_paged_attention_v2_partial assumes a 32-lane wavefront" + #endif + const float neg_inf = -__builtin_huge_valf(); + + uint32_t lane = threadIdx.x; // 0 .. 31 (within warp) + uint32_t sg = threadIdx.y; // 0 .. NumWarps-1 + uint32_t yblk = blockIdx.y; // kv_head * QGroups + qgrp + uint32_t chunk = blockIdx.z; // flat (request, tile) id + + uint32_t dim = (uint32_t)Dim; + uint32_t dpt = (uint32_t)DimsPerThread; + uint32_t d0 = lane * dpt; + uint32_t page_size = (uint32_t)PageSize; + + uint32_t hq_count = (uint32_t)NumQHeads; + uint32_t hkv_count = (uint32_t)NumKVHeads; + + uint32_t kv_head = yblk / (uint32_t)QGroups; + uint32_t q_group = yblk - kv_head * (uint32_t)QGroups; + if (kv_head >= hkv_count) { + kv_head = 0; + } + uint32_t q_base = kv_head * (uint32_t)NRep + q_group * (uint32_t)QHeads; + + int req_i = request_indices[chunk]; + int tile_i = kv_tile_indices[chunk]; + uint32_t r = req_i > 0 ? (uint32_t)req_i : 0u; + uint32_t tile = tile_i > 0 ? (uint32_t)tile_i : 0u; + + uint32_t page_begin = (uint32_t)indptr[r]; + uint32_t page_end = (uint32_t)indptr[r + 1]; + uint32_t npages = page_end > page_begin ? page_end - page_begin : 0u; + uint32_t fpo = (uint32_t)first_page_offset[r]; + uint32_t lpl = (uint32_t)last_page_len[r]; + uint32_t seq_len = 0u; + if (npages > 0u) { + uint32_t total = (npages - 1u) * page_size + lpl; + seq_len = total > fpo ? total - fpo : 0u; + } + + uint32_t ppc = (uint32_t)params[0]; + uint32_t chunk_page_end = page_begin + (tile + 1u) * ppc; + if (chunk_page_end > page_end) { + chunk_page_end = page_end; + } + uint32_t t_begin = tile * ppc * page_size; + t_begin = t_begin > fpo ? t_begin - fpo : 0u; + uint32_t t_end; + if (chunk_page_end >= page_end) { + t_end = seq_len; + } else { + uint32_t span = (chunk_page_end - page_begin) * page_size; + t_end = span > fpo ? span - fpo : 0u; + } + if (t_end > seq_len) { + t_end = seq_len; + } + + float q_reg[QHeads * DimsPerThread]; + for (uint32_t g = 0; g < (uint32_t)QHeads; g++) { + uint32_t h = q_base + g; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + q_reg[g * dpt + j] = (h < hq_count && d < dim) + ? (float)q[(r * hq_count + h) * dim + d] + : 0.0f; + } + } + + float m[QHeads]; + float l[QHeads]; + float acc[QHeads * DimsPerThread]; + for (uint32_t g = 0; g < (uint32_t)QHeads; g++) { + m[g] = neg_inf; + l[g] = 0.0f; + for (uint32_t j = 0; j < dpt; j++) { + acc[g * dpt + j] = 0.0f; + } + } + + float scale_v = (float)scale[0] * 1.4426950408889634f; + uint32_t stride_kv = hkv_count * dim; + + for (uint32_t t = t_begin + sg; t < t_end; t += (uint32_t)NumWarps) { + uint32_t abs_pos = fpo + t; + uint32_t page_off = abs_pos / page_size; + uint32_t entry = abs_pos - page_off * page_size; + uint32_t row = (uint32_t)indices[page_begin + page_off]; + uint32_t base = (row * page_size + entry) * stride_kv + kv_head * dim; + + float k_reg[DimsPerThread]; + float v_reg[DimsPerThread]; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + k_reg[j] = (d < dim) ? (float)k_pool[base + d] : 0.0f; + v_reg[j] = (d < dim) ? (float)v_pool[base + d] : 0.0f; + } + + for (uint32_t g = 0; g < (uint32_t)QHeads; g++) { + float partial = 0.0f; + for (uint32_t j = 0; j < dpt; j++) { + partial += q_reg[g * dpt + j] * k_reg[j]; + } + #pragma unroll + for (int o = 16; o > 0; o >>= 1) { + partial += __shfl_xor(partial, o, 32); + } + float score = partial * scale_v; + float m_new = fmaxf(m[g], score); + float corr = exp2f(m[g] - m_new); + float p = exp2f(score - m_new); + l[g] = l[g] * corr + p; + for (uint32_t j = 0; j < dpt; j++) { + acc[g * dpt + j] = acc[g * dpt + j] * corr + p * v_reg[j]; + } + m[g] = m_new; + } + } + + __shared__ float tg_m[NumWarps * QHeads]; + __shared__ float tg_l[NumWarps * QHeads]; + __shared__ float tg_acc[NumWarps * QHeads * Dim]; + + for (uint32_t g = 0; g < (uint32_t)QHeads; g++) { + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + tg_acc[(sg * (uint32_t)QHeads + g) * dim + d] = acc[g * dpt + j]; + } + } + if (lane == 0u) { + tg_m[sg * (uint32_t)QHeads + g] = m[g]; + tg_l[sg * (uint32_t)QHeads + g] = l[g]; + } + } + __syncthreads(); + + if (sg == 0u) { + for (uint32_t g = 0; g < (uint32_t)QHeads; g++) { + uint32_t h = q_base + g; + if (h >= hq_count) { + continue; + } + float m_g = tg_m[g]; + for (uint32_t s = 1; s < (uint32_t)NumWarps; s++) { + m_g = fmaxf(m_g, tg_m[s * (uint32_t)QHeads + g]); + } + float l_g = 0.0f; + if (m_g > neg_inf) { + for (uint32_t s = 0; s < (uint32_t)NumWarps; s++) { + uint32_t idx = s * (uint32_t)QHeads + g; + l_g += tg_l[idx] * exp2f(tg_m[idx] - m_g); + } + } + float inv_l = l_g > 0.0f ? (1.0f / l_g) : 0.0f; + uint32_t out_base = (chunk * hq_count + h) * dim; + for (uint32_t j = 0; j < dpt; j++) { + uint32_t d = d0 + j; + if (d < dim) { + float a = 0.0f; + if (l_g > 0.0f) { + for (uint32_t s = 0; s < (uint32_t)NumWarps; s++) { + uint32_t idx = s * (uint32_t)QHeads + g; + a += tg_acc[idx * dim + d] * exp2f(tg_m[idx] - m_g); + } + } + out_v[out_base + d] = a * inv_l; + } + } + if (lane == 0u) { + out_lse[chunk * hq_count + h] = + l_g > 0.0f ? (m_g + log2f(l_g)) : neg_inf; + } + } + } +)"; + +// v2 merge (port of `PAGED_ATTENTION_MERGE_CUDA_SOURCE`). Grid `(Dim, H, M)` +// over threadgroup `(Dim, 1, 1)` yields blocks `(1, H, M)`, so `gridDim.y` is +// the head count the CUDA body reads from the partials' shape. No barrier, so +// the early `return` is safe. +inline constexpr const char* PAGED_ATTENTION_MERGE_HIP_SOURCE = R"( + const float neg_inf = -__builtin_huge_valf(); + + uint32_t d = threadIdx.x; + uint32_t h = blockIdx.y; + uint32_t o = blockIdx.z; + + uint32_t dim = (uint32_t)Dim; + uint32_t heads = (uint32_t)gridDim.y; + if (d >= dim || h >= heads) { + return; + } + + uint32_t begin = (uint32_t)o_indptr[o]; + uint32_t end = (uint32_t)o_indptr[o + 1]; + + float m = neg_inf; + float l = 0.0f; + float acc = 0.0f; + for (uint32_t i = begin; i < end; i++) { + float s = (float)lse_in[i * heads + h]; + // Skips -inf (an empty chunk) and any NaN, so an empty partial can + // never poison the rescale with (-inf) - (-inf). + if (!(s > neg_inf)) { + continue; + } + float m_new = fmaxf(m, s); + float corr = exp2f(m - m_new); + float w = exp2f(s - m_new); + l = l * corr + w; + acc = acc * corr + w * (float)v_in[(i * heads + h) * dim + d]; + m = m_new; + } + + out_v[(o * heads + h) * dim + d] = l > 0.0f ? (acc / l) : 0.0f; + if (d == 0u) { + out_lse[o * heads + h] = l > 0.0f ? (m + log2f(l)) : neg_inf; + } +)"; + +// `fast::hip_kernel` declares `_shape`, `_strides` and +// `_ndim` arguments for every input whose name it finds followed by one +// of those suffixes anywhere in the source text, comments included, and the +// vendored overlay declares the first two with a type the launch does not pass +// (see the file comment). None of these bodies may name one. +namespace paged_attention_hip_detail { +constexpr bool names_no_shape_metadata(std::string_view src) { + return src.find("_shape") == std::string_view::npos && + src.find("_strides") == std::string_view::npos && + src.find("_ndim") == std::string_view::npos; +} +} // namespace paged_attention_hip_detail + +static_assert( + paged_attention_hip_detail::names_no_shape_metadata( + PAGED_ATTENTION_DECODE_HIP_SOURCE), + "the v1 HIP body must not name _shape/_strides/_ndim"); +static_assert( + paged_attention_hip_detail::names_no_shape_metadata( + PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE), + "the v2 partial HIP body must not name _shape/_strides/_ndim"); +static_assert( + paged_attention_hip_detail::names_no_shape_metadata( + PAGED_ATTENTION_MERGE_HIP_SOURCE), + "the merge HIP body must not name _shape/_strides/_ndim"); + +} // namespace mlxcel::turbo diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp index 568d1edc4..ac82c0e9c 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp @@ -16,6 +16,7 @@ #include #include "gpu_backend.h" #include "kernel_port.h" +#include "paged_attention_hip.h" #include #include @@ -455,7 +456,8 @@ constexpr const char* PAGED_ATTENTION_V2_PARTIAL_CUDA_SOURCE = R"( } )"; -// Apple Silicon SIMD width; the warp width on CUDA. +// Apple Silicon SIMD width; the warp width on CUDA, and the wavefront width of +// the wave32 targets the HIP port is built for (`paged_attention_hip.h`). constexpr int PAGED_V2_SIMD_WIDTH = 32; // Threadgroup-memory budget for `tg_acc`, matching the ~28 KB v1 ceiling. @@ -489,39 +491,69 @@ const std::vector& partial_output_names() { return names; } -// Thread-safe lazy-initialised holders for the two JIT-compiled bodies, using -// the `std::call_once` pattern of `paged_attention.cpp` / `sparse_v_sdpa.cpp`: -// the server reaches first use concurrently from per-request blocking workers, -// and `call_once` re-runs the initializer if MLX device lookup throws. +// Compiles the partial body for one backend. Reached only through a port +// table entry, so `backend` is always a backend that entry names; the `None` +// arm and the non-ROCm-build throw keep the switch exhaustive and are +// unreachable. +mlx::core::fast::CustomKernelFunction make_partial_kernel( + mlxcel::GpuKernelBackend backend) { + switch (backend) { + case mlxcel::GpuKernelBackend::Metal: + return mlx::core::fast::metal_kernel( + "mlxcel_paged_attention_v2_partial", + partial_input_names(), + partial_output_names(), + std::string(PAGED_ATTENTION_V2_PARTIAL_SOURCE)); + case mlxcel::GpuKernelBackend::Cuda: + return mlx::core::fast::cuda_kernel( + "mlxcel_paged_attention_v2_partial", + partial_input_names(), + partial_output_names(), + std::string(PAGED_ATTENTION_V2_PARTIAL_CUDA_SOURCE)); + case mlxcel::GpuKernelBackend::Rocm: +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + // Issue #2068: the CUDA body ported to HIP, in + // `paged_attention_hip.h`. + return mlx::core::fast::hip_kernel( + "mlxcel_paged_attention_v2_partial", + partial_input_names(), + partial_output_names(), + std::string(PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE)); +#else + throw std::runtime_error( + "[paged_attention_decode_v2_partial] this build has no ROCm backend"); +#endif + case mlxcel::GpuKernelBackend::None: + break; + } + throw std::logic_error( + "[paged_attention_decode_v2_partial] no kernel body for this backend"); +} + +// Thread-safe lazy-initialised holder for the JIT-compiled body, one per +// backend, using the `std::call_once` pattern of `paged_attention.cpp` / +// `sparse_v_sdpa.cpp`: the server reaches first use concurrently from +// per-request blocking workers, and `call_once` re-runs the initializer if MLX +// device lookup throws. The holder names its backend instead of carrying a +// `use_cuda` flag, whose false arm read as Metal on any third backend +// (issue #2068). struct PagedV2PartialHolder { std::optional kernel; std::once_flag init_flag; - bool cuda; + mlxcel::GpuKernelBackend backend; - explicit PagedV2PartialHolder(bool use_cuda) : cuda(use_cuda) {} + explicit PagedV2PartialHolder(mlxcel::GpuKernelBackend b) : backend(b) {} mlx::core::fast::CustomKernelFunction& get() { - std::call_once(init_flag, [this] { - kernel = cuda - ? mlx::core::fast::cuda_kernel( - "mlxcel_paged_attention_v2_partial", - partial_input_names(), - partial_output_names(), - std::string(PAGED_ATTENTION_V2_PARTIAL_CUDA_SOURCE)) - : mlx::core::fast::metal_kernel( - "mlxcel_paged_attention_v2_partial", - partial_input_names(), - partial_output_names(), - std::string(PAGED_ATTENTION_V2_PARTIAL_SOURCE)); - }); + std::call_once(init_flag, [this] { kernel = make_partial_kernel(backend); }); return *kernel; } }; -inline PagedV2PartialHolder& get_partial_kernel(bool cuda) { - static PagedV2PartialHolder metal_holder(false); - static PagedV2PartialHolder cuda_holder(true); - return cuda ? cuda_holder : metal_holder; +template +PagedV2PartialHolder& get_partial_kernel() { + static PagedV2PartialHolder holder(Backend); + return holder; } @@ -531,13 +563,14 @@ inline PagedV2PartialHolder& get_partial_kernel(bool cuda) { const mlxcel::KernelPorts& paged_v2_partial_ports() { static const mlxcel::KernelPorts ports{ .metal = +[]() -> mlx::core::fast::CustomKernelFunction& { - return get_partial_kernel(false).get(); + return get_partial_kernel().get(); }, .cuda = +[]() -> mlx::core::fast::CustomKernelFunction& { - return get_partial_kernel(true).get(); + return get_partial_kernel().get(); + }, + .rocm = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_partial_kernel().get(); }, - // No HIP port yet (#1814). - .rocm = nullptr, }; return ports; } @@ -607,6 +640,19 @@ std::vector paged_attention_decode_v2_partial( const auto& q_shape = q.shape(); // [B, Hq, 1, Dim] const auto& kp_shape = k_pool.shape(); // [num_blocks, PageSize, Hkv, Dim] + // Shapes the kernel bodies cannot index are refused here, on every + // backend, instead of being read out of bounds (issue #2068). + // `v_pool` may be narrower than `k_pool` in heads (MiniMax-M3 keeps an + // index key at head `Hkv` of the K allocation; the sparse launch accepts + // that layout at batch 1), so only its rank and head dimension are checked. + if (q.ndim() != 4 || k_pool.ndim() != 4 || v_pool.ndim() != 4 || + q_shape[2] != 1 || kp_shape[3] != q_shape[3] || + v_pool.shape(3) != q_shape[3] || kp_shape[1] < 1 || kp_shape[2] < 1) { + throw std::invalid_argument( + "[paged_attention_decode_v2_partial] expects q [B, Hq, 1, D] and " + "k_pool, v_pool [blocks, page_size, heads, D]"); + } + int hq = q_shape[1]; int dim = q_shape[3]; int page_size = kp_shape[1]; @@ -651,6 +697,14 @@ std::vector paged_attention_decode_v2_partial( // the `{"T", T}` the fused decode-MoE kernels have always passed, which is // why those kernels never had the defect. Do not drop these because the // kernel body does not reference them. + // + // `NumQHeads`/`NumKVHeads` carry the head counts the Metal and CUDA bodies + // read from `q_shape`/`k_pool_shape`, for the HIP body, which cannot read + // `_shape` (see `paged_attention_hip.h`). They are passed on every + // backend because a launch that branches on the backend kind is what + // `make verify-kernel-port-dispatch` rejects; the Metal and CUDA bodies + // never name them, so there they only extend the module name with values + // that are fixed per model (issue #2068). std::vector> template_args = { {"Dim", dim}, {"PageSize", page_size}, @@ -659,6 +713,8 @@ std::vector paged_attention_decode_v2_partial( {"QGroups", q_groups}, {"DimsPerThread", dims_per_thread}, {"NumWarps", num_warps}, + {"NumQHeads", hq}, + {"NumKVHeads", hkv}, {"QType", q.dtype()}, {"KVType", k_pool.dtype()}, {"VType", v_pool.dtype()}, diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2.h b/src/lib/mlx-cpp/turbo/paged_attention_v2.h index 4721c5435..7f042ba00 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2.h +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2.h @@ -164,7 +164,7 @@ std::vector paged_attention_merge_states( // True when the resolved GPU backend has a port of the partial kernel, and of // the merge kernel, respectively. Each reads the `KernelPorts` table its // launcher dispatches through, so a predicate and its dispatch cannot disagree. -// Metal and CUDA today; ROCm answers false until lablup/mlxcel#1814 ports them. +// True on Metal, CUDA and ROCm (the HIP ports, lablup/mlxcel#2068). bool paged_attention_v2_partial_available(); bool paged_attention_merge_states_available(); diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp index 6da0ca463..678087e49 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp @@ -24,6 +24,7 @@ #include #include "gpu_backend.h" #include "kernel_port.h" +#include "paged_attention_hip.h" #include #include @@ -154,35 +155,65 @@ const std::vector& merge_output_names() { return names; } +// Compiles the merge body for one backend. Reached only through a port table +// entry, so `backend` is always a backend that entry names; the `None` arm and +// the non-ROCm-build throw keep the switch exhaustive and are unreachable. +mlx::core::fast::CustomKernelFunction make_merge_kernel( + mlxcel::GpuKernelBackend backend) { + switch (backend) { + case mlxcel::GpuKernelBackend::Metal: + return mlx::core::fast::metal_kernel( + "mlxcel_paged_attention_merge_states", + merge_input_names(), + merge_output_names(), + std::string(PAGED_ATTENTION_MERGE_SOURCE)); + case mlxcel::GpuKernelBackend::Cuda: + return mlx::core::fast::cuda_kernel( + "mlxcel_paged_attention_merge_states", + merge_input_names(), + merge_output_names(), + std::string(PAGED_ATTENTION_MERGE_CUDA_SOURCE)); + case mlxcel::GpuKernelBackend::Rocm: +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + // Issue #2068: the CUDA body ported to HIP, in + // `paged_attention_hip.h`. + return mlx::core::fast::hip_kernel( + "mlxcel_paged_attention_merge_states", + merge_input_names(), + merge_output_names(), + std::string(PAGED_ATTENTION_MERGE_HIP_SOURCE)); +#else + throw std::runtime_error( + "[paged_attention_merge_states] this build has no ROCm backend"); +#endif + case mlxcel::GpuKernelBackend::None: + break; + } + throw std::logic_error( + "[paged_attention_merge_states] no kernel body for this backend"); +} + +// One lazily compiled holder per backend, under the `std::call_once` contract +// of `paged_attention_v2.cpp`. It names its backend instead of carrying a +// `use_cuda` flag, whose false arm read as Metal on any third backend +// (issue #2068). struct PagedMergeHolder { std::optional kernel; std::once_flag init_flag; - bool cuda; + mlxcel::GpuKernelBackend backend; - explicit PagedMergeHolder(bool use_cuda) : cuda(use_cuda) {} + explicit PagedMergeHolder(mlxcel::GpuKernelBackend b) : backend(b) {} mlx::core::fast::CustomKernelFunction& get() { - std::call_once(init_flag, [this] { - kernel = cuda - ? mlx::core::fast::cuda_kernel( - "mlxcel_paged_attention_merge_states", - merge_input_names(), - merge_output_names(), - std::string(PAGED_ATTENTION_MERGE_CUDA_SOURCE)) - : mlx::core::fast::metal_kernel( - "mlxcel_paged_attention_merge_states", - merge_input_names(), - merge_output_names(), - std::string(PAGED_ATTENTION_MERGE_SOURCE)); - }); + std::call_once(init_flag, [this] { kernel = make_merge_kernel(backend); }); return *kernel; } }; -inline PagedMergeHolder& get_merge_kernel(bool cuda) { - static PagedMergeHolder metal_holder(false); - static PagedMergeHolder cuda_holder(true); - return cuda ? cuda_holder : metal_holder; +template +PagedMergeHolder& get_merge_kernel() { + static PagedMergeHolder holder(Backend); + return holder; } // This kernel's ports, in one place. `has_kernel_port` and `select_kernel_port` @@ -191,13 +222,14 @@ inline PagedMergeHolder& get_merge_kernel(bool cuda) { const mlxcel::KernelPorts& paged_merge_ports() { static const mlxcel::KernelPorts ports{ .metal = +[]() -> mlx::core::fast::CustomKernelFunction& { - return get_merge_kernel(false).get(); + return get_merge_kernel().get(); }, .cuda = +[]() -> mlx::core::fast::CustomKernelFunction& { - return get_merge_kernel(true).get(); + return get_merge_kernel().get(); + }, + .rocm = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_merge_kernel().get(); }, - // No HIP port yet (#1814). - .rocm = nullptr, }; return ports; } @@ -217,6 +249,18 @@ std::vector paged_attention_merge_states( using mlx::core::fast::TemplateArg; const auto& v_shape = v_in.shape(); // [N, H, Dim] + + // Shapes the kernel bodies cannot index are refused here, on every + // backend, instead of being read out of bounds (issue #2068). The HIP body + // takes the head count from the grid, so `lse_in` must agree with `v_in` + // on it. + if (v_in.ndim() != 3 || lse_in.ndim() != 2 || o_indptr.ndim() != 1 || + lse_in.shape(0) != v_shape[0] || lse_in.shape(1) != v_shape[1]) { + throw std::invalid_argument( + "[paged_attention_merge_states] expects v_in [N, H, D], lse_in " + "[N, H] and o_indptr [M + 1]"); + } + int heads = v_shape[1]; int dim = v_shape[2]; int num_outputs = static_cast(o_indptr.size()) - 1; diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index bc3dfe8ae..68ad621dd 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/src/lib/mlxcel-core/build.rs @@ -294,6 +294,8 @@ fn main() { println!("cargo:rerun-if-changed=../mlx-cpp/turbo/paged_attention.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/paged_attention.cpp"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/paged_attention.metal"); + // HIP bodies of the three paged-attention kernels (#2068). + println!("cargo:rerun-if-changed=../mlx-cpp/turbo/paged_attention_hip.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/paged_attention_v2.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/paged_attention_v2.cpp"); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 0bd55baaf..7353612b4 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -2480,7 +2480,7 @@ void paged_attention_merge_states( // tables through `has_kernel_port`, so a predicate cannot disagree with the // dispatch: all three kernels (v1 decode, v2 partial, merge); the merge kernel // alone; the v1 decode kernel alone; and the v2 pair (partial and merge). -// Metal and CUDA today; ROCm answers false until lablup/mlxcel#1814. +// True on Metal, CUDA and ROCm (the HIP ports, lablup/mlxcel#2068). bool paged_attention_kernels_available(); bool paged_attention_merge_available(); bool paged_attention_decode_available(); diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index 08817e406..4433b1639 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/src/lib/mlxcel-core/src/cache/paged.rs @@ -2085,11 +2085,11 @@ impl PagedBlockPool { }; // Both fused paths end in the paged-attention kernels, which have - // Metal and CUDA ports only. On a backend without them the launch + // Metal, CUDA and HIP ports. On a backend without them the launch // would be refused, so decline before planning and let the caller's // gather-then-SDPA fallback serve the step (issue #1803). The - // predicate reads the kernels' own port tables, so a ROCm port - // (#1814) turns this path on without another edit here. + // predicate reads the kernels' own port tables, which is how the HIP + // ports (#2068) turned this path on for ROCm with no edit here. if !crate::ffi::paged_attention_kernels_available() { return Ok(( None, @@ -2315,8 +2315,13 @@ impl PagedBlockPool { // bare UniquePtr, so a failed launch would abort rather than error; // decline here and let the caller take the gather fallback (#634 // security review). Unreachable for the single-slab island's small - // batches, kept as a guard for external library callers. - if crate::layers::paged_decode_backend() == crate::layers::PagedDecodeBackend::Cuda { + // batches, kept as a guard for external library callers. The HIP port + // (#2068) uses the same grid, so ROCm keeps the CUDA bound rather than + // relying on a device-reported limit. + if matches!( + crate::layers::paged_decode_backend(), + crate::layers::PagedDecodeBackend::Cuda | crate::layers::PagedDecodeBackend::Rocm + ) { let q_shape = ffi::array_shape(q); let grid_z = i64::from(q_shape.first().copied().unwrap_or(0).max(0)) * i64::from(q_shape.get(1).copied().unwrap_or(0).max(0)); diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 56c1ee2ac..908149580 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -5587,8 +5587,9 @@ pub enum PagedDecodeDispatch { } /// Compute backend the pooled decode runs on. The fused kernel has a Metal JIT -/// body (ADR 0001, measured on Apple Silicon) and a CUDA JIT body (#634). Both -/// are native candidates; the selector applies backend-specific thresholds. +/// body (ADR 0001, measured on Apple Silicon), a CUDA JIT body (#634) and a HIP +/// JIT body (#2068). All three are native candidates; the selector applies +/// backend-specific thresholds. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PagedDecodeBackend { /// Apple Silicon Metal (the fused kernel's original home). @@ -5597,11 +5598,12 @@ pub enum PagedDecodeBackend { /// pool blocks with no gather pass, so its advantage grows with context; /// the Metal-measured batch/context ceilings do not apply. Cuda, + /// AMD ROCm (the HIP port of the CUDA kernel, #2068). It follows the CUDA + /// rule, since it is the same kernel: native wherever the kernel can serve + /// the layer, with the same `gridDim.z` guard in front of the launch. + Rocm, /// Any backend with no fused kernel port, which today means CPU-only - /// builds and ROCm (issue #1803). Always gather. ROCm is a real GPU here, - /// not a CPU, and gets its own variant when its kernels are ported - /// (issue #1814); until then the answer is the same, so it shares this one - /// rather than pretending the ports exist. + /// builds and machines with no usable GPU. Always gather. Other, } @@ -5669,6 +5671,10 @@ pub fn select_pooled_paged_dispatch( // Metal-measured batch and context ceilings do not apply because the // CUDA win grows with context rather than eroding at long context. PagedDecodeBackend::Cuda => slab_count <= NATIVE_MAX_SLABS, + // ROCm (#2068) runs the HIP port of the CUDA kernel, which reads the + // pool the same way, so it takes the CUDA rule. No ROCm-measured + // batch or context ceiling exists to justify a narrower island. + PagedDecodeBackend::Rocm => slab_count <= NATIVE_MAX_SLABS, PagedDecodeBackend::Other => false, }; if native { @@ -5791,14 +5797,15 @@ pub(crate) fn resolve_paged_v2_dispatch( /// The backend the fused kernel would run on, cached for the decode hot path. /// -/// The fused kernel has a Metal JIT body and a CUDA JIT body (#634) and no HIP -/// port yet (#1814). It is a native candidate only when +/// The fused kernel has a Metal JIT body, a CUDA JIT body (#634) and a HIP JIT +/// body (#2068). It is a native candidate only when /// [`crate::paged_attention_decode_available`] says the resolved backend has a /// port in the kernel's own table; then [`crate::metal_is_available`] / -/// [`crate::cuda_is_available`] name which one. A CPU-only build, a machine -/// with no usable GPU and a ROCm build all fall to gather, none of which a -/// compile-time `target_os` check would catch. Metal is probed first so the -/// Apple path is unchanged. Detection is process-static, so it is read once. +/// [`crate::cuda_is_available`] / [`crate::hardware::gpu_backend_kind`] name +/// which one. A CPU-only build and a machine with no usable GPU fall to gather, +/// neither of which a compile-time `target_os` check would catch. Metal is +/// probed first and CUDA second so those paths are unchanged. Detection is +/// process-static, so it is read once. pub(crate) fn paged_decode_backend() -> PagedDecodeBackend { use std::sync::OnceLock; static BACKEND: OnceLock = OnceLock::new(); @@ -5814,6 +5821,8 @@ pub(crate) fn paged_decode_backend() -> PagedDecodeBackend { PagedDecodeBackend::Metal } else if crate::cuda_is_available() { PagedDecodeBackend::Cuda + } else if crate::hardware::gpu_backend_kind() == crate::hardware::GpuBackendKind::Rocm { + PagedDecodeBackend::Rocm } else { PagedDecodeBackend::Other } @@ -5848,7 +5857,7 @@ impl PagedDispatchCache { /// Bit carrying the decision alongside the packed key: set means /// [`PagedDecodeDispatch::Native`], clear means [`PagedDecodeDispatch::Gather`]. /// The key uses bits `0..=49` ([`Self::pack_key`]: the backend tag now needs - /// two bits for the Metal/Cuda/Other trichotomy), so bit `50` is free. + /// two bits for the Metal/Cuda/Rocm/Other tags), so bit `50` is free. const DECISION_BIT: u64 = 1u64 << 50; /// Mask covering the packed-key bits (`0..=49`), used to compare a cell's /// key half against a freshly packed key while ignoring the decision bit. @@ -5878,6 +5887,7 @@ impl PagedDispatchCache { PagedDecodeBackend::Metal => 0u64, PagedDecodeBackend::Cuda => 1u64, PagedDecodeBackend::Other => 2u64, + PagedDecodeBackend::Rocm => 3u64, }; b | (v << 16) | (s << 32) | (k << 48) } @@ -9423,8 +9433,8 @@ mod tests { use super::{ NativePagedOverride, PAGED_DISPATCH_CACHE_EMPTY, PagedDecodeBackend, PagedDecodeDispatch, - PagedDispatchCache, parse_native_paged_override, resolve_dispatch_decision, - select_pooled_paged_dispatch, + PagedDispatchCache, paged_decode_backend, parse_native_paged_override, + resolve_dispatch_decision, select_pooled_paged_dispatch, }; #[test] @@ -9522,6 +9532,57 @@ mod tests { } } + #[test] + fn selector_rocm_backend_follows_the_cuda_rule() { + // ROCm (#2068) runs the HIP port of the CUDA kernel, so it takes the + // CUDA rule: native on every single-slab shape, including the b=1 and + // long-context regimes Metal routes to gather, and gather past one + // slab. Compared against CUDA shape by shape so the two cannot drift. + for &(b, ctx) in &[ + (1usize, 128usize), + (1, 65536), + (3, 4096), + (4, 4097), + (16, 131072), + ] { + for &slabs in &[0usize, 1, 2, 8] { + let rocm = select_pooled_paged_dispatch(b, ctx, slabs, PagedDecodeBackend::Rocm); + assert_eq!( + rocm, + select_pooled_paged_dispatch(b, ctx, slabs, PagedDecodeBackend::Cuda), + "ROCm must match CUDA at b={b} ctx={ctx} slabs={slabs}" + ); + let expected = if slabs <= 1 { + PagedDecodeDispatch::Native + } else { + PagedDecodeDispatch::Gather + }; + assert_eq!(rocm, expected, "ROCm b={b} ctx={ctx} slabs={slabs}"); + } + } + } + + #[test] + fn paged_decode_backend_names_the_resolved_backend() { + // The backend the fused kernel would run on agrees with the resolved + // GPU backend wherever that backend has the v1 port. Before #2068 a + // ROCm build answered `Other` even with a port in the table, so the + // native path was unreachable there. + use crate::hardware::{GpuBackendKind, gpu_backend_kind}; + let backend = paged_decode_backend(); + if !crate::paged_attention_decode_available() { + assert_eq!(backend, PagedDecodeBackend::Other); + return; + } + let expected = match gpu_backend_kind() { + GpuBackendKind::Metal => PagedDecodeBackend::Metal, + GpuBackendKind::Cuda => PagedDecodeBackend::Cuda, + GpuBackendKind::Rocm => PagedDecodeBackend::Rocm, + _ => PagedDecodeBackend::Other, + }; + assert_eq!(backend, expected); + } + #[test] fn selector_batch_and_context_boundaries() { // batch boundary: 3 -> gather, 4 -> native (NATIVE_MIN_BATCH). @@ -9693,6 +9754,18 @@ mod tests { PagedDecodeDispatch::Gather, "Other must be gather (no kernel)" ); + // ROCm (#2068) at the same shape: native like CUDA, so a ROCm tag that + // aliased Metal's or Other's cell would return their cached gather. + assert_eq!( + cache.select(1, 512, 1, PagedDecodeBackend::Rocm), + PagedDecodeDispatch::Native, + "ROCm single-slab b=1 must be native (not an aliased Metal/Other gather)" + ); + assert_eq!( + cache.select(1, 512, 1, PagedDecodeBackend::Other), + PagedDecodeDispatch::Gather, + "Other must be gather after a ROCm query (not the aliased ROCm native)" + ); // Re-query CUDA: still native, proving the Metal/Other writes did not // clobber the CUDA cell (each backend keeps its own last-key cell only // for the most recent distinct key, so this also confirms recompute @@ -9706,9 +9779,10 @@ mod tests { #[test] fn dispatch_cache_cuda_pack_key_roundtrips_without_collision() { - // The 2-bit backend tag (Metal=0, Cuda=1, Other=2) must produce three - // distinct packed keys for one identical shape, and each must stay clear - // of the decision bit and the empty sentinel. + // The 2-bit backend tag (Metal=0, Cuda=1, Other=2, Rocm=3) must + // produce four distinct packed keys for one identical shape, and each + // must stay clear of the decision bit and the empty sentinel, also at + // the saturated shape where every field is all ones. let shape = (4usize, 4096usize, 1usize); let k_metal = PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Metal); @@ -9716,10 +9790,24 @@ mod tests { PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Cuda); let k_other = PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Other); + let k_rocm = + PagedDispatchCache::pack_key(shape.0, shape.1, shape.2, PagedDecodeBackend::Rocm); + let k_rocm_saturated = PagedDispatchCache::pack_key( + usize::MAX, + usize::MAX, + usize::MAX, + PagedDecodeBackend::Rocm, + ); assert_ne!(k_metal, k_cuda, "Metal and CUDA keys must differ"); assert_ne!(k_cuda, k_other, "CUDA and Other keys must differ"); assert_ne!(k_metal, k_other, "Metal and Other keys must differ"); for k in [k_metal, k_cuda, k_other] { + assert_ne!( + k, k_rocm, + "ROCm's key must differ from every other backend's" + ); + } + for k in [k_metal, k_cuda, k_other, k_rocm, k_rocm_saturated] { assert_eq!( k & PagedDispatchCache::DECISION_BIT, 0, diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 06c0eb0c9..9805c35e1 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1514,8 +1514,7 @@ mod ffi { /// and the merge kernel. Answered from the kernels' own `KernelPorts` /// tables (`has_kernel_port`), the ones their launchers dispatch /// through, so this cannot say yes to a launch that would refuse. - /// Metal and CUDA today; ROCm answers `false` until lablup/mlxcel#1814 - /// ports them. + /// True on Metal, CUDA and ROCm (the HIP ports, #2068). fn paged_attention_kernels_available() -> bool; /// True when the resolved GPU backend has a port of the merge kernel diff --git a/src/lib/mlxcel-core/src/test_support/kernel_ports.rs b/src/lib/mlxcel-core/src/test_support/kernel_ports.rs index 577526e25..a5e10f64b 100644 --- a/src/lib/mlxcel-core/src/test_support/kernel_ports.rs +++ b/src/lib/mlxcel-core/src/test_support/kernel_ports.rs @@ -16,13 +16,19 @@ //! no port of that kernel (lablup/mlxcel#1809). //! //! The ROCm test gate (`make verify-test-rocm`) runs every GPU test on gfx1151. -//! The paged-attention kernels (v1 decode, v2 partial, merge) have Metal and -//! CUDA ports and no HIP port yet; that port is lablup/mlxcel#1814. Their tests -//! either launch the kernels directly and get the launcher's refusal, or go -//! through a production path that declines on the missing port before it -//! reaches the decision the test checks (the batched paged decode's floor and -//! multi-slab reports, for example). Either way they fail for a missing port, -//! which reads as a correctness failure. +//! The paged-attention kernels (v1 decode, v2 partial, merge) had Metal and +//! CUDA ports and no HIP port until lablup/mlxcel#2068 (part of #1814), and +//! these skips were written for that gap. Without a port the tests either +//! launch the kernels directly and get the launcher's refusal, or go through a +//! production path that declines on the missing port before it reaches the +//! decision the test checks (the batched paged decode's floor and multi-slab +//! reports, for example). Either way they fail for a missing port, which reads +//! as a correctness failure. +//! +//! With the HIP ports in the tables every predicate below is true on ROCm, so +//! nothing skips there any more. The macros stay, so that a kernel added to +//! this family, or a backend added later, skips visibly here too instead of +//! failing. //! //! A skip here is narrow on purpose, so it cannot hide a real defect: //! From 586e7e0928cec2465a82b98fbcb2f4c8e4de7d4f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 07:05:18 +0900 Subject: [PATCH 2/6] update(rocm): label the paged kernel bench with the ROCm backend The example picked its selector label from metal_is_available() and cuda_is_available() only, so on ROCm it printed the gather decision while production now dispatches the HIP port. It reads the decode port predicate and gpu_backend_kind() the way paged_decode_backend() does. Refs #2068 --- examples/paged_attention_kernel_bench.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/examples/paged_attention_kernel_bench.rs b/examples/paged_attention_kernel_bench.rs index 5b10fa6ac..8e4f5fba1 100644 --- a/examples/paged_attention_kernel_bench.rs +++ b/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 } From f4ca4b9a11b0f28425be6bf198c45f04317d286a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 07:08:41 +0900 Subject: [PATCH 3/6] update(rocm): read paged-attention geometry from _shape on HIP With #2100 merged the ROCm overlay passes _shape by value, so the HIP bodies now read the head counts, block size and merge head count from q_shape, k_pool_shape and v_in_shape exactly as the CUDA bodies do. This drops the NumQHeads, NumKVHeads and PoolBlockSize template arguments that the first commit added on every backend, so the Metal and CUDA launches are again identical to main. On gfx1151 the paged tests pass (0 skips); with the lane fold started at 8 and the merge in base e, 29 of them fail. Refs #2068 --- src/lib/mlx-cpp/turbo/paged_attention.cpp | 11 ---- src/lib/mlx-cpp/turbo/paged_attention_hip.h | 59 ++++---------------- src/lib/mlx-cpp/turbo/paged_attention_v2.cpp | 10 ---- 3 files changed, 12 insertions(+), 68 deletions(-) diff --git a/src/lib/mlx-cpp/turbo/paged_attention.cpp b/src/lib/mlx-cpp/turbo/paged_attention.cpp index a64a443fa..93b3078c2 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention.cpp @@ -570,22 +570,11 @@ mlx::core::array paged_attention_decode( // int-only args, two callers that share this geometry but differ in pool // dtype get one compiled module and the later one reads through the wrong // pointer type. Metal already folds the dtypes into its key. - // - // `NumQHeads`/`NumKVHeads`/`PoolBlockSize` carry the geometry the Metal and - // CUDA bodies read from `q_shape`/`k_pool_shape`, for the HIP body, which - // cannot read `_shape` (see `paged_attention_hip.h`). They are - // passed on every backend, because a launch that branches on the backend - // kind is what `make verify-kernel-port-dispatch` rejects; the Metal and - // CUDA bodies never name them, so there they only extend the module name - // with values that are fixed per model (issue #2068). std::vector> template_args = { {"Dim", dim}, {"NRep", n_rep}, {"DimsPerThread", dims_per_thread}, {"NumSplits", num_splits}, - {"NumQHeads", hq}, - {"NumKVHeads", hkv}, - {"PoolBlockSize", kp_shape[1]}, {"QType", q.dtype()}, {"KVType", k_pool.dtype()}, {"VType", v_pool.dtype()}, diff --git a/src/lib/mlx-cpp/turbo/paged_attention_hip.h b/src/lib/mlx-cpp/turbo/paged_attention_hip.h index cba8fb9b5..b3882d95f 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_hip.h +++ b/src/lib/mlx-cpp/turbo/paged_attention_hip.h @@ -39,16 +39,10 @@ // - The KV reads keep the CUDA text's explicit `(float)`: `hip_bfloat16`, the // type MLX substitutes for bfloat16 on ROCm, converts to float only through // an `explicit` operator. -// - The geometry the CUDA bodies read from `_shape` comes from template -// arguments instead (`NumQHeads`, `NumKVHeads`, `PoolBlockSize`, which the -// launchers pass on every backend) or, for the merge kernel's head count, -// from `gridDim.y`. The vendored `fast::hip_kernel` declares -// `_shape` as a pointer while the launch passes the shape by value, so -// a HIP body that names one faults the queue (fixed by lablup/mlxcel#2100, -// LOCAL_FIXES item 30). The `static_assert`s at the end of this file keep -// every such name out of these bodies, comments included, because the -// overlay decides whether to declare the argument by searching the source -// text. +// - The geometry is read from `_shape` exactly as the CUDA bodies read +// it. That relies on the overlay passing shapes by value, which it does +// since lablup/mlxcel#2100 (`patches-rocm/LOCAL_FIXES.md` item 30); before +// that fix a HIP body that named `_shape` faulted the queue. // // Wave32 guard. The v1 and v2-partial bodies fold a dot product across 32 // lanes with an XOR butterfly that starts at 16. They carry the `#error` guard @@ -63,8 +57,6 @@ // shuffle and no barrier), so it is correct for any wavefront size and carries // no guard: an `#error` there would only reject a correct kernel. -#include - namespace mlxcel::turbo { // v1 split-K flash-decoding (port of `PAGED_ATTENTION_DECODE_CUDA_SOURCE`). @@ -84,9 +76,9 @@ inline constexpr const char* PAGED_ATTENTION_DECODE_HIP_SOURCE = R"( uint32_t sg = threadIdx.y; // 0 .. NumSplits-1 uint32_t bhq = blockIdx.z; // 0 .. B*Hq-1 - uint32_t hq_count = (uint32_t)NumQHeads; // Hq - uint32_t block_size = (uint32_t)PoolBlockSize; // tokens per block - uint32_t hkv_count = (uint32_t)NumKVHeads; // Hkv + uint32_t hq_count = (uint32_t)q_shape[1]; // Hq + uint32_t block_size = (uint32_t)k_pool_shape[1]; // tokens per block + uint32_t hkv_count = (uint32_t)k_pool_shape[2]; // Hkv uint32_t dim = (uint32_t)Dim; uint32_t dpt = (uint32_t)DimsPerThread; // dims this lane owns uint32_t d0 = lane * dpt; // first dim of this lane @@ -229,8 +221,8 @@ inline constexpr const char* PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE = R"( uint32_t d0 = lane * dpt; uint32_t page_size = (uint32_t)PageSize; - uint32_t hq_count = (uint32_t)NumQHeads; - uint32_t hkv_count = (uint32_t)NumKVHeads; + uint32_t hq_count = (uint32_t)q_shape[1]; + uint32_t hkv_count = (uint32_t)k_pool_shape[2]; uint32_t kv_head = yblk / (uint32_t)QGroups; uint32_t q_group = yblk - kv_head * (uint32_t)QGroups; @@ -393,9 +385,8 @@ inline constexpr const char* PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE = R"( )"; // v2 merge (port of `PAGED_ATTENTION_MERGE_CUDA_SOURCE`). Grid `(Dim, H, M)` -// over threadgroup `(Dim, 1, 1)` yields blocks `(1, H, M)`, so `gridDim.y` is -// the head count the CUDA body reads from the partials' shape. No barrier, so -// the early `return` is safe. +// over threadgroup `(Dim, 1, 1)` yields blocks `(1, H, M)`. No barrier, so the +// early `return` is safe. inline constexpr const char* PAGED_ATTENTION_MERGE_HIP_SOURCE = R"( const float neg_inf = -__builtin_huge_valf(); @@ -404,7 +395,7 @@ inline constexpr const char* PAGED_ATTENTION_MERGE_HIP_SOURCE = R"( uint32_t o = blockIdx.z; uint32_t dim = (uint32_t)Dim; - uint32_t heads = (uint32_t)gridDim.y; + uint32_t heads = (uint32_t)v_in_shape[1]; if (d >= dim || h >= heads) { return; } @@ -436,30 +427,4 @@ inline constexpr const char* PAGED_ATTENTION_MERGE_HIP_SOURCE = R"( } )"; -// `fast::hip_kernel` declares `_shape`, `_strides` and -// `_ndim` arguments for every input whose name it finds followed by one -// of those suffixes anywhere in the source text, comments included, and the -// vendored overlay declares the first two with a type the launch does not pass -// (see the file comment). None of these bodies may name one. -namespace paged_attention_hip_detail { -constexpr bool names_no_shape_metadata(std::string_view src) { - return src.find("_shape") == std::string_view::npos && - src.find("_strides") == std::string_view::npos && - src.find("_ndim") == std::string_view::npos; -} -} // namespace paged_attention_hip_detail - -static_assert( - paged_attention_hip_detail::names_no_shape_metadata( - PAGED_ATTENTION_DECODE_HIP_SOURCE), - "the v1 HIP body must not name _shape/_strides/_ndim"); -static_assert( - paged_attention_hip_detail::names_no_shape_metadata( - PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE), - "the v2 partial HIP body must not name _shape/_strides/_ndim"); -static_assert( - paged_attention_hip_detail::names_no_shape_metadata( - PAGED_ATTENTION_MERGE_HIP_SOURCE), - "the merge HIP body must not name _shape/_strides/_ndim"); - } // namespace mlxcel::turbo diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp index ac82c0e9c..186e41062 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp @@ -697,14 +697,6 @@ std::vector paged_attention_decode_v2_partial( // the `{"T", T}` the fused decode-MoE kernels have always passed, which is // why those kernels never had the defect. Do not drop these because the // kernel body does not reference them. - // - // `NumQHeads`/`NumKVHeads` carry the head counts the Metal and CUDA bodies - // read from `q_shape`/`k_pool_shape`, for the HIP body, which cannot read - // `_shape` (see `paged_attention_hip.h`). They are passed on every - // backend because a launch that branches on the backend kind is what - // `make verify-kernel-port-dispatch` rejects; the Metal and CUDA bodies - // never name them, so there they only extend the module name with values - // that are fixed per model (issue #2068). std::vector> template_args = { {"Dim", dim}, {"PageSize", page_size}, @@ -713,8 +705,6 @@ std::vector paged_attention_decode_v2_partial( {"QGroups", q_groups}, {"DimsPerThread", dims_per_thread}, {"NumWarps", num_warps}, - {"NumQHeads", hq}, - {"NumKVHeads", hkv}, {"QType", q.dtype()}, {"KVType", k_pool.dtype()}, {"VType", v_pool.dtype()}, From 0d1afcb3c0f3a255db2c8100888776a9cf8c86c4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 10:51:41 +0900 Subject: [PATCH 4/6] fix(rocm): tighten paged launcher shape checks and label ROCm tactics Review of #2103 found the V-pool check too loose: the bodies address V with K's block size and head stride, so a V that differs from K in axes 1 to 3 would be read out of bounds. MiniMax-M3's sparse launch reshapes both pools to [rows, 1, 1, D] and differs only in axis 0, which stays free. The v1 and v2 checks now also refuse D = 0, and the merge check refuses a head dim that is not a launchable block width (1 to 1024). The paged autotune ops label ROCm tactics "rocm" instead of "cuda", and stale comments in mla/mod.rs, cache/paged.rs and the merge launcher are corrected. Paged, MLA, autotune and layers tests on gfx1151: 550 passed. clippy -p mlxcel-core --lib --tests clean. Refs #2068 --- src/lib/mlx-cpp/turbo/paged_attention.cpp | 12 +++++++----- src/lib/mlx-cpp/turbo/paged_attention_v2.cpp | 12 +++++++----- src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp | 10 ++++++---- src/lib/mlxcel-core/src/autotune/ops/mod.rs | 4 ++-- .../src/autotune/ops/paged_decode_splits.rs | 7 +++++-- .../src/autotune/ops/paged_decode_v2_chunk.rs | 7 +++++-- src/lib/mlxcel-core/src/cache/paged.rs | 8 ++++---- src/lib/mlxcel-core/src/mla/mod.rs | 2 +- 8 files changed, 37 insertions(+), 25 deletions(-) diff --git a/src/lib/mlx-cpp/turbo/paged_attention.cpp b/src/lib/mlx-cpp/turbo/paged_attention.cpp index 93b3078c2..c0c42a387 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention.cpp @@ -514,12 +514,14 @@ mlx::core::array paged_attention_decode( // backend, instead of being read out of bounds (issue #2068; #2067 did the // same for the SSM update kernel). The bridge declares this function // `Result`, so the throw reaches Rust as an `Err`. - // `v_pool` may be narrower than `k_pool` in heads (MiniMax-M3 keeps an - // index key at head `Hkv` of the K allocation; the sparse launch accepts - // that layout at batch 1), so only its rank and head dimension are checked. + // `v_pool` may hold fewer rows than `k_pool` (the MiniMax-M3 sparse launch + // reshapes both allocations to `[rows, 1, 1, D]`, and K carries the + // index-key side head), so axis 0 is not compared. The bodies address V + // with K's block size and head stride, so axes 1 to 3 must match. if (q.ndim() != 4 || k_pool.ndim() != 4 || v_pool.ndim() != 4 || - q_shape[2] != 1 || kp_shape[3] != q_shape[3] || - v_pool.shape(3) != q_shape[3] || kp_shape[1] < 1 || kp_shape[2] < 1) { + q_shape[2] != 1 || q_shape[3] < 1 || kp_shape[3] != q_shape[3] || + kp_shape[1] < 1 || kp_shape[2] < 1 || v_pool.shape(1) != kp_shape[1] || + v_pool.shape(2) != kp_shape[2] || v_pool.shape(3) != kp_shape[3]) { throw std::invalid_argument( "[paged_attention_decode] expects q [B, Hq, 1, D] and k_pool, v_pool " "[blocks, block_size, heads, D]"); diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp index 186e41062..22f7f6643 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2.cpp @@ -642,12 +642,14 @@ std::vector paged_attention_decode_v2_partial( // Shapes the kernel bodies cannot index are refused here, on every // backend, instead of being read out of bounds (issue #2068). - // `v_pool` may be narrower than `k_pool` in heads (MiniMax-M3 keeps an - // index key at head `Hkv` of the K allocation; the sparse launch accepts - // that layout at batch 1), so only its rank and head dimension are checked. + // `v_pool` may hold fewer rows than `k_pool` (the MiniMax-M3 sparse launch + // reshapes both allocations to `[rows, 1, 1, D]`, and K carries the + // index-key side head), so axis 0 is not compared. The bodies address V + // with K's block size and head stride, so axes 1 to 3 must match. if (q.ndim() != 4 || k_pool.ndim() != 4 || v_pool.ndim() != 4 || - q_shape[2] != 1 || kp_shape[3] != q_shape[3] || - v_pool.shape(3) != q_shape[3] || kp_shape[1] < 1 || kp_shape[2] < 1) { + q_shape[2] != 1 || q_shape[3] < 1 || kp_shape[3] != q_shape[3] || + kp_shape[1] < 1 || kp_shape[2] < 1 || v_pool.shape(1) != kp_shape[1] || + v_pool.shape(2) != kp_shape[2] || v_pool.shape(3) != kp_shape[3]) { throw std::invalid_argument( "[paged_attention_decode_v2_partial] expects q [B, Hq, 1, D] and " "k_pool, v_pool [blocks, page_size, heads, D]"); diff --git a/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp b/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp index 678087e49..51417a279 100644 --- a/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp +++ b/src/lib/mlx-cpp/turbo/paged_attention_v2_merge.cpp @@ -251,11 +251,13 @@ std::vector paged_attention_merge_states( const auto& v_shape = v_in.shape(); // [N, H, Dim] // Shapes the kernel bodies cannot index are refused here, on every - // backend, instead of being read out of bounds (issue #2068). The HIP body - // takes the head count from the grid, so `lse_in` must agree with `v_in` - // on it. + // backend, instead of being read out of bounds (issue #2068). Every body + // indexes `lse_in` with the head count it reads from `v_in`, so `lse_in` + // must be `[N, H]` with `v_in`'s N and H. The threadgroup is `(D, 1, 1)`, + // so D must be a launchable block width. if (v_in.ndim() != 3 || lse_in.ndim() != 2 || o_indptr.ndim() != 1 || - lse_in.shape(0) != v_shape[0] || lse_in.shape(1) != v_shape[1]) { + lse_in.shape(0) != v_shape[0] || lse_in.shape(1) != v_shape[1] || + v_shape[2] < 1 || v_shape[2] > 1024) { throw std::invalid_argument( "[paged_attention_merge_states] expects v_in [N, H, D], lse_in " "[N, H] and o_indptr [M + 1]"); diff --git a/src/lib/mlxcel-core/src/autotune/ops/mod.rs b/src/lib/mlxcel-core/src/autotune/ops/mod.rs index 359621463..f9efd68ab 100644 --- a/src/lib/mlxcel-core/src/autotune/ops/mod.rs +++ b/src/lib/mlxcel-core/src/autotune/ops/mod.rs @@ -16,8 +16,8 @@ //! //! | Op | Backend | Knob | Status | //! |----|---------|------|--------| -//! | [`paged_decode_splits`] | Metal + CUDA | v1 paged-decode `NumSplits` launch shape | Wired end to end; the Apple-Silicon-tunable op. | -//! | [`paged_decode_v2_chunk`] | Metal + CUDA | v2 paged-decode `pages_per_chunk` | Wired end to end (issue #898); reachable only when `MLXCEL_PAGED_ATTENTION_V2=1` selects the v2 path. | +//! | [`paged_decode_splits`] | Metal + CUDA + ROCm | v1 paged-decode `NumSplits` launch shape | Wired end to end; the Apple-Silicon-tunable op. | +//! | [`paged_decode_v2_chunk`] | Metal + CUDA + ROCm | v2 paged-decode `pages_per_chunk` | Wired end to end (issue #898); reachable only when `MLXCEL_PAGED_ATTENTION_V2=1` selects the v2 path. | //! | [`cuda_kernel_knobs::QmmTileOp`] | CUDA | Blackwell qmm CTA `tile_m` | Wired, **unvalidated** (no CUDA host was available). | //! | [`cuda_kernel_knobs::QmvMultirowOp`] | CUDA | multirow-qmv row-window ceiling | Wired, **unvalidated** (no CUDA host was available). | //! diff --git a/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs b/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs index b50d55bf3..0343dd5cb 100644 --- a/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs +++ b/src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs @@ -202,10 +202,13 @@ impl TunableOp for PagedDecodeSplitsOp<'_> { } fn runner_id(&self) -> String { - // The launcher picks the Metal or CUDA JIT body from the live backend, - // and the two are different kernels; keep their tactics apart. + // The launcher picks the Metal, CUDA or HIP (#2068) JIT body from the + // live backend, and they are different kernels; keep their tactics + // apart. if crate::metal_is_available() { "metal".to_string() + } else if crate::hardware::gpu_backend_kind() == crate::hardware::GpuBackendKind::Rocm { + "rocm".to_string() } else { "cuda".to_string() } diff --git a/src/lib/mlxcel-core/src/autotune/ops/paged_decode_v2_chunk.rs b/src/lib/mlxcel-core/src/autotune/ops/paged_decode_v2_chunk.rs index 0e3d276ed..1afbef8cb 100644 --- a/src/lib/mlxcel-core/src/autotune/ops/paged_decode_v2_chunk.rs +++ b/src/lib/mlxcel-core/src/autotune/ops/paged_decode_v2_chunk.rs @@ -193,10 +193,13 @@ impl TunableOp for PagedDecodeV2ChunkOp<'_> { } fn runner_id(&self) -> String { - // The launcher picks the Metal or CUDA JIT body from the live backend, - // and the two are different kernels; keep their tactics apart. + // The launcher picks the Metal, CUDA or HIP (#2068) JIT body from the + // live backend, and they are different kernels; keep their tactics + // apart. if crate::metal_is_available() { "metal".to_string() + } else if crate::hardware::gpu_backend_kind() == crate::hardware::GpuBackendKind::Rocm { + "rocm".to_string() } else { "cuda".to_string() } diff --git a/src/lib/mlxcel-core/src/cache/paged.rs b/src/lib/mlxcel-core/src/cache/paged.rs index 4433b1639..af7f9af27 100644 --- a/src/lib/mlxcel-core/src/cache/paged.rs +++ b/src/lib/mlxcel-core/src/cache/paged.rs @@ -2311,10 +2311,10 @@ impl PagedBlockPool { }; // The CUDA launch places batch * query_heads in gridDim.z, which CUDA - // caps at 65535 (Metal has no such cap). The bridge entry returns a - // bare UniquePtr, so a failed launch would abort rather than error; - // decline here and let the caller take the gather fallback (#634 - // security review). Unreachable for the single-slab island's small + // caps at 65535 (Metal has no such cap). The limit is hit when the + // graph is evaluated, not when the bridge call returns, so the + // launcher's `Result` cannot report it; decline here and let the + // caller take the gather fallback (#634 security review). Unreachable for the single-slab island's small // batches, kept as a guard for external library callers. The HIP port // (#2068) uses the same grid, so ROCm keeps the CUDA bound rather than // relying on a device-reported limit. diff --git a/src/lib/mlxcel-core/src/mla/mod.rs b/src/lib/mlxcel-core/src/mla/mod.rs index 8e84c2ffc..b6534ef0a 100644 --- a/src/lib/mlxcel-core/src/mla/mod.rs +++ b/src/lib/mlxcel-core/src/mla/mod.rs @@ -140,7 +140,7 @@ pub fn split_kv_enabled() -> bool { *ENABLED.get_or_init(|| parse_enabled(std::env::var(SPLIT_KV_ENV).ok().as_deref())) && absorbed_enabled() // The split path ends in `paged_attention_merge_states`, a custom - // kernel with Metal and CUDA ports only. On a backend without one the + // kernel with Metal, CUDA and HIP (#2068) ports. On a backend without one the // launch would be refused. Declining here takes the `absorbed_decode` // fallback the caller already has for a plan that declines (issue // #1803). The predicate reads the merge kernel's own port table. From 0bf1c97c4f5a2352e09c4a881cd1530d3e5f5503 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 13:53:07 +0900 Subject: [PATCH 5/6] docs(rocm): record the paged-attention HIP port measurement Server paged decode on gfx1151 with Meta-Llama-3.1-8B-Instruct-4bit, fused v2 against MLXCEL_PAGED_ATTENTION_NATIVE=0 on the same binary, each pair inside one rocm_gpu_guard.sh window: per-request decode 1.28x at batch 4 and ~16K tokens (three runs each, spread under 0.5 tok/s), 1.40x and 1.77x for one sequence at ~16K and ~32K (one run each), 1.09x at ~1K, and no measurable change at ~4K. Time to first token is unchanged. The installation status table gets a paged-attention row and drops paged attention from the list of kernels that still run as graph fallbacks. Refs #2068 --- ...rocm-paged-attention-gfx1151-2026-10-05.md | 40 +++++++++++++++++++ docs/installation.md | 3 +- 2 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 docs/benchmark_results/rocm-paged-attention-gfx1151-2026-10-05.md diff --git a/docs/benchmark_results/rocm-paged-attention-gfx1151-2026-10-05.md b/docs/benchmark_results/rocm-paged-attention-gfx1151-2026-10-05.md new file mode 100644 index 000000000..6340cf6f6 --- /dev/null +++ b/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. diff --git a/docs/installation.md b/docs/installation.md index 563cb1805..defef4307 100644 --- a/docs/installation.md +++ b/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. | From 63f770b9c6b30d27ec0f5ea1607f151df7ce8cbd Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 14:18:24 +0900 Subject: [PATCH 6/6] docs: add technical report for PR #2103 --- ...cm-paged-attention-hip-port-20261005.en.md | 177 ++++++++++++++++++ ...cm-paged-attention-hip-port-20261005.ko.md | 177 ++++++++++++++++++ 2 files changed, 354 insertions(+) create mode 100644 TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.en.md create mode 100644 TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.ko.md diff --git a/TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.en.md b/TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.en.md new file mode 100644 index 000000000..8c30f4a2a --- /dev/null +++ b/TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.en.md @@ -0,0 +1,177 @@ +# Technical Report: PR #2103 - Port paged attention (v1, v2 partial, merge) to HIP + +**Date**: 2026-10-05 + +**Status**: Implemented and measured on the gfx1151 host; head `0bf1c97c` on origin/main `c05d5438`, pending merge. + +**Languages**: C++ (HIP kernel sources, kernel holders, port tables, launcher shape checks), Rust (`PagedDecodeBackend::Rocm`, dispatch selector and memo key, `gridDim.z` guard, autotune runner labels, kernel bench example, unit tests), CMake and `build.rs` (header registration), Markdown (environment variables, installation, benchmark results), Makefile (comment) + +**Risk Level**: Medium (it changes the default paged decode path on ROCm for the server's batched decode, MLA split-KV and sparse paged decode, and it adds host-side shape refusals and a holder refactor on Metal and CUDA, neither of which could be run on the development host) + +## Executive Summary + +Issue #2068 (part of #1814, epic #1801) asked for HIP ports of the three paged-attention kernels: the v1 split-K decode, the v2 partial and the v2 merge. Before this PR the three port tables 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, and 36 tests in the ROCm gate skipped with a `skipping ... lablup/mlxcel#1814` line. + +The PR adds `paged_attention_hip.h` with hipRTC ports of the three CUDA bodies, fills the three `.rocm` slots, and makes `paged_decode_backend()` return a new `PagedDecodeBackend::Rocm` so the v1 dispatch selector can pick the native kernel on ROCm. The HIP bodies read their geometry from `_shape` exactly as the CUDA bodies do. That became possible when #2100 merged its overlay fix; the PR's first commit had carried a workaround for the missing fix, and a later commit removed it, so the Metal and CUDA kernel bodies and template arguments are byte-identical to main. The three launchers also gained shape refusals on every backend. + +On gfx1151 the 36 previously skipped tests run and pass with no test edits, and a deliberate mutation of the HIP bodies makes 29 of them fail. Server paged decode with Meta-Llama-3.1-8B-Instruct-4bit improved per-request decode by 1.28x at batch 4 and ~16K tokens (three runs per arm), and by 1.40x and 1.77x for a single sequence at ~16K and ~32K (one run per arm). The ~4K case shows no measurable change. + +## 1. Problem Statement + +### 1.1 Three empty slots + +The three kernels each had a `KernelPorts` table (`paged_attention_ports()`, `paged_v2_partial_ports()`, `paged_merge_ports()`) with Metal and CUDA entries and `.rocm = nullptr`. Their support predicates read `has_kernel_port`, so on ROCm they answered false, and every production gate in front of the kernels (`cache/paged.rs`, `mla/mod.rs`, `paged_v2/sparse.rs`) declined to the gather fallback. PR #2059 had already added the skip macros in `test_support/kernel_ports.rs`, so the tests skipped visibly instead of failing; the issue's goal was to make them run with no edits once the tables were filled. + +### 1.2 A gate that did not follow the predicate + +`paged_decode_backend()` returned `Other` unless `metal_is_available()` or `cuda_is_available()` was true. On ROCm a filled v1 table would still have mapped to `Other`, and `select_pooled_paged_dispatch` would never have chosen native. Filling the table alone was not enough to reach the v1 kernel from the library entry point. + +### 1.3 Two holders keyed on a boolean + +`PagedV2PartialHolder` and `PagedMergeHolder` took `bool use_cuda`, whose false arm compiled the Metal body. On a third backend that flag reads as Metal, so the holders had to name their backend before a ROCm entry could be added. + +## 2. Change Summary + +| Area | Change | +|---|---| +| `paged_attention_hip.h` (new, 430 lines) | `PAGED_ATTENTION_DECODE_HIP_SOURCE`, `PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE`, `PAGED_ATTENTION_MERGE_HIP_SOURCE`: the CUDA bodies with spelling changes only | +| `paged_attention.cpp` | `PagedAttentionKernelHolderHip` calling `fast::hip_kernel` under `MLXCEL_BRIDGE_ROCM_BACKEND`; `.rocm` getter; host shape refusal | +| `paged_attention_v2.cpp` | `make_partial_kernel(GpuKernelBackend)`, `PagedV2PartialHolder` keyed on `GpuKernelBackend`, one holder per backend via a template getter; `.rocm` getter; host shape refusal | +| `paged_attention_v2_merge.cpp` | Same refactor as the partial (`make_merge_kernel`, `PagedMergeHolder`); `.rocm` getter; host shape refusal | +| `layers.rs` | `PagedDecodeBackend::Rocm`, returned by `paged_decode_backend()`; CUDA rule in `select_pooled_paged_dispatch`; memo tag 3; three new or extended unit tests | +| `cache/paged.rs` | `gridDim.z <= 65535` guard applies to `Rocm` as to `Cuda`; comments | +| Autotune ops | `paged_decode_splits` and `paged_decode_v2_chunk` label ROCm tactics `rocm` instead of `cuda` | +| `examples/paged_attention_kernel_bench.rs` | Selector label reads the decode port predicate and `gpu_backend_kind()` as production does | +| Build | Header registered in `turbo/CMakeLists.txt` and `mlxcel-core/build.rs` | +| Comments and docs | `kernel_ports.rs`, `lib.rs`, `mlx_cxx_bridge.h`, `mla/mod.rs`, Makefile comment, `environment-variables.md`, `installation.md`, new `rocm-paged-attention-gfx1151-2026-10-05.md` | + +The branch has five commits: the ports (`16b6b25d`), the bench label (`586e7e09`), the `_shape` switch after #2100 (`f4ca4b9a`), the review fixes (`0d1afcb3`), and the results doc (`0bf1c97c`). The diff against origin/main is 22 files, 840 insertions and 109 deletions. + +## 3. The Ports + +### 3.1 Same kernels, different spelling + +Each HIP body keeps the CUDA body's thread mapping, kernel name, inputs, outputs, grid and template arguments. The header comment notes that the ROCm `CustomKernel::eval_gpu` ceil-divides the Metal-style total-thread grid by the threadgroup tuple the same way CUDA does, so every `blockIdx` and `threadIdx` means what it means in the CUDA body. The differences are: + +- the 32-lane all-reduce is `__shfl_xor(v, o, 32)`, with the width stated, because `__shfl_xor_sync` exists in HIP only as a shim that ignores its mask; +- negative infinity is `-__builtin_huge_valf()`, because hipRTC compiles the body with only the headers `fast::hip_kernel` prepends and `INFINITY` comes from ``; +- KV reads keep the explicit `(float)`, since `hip_bfloat16` converts to float only through an `explicit` operator. + +The launches, and the `fast::hip_kernel` call that compiles each string, stay in the `.cpp` file that already holds the CUDA launch. The header holds data only. That keeps `make verify-kernel-dtype-keys` at its pinned scope: `EXPECTED_IN_SCOPE` and the `9 in scope` count are unchanged. + +### 3.2 The wave32 guard, and why the merge kernel has none + +The #1814 port requirements tell every port that reduces or shuffles across lanes to carry two `#error` checks on `__AMDGCN_WAVEFRONT_SIZE__` and `__AMDGCN_WAVEFRONT_SIZE`, because a fold that starts at lane 16 on a 64-lane CDNA wave drops half the lanes and still returns a finite, wrong result. + +- **v1 decode and v2 partial** fold a dot product across 32 lanes with an XOR butterfly starting at 16. They carry the guard. As #2067 found, it is inert with ROCm 10's AMD clang, which defines neither macro for gfx1151, gfx942 or gfx90a, so it documents the assumption rather than enforcing it. The header comment adds a second line of reasoning: the 32 lanes that share a `threadIdx.y` are consecutive in the block's linear order, so on a wave64 target the explicit width of 32 should keep each butterfly inside its own row. That is stated as not run, since no wave64 device was available. +- **merge** has no lane-level operation: one thread per output element, no shuffle and no barrier. It is correct for any wavefront size, so it carries no guard; an `#error` there would only reject a correct kernel on wave64. The issue's acceptance checkbox said "all three tables ... with the wave32 guard". The PR deviates from that wording on purpose and records the reason next to the ticked box, in line with the port requirement's own scope ("every ported kernel that reduces or shuffles across lanes"). + +### 3.3 Holders that name their backend + +`PagedV2PartialHolder` and `PagedMergeHolder` now take a `mlxcel::GpuKernelBackend`. A `make_partial_kernel` / `make_merge_kernel` switch compiles the Metal, CUDA or HIP body, and a template getter `get_partial_kernel()` gives one static holder per backend. `std::call_once` is kept, because the server reaches first use concurrently from per-request workers and `call_once` re-runs the initializer if MLX device lookup throws. The `None` arm and the non-ROCm-build throw keep the switch exhaustive; they are unreachable because a holder is only reached through a table entry that names its backend. + +## 4. Removing the #2100 Workaround + +The CUDA bodies read head counts, block size and the merge head count from `q_shape`, `k_pool_shape` and `v_in_shape`. On ROCm that faulted until #2100: the vendored `fast::hip_kernel` declared `_shape` as a pointer while the launch passed it by value (LOCAL_FIXES item 30). + +The first commit (`16b6b25d`) worked around that. It passed the geometry as new `NumQHeads`, `NumKVHeads` and `PoolBlockSize` template arguments on every backend, read the merge head count from `gridDim.y`, and used `static_assert`s to keep `_shape`, `_strides` and `_ndim` out of the HIP text. The cost was a change to the Metal and CUDA launches: extra template arguments that those bodies never read, which also change their JIT cache keys, on two backends the host could not run. + +Once #2100 merged, commit `f4ca4b9a` dropped the workaround (12 insertions, 68 deletions across three files). The HIP bodies now read `_shape` exactly as the CUDA bodies do, and the Metal and CUDA launches match main again. Checked for this report against origin/main `c05d5438`: all six Metal and CUDA source strings in the three launcher files are byte-identical, no `.metal` file changed, and the diff touches no `TemplateArg` list. The removal leaves one way to read the geometry across three backends and keeps the untestable backends' kernels exactly as they were. + +## 5. Host-Side Shape Refusals on Every Backend + +The three launchers now refuse shapes the bodies cannot index, before building the launch, on every backend. They throw `std::invalid_argument`; the bridge declares these functions `Result`, so the throw reaches Rust as an `Err`. This follows what #2067 did for the SSM update kernel. + +- **v1 decode and v2 partial**: `q`, `k_pool` and `v_pool` must be rank 4; `q` must be `[B, Hq, 1, D]` with `D >= 1`; `k_pool`'s block size and head count must be at least 1 and its D must match `q`'s; and `v_pool` must match `k_pool` in axes 1 to 3. Axis 0 is not compared, because MiniMax-M3's sparse launch reshapes both pools to `[rows, 1, 1, D]` and K carries an extra index-key side head, so K has more rows than V. +- **merge**: `v_in` rank 3, `lse_in` rank 2 with `v_in`'s N and H, `o_indptr` rank 1, and a head dim from 1 to 1024, since the threadgroup is `(D, 1, 1)` and D must be a launchable block width. + +The first version of these checks compared only rank and head dim. Review of #2103 found the V-pool check too loose: the bodies address V with K's block size and head stride, so a V differing in axes 1 to 3 would be read out of bounds. Commit `0d1afcb3` tightened it and added the `D = 0` and merge-width refusals. + +These refusals run on Metal and CUDA too. A call that used to launch and read out of bounds now returns an error. Production callers on ROCm pass conforming shapes (the paged, MLA, autotune and layers selectors, 550 tests, pass after the change), but the Metal and CUDA paths were not run with them. + +## 6. `PagedDecodeBackend::Rocm` + +- **Detection.** `paged_decode_backend()` returns `Rocm` when the v1 predicate `paged_attention_decode_available()` is true, neither Metal nor CUDA is available, and `crate::hardware::gpu_backend_kind()` is `Rocm`. Metal is still probed first and CUDA second, so those answers are unchanged. The `Other` doc now covers only CPU-only builds and machines with no usable GPU. +- **Selector.** In `select_pooled_paged_dispatch`, `Rocm` takes the CUDA rule, `slab_count <= NATIVE_MAX_SLABS`: native on any single-slab layer, including the batch-1 and long-context regimes that Metal routes to gather. The HIP kernel is the CUDA kernel's port and reads the pool the same way, and no ROCm-measured batch or context ceiling exists to justify a narrower island. +- **Memo key.** `PagedDispatchCache::pack_key` gives `Rocm` tag 3, after Metal 0, CUDA 1 and Other 2. The 2-bit tag field is now full. +- **`gridDim.z` guard.** The CUDA launch puts `batch * query_heads` in `gridDim.z`, which CUDA caps at 65535. That limit is hit when the graph is evaluated, not when the bridge call returns, so the launcher's `Result` cannot report it, and `cache/paged.rs` declines to gather beforehand. The HIP port uses the same grid, so the guard now matches `Cuda | Rocm`. ROCm keeps the CUDA bound instead of relying on a device-reported limit. +- **Tests.** `selector_rocm_backend_follows_the_cuda_rule` compares `Rocm` with `Cuda` shape by shape over five (batch, context) pairs and four slab counts. `paged_decode_backend_names_the_resolved_backend` checks the answer against `gpu_backend_kind()` wherever the v1 port exists. The memo test checks that a ROCm query is not served from an aliased Metal or Other cell, and the pack-key test checks four distinct keys, including a saturated ROCm key, clear of the decision bit and the empty sentinel. + +Two smaller consumers follow the same detection. The kernel bench example now labels its `select=` column `Rocm` on ROCm instead of printing the gather decision production no longer makes. The two paged autotune ops return `rocm` as their runner id, so ROCm tactics are not stored under `cuda` next to a different kernel's. + +## 7. Correctness Evidence + +### 7.1 The 36 skipped tests + +The 36 mlxcel-core tests that skipped on ROCm for lack of these ports (36 `skipping ... lablup/mlxcel#1814` lines on `57d8ed29`) now run and pass on gfx1151, with none of them edited. The issue text counted 34; the gate printed 36. They cover 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. + +### 7.2 Mutation check + +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: v2 launch, cascade, sparse, split-KV, batched decode, and `test_fused_paged_decode_native_vs_fallback_matrix` for v1. The two v1 tests at head dim 8 do not see the fold change, because 8 dims fit in lanes 0 to 7 and the dropped fold stage only combined zeros. The tests therefore detect a port that folds or rescales incorrectly, not only one that fails to launch. + +### 7.3 Gates + +From the PR, on gfx1151 (Radeon 8060S, ROCm 10) at `f4ca4b9a` rebased on `c05d5438`: `make verify-rocm` OK, 11869 passed, 0 failed, 378 ignored across 147 test binaries, ROCm smoke OK, no `#1814` skip line. After `0d1afcb3`: paged, MLA, autotune and layers selectors pass (550 tests); `cargo clippy -p mlxcel-core --features rocm --lib --tests -- -D warnings`, `cargo fmt --check` and `dead_doc_pointers` pass. `make verify-kernel-dtype-keys verify-kernel-port-dispatch` pass with 9 in scope and `EXPECTED_IN_SCOPE` unchanged. + +Orchestrator verification on head `0bf1c97c`, up to date with origin/main `c05d5438`: `make verify-rocm` passed every step, 11,869 passed, 0 failed, 378 ignored; smoke OK; zero `skipping ... #1814` lines in the log. + +## 8. Measured Results + +### 8.1 Why the server, not `bench_decode.sh` + +The #1814 port requirements name `scripts/bench_decode.sh`. The decode profile put paged attention at 0% of decode because `bench_decode.sh` decodes into a dense KV cache, so it never reaches these kernels. Issue #2068's own acceptance criterion names `scripts/benchmark_paged_decode_production.sh` (issue #899), which drives the server's paged decode, and that is what was run. + +### 8.2 Method + +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, at `f4ca4b9a`. One binary for both arms: before is `MLXCEL_PAGED_ATTENTION_NATIVE=0` (gather-then-SDPA, which every ROCm build ran before this change), after is unset (fused v2). Each case started a fresh server per arm, and the before and after arms of one case ran back to back inside one `scripts/rocm_gpu_guard.sh` window of 55 s with `/sys/class/kfd/kfd/proc` empty and no compiler. Attempts that saw a foreign GPU process or a compiler were rejected and rerun. Every after-arm server log announces `paged decode v2: fused v2 launch`, and no before-arm log does. The later commits change host shape checks, comments and the autotune label, not the kernel bodies or dispatch, and were not re-measured. + +### 8.3 Per-request decode tok/s + +| Case | Runs per arm | Before | After | After / before (medians) | +|---|---|---|---|---| +| batch 4, ~1K prompt | 3 | 5.7 / 5.5 / 5.6 | 6.1 / 6.3 / 6.0 | 1.09x | +| batch 4, ~4K prompt | 3 | 10.1 / 10.6 / 10.2 | 10.4 / 10.6 / 10.5 | 1.03x, within run spread | +| batch 4, ~16K prompt | 3 | 10.9 / 10.9 / 10.9 | 13.9 / 13.9 / 13.9 | 1.28x | +| batch 1, ~16K prompt | 1 | 17.2 | 24.1 | 1.40x | +| batch 1, ~32K prompt | 1 | 11.2 | 19.8 | 1.77x | + +The batch-4 cases ran three times per arm, with run-to-run spread of at most 0.5 tok/s. The ~1K and ~16K gains are outside that spread; the ~4K gain is not and is reported as no measurable change. The two single-sequence cases ran once per arm, because a guarded window long enough for them was rare on the shared host, so their ratios have no spread estimate and rest on one pair each. The direction matches the batch-4 trend (the gain grows with context), but the 1.40x and 1.77x figures should be read as single observations. + +Time to first token is unchanged within a few percent in every case, as expected since prefill does not use these kernels: for example 155.1 s before and 151.1 s after (medians) at batch 4 and ~16K. + +### 8.4 Why the aggregate column is not used + +The script also reports aggregate tok/s, all completion tokens over the level's wall-clock span. At long context it moves the other way (batch 4, ~16K: 1.9 / 2.0 / 2.0 before, 1.7 / 1.7 / 1.7 after) because it is dominated by prefill, about 150 s of time to first token at that point, and because it depends on how many tokens each request generated before stopping, which differs between the two numeric paths. It does not measure the kernel. The results page lists it and explains this; the PR and this report do not use it for any claim. + +## 9. Technical Decisions + +- **Port the CUDA bodies line for line.** CUDA and HIP share grid semantics, `template_args` and runtime compilation. Identical bodies leave one place to look when backends disagree. +- **Guard only where lanes talk.** v1 and v2 partial carry the guard; merge does not. This departs from the checkbox's wording and follows the port requirement's scope. +- **Remove the template-argument workaround once #2100 landed.** Reading `_shape` keeps one geometry contract across backends and leaves Metal and CUDA byte-identical to main, which matters because neither could be run here. +- **Keep launches next to the CUDA launches.** The dtype-key checker's pinned scope and count stay unchanged. +- **Refuse bad shapes on the host for every backend.** An `Err` at the bridge is better than an out-of-bounds read on any backend; the axis-0 exemption keeps MiniMax-M3's sparse launch working. +- **ROCm takes the CUDA selector rule and the CUDA grid bound.** Same kernel, same grid; no ROCm-specific ceiling has been measured, and a device-reported limit was not trusted over the known bound. +- **Name the backend in holders, tags and labels.** A boolean `use_cuda`, a `cuda` runner id and an `Other` fallback each read wrong once a third backend exists. + +## 10. Residual Risks and What Was Not Verified + +- **Metal and CUDA were not run.** Their kernel bodies, template arguments and table entries are unchanged, but they are touched by the holder refactor (same `metal_kernel`/`cuda_kernel` calls, one holder per backend) and by the new shape refusals. +- **No wave64 device.** The guard is inert on ROCm 10 clang, so a wave64 build would not be stopped at compile time. Correctness of the v1 and v2 partial fold on wave64 is reasoned from the explicit shuffle width and lane layout, not run. +- **Single-run long-context points.** The 1.40x and 1.77x batch-1 figures are one pair each. +- **Guard window shorter than the port requirement's.** The runs used a 55 s idle window; the #1814 requirements name 90 s for `bench_decode.sh`. +- **The memo tag field is full.** Four backend tags use both bits; a fifth backend needs the key layout widened. + +## 11. Follow-ups + +- **32-bit index math past 2^32 pool elements.** All three backends compute the KV address as `(row * block_size + slot) * stride_kv + kv_head * dim` in 32-bit unsigned arithmetic (`uint` in the Metal bodies, `uint32_t` in the CUDA and HIP bodies). A pool with more than 2^32 elements wraps that offset and reads the wrong block. For a pool with 8 KV heads of dimension 128, that is 4,194,304 token slots. This predates the PR, is shared by every backend, and is not caught by the new shape refusals, which check rank and per-axis agreement, not total size. Fixing it means 64-bit offsets in all three bodies (or a host refusal past the limit), which changes Metal and CUDA and so needs those hosts. +- **Slow batched server decode is a scheduling issue.** Batch-4 per-request decode on this server (5.6 to 13.9 tok/s) is far below the single-stream `bench_decode.sh` rate for the same model, 35 tok/s, because the per-request rate includes time spent waiting while the other three requests prefill. That is the same in both arms and is not addressed by these kernels. + +## 12. Learning Points + +- **Filling a port table is not always enough.** `paged_decode_backend()` enumerated backends instead of reading the predicate's backend, so the v1 kernel stayed unreachable until it learned the new variant. +- **Remove a workaround as soon as its cause is fixed.** The template-argument detour worked, but it changed two backends that could not be tested. Dropping it after #2100 restored byte-identity with main. +- **A guard belongs only where lanes communicate.** The merge kernel's one-thread-per-element design is wave-size independent, and adding a guard would reject a correct kernel. +- **Mutation checks show what a suite can see.** The fold mutation went unnoticed by the two v1 tests at head dim 8; the matrix test caught it. Without the mutation the suite's blind spot at small D would not be known. +- **State run counts with the numbers.** A table that mixes three-run medians and single runs should say which is which, and a column dominated by something other than the kernel should be explained and set aside, not quoted. + +Refs: #2068, #1814, #1801, #2100, #2067, #2059, #2061, #1803, #899, #898, #634. diff --git a/TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.ko.md b/TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.ko.md new file mode 100644 index 000000000..6cdaefdb3 --- /dev/null +++ b/TECHNICAL_REPORTS/2103-rocm-paged-attention-hip-port-20261005.ko.md @@ -0,0 +1,177 @@ +# 기술 보고서: PR #2103 - Paged attention (v1, v2 partial, merge)을 HIP으로 port + +**날짜**: 2026-10-05 + +**상태**: gfx1151 호스트에서 구현 및 측정 완료. head `0bf1c97c`, origin/main `c05d5438` 위, 머지 대기 중. + +**언어**: C++ (HIP kernel source, kernel holder, port table, launcher shape 검사), Rust (`PagedDecodeBackend::Rocm`, dispatch selector와 memo key, `gridDim.z` guard, autotune runner label, kernel bench 예제, unit test), CMake와 `build.rs` (header 등록), Markdown (환경 변수, 설치, benchmark 결과), Makefile (주석) + +**위험도**: 중간 (ROCm에서 server의 batched decode, MLA split-KV, sparse paged decode의 기본 paged decode 경로가 바뀝니다. Metal과 CUDA에도 host 쪽 shape 거부와 holder refactor가 들어가며, 두 backend는 개발 호스트에서 실행할 수 없었습니다) + +## 요약 + +Issue #2068 (#1814의 일부, epic #1801)은 세 paged-attention kernel, 즉 v1 split-K decode, v2 partial, v2 merge의 HIP port를 요청했습니다. 이 PR 이전에는 세 port table에 `.rocm` 항목이 없었으므로, ROCm에서 server의 batched paged decode, MLA split-KV, sparse paged decode가 모두 gather-then-SDPA로 실행되었고, ROCm gate의 test 36개가 `skipping ... lablup/mlxcel#1814` 줄을 남기고 skip되었습니다. + +이 PR은 세 CUDA 본문을 hipRTC로 옮긴 `paged_attention_hip.h`를 추가하고, 세 `.rocm` slot을 채우며, `paged_decode_backend()`가 새 `PagedDecodeBackend::Rocm`을 반환하게 하여 v1 dispatch selector가 ROCm에서 native kernel을 고를 수 있게 합니다. HIP 본문은 CUDA 본문과 똑같이 `_shape`에서 geometry를 읽습니다. 이는 #2100의 overlay 수정이 머지되면서 가능해졌습니다. PR의 첫 commit은 그 수정이 없는 상태를 우회하는 workaround를 담았고, 이후 commit이 이를 제거했으므로 Metal과 CUDA의 kernel 본문과 template argument는 main과 byte 단위로 같습니다. 세 launcher에는 모든 backend에서 동작하는 shape 거부도 추가되었습니다. + +gfx1151에서 이전에 skip되던 test 36개가 test 수정 없이 실행되어 통과하며, HIP 본문을 의도적으로 변형하면 그중 29개가 실패합니다. Meta-Llama-3.1-8B-Instruct-4bit로 server paged decode를 측정한 결과, 요청당 decode는 batch 4, ~16K token에서 1.28x (arm당 3회), 단일 sequence ~16K와 ~32K에서 1.40x와 1.77x (arm당 1회) 빨라졌습니다. ~4K는 측정 가능한 변화가 없습니다. + +## 1. 문제 정의 + +### 1.1 빈 slot 세 개 + +세 kernel은 각각 `KernelPorts` table (`paged_attention_ports()`, `paged_v2_partial_ports()`, `paged_merge_ports()`)을 가지고 있었고, Metal과 CUDA 항목만 있고 `.rocm = nullptr`였습니다. Support predicate가 `has_kernel_port`를 읽으므로 ROCm에서는 false를 답했고, kernel 앞의 production gate (`cache/paged.rs`, `mla/mod.rs`, `paged_v2/sparse.rs`)는 모두 gather fallback으로 내려갔습니다. PR #2059가 이미 `test_support/kernel_ports.rs`에 skip macro를 넣어 두었으므로 test는 실패 대신 눈에 보이게 skip했고, issue의 목표는 table을 채우면 test를 수정 없이 실행되게 하는 것이었습니다. + +### 1.2 Predicate를 따르지 않는 gate + +`paged_decode_backend()`는 `metal_is_available()`이나 `cuda_is_available()`이 true가 아니면 `Other`를 반환했습니다. ROCm에서 v1 table을 채워도 `Other`로 매핑되어 `select_pooled_paged_dispatch`가 native를 고를 수 없었습니다. Library entry point에서 v1 kernel에 도달하려면 table을 채우는 것만으로는 부족했습니다. + +### 1.3 Boolean으로 구분하는 holder 두 개 + +`PagedV2PartialHolder`와 `PagedMergeHolder`는 `bool use_cuda`를 받았고, false 쪽은 Metal 본문을 compile했습니다. 세 번째 backend에서는 이 flag가 Metal로 읽히므로, ROCm 항목을 넣기 전에 holder가 자기 backend를 명시해야 했습니다. + +## 2. 변경 요약 + +| 영역 | 변경 | +|---|---| +| `paged_attention_hip.h` (신규, 430줄) | `PAGED_ATTENTION_DECODE_HIP_SOURCE`, `PAGED_ATTENTION_V2_PARTIAL_HIP_SOURCE`, `PAGED_ATTENTION_MERGE_HIP_SOURCE`: 철자만 바꾼 CUDA 본문 | +| `paged_attention.cpp` | `MLXCEL_BRIDGE_ROCM_BACKEND` 아래에서 `fast::hip_kernel`을 호출하는 `PagedAttentionKernelHolderHip`, `.rocm` getter, host shape 거부 | +| `paged_attention_v2.cpp` | `make_partial_kernel(GpuKernelBackend)`, `GpuKernelBackend`로 구분하는 `PagedV2PartialHolder`, template getter로 backend마다 holder 하나, `.rocm` getter, host shape 거부 | +| `paged_attention_v2_merge.cpp` | partial과 같은 refactor (`make_merge_kernel`, `PagedMergeHolder`), `.rocm` getter, host shape 거부 | +| `layers.rs` | `paged_decode_backend()`가 반환하는 `PagedDecodeBackend::Rocm`, `select_pooled_paged_dispatch`의 CUDA 규칙, memo tag 3, 신규 또는 확장 unit test 세 개 | +| `cache/paged.rs` | `gridDim.z <= 65535` guard를 `Cuda`처럼 `Rocm`에도 적용, 주석 | +| Autotune op | `paged_decode_splits`와 `paged_decode_v2_chunk`가 ROCm tactic을 `cuda` 대신 `rocm`으로 label | +| `examples/paged_attention_kernel_bench.rs` | Selector label이 production처럼 decode port predicate와 `gpu_backend_kind()`를 읽음 | +| Build | `turbo/CMakeLists.txt`와 `mlxcel-core/build.rs`에 header 등록 | +| 주석과 문서 | `kernel_ports.rs`, `lib.rs`, `mlx_cxx_bridge.h`, `mla/mod.rs`, Makefile 주석, `environment-variables.md`, `installation.md`, 신규 `rocm-paged-attention-gfx1151-2026-10-05.md` | + +Branch의 commit은 다섯 개입니다: port (`16b6b25d`), bench label (`586e7e09`), #2100 이후 `_shape` 전환 (`f4ca4b9a`), review 수정 (`0d1afcb3`), 결과 문서 (`0bf1c97c`). origin/main 대비 diff는 22개 파일, 840줄 추가, 109줄 삭제입니다. + +## 3. Port + +### 3.1 같은 kernel, 다른 철자 + +각 HIP 본문은 CUDA 본문의 thread mapping, kernel 이름, input, output, grid, template argument를 그대로 유지합니다. Header 주석에 적힌 대로, ROCm `CustomKernel::eval_gpu`는 Metal 방식의 전체 thread grid를 threadgroup tuple로 CUDA와 같은 방식으로 올림 나눗셈하므로, 모든 `blockIdx`와 `threadIdx`는 CUDA 본문에서와 같은 의미입니다. 차이는 다음과 같습니다. + +- 32-lane all-reduce는 width를 명시한 `__shfl_xor(v, o, 32)`입니다. HIP의 `__shfl_xor_sync`는 mask를 무시하는 호환 shim일 뿐이기 때문입니다. +- 음의 무한대는 `-__builtin_huge_valf()`입니다. hipRTC는 `fast::hip_kernel`이 앞에 붙이는 header만으로 본문을 compile하는데, `INFINITY`는 ``에서 오기 때문입니다. +- KV 읽기는 명시적 `(float)`를 유지합니다. `hip_bfloat16`은 `explicit` operator로만 float로 변환됩니다. + +Launch와 각 문자열을 compile하는 `fast::hip_kernel` 호출은 이미 CUDA launch가 있는 `.cpp` 파일에 남습니다. Header는 data만 담습니다. 그래서 `make verify-kernel-dtype-keys`의 pin 범위가 유지되고, `EXPECTED_IN_SCOPE`와 `9 in scope` 개수는 바뀌지 않습니다. + +### 3.2 Wave32 guard, 그리고 merge kernel에 guard가 없는 이유 + +#1814 port 요구 사항은 lane 사이에서 reduction이나 shuffle을 하는 모든 port에 `__AMDGCN_WAVEFRONT_SIZE__`와 `__AMDGCN_WAVEFRONT_SIZE`에 대한 `#error` 검사 두 개를 넣으라고 합니다. Lane 16에서 시작하는 fold는 64-lane CDNA wave에서 lane 절반을 버리고도 유한하지만 틀린 결과를 내기 때문입니다. + +- **v1 decode와 v2 partial**은 16에서 시작하는 XOR butterfly로 32 lane에 걸쳐 dot product를 fold합니다. 이 둘에는 guard가 있습니다. #2067에서 확인했듯이 ROCm 10의 AMD clang은 gfx1151, gfx942, gfx90a 어디에서도 두 macro를 정의하지 않으므로 guard는 동작하지 않고, 가정을 강제하기보다 기록하는 역할입니다. Header 주석은 두 번째 근거를 덧붙입니다. 같은 `threadIdx.y`를 공유하는 32 lane은 block의 선형 순서에서 연속이므로, wave64 target에서도 width 32를 명시한 shuffle이 각 butterfly를 자기 행 안에 가둘 것이라는 점입니다. Wave64 장치가 없었으므로 이는 실행하지 않은 추론으로 적혀 있습니다. +- **merge**에는 lane 수준 연산이 없습니다. Output element마다 thread 하나이고, shuffle도 barrier도 없습니다. 어떤 wavefront 크기에서도 맞으므로 guard를 두지 않습니다. 거기에 `#error`를 넣으면 wave64에서 올바른 kernel을 거부할 뿐입니다. Issue의 acceptance checkbox는 "세 table 모두 wave32 guard와 함께"라고 적었습니다. PR은 port 요구 사항 자체의 범위 ("lane 사이에서 reduction이나 shuffle을 하는 모든 port")를 따라 이 문구와 의도적으로 다르게 했고, 체크한 항목 옆에 이유를 적었습니다. + +### 3.3 자기 backend를 명시하는 holder + +`PagedV2PartialHolder`와 `PagedMergeHolder`는 이제 `mlxcel::GpuKernelBackend`를 받습니다. `make_partial_kernel` / `make_merge_kernel` switch가 Metal, CUDA, HIP 본문 중 하나를 compile하고, template getter `get_partial_kernel()`가 backend마다 static holder 하나를 줍니다. Server가 요청별 worker에서 동시에 처음 사용에 도달하고, MLX device 조회가 throw하면 `call_once`가 initializer를 다시 실행하므로 `std::call_once`는 유지됩니다. `None` arm과 ROCm이 아닌 build의 throw는 switch를 exhaustive하게 만들 뿐이며, holder는 자기 backend를 명시한 table 항목을 통해서만 도달하므로 실행되지 않습니다. + +## 4. #2100 workaround 제거 + +CUDA 본문은 head 수, block 크기, merge head 수를 `q_shape`, `k_pool_shape`, `v_in_shape`에서 읽습니다. ROCm에서는 #2100 전까지 이것이 fault를 냈습니다. Vendored `fast::hip_kernel`이 `_shape`를 pointer로 선언했지만 launch는 값으로 넘겼기 때문입니다 (LOCAL_FIXES item 30). + +첫 commit (`16b6b25d`)은 이를 우회했습니다. Geometry를 새 template argument `NumQHeads`, `NumKVHeads`, `PoolBlockSize`로 모든 backend에 넘기고, merge head 수를 `gridDim.y`에서 읽고, `static_assert`로 `_shape`, `_strides`, `_ndim`이 HIP 문자열에 들어가지 않게 했습니다. 대가는 Metal과 CUDA launch의 변경이었습니다. 두 본문이 읽지도 않는 template argument가 추가되어 JIT cache key까지 바뀌었고, 두 backend는 호스트에서 실행할 수 없었습니다. + +#2100이 머지된 뒤 commit `f4ca4b9a`가 workaround를 제거했습니다 (세 파일, 12줄 추가, 68줄 삭제). HIP 본문은 이제 CUDA 본문과 똑같이 `_shape`를 읽고, Metal과 CUDA launch는 다시 main과 같습니다. 이 보고서를 위해 origin/main `c05d5438`과 비교한 결과, 세 launcher 파일의 Metal과 CUDA source 문자열 여섯 개가 모두 byte 단위로 같고, `.metal` 파일은 바뀌지 않았으며, diff는 어떤 `TemplateArg` 목록도 건드리지 않습니다. 제거로 세 backend가 geometry를 읽는 방법이 하나로 남고, 테스트할 수 없는 backend의 kernel은 그대로 유지됩니다. + +## 5. 모든 backend의 host 쪽 shape 거부 + +세 launcher는 이제 launch를 만들기 전에, 모든 backend에서, 본문이 index할 수 없는 shape를 거부합니다. `std::invalid_argument`를 throw하며, bridge가 이 함수들을 `Result`로 선언하므로 throw는 Rust에 `Err`로 전달됩니다. #2067이 SSM update kernel에 한 방식을 따른 것입니다. + +- **v1 decode와 v2 partial**: `q`, `k_pool`, `v_pool`은 rank 4여야 합니다. `q`는 `D >= 1`인 `[B, Hq, 1, D]`, `k_pool`의 block 크기와 head 수는 1 이상, D는 `q`와 같아야 하며, `v_pool`은 axis 1부터 3까지 `k_pool`과 같아야 합니다. Axis 0은 비교하지 않습니다. MiniMax-M3의 sparse launch가 두 pool을 `[rows, 1, 1, D]`로 reshape하고 K에는 index-key side head가 더 있어서 K의 행이 V보다 많기 때문입니다. +- **merge**: `v_in` rank 3, `v_in`의 N과 H를 가진 rank 2 `lse_in`, rank 1 `o_indptr`, 그리고 1부터 1024까지의 head dim입니다. Threadgroup이 `(D, 1, 1)`이므로 D는 launch 가능한 block 폭이어야 합니다. + +처음 버전의 검사는 rank와 head dim만 비교했습니다. #2103 review에서 V-pool 검사가 너무 느슨하다는 점이 발견되었습니다. 본문은 K의 block 크기와 head stride로 V를 주소 계산하므로, axis 1부터 3이 다른 V는 범위 밖에서 읽힙니다. Commit `0d1afcb3`이 이를 조이고 `D = 0`과 merge 폭 거부를 추가했습니다. + +이 거부는 Metal과 CUDA에서도 실행됩니다. 이전에 launch되어 범위 밖을 읽던 호출은 이제 error를 반환합니다. ROCm의 production caller는 조건에 맞는 shape를 넘기지만 (변경 후 paged, MLA, autotune, layers selector test 550개 통과), Metal과 CUDA 경로는 이 검사와 함께 실행되지 않았습니다. + +## 6. `PagedDecodeBackend::Rocm` + +- **검출.** `paged_decode_backend()`는 v1 predicate `paged_attention_decode_available()`이 true이고, Metal과 CUDA가 모두 없고, `crate::hardware::gpu_backend_kind()`가 `Rocm`일 때 `Rocm`을 반환합니다. Metal을 먼저, CUDA를 두 번째로 확인하는 순서는 그대로이므로 두 backend의 답은 바뀌지 않습니다. `Other` 문서는 이제 CPU 전용 build와 사용 가능한 GPU가 없는 기계만 다룹니다. +- **Selector.** `select_pooled_paged_dispatch`에서 `Rocm`은 CUDA 규칙 `slab_count <= NATIVE_MAX_SLABS`를 따릅니다. Metal이 gather로 보내는 batch 1과 긴 context 구간을 포함해 single-slab layer라면 모두 native입니다. HIP kernel은 CUDA kernel의 port로 pool을 같은 방식으로 읽고, 더 좁은 구간을 정당화할 ROCm 측정 batch나 context 상한이 없습니다. +- **Memo key.** `PagedDispatchCache::pack_key`는 Metal 0, CUDA 1, Other 2 다음으로 `Rocm`에 tag 3을 줍니다. 2-bit tag 필드는 이제 가득 찼습니다. +- **`gridDim.z` guard.** CUDA launch는 `batch * query_heads`를 `gridDim.z`에 넣고, CUDA는 이를 65535로 제한합니다. 이 한도는 bridge 호출이 반환될 때가 아니라 graph가 evaluate될 때 걸리므로 launcher의 `Result`로는 보고할 수 없고, `cache/paged.rs`가 미리 gather로 내려갑니다. HIP port는 같은 grid를 쓰므로 guard는 이제 `Cuda | Rocm`에 걸립니다. ROCm은 장치가 보고하는 한도에 기대지 않고 CUDA 한도를 그대로 씁니다. +- **Test.** `selector_rocm_backend_follows_the_cuda_rule`은 다섯 개의 (batch, context) 쌍과 네 개의 slab 수에 대해 `Rocm`을 `Cuda`와 shape마다 비교합니다. `paged_decode_backend_names_the_resolved_backend`는 v1 port가 있는 곳에서 답을 `gpu_backend_kind()`와 대조합니다. Memo test는 ROCm 질의가 Metal이나 Other의 cell을 잘못 공유해 답을 받지 않는지 확인하고, pack-key test는 포화된 ROCm key를 포함해 네 key가 서로 다르고 decision bit와 empty sentinel과 겹치지 않는지 확인합니다. + +작은 소비자 두 곳도 같은 검출을 따릅니다. Kernel bench 예제는 ROCm에서 production이 더는 내리지 않는 gather 결정을 출력하는 대신 `select=` 열을 `Rocm`으로 표시합니다. 두 paged autotune op은 runner id로 `rocm`을 반환하므로, ROCm tactic이 다른 kernel의 tactic 옆에 `cuda`로 저장되지 않습니다. + +## 7. 정확성 근거 + +### 7.1 Skip되던 test 36개 + +이 port가 없어서 ROCm에서 skip되던 mlxcel-core test 36개 (`57d8ed29`에서 `skipping ... lablup/mlxcel#1814` 줄 36개)가 이제 gfx1151에서 실행되어 통과하며, 그중 어느 것도 수정하지 않았습니다. Issue 본문은 34개로 셌지만 gate는 36개를 출력했습니다. 범위는 200 step에 걸친 gather 경로 대비 v1 decode와 GQA 및 batch matrix, host reference와 gather 경로 대비 v2 partial과 merge (f32와 f16 pool, 빈 요청, 잘린 window, GQA head mapping, 한 geometry에서 두 pool dtype), cascade, sparse, MLA split-KV입니다. + +### 7.2 Mutation 검사 + +HIP lane fold를 16 대신 8에서 시작하고 HIP merge를 base 2 대신 base e로 계산하면 이 test 중 29개가 실패합니다. v2 launch, cascade, sparse, split-KV, batched decode, 그리고 v1에 대한 `test_fused_paged_decode_native_vs_fallback_matrix`입니다. Head dim 8인 v1 test 두 개는 fold 변경을 보지 못합니다. 8개 dim이 lane 0부터 7에 들어가고, 빠진 fold 단계는 0끼리만 합쳤을 것이기 때문입니다. 즉 test는 launch 실패만이 아니라 fold나 rescale이 틀린 port도 잡아냅니다. + +### 7.3 Gate + +PR 기준, gfx1151 (Radeon 8060S, ROCm 10)에서 `c05d5438` 위로 rebase한 `f4ca4b9a`: `make verify-rocm` OK, test binary 147개에서 11869 passed, 0 failed, 378 ignored, ROCm smoke OK, `#1814` skip 줄 없음. `0d1afcb3` 이후: paged, MLA, autotune, layers selector 통과 (test 550개), `cargo clippy -p mlxcel-core --features rocm --lib --tests -- -D warnings`, `cargo fmt --check`, `dead_doc_pointers` 통과. `make verify-kernel-dtype-keys verify-kernel-port-dispatch`는 9 in scope로 통과하고 `EXPECTED_IN_SCOPE`는 바뀌지 않았습니다. + +Orchestrator 검증, origin/main `c05d5438`과 최신 상태인 head `0bf1c97c`: `make verify-rocm`의 모든 단계 통과, 11,869 passed, 0 failed, 378 ignored, smoke OK, log에 `skipping ... #1814` 줄 0개. + +## 8. 측정 결과 + +### 8.1 `bench_decode.sh`가 아니라 server를 측정한 이유 + +#1814 port 요구 사항은 `scripts/bench_decode.sh`를 명시합니다. 그러나 decode profile은 `bench_decode.sh`가 dense KV cache로 decode하기 때문에 paged attention을 decode의 0%로 잡았고, 이 script는 이 kernel에 도달하지 않습니다. Issue #2068의 acceptance criterion 자체가 server의 paged decode를 구동하는 `scripts/benchmark_paged_decode_production.sh` (issue #899)를 명시하며, 이것을 실행했습니다. + +### 8.2 방법 + +Meta-Llama-3.1-8B-Instruct-4bit, `mlxcel-server --parallel 4 --ctx-size 131072`, 요청당 decode token 128개로 `scripts/bench_serving_concurrency.py`, `f4ca4b9a`에서 측정했습니다. 두 arm은 같은 binary를 씁니다. Before는 `MLXCEL_PAGED_ATTENTION_NATIVE=0` (이 변경 전 모든 ROCm build가 실행하던 gather-then-SDPA), after는 미설정 (fused v2)입니다. 각 case는 arm마다 server를 새로 띄웠고, 한 case의 before와 after arm은 `/sys/class/kfd/kfd/proc`이 비어 있고 compiler가 없는 55초짜리 `scripts/rocm_gpu_guard.sh` window 하나 안에서 연달아 실행했습니다. 외부 GPU process나 compiler를 본 시도는 버리고 다시 실행했습니다. 모든 after arm의 server log에는 `paged decode v2: fused v2 launch`가 찍히고, before arm log에는 없습니다. 이후 commit은 host shape 검사, 주석, autotune label만 바꾸고 kernel 본문이나 dispatch는 바꾸지 않으며, 다시 측정하지 않았습니다. + +### 8.3 요청당 decode tok/s + +| Case | Arm당 실행 횟수 | Before | After | After / before (중앙값) | +|---|---|---|---|---| +| batch 4, ~1K prompt | 3 | 5.7 / 5.5 / 5.6 | 6.1 / 6.3 / 6.0 | 1.09x | +| batch 4, ~4K prompt | 3 | 10.1 / 10.6 / 10.2 | 10.4 / 10.6 / 10.5 | 1.03x, 실행 간 편차 안 | +| batch 4, ~16K prompt | 3 | 10.9 / 10.9 / 10.9 | 13.9 / 13.9 / 13.9 | 1.28x | +| batch 1, ~16K prompt | 1 | 17.2 | 24.1 | 1.40x | +| batch 1, ~32K prompt | 1 | 11.2 | 19.8 | 1.77x | + +Batch 4 case는 arm마다 세 번 실행했고 실행 간 편차는 최대 0.5 tok/s입니다. ~1K와 ~16K의 향상은 이 편차 밖이고, ~4K의 향상은 편차 안이므로 측정 가능한 변화 없음으로 보고합니다. 두 단일 sequence case는 공유 호스트에서 그만큼 긴 guard window를 얻기가 드물어 arm마다 한 번만 실행했으므로, 비율에 편차 추정이 없고 각각 한 쌍에 기댑니다. 방향은 batch 4의 경향 (context가 길수록 향상이 커짐)과 맞지만, 1.40x와 1.77x는 단일 관측으로 읽어야 합니다. + +첫 token까지의 시간은 모든 case에서 몇 퍼센트 안으로 변하지 않았습니다. Prefill은 이 kernel을 쓰지 않으므로 예상대로입니다. 예를 들어 batch 4, ~16K에서 before 155.1초, after 151.1초 (중앙값)입니다. + +### 8.4 Aggregate 열을 쓰지 않는 이유 + +Script는 aggregate tok/s, 즉 level의 wall-clock 구간 동안의 전체 completion token도 보고합니다. 긴 context에서 이 값은 반대로 움직입니다 (batch 4, ~16K: before 1.9 / 2.0 / 2.0, after 1.7 / 1.7 / 1.7). 그 지점에서 첫 token까지 약 150초가 걸리는 prefill이 지배하고, 각 요청이 멈추기 전에 생성한 token 수에 따라 달라지는데 그 수가 두 수치 경로 사이에서 다르기 때문입니다. Kernel을 측정하는 값이 아닙니다. 결과 페이지는 이 값을 싣고 이유를 설명하며, PR과 이 보고서는 어떤 주장에도 이 값을 쓰지 않습니다. + +## 9. 기술적 선택과 그 이유 + +- **CUDA 본문을 줄 단위로 옮긴다.** CUDA와 HIP은 grid 의미, `template_args`, runtime compile을 공유합니다. 본문이 같으면 backend가 어긋날 때 볼 곳이 하나입니다. +- **Lane끼리 통신하는 곳에만 guard를 둔다.** v1과 v2 partial에는 guard가 있고 merge에는 없습니다. Checkbox 문구와는 다르지만 port 요구 사항의 범위를 따릅니다. +- **#2100이 들어오자마자 template argument workaround를 제거한다.** `_shape`를 읽으면 backend 사이에 geometry 계약이 하나로 유지되고 Metal과 CUDA가 main과 byte 단위로 같게 남습니다. 둘 다 여기서 실행할 수 없었으므로 중요합니다. +- **Launch를 CUDA launch 옆에 둔다.** dtype-key checker의 pin 범위와 개수가 바뀌지 않습니다. +- **잘못된 shape를 모든 backend에서 host에서 거부한다.** 어느 backend에서든 범위 밖 읽기보다 bridge의 `Err`가 낫습니다. Axis 0 예외로 MiniMax-M3의 sparse launch가 계속 동작합니다. +- **ROCm은 CUDA selector 규칙과 CUDA grid 한도를 따른다.** 같은 kernel, 같은 grid이고, ROCm 고유 상한은 측정된 적이 없으며, 알려진 한도 대신 장치가 보고하는 한도를 믿지 않았습니다. +- **Holder, tag, label에 backend 이름을 쓴다.** Boolean `use_cuda`, `cuda` runner id, `Other` fallback은 세 번째 backend가 생기면 모두 틀리게 읽힙니다. + +## 10. 남은 위험과 검증하지 않은 것 + +- **Metal과 CUDA는 실행하지 않았습니다.** Kernel 본문, template argument, table 항목은 바뀌지 않았지만, holder refactor (같은 `metal_kernel`/`cuda_kernel` 호출, backend마다 holder 하나)와 새 shape 거부의 영향을 받습니다. +- **Wave64 장치가 없습니다.** ROCm 10 clang에서 guard가 동작하지 않으므로 wave64 build는 compile 시점에 막히지 않습니다. Wave64에서 v1과 v2 partial fold의 정확성은 명시한 shuffle 폭과 lane 배치로부터 추론한 것이며 실행하지 않았습니다. +- **긴 context 지점은 단일 실행입니다.** Batch 1의 1.40x와 1.77x는 각각 한 쌍입니다. +- **Guard window가 port 요구 사항보다 짧습니다.** 측정은 55초 idle window를 썼고, #1814 요구 사항은 `bench_decode.sh`에 대해 90초를 명시합니다. +- **Memo tag 필드가 가득 찼습니다.** Backend tag 네 개가 두 bit를 모두 쓰므로, 다섯 번째 backend는 key 배치를 넓혀야 합니다. + +## 11. 후속 작업 + +- **Pool element가 2^32를 넘으면 32-bit index 계산이 넘칩니다.** 세 backend 모두 KV 주소를 `(row * block_size + slot) * stride_kv + kv_head * dim`으로 32-bit unsigned 산술 (Metal 본문은 `uint`, CUDA와 HIP 본문은 `uint32_t`)로 계산합니다. Element가 2^32개를 넘는 pool에서는 이 offset이 wrap되어 잘못된 block을 읽습니다. KV head 8개, 차원 128인 pool이라면 token slot 4,194,304개입니다. 이 PR 이전부터 있던 문제이고 모든 backend가 공유하며, rank와 axis별 일치만 보고 전체 크기는 보지 않는 새 shape 거부로는 잡히지 않습니다. 고치려면 세 본문 모두 64-bit offset으로 바꾸거나 (또는 한도를 넘으면 host에서 거부하거나) 해야 하며, Metal과 CUDA가 바뀌므로 그 호스트가 필요합니다. +- **느린 batched server decode는 scheduling 문제입니다.** 이 server에서 batch 4의 요청당 decode (5.6에서 13.9 tok/s)는 같은 모델의 단일 stream `bench_decode.sh` 속도 35 tok/s보다 훨씬 낮습니다. 요청당 속도에 나머지 세 요청이 prefill하는 동안 기다린 시간이 포함되기 때문입니다. 두 arm에서 같으며, 이 kernel로 해결되는 문제가 아닙니다. + +## 12. 학습 포인트 + +- **Port table을 채우는 것만으로 충분하지 않을 수 있습니다.** `paged_decode_backend()`가 predicate의 backend를 읽지 않고 backend를 나열했기 때문에, 새 variant를 알기 전까지 v1 kernel에 도달할 수 없었습니다. +- **원인이 고쳐지면 workaround는 바로 걷어냅니다.** Template argument 우회는 동작했지만 테스트할 수 없는 두 backend를 바꿨습니다. #2100 이후 제거하면서 main과 byte 단위 동일성이 돌아왔습니다. +- **Guard는 lane끼리 통신하는 곳에만 둡니다.** Element마다 thread 하나인 merge kernel은 wave 크기와 무관하며, guard를 넣으면 올바른 kernel을 거부합니다. +- **Mutation 검사는 suite가 무엇을 볼 수 있는지 보여 줍니다.** Fold mutation은 head dim 8인 v1 test 두 개에서는 드러나지 않았고 matrix test가 잡았습니다. Mutation이 없었다면 작은 D에서의 사각지대를 알 수 없었습니다. +- **수치와 함께 실행 횟수를 적습니다.** 세 번의 중앙값과 단일 실행이 섞인 표는 어느 쪽인지 밝혀야 하고, kernel이 아닌 다른 것이 지배하는 열은 인용하지 말고 설명한 뒤 제쳐 두어야 합니다. + +Refs: #2068, #1814, #1801, #2100, #2067, #2059, #2061, #1803, #899, #898, #634.