From a24b1e494152e14ecd2ed9c47e9e8dcc3a302ad4 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:34:30 +0800 Subject: [PATCH] fix: honor input CUDA device and stream in topk --- csrc/api.cpp | 17 ++++++++-- deep_select/interface.py | 1 + tests/test_devices.py | 68 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 tests/test_devices.py diff --git a/csrc/api.cpp b/csrc/api.cpp index b1d9285..ac8290c 100644 --- a/csrc/api.cpp +++ b/csrc/api.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #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; diff --git a/deep_select/interface.py b/deep_select/interface.py index 06081f2..dc43248 100644 --- a/deep_select/interface.py +++ b/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. diff --git a/tests/test_devices.py b/tests/test_devices.py new file mode 100644 index 0000000..056414c --- /dev/null +++ b/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()