Skip to content

Fix HIP memory leaks in RNN kernels#728

Draft
magaonka-amd wants to merge 1 commit intoROCm:rocm-jaxlib-v0.8.2from
magaonka-amd:fix/rnn-hip-memleak-v0.8.2
Draft

Fix HIP memory leaks in RNN kernels#728
magaonka-amd wants to merge 1 commit intoROCm:rocm-jaxlib-v0.8.2from
magaonka-amd:fix/rnn-hip-memleak-v0.8.2

Conversation

@magaonka-amd
Copy link

Free dropout_states_dev GPU memory allocated via gpuMalloc in the HIP paths of DoRnnComputeWorkspaceReserveSpaceSizes, DnnRNNForward_, and DnnRNNBackward_. Also destroy the leaked miopenTensorDescriptor in the forward and backward functions.

Upstream PR: jax-ml#35638

Free dropout_states_dev GPU memory allocated via gpuMalloc in the HIP
paths of DoRnnComputeWorkspaceReserveSpaceSizes, DnnRNNForward_, and
DnnRNNBackward_. Also destroy the leaked miopenTensorDescriptor in
the forward and backward functions.
@magaonka-amd magaonka-amd force-pushed the fix/rnn-hip-memleak-v0.8.2 branch from e5c0398 to 03a0e13 Compare March 5, 2026 16:50
@magaonka-amd magaonka-amd marked this pull request as draft March 5, 2026 17:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant