Preserve the CUDA default stream in DSA wrappers

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

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