Repository navigation
fix(rocm): stop leaking alpha/beta pairs in the GEMM pointer wrappers #2243
Description
Part of #1801
Problem / Background
Three ROCm GEMM wrappers heap-allocate an alpha/beta pair with new float[2]{alpha, beta} for any scale pair other than (1, 0) and (1, 1) and never free it, so every such call leaks 8 bytes. Found while reviewing PR #2208. References are against origin/main 1b4e3657, paths relative to src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/.
Current Behavior
gemms/hipblaslt_gemm.cpp:1204inhipblaslt_gemm_ptrs, passed into thelaunch_kernellambda at:1209.gemms/hipblaslt_gemm.cpp:1272inhipblaslt_gemm_rowmajor_on_stream, which callshipblaslt_gemm_impldirectly on the caller's stream.gemms/rocblas_gemm.cpp:290inrocblas_gemm_ptrs, passed torocblas_gemm_exinside itslaunch_kernellambda.
The comments (hipblaslt_gemm.cpp:1187-1191, :1257-1259) justify the leak with "hipBLASLt may retain host alpha/beta pointers until the stream kernel finishes" and a train-graph failure. The code does not support that: CommandEncoder::launch_kernel (device.h:629-690) invokes the functor synchronously on every path (eager func(stream_), and under use_hip_graphs(), which is constant false, either inside a stream capture or right after commit()), no pointer mode is set (no HIPBLASLT_MATMUL_DESC_POINTER_MODE anywhere), so the host scalars are read during the hipblasLtMatmul call. The production paths already rely on that: hipblaslt_gemm_epilogue (hipblaslt_gemm.cpp:1393-1416, every bf16/f16 dense GEMM) and hipblaslt_gemm_batched (:1453-1470) pass &alpha and &beta of the lambda's by-value captures, and the f16 arm of rocblas_gemm_ptrs passes addresses of lambda locals (rocblas_gemm.cpp:338-350).
The leak is latent: the current callers (matmul.cpp:79, :1622, :1796) pass 1.0/0.0, and moe_swiglu.cpp:993 forwards a beta from its local gemm helper. Any new caller with another scale leaks per call.
Proposed Solution
In all three functions pass the addresses of the scalars that live for the duration of the library call, as the epilogue and batched paths do: &alpha/&beta of the lambda's by-value captures in hipblaslt_gemm_ptrs and rocblas_gemm_ptrs, and of the function parameters in hipblaslt_gemm_rowmajor_on_stream (synchronous). Delete the static kOne/kZero branches and the new float[2], and replace the stale comments with one line stating that host-mode scalars are read at enqueue and launch_kernel runs its functor synchronously. Record the change as a new patches-rocm/LOCAL_FIXES.md entry (mlxcelverse only, not proposed to any fork upstream).
Rejected: keeping a process-lifetime interned table per (alpha, beta). It bounds the leak but keeps the false premise that the library reads the pointer after enqueue.
Edge case: if use_hip_graphs() is ever re-enabled with manual graph nodes that defer the functor, this must be revisited; note it in the use_hip_graphs() comment (device.cpp:46-57).
Acceptance Criteria
-
git grep 'new float\[' src/lib/mlx-cpp/patches-rocmreturns nothing. -
tests/rocm_hipblaslt_concurrency.rs,tests/rocm_rocblas_handle_concurrency.rsand the MoE GEMM tests pass unchanged. - If a Rust-reachable path with a scale other than (1, 0) exists (for example
addmmwithalpha != 1), a test checks its output against a CPU reference; otherwise the PR says none exists.
Verification
make verify-rocm-overlay
cargo test --profile test-fast --features rocm --test rocm_hipblaslt_concurrency --test rocm_rocblas_handle_concurrency -- --test-threads=1
make verify-rocm
Activity
fix(rocm): stop leaking alpha/beta pairs in the GEMM pointer wrappers
fix(rocm): stop leaking alpha/beta pairs in the GEMM pointer wrappers (…