Multiple updates and refactorings (#280)

2026-01-16 17:06:52 +08:00
parent 3ccf40c53a
commit 0f5f266202
55 changed files with 2706 additions and 891 deletions
@@ -1,10 +1,16 @@
-from .._C import (
-    get_tma_aligned_size,
-    get_mk_alignment_for_contiguous_layout,
-    get_mn_major_tma_aligned_tensor,
-    get_mn_major_tma_aligned_packed_ue8m0_tensor,
-    get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
-)
+try:
+    from .._C import (
+        get_tma_aligned_size,
+        get_mn_major_tma_aligned_tensor,
+        get_mn_major_tma_aligned_packed_ue8m0_tensor,
+        get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
+    )
+except ImportError:
+    # Expected behavior for CUDA runtime version before 12.1
+    pass
+
+# Valid for all CUDA versions
+from .._C import get_mk_alignment_for_contiguous_layout

 # Some alias
 get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout