feat: unify attention backend with multi-dim mask support

- Add attention() functional entry delegating to active backend
- GQA/MLA forward calls attention() instead of inline cache/SDPA
- CUDA kernels support 2D/3D/4D mask via mask_h_stride field
- CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens
- KVCache.max_len precomputed in bind_tasks to avoid GPU sync
- batch==1 decode short-circuits mask=None
- Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask
- 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native

tests/extension/test_kernel_mask.py

@@ -0,0 +1,74 @@
+"""Kernel-level mask dimension support (2D, 3D, 4D)."""
+
+import torch
+
+from tests.extension.conftest import D, skip_no_cuda
+
+
+@skip_no_cuda
+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)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(batch, kv_len, dtype=torch.bool, device="cuda")
+    mask[:, 4:] = False
+
+    out = attn_prefill(q, k, v, mask=mask, is_causal=False)
+    assert out.shape == (batch, q_len, n_heads, D)
+
+
+@skip_no_cuda
+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)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(batch, q_len, kv_len, dtype=torch.bool, device="cuda")
+
+    out = attn_prefill(q, k, v, mask=mask, is_causal=False)
+    assert out.shape == (batch, q_len, n_heads, D)
+
+
+@skip_no_cuda
+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)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(batch, 1, q_len, kv_len, dtype=torch.bool, device="cuda")
+    mask[:, :, :, 4:] = False
+
+    out = attn_prefill(q, k, v, mask=mask, is_causal=False)
+    assert out.shape == (batch, q_len, n_heads, D)
+
+
+@skip_no_cuda
+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)
+    k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
+
+    out_no_mask = attn_prefill(q, k, v, mask=None, is_causal=False)
+    mask = torch.ones(batch, 1, q_len, kv_len, dtype=torch.bool, device="cuda")
+    out_with_mask = attn_prefill(q, k, v, mask=mask, is_causal=False)
+
+    diff = (out_no_mask.float() - out_with_mask.float()).abs().max().item()
+    assert diff == 0.0, f"4D all-True mask diff: {diff}"