This is an automated email from the ASF dual-hosted git repository.
jinhongyii 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 b4e62497dc fix(tirx): stabilize multi-GPU correctness tests (#20213)
b4e62497dc is described below
commit b4e62497dc76461f1365e71df13bd6da2e90fb66
Author: Bohan Hou <[email protected]>
AuthorDate: Fri Aug 28 00:08:05 2026 -0400
fix(tirx): stabilize multi-GPU correctness tests (#20213)
## Summary
- preserve the caller current CUDA device while unloading per-device
modules
- add a multi-GPU regression test for module cleanup
- align the TensorMap TFLOAT32 validation expectation with the current
diagnostic
- skip MSA registry cases that require an isolated CuTeDSL 4.5.3 process
## Testing
- python -m pytest -q tests/python/tirx/codegen/test_codegen_cuda.py
- targeted TensorMap and registry correctness tests
- pre-commit run --files src/backend/cuda/runtime/cuda_module.cc
tests/python/tirx/codegen/test_codegen_cuda.py
tests/python/tirx/codegen/test_codegen_hopper.py
tests/python/tirx/test_tirx_kernels_registry_correctness.py
- python -m pytest tests/python/tirx -n16 --tb=short -ra: 2845 passed,
91 skipped, 3 xpassed
---
src/backend/cuda/runtime/cuda_module.cc | 6 ++++
tests/python/tirx/codegen/test_codegen_cuda.py | 32 ++++++++++++++++++++++
tests/python/tirx/codegen/test_codegen_hopper.py | 2 +-
.../tirx/test_tirx_kernels_registry_correctness.py | 2 ++
4 files changed, 41 insertions(+), 1 deletion(-)
diff --git a/src/backend/cuda/runtime/cuda_module.cc
b/src/backend/cuda/runtime/cuda_module.cc
index 604a7de33c..3984bfde39 100644
--- a/src/backend/cuda/runtime/cuda_module.cc
+++ b/src/backend/cuda/runtime/cuda_module.cc
@@ -74,6 +74,8 @@ class CUDAModuleNode : public ffi::ModuleObj {
}
// destructor
~CUDAModuleNode() {
+ int previous_device = -1;
+ cudaError_t get_device_err = cudaGetDevice(&previous_device);
for (size_t i = 0; i < module_.size(); ++i) {
if (module_[i] != nullptr) {
cudaError_t set_err = cudaSetDevice(static_cast<int>(i));
@@ -85,6 +87,10 @@ class CUDAModuleNode : public ffi::ModuleObj {
(void)result;
}
}
+ if (get_device_err == cudaSuccess) {
+ // Preserve the caller's current device after unloading per-device
modules.
+ (void)cudaSetDevice(previous_device);
+ }
}
const char* kind() const final { return "cuda"; }
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index 93993e11b9..f9b47645e8 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -15,6 +15,7 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=missing-function-docstring
+import gc
import re
import numpy as np
@@ -71,6 +72,37 @@ 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_module_destructor_preserves_current_device():
+ torch = pytest.importorskip("torch")
+
+ @T.prim_func
+ def main(A: T.Buffer((1,), "int32")):
+ T.device_entry()
+ tx = T.thread_id([1])
+ if tx == 0:
+ A[0] = A[0] + 1
+
+ _, mod = _get_source(main)
+ 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))
+ mod["main"](data)
+ tvm.cuda(0).sync()
+ del data
+ gc.collect()
+
+ torch.cuda.set_device(1)
+ del mod
+ gc.collect()
+
+ assert torch.cuda.current_device() == 1
+ finally:
+ torch.cuda.set_device(original_device)
+
+
def test_vector_access_ptr_preserves_packed_offset(monkeypatch):
buffer = tvm.tirx.decl_buffer((8,), "int4x4", name="A")
data = tvm.tirx.Var("A_data", tvm.tirx.buffer_data_pointer_type(buffer))
diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py
b/tests/python/tirx/codegen/test_codegen_hopper.py
index e25c5a0027..1ac6c43cd8 100644
--- a/tests/python/tirx/codegen/test_codegen_hopper.py
+++ b/tests/python/tirx/codegen/test_codegen_hopper.py
@@ -583,7 +583,7 @@ def
test_cp_async_bulk_tensor_global_to_shared_unicast(dtype, inputs):
(16, 16),
"float32",
[16, 16, 64, 4, 16, 1, 1, 0, 0, 0, 0, 7],
- r"force_cu_dtype only supports CU_TENSOR_MAP_DATA_TYPE_TFLOAT32",
+ r"force_cu_dtype accepts CU_TENSOR_MAP_DATA_TYPE_TFLOAT32",
),
(
(16, 16),
diff --git a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
index 43b8a71b98..e4d4f0572a 100644
--- a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
+++ b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
@@ -187,5 +187,7 @@ def test_manifest_tirx_kernel_correctness(kernel_name,
config):
"MegaMoE requires its dedicated multi-process scheduler; this
suite's "
"processes own CUDA contexts that its physical-device assignment
rejects"
)
+ if getattr(_KERNELS[kernel_name], "KERNEL_META", {}).get("category") ==
"msa":
+ pytest.skip("MSA references require an isolated CuTeDSL 4.5.3 process")
with _registry_gpu_lock(kernel_name, config):
kernel_runner.run_kernel_test(kernel_name, config, registry=_KERNELS)
