refactor: separate extension ops and backends

tests/extension/test_backend.py

@@ -10,6 +10,8 @@ from astrai.extension import (
     ATTN_BACKEND,
     AttentionBackendFactory,
     CudaBackend,
+    FlashAttnBackend,
+    TorchNativeBackend,
     attn_backend,
     get_backend,
 )
@@ -17,13 +19,6 @@ from astrai.extension import (
 
 def test_default_backend_resolves_to_available():
     """Default backend is the first available in cuda > flash > torch order."""
-    from astrai.extension.attention_backend import (
-        CudaBackend,
-        FlashAttnBackend,
-        TorchNativeBackend,
-        _resolve_default_backend,
-    )
-
     backend = get_backend()
     assert isinstance(backend, (CudaBackend, FlashAttnBackend, TorchNativeBackend))
 

tests/extension/test_kernel_mask.py

@@ -2,14 +2,13 @@
 
 import torch
 
+from astrai.extension.ops.attention import attn_prefill
 from tests.extension.conftest import D, skip_no_kernel
 
 
 @skip_no_kernel
 def test_kernel_accepts_2d_mask():
     """Kernel should accept 2D mask [batch, kv_len]."""
-    from astrai.extension.attention_ops import attn_prefill
-
     batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
     kv_len = 8
     q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
@@ -25,8 +24,6 @@ def test_kernel_accepts_2d_mask():
 @skip_no_kernel
 def test_kernel_accepts_3d_mask():
     """Kernel should accept 3D mask [batch, q_len, kv_len]."""
-    from astrai.extension.attention_ops import attn_prefill
-
     batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
     kv_len = 8
     q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
@@ -41,8 +38,6 @@ def test_kernel_accepts_3d_mask():
 @skip_no_kernel
 def test_kernel_accepts_4d_mask():
     """Kernel should accept 4D mask [batch, n_heads, q_len, kv_len]."""
-    from astrai.extension.attention_ops import attn_prefill
-
     batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
     kv_len = 8
     q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
@@ -58,8 +53,6 @@ def test_kernel_accepts_4d_mask():
 @skip_no_kernel
 def test_4d_mask_matches_no_mask_when_all_true():
     """A 4D all-True mask should produce the same output as no mask."""
-    from astrai.extension.attention_ops import attn_prefill
-
     batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
     kv_len = 8
     q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)