Repository navigation
fix(tirx): stabilize multi-GPU correctness tests - #20213
Merged
jinhongyii merged 3 commits intoAug 28, 2026
Merged
Conversation
jinhongyii
approved these changes
Aug 28, 2026
spectrometerHBH
added a commit
that referenced
this pull request
Aug 29, 2026
…p-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
tqchen
pushed a commit
that referenced
this pull request
Oct 16, 2026
## 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
tqchen
pushed a commit
that referenced
this pull request
Oct 16, 2026
…p-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
tqchen
pushed a commit
that referenced
this pull request
Oct 18, 2026
## 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
tqchen
pushed a commit
that referenced
this pull request
Oct 18, 2026
…p-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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Testing