Thank you very much for the quick response @glaringlee. I can create a CUDA stream and use it here in the extension. But it will not make sense because the CUDA stream I used for the extension is not the same as the one used in the neural network backbone. Therefore, the execution of the extension will not be truly async. Ideally, it is PyTorch’s responsibility to expose the PyTorch CUDA stream to the user.
I think TensorFlow exposes CUDA stream to the user:

  
      
      TensorFlow
  
  
    

Create an op  |  TensorFlow Core



  
  
    
    
  
  


Please ctrl + F and search “stream”.