Repository navigation
launch_host_func doesn't work with CUDA Graph #9274
Description
Description
Hi,
I tried launch_host_func, which I assume exposes cudaLaunchHostFunc to Python. However, unlike cudaLaunchHostFunc, launch_host_func doesn't work with CUDA Graph.
The host function can be captured by CUDA graph, but when replying CUDA graph, only the first reply can call the host function, while more replays cannot. Is this expected?
The cupy version is 13.3.0.
Thanks!
To Reproduce
import time
import cupy as cp
import torch
def device_func(x, y):
torch.matmul(x, x, out=y)
def my_callback(data):
print("Host callback says:", data)
def main():
x = torch.randn(4096, 4096, device='cuda')
y = torch.empty_like(x)
device_func(x, y)
data = "Hello Host Node"
stream = cp.cuda.Stream(non_blocking=True)
stream_pytorch = torch.cuda.ExternalStream(stream.ptr)
with stream, stream_pytorch:
stream.begin_capture()
device_func(x, y)
stream.launch_host_func(my_callback, data)
g = stream.end_capture()
torch.cuda.synchronize()
print("Graph captured", flush=True)
for i in range(5):
with stream, stream_pytorch:
x.copy_(torch.ones_like(x))
y.copy_(torch.ones_like(y))
g.launch(stream=stream)
torch.cuda.synchronize()
print(y)
if __name__ == "__main__":
main()Installation
Wheel (pip install cupy-***)
Environment
Graph captured
Host callback says: Hello Host Node
tensor([[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
...,
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.]], device='cuda:0')
tensor([[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
...,
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.]], device='cuda:0')
tensor([[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
...,
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.]], device='cuda:0')
tensor([[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
...,
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.]], device='cuda:0')
tensor([[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
...,
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.],
[4096., 4096., 4096., ..., 4096., 4096., 4096.]], device='cuda:0')
You can see that the kernel is replayed correctly, while the host function is replayed only once.
Additional Information
No response
Activity
kmaehashi commentedon Aug 1, 2025
Thanks for the feedback! Reproduced on my end.
This issue looks happening because this API was designed before the CUDA graph era, and it destroys arguments passed in the first call.
cupy/cupy_backends/cuda/api/runtime.pyx
Lines 890 to 894 in c7b45b9
syuoni commentedon Aug 1, 2025
Hi @kmaehashi , thanks for replying! May I have an expectation when this issue can be fixed?
kmaehashi commentedon Aug 2, 2025
I've made a PR for this - #9279. Will try to ship the fix in the next release.