Skip to content

fix(rocm): stop leaking alpha/beta pairs in the GEMM pointer wrappers #2243

Closed

Description

@inureyes

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:1204 in hipblaslt_gemm_ptrs, passed into the launch_kernel lambda at :1209.
  • gemms/hipblaslt_gemm.cpp:1272 in hipblaslt_gemm_rowmajor_on_stream, which calls hipblaslt_gemm_impl directly on the caller's stream.
  • gemms/rocblas_gemm.cpp:290 in rocblas_gemm_ptrs, passed to rocblas_gemm_ex inside its launch_kernel lambda.

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-rocm returns nothing.
  • tests/rocm_hipblaslt_concurrency.rs, tests/rocm_rocblas_handle_concurrency.rs and the MoE GEMM tests pass unchanged.
  • If a Rust-reachable path with a scale other than (1, 0) exists (for example addmm with alpha != 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

Related: PR #2208 (#2200), PR #2193.

Activity

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    platform:linuxLinux (CUDA / packaging) specific
    on Oct 8, 2026
  2. added and removed on Oct 8, 2026
  3. added a commit that references this issue on Oct 8, 2026
    9eb02c9
  4. added a commit that references this issue on Oct 12, 2026
    25baebc
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificpriority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions