diff --git a/python/cudnn/deepseek_sparse_attention/utils/runtime.py b/python/cudnn/deepseek_sparse_attention/utils/runtime.py index d9984a258..145fba12c 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/runtime.py +++ b/python/cudnn/deepseek_sparse_attention/utils/runtime.py @@ -56,5 +56,5 @@ def torch_stream_context(current_stream: Optional[cuda.CUstream] = None) -> Iter if current_stream is None: yield return - with torch.cuda.stream(torch.cuda.ExternalStream(int(current_stream))): + with torch.cuda.stream(torch.cuda.get_stream_from_external(int(current_stream))): yield diff --git a/test/python/fe_api/dsa/test_DSA_runtime.py b/test/python/fe_api/dsa/test_DSA_runtime.py new file mode 100644 index 000000000..1b4611f21 --- /dev/null +++ b/test/python/fe_api/dsa/test_DSA_runtime.py @@ -0,0 +1,17 @@ +import pytest +import torch + + +@pytest.mark.L0 +def test_torch_stream_context_preserves_legacy_default_stream(): + try: + from cuda.bindings import driver as cuda + from cudnn.deepseek_sparse_attention.utils.runtime import torch_stream_context + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + outer_stream = torch.cuda.Stream() + with torch.cuda.stream(outer_stream): + with torch_stream_context(cuda.CUstream(0)): + assert torch.cuda.current_stream().cuda_stream == 0 + assert torch.cuda.current_stream().cuda_stream == outer_stream.cuda_stream