Relative import first

ops/__init__.py

@@ -0,0 +1,10 @@
+# ops/__init__.py
+from .pooling import mean_pooling
+from .compressed_attention import compressed_attention
+from .topk_sparse_attention import topk_sparse_attention
+
+__all__ = [
+    "mean_pooling",
+    "compressed_attention",
+    "topk_sparse_attention",
+]

ops/compressed_attention.py

@@ -20,9 +20,9 @@ import triton
 import triton.language as tl
 
 try:
-    from ops.utils import get_num_warps_stages, is_hopper_gpu
-except ImportError:
     from .ops.utils import get_num_warps_stages, is_hopper_gpu
+except ImportError:
+    from ops.utils import get_num_warps_stages, is_hopper_gpu
 
 IS_HOPPER_GPU = is_hopper_gpu()
 

ops/topk_sparse_attention.py

@@ -19,9 +19,9 @@ import triton
 import triton.language as tl
 
 try:
-    from ops.utils import get_num_warps_stages, is_hopper_gpu
-except ImportError:
     from .ops.utils import get_num_warps_stages, is_hopper_gpu
+except ImportError:
+    from ops.utils import get_num_warps_stages, is_hopper_gpu
 
 IS_HOPPER_GPU = is_hopper_gpu()