fix: enable mempool by default on HIP when graph capture is supported (for PR #15)#16
Merged
rtmadduri merged 1 commit intoJul 20, 2026
Conversation
PR AMD-Ecosystem#15 enables graph capture on HIP (supports_graph_capture=True). However, the device initialization still disables mempool by default (enable_mempools_at_init=False), which causes hipGraph allocation to fail: Failed to allocate memory during graph capture because memory pools are not enabled on device. Try calling wp.set_mempool_enabled(True). Fix: when a HIP device supports both mempool and graph capture, enable mempool by default. This is safe on ROCm 7.2+ where graph capture has been validated. Users can still disable mempool explicitly via wp.set_mempool_enabled(device, False) or by unsetting warp.config.enable_mempools_at_init before warp.init(). Without this fix, users must call wp.config.enable_mempools_at_init=True before wp.init() as a workaround — which is not discoverable.
|
LGTM! I will merge this into PR#15 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
PR #15 enables graph capture on HIP (
supports_graph_capture=True), but the device initialization still disables mempool by default (enable_mempools_at_init=False). This causes hipGraph allocation to fail during capture:The
[WARNING] CUDA Graphs disabled: mempool disabledwarning then causes mjlab (and other frameworks) to disable graph capture entirely, negating PR #15's benefit.Fix
When a HIP device supports both mempool AND graph capture, enable mempool by default at device initialization. This is safe on ROCm 7.2+ where graph capture has been validated by PR #15.
Users can still disable mempool explicitly:
wp.set_mempool_enabled(device, False)at runtimewarp.config.enable_mempools_at_init = Falsebeforewp.init()Without this fix
Users must set
wp.config.enable_mempools_at_init = Truebeforewp.init()as a workaround — which is not documented or discoverable.Testing
Validated on AMD MI325X (gfx942, ROCm 7.2) with mujoco_warp physics simulation using mjlab RL framework. Graph capture now works end-to-end without the mempool workaround.
Related: AIOSS-5496, AIOSS-5700