deepseek-ai / deepseek-ai/DeepSelect

Calling Top-K with input on a non-current GPU uses the wrong device

Open
#3 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Cuda
Stars
343
Forks
21
Avg merge
10m
Merged PRs (30d)
2

Description

`topk()` builds its launch arguments from the current CUDA device and its current stream, even when `input` belongs to another GPU. The binding also accepts `end`, output tensors, and `output_idx_offset` from different CUDA devices because its device checks only require CUDA tensors.

For example, this call should execute on device 1 and leave device 0 current (requires a supported GPU at device 1):

```python
import torch
import deep_select

torch.cuda.set_device(0)
x = torch.arange(256, device="cuda:1", dtype=torch.float32).reshape(1, 256)
values, indices = deep_select.topk(x, 8, sorted=True)
torch.cuda.synchronize(1)
assert torch.cuda.current_device() == 0
torch.testing.assert_close(values, torch.topk(x, 8).values)
```

This is a source-derived reproduction, not a recorded multi-GPU run. In `csrc/api.cpp`, `getDeviceProperties(current_device())` and `getCurrentCUDAStream()` select device 0 while the launch receives device-1 pointers. Depending on device access and configuration, that can fail or access memory through an unintended device and stream. No input-device guard is installed before dispatch.

The plain pybind tensor conversion does not supply that guard: [PyTorch's tensor caster](https://github.com/pytorch/pytorch/blob/main/torch/csrc/utils.cpp) only unpacks the tensor. [CUDAGuard](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDAGuard.h) provides scoped device selection; [CUDAStream.h](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDAStream.h) documents that current streams are maintained per device.

Expected behavior: select the input device for launch setup and execution, use that device's current stream, restore the caller's current device, and reject mixed-device tensor arguments before launch.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in csrc/api.cpp, focusing on how device properties, CUDA streams, and tensor arguments are selected before the Top-K launch. Use the issue's multi-GPU reproduction as the first check, then verify mixed-device arguments are rejected and that the caller's current device is restored after execution. Done means launch setup follows the input device and its stream without changing the caller's current device.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.