From e85abb151e06de5505b26a522ad02b31a638e5c0 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Mon, 17 Aug 2026 20:07:01 -0400 Subject: [PATCH 1/2] fix(kernel): gate CUDA 12.8 tensor-map enums by version instead of #ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B and friends are enumerators, not macros, so every #ifdef guard around them was always false and the features they gate were unconditionally rejected -- requesting the 128B_ATOM_32B swizzle failed host validation with "Unsupported swizzle enum value: 4" even on CUDA 13.2. Use CUDA_VERSION >= 12080 checks for the CUDA 12.8 additions and drop the dead guards around the base-enum TFLOAT32/FLOAT32_FTZ members, which exist wherever the tensor-map API does. --- src/backend/cuda/runtime/cuda_device_api.cc | 52 +++++++-------------- 1 file changed, 16 insertions(+), 36 deletions(-) diff --git a/src/backend/cuda/runtime/cuda_device_api.cc b/src/backend/cuda/runtime/cuda_device_api.cc index 01222b408588..d7728563833a 100644 --- a/src/backend/cuda/runtime/cuda_device_api.cc +++ b/src/backend/cuda/runtime/cuda_device_api.cc @@ -602,15 +602,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto is_valid_swizzle = swizzle_kind == CU_TENSOR_MAP_SWIZZLE_NONE || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_32B || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_64B || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B - is_valid_swizzle = is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B - is_valid_swizzle = - is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B - is_valid_swizzle = is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B; +#if (CUDA_VERSION >= 12080) + is_valid_swizzle = is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B || + swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B || + swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B; #endif TVM_FFI_ICHECK(is_valid_swizzle) << "Unsupported swizzle enum value: " << static_cast(swizzle_kind); @@ -628,15 +623,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { << "Unsupported oobFill enum value: " << static_cast(oob_fill_kind); bool is_packed_16u4_align8 = false; -#ifdef CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B - is_packed_16u4_align8 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; -#endif bool is_packed_16u4_align16 = false; -#ifdef CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B - is_packed_16u4_align16 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; -#endif bool is_packed_16u6_align16 = false; -#ifdef CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B +#if (CUDA_VERSION >= 12080) + is_packed_16u4_align8 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; + is_packed_16u4_align16 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; is_packed_16u6_align16 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B; #endif auto is_packed_align16 = is_packed_16u4_align16 || is_packed_16u6_align16; @@ -644,27 +635,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto is_floating_dtype = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT16 || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT32 || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT64 || - cu_dtype == CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; -#ifdef CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ - is_floating_dtype = is_floating_dtype || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ; -#endif -#ifdef CU_TENSOR_MAP_DATA_TYPE_TFLOAT32 - is_floating_dtype = is_floating_dtype || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; -#endif -#ifdef CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ - is_floating_dtype = is_floating_dtype || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ; -#endif + cu_dtype == CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 || + cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ || + cu_dtype == CU_TENSOR_MAP_DATA_TYPE_TFLOAT32 || + cu_dtype == CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ; auto is_128b_swizzle = swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B - is_128b_swizzle = is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B - is_128b_swizzle = - is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B - is_128b_swizzle = is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B; +#if (CUDA_VERSION >= 12080) + is_128b_swizzle = is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B || + swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B || + swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B; #endif // Host-side validation for documented cuTensorMapEncodeTiled requirements. @@ -702,7 +682,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { if (is_packed_16u4_align16) { bool supported_swizzle = swizzle_kind == CU_TENSOR_MAP_SWIZZLE_NONE || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B +#if (CUDA_VERSION >= 12080) supported_swizzle = supported_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; #endif TVM_FFI_ICHECK(supported_swizzle) From 9c5cf2aef95e289fad7826766c3507c4b19f3a26 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Mon, 17 Aug 2026 20:07:01 -0400 Subject: [PATCH 2/2] test(kernel): release registry GPU locks on teardown failure and skip MegaMoE Release the per-device flocks in a nested finally so a poisoned CUDA context (empty_cache raising after a kernel fault) cannot leak a held lock: pytest keeps the failure's traceback frame alive, so the leaked lock's fd never closes and the worker deadlocks against itself when its next test reopens the same lock file. Skip sm100_fp8_fp4_mega_moe here: its dedicated multi-process scheduler validates physical-device assignments that reject any process already owning a CUDA context, which every process in this suite does (the free-memory probe alone creates one). --- .../test_tirx_kernels_registry_correctness.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/python/tirx/test_tirx_kernels_registry_correctness.py b/tests/python/tirx/test_tirx_kernels_registry_correctness.py index fec2016789cc..548547718873 100644 --- a/tests/python/tirx/test_tirx_kernels_registry_correctness.py +++ b/tests/python/tirx/test_tirx_kernels_registry_correctness.py @@ -53,6 +53,7 @@ def _load_workloads(): _DISTRIBUTED_KERNELS = frozenset( {"allgather_gemm", "deepgemm_fp8_fp4_mega_moe", "gemm_reduce_scatter"} ) +_MEGA_MOE_KERNELS = frozenset({"deepgemm_fp8_fp4_mega_moe", "sm100_fp8_fp4_mega_moe"}) _XDIST_CUDA_DEVICE = None @@ -157,10 +158,12 @@ def _registry_gpu_lock(kernel_name, config): ) yield finally: - torch.cuda.empty_cache() - for lock_file in reversed(lock_files): - fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) - lock_file.close() + try: + torch.cuda.empty_cache() + finally: + for lock_file in reversed(lock_files): + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + lock_file.close() @pytest.mark.parametrize(("kernel_name", "config"), _manifest_kernel_config_cases()) @@ -172,5 +175,10 @@ def test_manifest_tirx_kernel_correctness(kernel_name, config): pytest.skip( f"requires {required_devices} CUDA devices, but only {visible_devices} are visible" ) + if kernel_name in _MEGA_MOE_KERNELS: + pytest.skip( + "MegaMoE requires its dedicated multi-process scheduler; this suite's " + "processes own CUDA contexts that its physical-device assignment rejects" + ) with _registry_gpu_lock(kernel_name, config): kernel_runner.run_kernel_test(kernel_name, config, registry=_KERNELS)