[torchAPI] Inherit cuda stream from torch (#618)

vadiklyutiy · vadiklyutiy · commit ad4e00a0373e · 2024-12-20T20:28:38.000+04:00
When we work through `torch.compile` we didn't inherit the stream from torch. Fix it.

PR fixes issue #563  

Tested with the main branch of vllm cloned 28th of Nov. Options 
`--model meta-llama/Llama-3.1-8B-Instruct --max_model_len 4000`

Before output was 
``` 
about!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
```

after 
```
 about a cat named Whiskers
Once upon a time, in a cozy little house on a quiet street, there lived a beautiful and elegant cat named Whiskers. Whiskers was a stunning feline with shimmering grey fur, bright green eyes, and a delicate pink nose. She was a gentle soul, with a soft and soothing purr that could calm even the most anxious of hearts.
Whiskers lived with her loving owner, a kind old lady named Mrs. Jenkins, who adored her dearly. Mrs. Jenkins had rescued Whiskers from a shelter when she was just a tiny kitten, and the two had been inseparable ever since. Whiskers spent her days lounging in the sunbeams that streamed through the windows, chasing the occasional fly, and purring contentedly as Mrs. Jenkins stroked her soft fur.
But Whiskers was more than just a lazy cat. She had a secret life, one that few people knew about. You```
```
as expected.
diff --git a/python/hidet/graph/frontend/torch/dynamo_backends.py b/python/hidet/graph/frontend/torch/dynamo_backends.py
@@ -164,8 +164,11 @@ def __call__(self, *args):
             else:
                 # ignore constant
                 pass
-
+        # Inherited cuda stream from torch
+        runtime_api.set_current_stream(torch.cuda.current_stream().cuda_stream)
+        # Prepare inputs
         tensor_args = preprocess_inputs(tensor_args)
+        # Run graph/model
         outputs = self.cgraph.run_async(tensor_args, output_to_torch_tensor=True)
         outputs: Sequence[torch.Tensor] = [
             tensor.torch() if isinstance(tensor, hidet.Tensor) else tensor for tensor in outputs
