perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814

Part of #1801. Phase 4. Depends on #1803, #1809.

Problem / Background

Correctness comes first in this epic; this issue collects the performance work the spike exposed, plus the ROCm ports of mlxcel's fused kernels now that the launchers refuse cleanly instead of aborting.

• Memory footprint. mlx-lm reported a peak of 20.2 GB for Meta-Llama-3.1-8B-Instruct-4bit (about 4.5 GB of weights) and 21.7 GB for Qwen3-30B-A3B-4bit (17 GB) on ROCm. The fork's allocator (arena, decode free list, GTT/managed fallbacks) may be caching aggressively; on a UMA host this competes with the OS and other GPU tenants.
• Decode bandwidth. Decode-shaped q4 GEMV reached about 150 GB/s including per-call sync. Llama-3.1-8B-4bit decodes at 32 tok/s, below the ceiling that weight traffic alone implies at that bandwidth.
• Fused kernels. Every launcher now resolves its port through one helper and refuses with a typed error on a backend that has no port, and its callers take the MLX graph path.
• MoE prefill (added 2026-09-29). Commit b3a2226e (perf(rocm): route f16 gather_qmm to the warp-shared kernel, Refs perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814) fixed f16 decode through gather_qmm (about 205 ms to 1.44 ms per call at Mixtral's expert shape; LOCAL_FIXES.md item 12) and left one item here: the sorted large-row path used by prefill stays slow for both f16 and bf16, because the expert-batched kernel is opt-in (MLX_ROCM_GATHER_QMV_EXPERT_BATCHED=1) after it returned wrong results for bf16 activations with sorted rhs indices (LOCAL_FIXES.md item 9).

Current Behavior

The earlier version of this issue claimed that #1803 already made ROCm run mlxcel's fused paths as graph fallbacks. That was false. #1803 (PR #1877) replaced the !metal::is_available() idiom and fixed the CUDA arm, but the false arm of each two-way dispatch still meant "Metal", so on ROCm six launchers reached fast::metal_kernel and aborted the process rather than falling back, and the bridge declarations were not Result, so the throw crossed a noexcept extern into std::terminate. #1885 (PR #2010) fixed the two samplers; PR #2018 fixed the other four. As of PR #2018 the premise finally holds.

That is a prerequisite for this issue, not history. A clean refusal is what distinguishes "no port for this backend" from "the port is wrong", so a ported kernel can be compared against a known-good graph result instead of against an abort. It is also what gives each dispatch a named rocm slot to fill.

The port surface

PR #2026 replaced the nine hand-written two-way choices with one helper, mlxcel::select_kernel_port (src/lib/mlx-cpp/turbo/kernel_port.h:94, kernel_port.cpp:57), reading a per-kernel KernelPorts table of three getters (kernel_port.h:62, struct KernelPorts), with has_kernel_port (kernel_port.h:74, kernel_port.cpp:53) answering the support predicate from the same table. An absent port is the value nullptr in a named field rather than a case a dispatch falls through to, so this issue's unit of work is "fill the .rocm field", not "restructure a dispatch". scripts/ci/check_kernel_port_dispatch.py enforces the pattern. The reasoning is in TECHNICAL_REPORTS/2026-kernel-port-dispatch-standard-20260929.en.md and TECHNICAL_REPORTS/2029-finish-kernel-port-standardization-20260929.en.md; read those rather than a restatement here.

PR #2029 (merged as d8d34e2b) extended this to every remaining launcher. Every line number below is origin/main at d8d34e2b, re-verified on 2026-09-29. There are 19 KernelPorts tables; 18 carry .rocm = nullptr; exactly one has a ROCm port, bitlinear_ports() (#1862, PR #1870). The checker reports 10 launcher file(s) route through select_kernel_port, 0 exempt. Fourteen of the eighteen are in scope here and four are out of scope; all eighteen are listed by accessor so the count is checkable with grep -rnE '^(const )?(mlxcel::)?KernelPorts& [a-z_0-9]+\(\)' --include=*.cpp src/.

Metal + CUDA, .rocm = nullptr (10 tables, in scope). These are the cheap ports, because an existing CUDA implementation is a far closer starting point for HIP than Metal source is.

run_fused_moe_two_kernel (mlx_cxx_kernels.cpp:2306, was :2260) is reached through two bridge entry points, fused_moe_expert_kernel (:2552, was :2506) and fused_moe_geglu_kernel (:2570, was :2524), so a port must satisfy both callers, and it needs both of its tables filled to run natively. moe_down_ports() is also read by moe_down_kernel_fn() (:2297) on the Nemotron-H fused_moe_forward path, so a HIP down kernel has three callers, not two.

Metal only, .cuda = nullptr, .rocm = nullptr (4 tables, in scope). Every earlier survey of this surface missed these, because the survey looked for two-way dispatch sites and these had none: each reached its kernel holder directly, which neither of the checker's original two rules could see. PR #2029 added rule 3 and moved them into tables. Both ports have to be written from Metal source, so they are the expensive half of the list.

Metal only, out of scope (4 tables). turbo4_delegated_ports() (src/lib/mlx-cpp/turbo/turbo4_delegated_sdpa.cpp:193, dispatch :700), turbo4_bulk_dequant_rotated_ports() (:264, dispatch :760), turbo4_delegated_steel_ports() (:666, dispatch :828), and sparse_v_ports() (src/lib/mlx-cpp/turbo/sparse_v_sdpa.cpp:223, dispatch :257). Why they are out of scope is under the SDPA heading below.

bitlinear_matmul is the one kernel with all three ports already filled (#1862, PR #1870), so it is the worked example for what filling a slot looks like end to end: bitlinear_ports() at mlx_cxx_kernels.cpp:240 with dispatch at :277 (merged in PR #2029; before that the same three ports were chosen by a hand-written switch), HIP source and holder at :220 (fast::hip_kernel) and :234 (get_bitlinear_kernel_hip()), and the wavefront guard at :170-175.

Filling .rocm is not always sufficient on its own. Several Rust-side and C++ support predicates still answer with backend availability rather than with the table, so a filled slot can stay unreachable. Verified 2026-09-29:

• ssm_kernel_available() (mlx_cxx_kernels.cpp:1576, was :1561) returns cu::is_available() off Apple.
• mamba1_scan_kernel_available() (:1751, was :1722) returns false off Apple outright.
• residual_add3_layer_norm gates on ffi::metal_is_available() (src/lib/mlxcel-core/src/layers.rs:1110).
• Added 2026-09-29: fused_xielu returns its elementwise fallback when !metal::is_available() (mlx_cxx_kernels.cpp:463) before it resolves a port, so a HIP xielu port is unreachable until that early return reads the table.
• Added 2026-09-29: custom_kernels_available() is itself backend-enumerating (custom_kernels_available_for in src/lib/mlx-cpp/turbo/gpu_backend.h:50-53 returns true for Metal or CUDA only) and still answers one C++ support predicate, rejection_sample_supported() (src/lib/mlx-cpp/turbo/sampling_rejection.cpp:765, returning custom_kernels_available() at :769 rather than has_kernel_port(rejection_ports())), and gates five Rust paths: fused MoE (src/models/switch_layers.rs:103), Nemotron-H fused MoE (src/models/nemotron_h.rs:1359), MLA split-KV merge (src/lib/mlxcel-core/src/mla/mod.rs:148), batched paged decode (src/lib/mlxcel-core/src/cache/paged.rs:2006), and paged_decode_backend() (src/lib/mlxcel-core/src/layers.rs:5543). Checker rule 4 does not catch these because they are not spelled metal_is_available() || cuda_is_available(). A HIP MoE, merge or paged port behind them stays unreachable.

Each of those becomes mlxcel::has_kernel_port(...) on the kernel's own table when its port lands, which is the shape fused_norm.cpp:347, fused_rope_append.cpp:394 and sampling.cpp:377 already use.

Rule 4 exemptions owned here (added 2026-09-29)

PR #2029 added checker rule 4 (no Rust gate spelled metal_is_available() || cuda_is_available()). Its only two exemptions sit in BACKEND_ENUMERATION_TODO (scripts/ci/check_kernel_port_dispatch.py:92-103), both tracked to this issue, and both are blocked on a missing predicate rather than on unapplied convention:

• src/lib/mlxcel-core/src/ffi_tests.rs:815 and :919 gate the fused paged-decode kernel. paged_attention.cpp has a table (paged_attention_ports(), :427) but exports no *_available() for the cxx bridge to call.
• src/lib/mlxcel-core/src/grouped_gemm_numeric_tests.rs:187, :286 and :353 gate MLX's own gather_mm, which is not an mlxcel port, so no table applies. gpu_backend_available() would be the honest predicate, but whether MLX's ROCm backend implements the grouped-GEMM path is unverified, and widening a gate on an assumption is how the feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 abort was reached.

The two SDPA launchers are safe; the earlier conclusion here was wrong

An earlier version of this section said the two Metal-only SDPA launchers, src/lib/mlx-cpp/turbo/turbo4_delegated_sdpa.cpp and src/lib/mlx-cpp/turbo/sparse_v_sdpa.cpp, were not confirmed safe on ROCm and needed checking before any porting, and that whether sparse_v_available() needs a backend term was an unresolved correctness question deserving its own issue. That item is resolved and the conclusion was wrong.

Both launchers are reached only through cache::turbo::sparse_v::kernel_enabled() (src/lib/mlxcel-core/src/cache/turbo/sparse_v.rs:751, :929, :1243, :1318; previously cited as :928, :1241, :1313), defined at :151, which returns false unconditionally when !cfg!(target_os = "macos"). Its own comment gives the reason: on non-Apple builds the cxx symbol exists via MLX upstream's no-op Metal stubs, but dispatching the JIT kernel without a Metal device hangs, so the graph fallback is forced.

Verified by execution on the gfx1151 host rather than by reading again: --kv-cache-mode turbo4-asym applies to 24 of 28 layers on Qwen3-0.6B-4bit, and forcing the sparse-V branch with MLXCEL_TURBO4_ASYM_DEQUANT_SDPA=0 still generates normally instead of aborting.

The observation that sparse_v_available() (src/lib/mlxcel-core/src/cache.rs:3565-3570) carries no backend term was correct; the conclusion drawn from it was not, because the kernel-level gate sits one layer below it. Two predicates answering different questions, and only the higher one was grepped for. No separate correctness issue is needed.

Both launchers are now in the port table (the four accessors listed above, merged in PR #2029), which makes them ordinary entries rather than an open question, and puts them out of this issue's scope entirely: kernel_enabled() forces the graph fallback off macOS whatever the table says, so a HIP port there would be unreachable until that gate gains a backend term. metal_gated_delta_forward_impl (mlx_cxx_kernels.cpp:1125, was :1110) is out of scope for the same kind of reason and is likewise safe: gated_delta_kernel_available() (:1081, was :1066) returns false off Apple through an #ifdef __APPLE__, and the launcher picks among five holders by shape rather than by backend, so it has no port table.

Proposed Solution

Measure first, then port the fused kernels with the largest decode impact and tune the allocator policy. Each change is gated by the #1809 correctness matrix. Three verified facts shape the approach.

1. The path is already open. fast::hip_kernel is declared in the ROCm overlay at src/lib/mlx-cpp/patches-rocm/mlx/fast.h:96, and #1862 (PR #1870) proved the three-arm choice end to end with the BitNet kernel: a HIP source string, a get_bitlinear_kernel_hip() holder (mlx_cxx_kernels.cpp:234, calling fast::hip_kernel at :220), and a filled .rocm entry in bitlinear_ports() (:240). Porting is per-kernel work against an established pattern, not new infrastructure.

2. The ahead-of-time compile cost from #1853 does not apply here. fast::hip_kernel compiles through hipRTC at runtime (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/jit_module.cpp:345, hiprtcCompileProgram), the same model metal_kernel and cuda_kernel use. The 26s-to-348s blowup that made per-dtype instantiation expensive in #1853 was in the overlay's own .hip files, which are compiled ahead of time. So these fused kernels follow the same per-dtype template-argument rule as Metal and CUDA: the ROCm custom kernel names its cache entry "custom_kernel_" + name + template_arguments_hash(template_args) (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp:230-231), exactly the CUDA scheme, and the dtype-key rule enforced by make verify-kernel-dtype-keys applies unchanged. Caveat: the checker scopes a file by the token cuda_kernel(, not hip_kernel(, so a HIP launch is checked only when it shares a file with a CUDA launch (see #1875). Every table in scope here lives in a file that already has one.

3. The binding constraint is wavefront width, and it fails silently. gfx1151 is wave32 while CDNA parts (gfx90a, gfx942) are wave64. A reduction fold that starts at lane 16 is correct only on wave32; on wave64 it folds half the lanes and the missing half is invisible, so the result stays finite and plausible and is wrong. static_assert(warpSize == 32) does not compile in HIP, because warpSize is an object with an operator int() rather than a constant expression. #1862 guards it with a preprocessor #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32 plus #error, and again for the non-underscored spelling (mlx_cxx_kernels.cpp:170-175). Every ported reduction kernel carries the same guard; this is a requirement, not a tip.

Scope

In scope: allocator policy defaults for UMA hosts; the .rocm entries of the fourteen tables listed above, in profile-determined order, together with the support predicate of each kernel ported; the two rule 4 exemptions; and the MoE prefill path left here by b3a2226e.

Out of scope: the four SDPA tables and metal_gated_delta_forward_impl, which cannot be reached off Apple whatever their table says. The sparse_v_available() backend-term question is closed, not deferred; do not file a follow-up for it.

Implementation plan

1. Profile decode on the test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809 models with rocprofv3 and attribute time to kernels, graph fallbacks and host overhead.
2. Explain and bound the allocator peak (compare MLX_* cache limits against the fork's allocator knobs); set sensible defaults for UMA hosts.
3. Port the top fused launchers to fast::hip_kernel, filling the .rocm entry of each table and keeping every input dtype in the template arguments so cached kernels are not reused across dtypes. Switch that kernel's support predicate, and any custom_kernels_available() or metal::is_available() gate in front of it from the list above, to mlxcel::has_kernel_port on its table in the same change, so the predicate and the dispatch cannot disagree.
4. Report before/after with the perf(bench): ROCm support in the benchmark harness and a published gfx1151 baseline #1810 harness.
5. (Added 2026-09-29) Export a paged-decode support predicate from paged_attention.cpp (a *_available() reading has_kernel_port(paged_attention_ports()), exposed through the cxx bridge), switch the gates at ffi_tests.rs:815 and :919 to it, and remove ffi_tests.rs from BACKEND_ENUMERATION_TODO. This is independent of any HIP port and can land first.
6. (Added 2026-09-29) Establish whether MLX's ROCm backend implements gather_mm (the grouped-GEMM path), by running the three grouped_gemm_numeric_tests.rs tests directly on gfx1151 rather than through the suite, since the feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 abort stops the test binary before later-sorting tests run. If they pass, switch the gates at :187, :286 and :353 to gpu_backend_available() and remove the file from BACKEND_ENUMERATION_TODO; if MLX's ROCm backend does not implement it, record that in the checker entry and keep the exemption.
7. (Added 2026-09-29) Profile the sorted large-row gather_qmm prefill path left by b3a2226e and decide whether the expert-batched kernel can be made correct for bf16 and enabled by default.

The profile decides the port order. The per-layer and per-token frequency below is a starting hypothesis for the reviewer to check against the profile, not a conclusion to assume: fused_add_rms_norm and fused_rope_qk_append run per layer per token on every model; the two samplers run once per token on every model; run_fused_moe_two_kernel and ssm_update_kernel run per token only on their model families; the three paged-attention launchers run only on the paged decode path; the four Metal-only entries run only on the model families named in their table row, which makes them lower priority than their per-token frequency alone suggests.

Acceptance criteria

• A profile-based breakdown of decode time on ROCm is published.
• Peak memory for an 8B 4-bit model is explained, and reduced or bounded by a documented default.
• At least the top two fused launchers by decode impact run natively on ROCm, with the correctness matrix unchanged and measured speedups.
• Every ported kernel that performs a cross-lane reduction carries the wavefront-width guard, so a wave64 target fails to compile rather than returning a plausible wrong answer.
• Each ported kernel is validated against the graph fallback it replaces, on the same checkpoint, not only against the CUDA arm.
• Each ported kernel's support predicate reads its KernelPorts table, so the newly filled port is actually reachable from Rust. In particular, no ported kernel is left behind ssm_kernel_available()'s cu::is_available(), mamba1_scan_kernel_available()'s off-Apple false, residual_add3_layer_norm's metal_is_available(), fused_xielu's metal::is_available() early return, or a custom_kernels_available() gate.
• make verify-kernel-dtype-keys and make verify-kernel-port-dispatch pass for every new launch site.
• (Added 2026-09-29) paged_attention.cpp exports a paged-decode *_available() predicate through the bridge, ffi_tests.rs gates on it, and ffi_tests.rs is gone from BACKEND_ENUMERATION_TODO.
• (Added 2026-09-29) Whether MLX's ROCm backend implements gather_mm is established by running the three grouped_gemm_numeric_tests.rs tests directly on gfx1151, and the checker entry is either removed (gates switched to gpu_backend_available()) or updated with the finding.
• (Added 2026-09-29) grep -rn '\.rocm = nullptr' --include=*.cpp src/ drops by one for every port landed, and python3 scripts/ci/check_kernel_port_dispatch.py reports 0 awaiting a predicate once both exemptions are resolved.

References

• One-helper dispatch and the KernelPorts table: PR refactor(core): choose custom-kernel ports through one helper #2026, TECHNICAL_REPORTS/2026-kernel-port-dispatch-standard-20260929.en.md
• Checker rules 3 and 4, the last five launchers, and the two rule 4 exemptions: PR refactor: finish the kernel port standardization #2029 (d8d34e2b), TECHNICAL_REPORTS/2029-finish-kernel-port-standardization-20260929.en.md
• Refusal guards that preceded it: fix(rocm): the Gumbel-max and rejection sampler entry points have no ROCm arm and abort on metal_kernel #1885 (PR fix(rocm): refuse instead of aborting at the sampler entry points #2010), PR fix(rocm): guard the four remaining custom-kernel launchers #2018 (TECHNICAL_REPORTS/2018-rocm-kernel-launcher-guards-20260928.en.md)
• Three-arm precedent and the wavefront guard: feat(rocm): port the BitNet bitlinear_matmul kernel to mlxcelverse #1862 (PR feat(rocm): port the BitNet bitlinear_matmul kernel #1870)
• Dtype-key rule and its blind spot: make verify-kernel-dtype-keys, chore(ci): kernel dtype-key checker cannot see launches that move into headers #1875
• ROCm CI and local gate: PR ci(rocm): build, link and generate on a self-hosted gfx1151 runner #1989, PR test(rocm): add a local verify-rocm gate with a shared build-and-generate smoke #2008, make verify-rocm (Makefile:804)
• f16 gather_qmm decode fix and the prefill item left here: commit b3a2226e

Refresh log

2026-09-29

• Replaced every "(in flight)" citation with its merged location: PR refactor: finish the kernel port standardization #2029 landed as d8d34e2b, so the Metal-only tables, the SDPA tables and the bitlinear_ports() dispatch are cited at origin/main, each line re-verified. Corrected lines that moved (kernel_port.h:94, the SSM and MoE tables and dispatches, run_fused_moe_two_kernel and its two entry points, the three stale predicates, gated_delta, sparse_v.rs).
• Made the port surface match the verified count: 19 tables, 18 with .rocm = nullptr (10 Metal + CUDA and 4 Metal-only in scope, 4 SDPA out of scope), one ROCm port (bitlinear_ports(), PR feat(rocm): port the BitNet bitlinear_matmul kernel #1870), all enumerated by accessor.
• Added the two rule 4 exemptions PR refactor: finish the kernel port standardization #2029 assigned here as work items (paged-decode *_available() predicate; whether MLX's ROCm backend implements gather_mm).
• Extended the stale-predicate criterion with two classes found while verifying: fused_xielu's metal::is_available() early return, and the rejection-sampler predicate plus five Rust paths gated on custom_kernels_available(), which is Metal-or-CUDA by definition and not caught by rule 4.
• Recorded the MoE prefill item that commit b3a2226e left to this issue, and the dtype-key checker's hip_kernel( scoping caveat (chore(ci): kernel dtype-key checker cannot see launches that move into headers #1875).

Split (2026-09-30)

This issue is now the tracking issue for the nine issues below, which replace its implementation plan; each stands on its own and is a native sub-issue here. State at the split (d1128266): plan item 5 is done (PR #2059 added the paged-attention predicates and ffi_tests.rs left BACKEND_ENUMERATION_TODO), and bitlinear_ports() is still the only table with a .rocm port.

1. perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061 Decode profile on gfx1151 and whether MLX's ROCm gather_mm works (plan items 1 and 6). No dependency.
2. perf(rocm): explain and bound allocator peak memory on UMA hosts #2062 Allocator peak memory on UMA and a bounded default (plan item 2). No dependency.
3. perf(rocm): port fused_add_rms_norm and fused_rope_qk_append to HIP #2063 fused_add_rms_norm and fused_rope_qk_append HIP ports. Depends on perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061.
4. perf(rocm): port gumbel_max_sample and rejection_sample to HIP #2064 gumbel_max_sample and rejection_sample HIP ports, plus rejection_sample_supported() reading its table. Depends on perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061.
5. perf(rocm): port the fused MoE decode kernels (moe_gateup, moe_down) #2065 Fused MoE decode ports (moe_gateup, moe_down) and the switch_layers.rs:103 and nemotron_h.rs:1359 gates. Depends on perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061.
6. perf(rocm): fix expert-batched gather_qmm for bf16 MoE prefill #2066 Expert-batched gather_qmm correct for bf16 MoE prefill, and whether it becomes the default (plan item 7, LOCAL_FIXES.md item 9, upstreaming candidate for chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #1813). No dependency.
7. perf(rocm): port ssm_update_kernel to HIP and fix ssm_kernel_available #2067 ssm_update_kernel HIP port and ssm_kernel_available() reading its table. Depends on perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061.
8. perf(rocm): port paged attention (v1 decode, v2 partial, merge) to HIP #2068 Paged-attention HIP ports (v1 decode, v2 partial, merge), which clear the 36 ROCm test skips. Depends on perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061.
9. perf(rocm): port Metal-only kernels (xielu, add3 norm, mamba1, relu2) #2069 Metal-only kernels (fused_xielu, fused_add3_layer_norm, mamba1_selective_scan, moe_fc1_relu2), lowest priority. The moe_fc1_relu2 part depends on perf(rocm): port the fused MoE decode kernels (moe_gateup, moe_down) #2065.

#2061's ranking may reorder items 3, 4, 5, 7 and 8. This issue closes when all nine are closed.