perf(rocm): profile decode on gfx1151 and settle MLX ROCm gather_mm #2061

Part of #1814 (epic #1801)

Problem / Background

The ROCm port work split out of #1814 (fused-kernel ports, one issue per kernel family) needs an order, and #1814 stated the order only as a hypothesis from per-layer and per-token call frequency. No per-kernel measurement of ROCm decode exists yet: the gfx1151 throughput baseline (docs/benchmark_results/rocm-baseline-gfx1151-2026-09-30.md, PR #2056) has end-to-end tok/s only, for example Llama-3.1-8B-4bit decoding at 35.41 tok/s, below what weight traffic alone implies at the roughly 150 GB/s q4 GEMV bandwidth measured in #1814.

A second open question sits in the same measurement session. src/lib/mlxcel-core/src/grouped_gemm_numeric_tests.rs gates its three tests (gather_mm_matches_dense_per_expert_reference at :186, gather_mm_selects_the_indexed_expert at :285, gather_mm_half_precision_matches_reference at :352) on !crate::metal_is_available() && !crate::cuda_is_available(). That is MLX's own gather_mm, not an mlxcel port, so no KernelPorts table applies, and it is the last entry in BACKEND_ENUMERATION_TODO in scripts/ci/check_kernel_port_dispatch.py (the checker currently prints 1 awaiting a predicate). Whether MLX's ROCm overlay implements the grouped-GEMM path was never run.

Proposed Solution

Profile. On the gfx1151 host, build cargo build --release --features rocm --bin mlxcel-bench-decode and run each model under rocprofv3 --kernel-trace --hip-graph-trace --stats -f csv with the default pp512/tg128 shape scripts/bench_decode.sh uses. Models: Meta-Llama-3.1-8B-Instruct-4bit, Qwen3-30B-A3B-4bit, granite-4.0-h-tiny-4bit, NVIDIA-Nemotron-3-Nano-30B-A3B-4bit. Use the idle-GPU guard from the #2056 baseline: 90 s with /sys/class/kfd/kfd/proc empty and no cargo or compiler process before the run, a 1 Hz monitor during it, and a rerun if any sample shows another GPU tenant. Separate the decode steps from prefill (by dispatch timestamp after the last prefill kernel) and report, per model: total decode GPU time per token, host gap per token (wall time minus GPU busy time), and the top kernels by share of decode GPU time.

Attribution to ports. For each port issue split from #1814 (units 3, 4, 5, 7, 8), name the MLX graph-fallback kernels that the port would replace and sum their share of decode time per model: fused_add_rms_norm and fused_rope_qk_append (the graph ops their Rust callers run when fused_add_rms_norm_available() / fused_rope_qk_append_available() is false, src/lib/mlxcel-core/src/layers.rs:996 and :1003), the samplers (gumbel_max_sample, rejection_sample), MoE (gather_qmm and the SwitchGLU elementwise ops), SSM update (the SSD graph path of Mamba2 layers, which ROCm runs at every width because ssm_kernel_available() is false there), paged attention (only when the paged decode path is active; state whether it is). Close the report with the resulting ranked order.

gather_mm. Run the three tests directly, not through the suite, because an abort in an earlier-sorting test stops the binary before these run. If all three pass, replace the three gates with if !crate::gpu_backend_available() { return; } and delete the file's entry from BACKEND_ENUMERATION_TODO. If any fails or aborts, keep the gates, and rewrite the entry's comment to record the result (which test, the error text, and the overlay commit).

Scope

In scope: the profile report and its raw stats; the grouped_gemm_numeric_tests.rs gates and its checker entry.

Out of scope: writing any port; allocator memory (a sibling issue split from #1814); MoE prefill gather_qmm (a sibling issue split from #1814).

Acceptance Criteria

• docs/benchmark_results/rocm-decode-profile-gfx1151-<YYYY-MM-DD>.md publishes, for all four models, decode GPU time per token, host gap per token, the top kernels by decode share, and the share attributable to each port unit, ending in the ranked port order; it records the mlxcel commit, overlay commit, ROCm and HIP versions, and the idle-GPU evidence.
• The rocprofv3 stats CSVs are committed under benchmarks/rocm_profiles/gfx1151_<commit>/.
• The three grouped_gemm_numeric_tests.rs tests were run on gfx1151 and the outcome is in the report.
• Either the gates read gpu_backend_available() and python3 scripts/ci/check_kernel_port_dispatch.py prints 0 awaiting a predicate, or the BACKEND_ENUMERATION_TODO entry names the observed failure.
• perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814's port issues are reordered (or confirmed) from the ranking in a comment on perf(rocm): allocator footprint and ROCm ports of mlxcel fused kernels #1814.

Verification

cargo test --release --features rocm -p mlxcel-core --lib grouped_gemm_numeric_tests -- --test-threads=1 --nocapture
python3 scripts/ci/check_kernel_port_dispatch.py
make verify-kernel-port-dispatch
make verify-clippy-rocm