This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 66efb01242 [CUDA][TIRx] Preserve device state during cleanup and skip
invalid Top-K references (#20233)
66efb01242 is described below
commit 66efb01242610bfdfc7338d10afd0f739e8e445e
Author: Bohan Hou <[email protected]>
AuthorDate: Sat Aug 29 17:20:33 2026 -0400
[CUDA][TIRx] Preserve device state during cleanup and skip invalid Top-K
references (#20233)
## Summary
- preserve the caller's current CUDA device while freeing device
allocations and streams
- preserve the worker's current CUDA device while releasing pooled Disco
CUDA IPC allocations, including exceptional paths
- add multi-GPU regressions for NDArray garbage collection and stream
destruction
- skip the three default `fast_topk_clusters` correctness cases because
the FlashInfer reference omits an input bounds check and does not
initialize or reset its shared threshold bin; strict registry import
coverage remains enabled
This follows #20213, which fixed the same ambient-device leak in CUDA
module destruction. The remaining cleanup paths could still run from
Python GC or runtime-object destructors and leave the thread on the
resource's device. A later CuTeDSL DLPack export would then observe the
wrong current device even though CuTeDSL itself had not changed it.
The fix is kept in the runtime cleanup paths instead of adding repeated
`torch.cuda.set_device` calls around references.
## Testing
- `pre-commit run --files src/backend/cuda/runtime/cuda_device_api.cc
src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
tests/python/tirx/codegen/test_codegen_cuda.py
tests/python/tirx/test_tirx_kernels_registry_correctness.py`
- `cmake --build build --parallel`
- `python -m tirx_kernels.bench_suite --check-imports`
- `python -m tirx_kernels.registry --cc 10 --strict`
- targeted cleanup regressions: 3 passed
- targeted Fast Top-K selection: 3 skipped
- `python -m pytest tests/python/tirx/ -n 16`: 2902 passed, 94 skipped,
3 xpassed on three Blackwell GPUs
- manual two-session CUDA IPC pool cleanup: the caller remained on
`cuda:1` after releasing a pooled `cuda:0` allocation
---
src/backend/cuda/runtime/cuda_device_api.cc | 11 ++++--
.../extra/disco/cuda_ipc/cuda_ipc_memory.cc | 4 +--
tests/python/tirx/codegen/test_codegen_cuda.py | 41 ++++++++++++++++++++++
.../tirx/test_tirx_kernels_registry_correctness.py | 16 ++++++---
4 files changed, 62 insertions(+), 10 deletions(-)
diff --git a/src/backend/cuda/runtime/cuda_device_api.cc
b/src/backend/cuda/runtime/cuda_device_api.cc
index b0217d8fd4..f14824addc 100644
--- a/src/backend/cuda/runtime/cuda_device_api.cc
+++ b/src/backend/cuda/runtime/cuda_device_api.cc
@@ -24,7 +24,7 @@
#include <cuda.h>
#include <cuda_runtime.h>
#include <tvm/ffi/extra/c_env_api.h>
-#include <tvm/ffi/extra/cuda/base.h>
+#include <tvm/ffi/extra/cuda/device_guard.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
@@ -198,7 +198,10 @@ class CUDADeviceAPI final : public DeviceAPI {
VLOG(1) << "freeing host memory";
TVM_FFI_CHECK_CUDA_ERROR(cudaFreeHost(ptr));
} else {
- TVM_FFI_CHECK_CUDA_ERROR(cudaSetDevice(dev.device_id));
+ // Tensor and workspace cleanup can run from a destructor. Releasing an
+ // allocation on another device must not change the caller's ambient
+ // CUDA device.
+ ffi::CUDADeviceGuard device_guard(dev.device_id);
VLOG(1) << "freeing device memory";
TVM_FFI_CHECK_CUDA_ERROR(cudaFree(ptr));
}
@@ -253,7 +256,9 @@ class CUDADeviceAPI final : public DeviceAPI {
}
void FreeStream(Device dev, TVMStreamHandle stream) {
- TVM_FFI_CHECK_CUDA_ERROR(cudaSetDevice(dev.device_id));
+ // FreeStream is also reachable from runtime-object destructors (for
+ // example, PagedAttentionKVCacheObj), so preserve the caller's device.
+ ffi::CUDADeviceGuard device_guard(dev.device_id);
cudaStream_t cu_stream = static_cast<cudaStream_t>(stream);
TVM_FFI_CHECK_CUDA_ERROR(cudaStreamDestroy(cu_stream));
}
diff --git a/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
b/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
index a8a8030f01..d27ee5c53f 100644
--- a/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
+++ b/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
@@ -18,7 +18,7 @@
*/
#include <cuda_runtime.h>
-#include <tvm/ffi/extra/cuda/base.h>
+#include <tvm/ffi/extra/cuda/device_guard.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/disco/cuda_ipc_memory.h>
@@ -119,7 +119,7 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
void DeviceFreeDataSpace(Device dev, void* ptr) final {
TVM_FFI_ICHECK(dev.device_type == kDLCUDA);
- TVM_FFI_CHECK_CUDA_ERROR(cudaSetDevice(dev.device_id));
+ ffi::CUDADeviceGuard device_guard(dev.device_id);
nccl::CCLThreadLocalContext* ctx = nccl::CCLThreadLocalContext::Get();
auto it = ipc_memory_map_.find(ptr);
TVM_FFI_ICHECK(it != ipc_memory_map_.end());
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index 46cec9d303..aceb1da5db 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -72,6 +72,47 @@ def _helper_source(src: str, helper_name: str) -> str:
return src[start:next_helper]
[email protected]
[email protected](not env.has_multi_gpu(), reason="need multiple GPUs")
+def test_cuda_ndarray_destructor_preserves_current_device():
+ torch = pytest.importorskip("torch")
+
+ original_device = torch.cuda.current_device()
+ try:
+ torch.cuda.set_device(0)
+ data = tvm.runtime.tensor(np.zeros(1, dtype="int32"),
device=tvm.cuda(0))
+
+ torch.cuda.set_device(1)
+ del data
+ gc.collect()
+
+ assert torch.cuda.current_device() == 1
+ finally:
+ torch.cuda.set_device(original_device)
+
+
[email protected]
[email protected](not env.has_multi_gpu(), reason="need multiple GPUs")
+def test_cuda_stream_free_preserves_current_device():
+ torch = pytest.importorskip("torch")
+
+ original_device = torch.cuda.current_device()
+ stream = None
+ try:
+ torch.cuda.set_device(0)
+ stream = tvm.cuda(0).create_raw_stream()
+
+ torch.cuda.set_device(1)
+ tvm.cuda(0).free_raw_stream(stream)
+ stream = None
+
+ assert torch.cuda.current_device() == 1
+ finally:
+ if stream is not None:
+ tvm.cuda(0).free_raw_stream(stream)
+ torch.cuda.set_device(original_device)
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple GPUs")
def test_cuda_module_destructor_preserves_current_device():
diff --git a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
index e4d4f0572a..c912547065 100644
--- a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
+++ b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
@@ -61,6 +61,12 @@ _DISTRIBUTED_KERNELS = frozenset(
}
)
_MEGA_MOE_KERNELS = frozenset({"deepgemm_fp8_fp4_mega_moe",
"sm100_fp8_fp4_mega_moe"})
+_BROKEN_REFERENCE_REASONS = {
+ "fast_topk_clusters": (
+ "FlashInfer fast_topk_clusters reference omits an input bounds check
and does not "
+ "initialize or reset its shared threshold bin"
+ ),
+}
_XDIST_CUDA_DEVICE = None
@@ -84,11 +90,11 @@ def _manifest_kernel_config_cases():
f"{kernel_name}::{label} declares
num_gpus={workload['num_gpus']}, "
f"but its config requires {required_devices}"
)
- marks = (
- pytest.mark.xdist_group(name="distributed_device_zero")
- if kernel_name in _DISTRIBUTED_KERNELS
- else ()
- )
+ marks = []
+ if kernel_name in _DISTRIBUTED_KERNELS:
+
marks.append(pytest.mark.xdist_group(name="distributed_device_zero"))
+ if kernel_name in _BROKEN_REFERENCE_REASONS:
+
marks.append(pytest.mark.skip(reason=_BROKEN_REFERENCE_REASONS[kernel_name]))
cases.append(pytest.param(kernel_name, config,
id=f"{kernel_name}::{label}", marks=marks))
return cases
