Honor the input CUDA device and stream in Top-K

csrc/api.cpp

@@ -3,6 +3,7 @@
 #include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/CUDAEvent.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #include "kerutils/supplemental/torch_tensors.h"
 
@@ -50,6 +51,17 @@ void topk(
     KU_CHECK_DEVICE(output_value);
     KU_CHECK_DEVICE(output_index);
     KU_CHECK_DEVICE(output_idx_offset);
+
+    TORCH_CHECK(output_index.device() == input.device(), "`output_index` must be on the same device as `input`");
+    if (end.has_value()) {
+        TORCH_CHECK(end->device() == input.device(), "`end` must be on the same device as `input`");
+    }
+    if (output_value.has_value()) {
+        TORCH_CHECK(output_value->device() == input.device(), "`output_value` must be on the same device as `input`");
+    }
+    if (output_idx_offset.has_value()) {
+        TORCH_CHECK(output_idx_offset->device() == input.device(), "`output_idx_offset` must be on the same device as `input`");
+    }
 
     KU_CHECK_SHAPE(input, batch_size, vocab_size);
     KU_CHECK_SHAPE(begin, batch_size);
@@ -88,7 +100,8 @@ void topk(
         check_dim0_stride("value", *output_value, OUTPUT_STRIDE_ALIGNMENT_REQUIREMENT);
     }
 
-    cudaDeviceProp* device_prop = at::cuda::getDeviceProperties(at::cuda::current_device());
+    const c10::cuda::CUDAGuard device_guard(input.device());
+    cudaDeviceProp* device_prop = at::cuda::getDeviceProperties(input.get_device());
     TORCH_CHECK(device_prop != nullptr);
     TopkSelectArgs args = {
         (uint32_t)batch_size,
@@ -114,7 +127,7 @@ void topk(
         abort_when_nan_found,
 
         device_prop->sharedMemPerBlockOptin,
-        at::cuda::getCurrentCUDAStream().stream()
+        at::cuda::getCurrentCUDAStream(input.get_device()).stream()
     };
 
     uint32_t num_sm = device_prop->multiProcessorCount;

deep_select/interface.py

@@ -32,6 +32,7 @@ def topk(
 ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
     """
     Arguments:
+        All tensor arguments must be on the same CUDA device as input. Execution uses that device's current stream.
         input: (b, vocab_size), dtype=torch.bfloat16/torch.float. stride(0) must be a multiple of `deep_select.get_stride_requirement()[0]` bytes, and stride(1) must be 1.
         topk: int. Select topk elements for each row.
         sorted: bool. Whether to return sorted **output_val**. Only supports fp32.

tests/test_devices.py

@@ -0,0 +1,68 @@
+"""Device-boundary regressions: python -m unittest tests.test_devices -v."""
+
+import unittest
+
+import torch
+
+
+class TestDevices(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if torch.cuda.device_count() < 2:
+            raise unittest.SkipTest("requires two CUDA devices")
+        import deep_select
+        from deep_select import deep_select_cuda
+
+        cls.api = deep_select
+        cls.backend = deep_select_cuda
+
+    def test_rejects_mixed_devices(self):
+        tensors = {
+            "input": torch.zeros((1, 256), device="cuda:0"),
+            "end": torch.full((1,), 256, dtype=torch.int32, device="cuda:0"),
+            "output_value": torch.empty((1, 8), device="cuda:0"),
+            "output_index": torch.empty((1, 8), dtype=torch.int32, device="cuda:0"),
+            "output_idx_offset": torch.zeros((1,), dtype=torch.int32, device="cuda:0"),
+        }
+        for name in ("end", "output_value", "output_index", "output_idx_offset"):
+            with self.subTest(tensor=name):
+                args = dict(tensors)
+                args[name] = args[name].to("cuda:1")
+                with self.assertRaisesRegex(RuntimeError, f"`{name}` must be on the same device as `input`"):
+                    self.backend.topk(
+                        args["input"], 8, None, args["end"], False, False,
+                        args["output_value"], args["output_index"], args["output_idx_offset"],
+                        2147483647, float("-inf"), True, True,
+                    )
+
+    def test_noncurrent_device_and_stream(self):
+        supported = [i for i in range(torch.cuda.device_count())
+                     if torch.cuda.get_device_capability(i) in ((10, 0), (10, 3))]
+        if not supported:
+            self.skipTest("requires an SM100 or SM103 device")
+        target = supported[0]
+        other = next(i for i in range(torch.cuda.device_count()) if i != target)
+        stream = torch.cuda.Stream(device=target)
+        with torch.cuda.stream(stream):
+            x = torch.arange(256, device=f"cuda:{target}", dtype=torch.float32).reshape(1, 256)
+            # Warm up lazy setup before capture.
+            self.api.topk(x, 8, sorted=True)
+        stream.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            with torch.cuda.device(other):
+                values, indices = self.api.topk(x, 8, sorted=True)
+                self.assertEqual(torch.cuda.current_device(), other)
+        # Capture proves the launch uses target's current stream: graph replay
+        # must recompute the result after changing input on that same stream.
+        with torch.cuda.stream(stream):
+            x.neg_()
+            graph.replay()
+        stream.synchronize()
+        expected = torch.topk(x, 8)
+        torch.testing.assert_close(values, expected.values)
+        torch.testing.assert_close(indices.to(torch.int64), expected.indices)
+
+
+if __name__ == "__main__":
+    unittest.main()