From ae0f6f6797bfef564816b4f9ec28b65ce025eebe Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Sun, 19 Jul 2026 11:43:52 -0400 Subject: [PATCH] fix: enable mempool by default on HIP when graph capture is supported MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #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. --- warp/_src/context.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/warp/_src/context.py b/warp/_src/context.py index 97afe7cd2c..15159846ee 100644 --- a/warp/_src/context.py +++ b/warp/_src/context.py @@ -3858,6 +3858,12 @@ def __init__(self, runtime, alias, ordinal=-1, is_primary=False, context=None): if warp.config.enable_mempools_at_init: # enable if supported self.is_mempool_enabled = self.is_mempool_supported + elif self.is_hip and self.is_mempool_supported: + # HIP/ROCm: enable mempool by default when graph capture is supported + # (hipGraph requires mempool for in-capture allocations). + # Previously disabled due to hipMemsetAsync unreliability, but + # graph capture is now validated on ROCm 7.2 (Warp PR #15). + self.is_mempool_enabled = True else: # disable by default self.is_mempool_enabled = False