test: deduplicate suites and prune low-value cases

- extract shared helpers for dataset writers, scheduler construction, thread interleaving, hf roundtrips, and moe configs
- remove about 20 cases whose only assertions were format checks, restated declarations, fake-taxonomy duplicates, or test-local scaffolding
- strengthen weak cases into exact reference comparisons, positional mask checks, and deterministic outcomes
- replace two schedule factory smoke tests with cosine/sgdr formula assertions
- delete root-level CLI tests whose merge-priority facts are covered by tests/config/test_cli.py
- suite shrinks from 857 to 826 items; ruff format, import order, and pytest all green

tests/extension/test_kernel_mask.py

@@ -1,5 +1,7 @@
 """Kernel-level mask dimension support (2D, 3D, 4D)."""
 
+import math
+
 import torch
 
 from astrai.extension.ops.attention import attn_prefill
@@ -7,9 +9,27 @@ from tests.conftest import skip_no_kernel
 from tests.extension.conftest import D
 
 
+def _reference(q, k, v, mask):
+    """fp32 masked GQA attention reference (True=keep)."""
+    b, s_q, h, d = q.shape
+    rep = h // k.shape[2]
+    qf = q.float().transpose(1, 2)
+    kf = k.float().repeat_interleave(rep, dim=2).transpose(1, 2)
+    vf = v.float().repeat_interleave(rep, dim=2).transpose(1, 2)
+    scores = qf @ kf.transpose(-1, -2) / math.sqrt(d)
+    if mask.dim() == 2:
+        mask = mask[:, None, None, :]
+    elif mask.dim() == 3:
+        mask = mask[:, None, :, :]
+    # 4D [batch, 1, q_len, kv_len] broadcasts over heads as-is
+    scores = scores.masked_fill(~mask, float("-inf"))
+    return (scores.softmax(dim=-1) @ vf).transpose(1, 2).to(q.dtype)
+
+
 @skip_no_kernel
 def test_kernel_accepts_2d_mask():
-    """Kernel should accept 2D mask [batch, kv_len]."""
+    """2D mask [batch, kv_len] gates the softmax, not just parses."""
+    torch.manual_seed(11)
     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)
@@ -20,25 +40,31 @@ def test_kernel_accepts_2d_mask():
 
     out = attn_prefill(q, k, v, mask=mask, is_causal=False)
     assert out.shape == (batch, q_len, n_heads, D)
+    torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05)
 
 
 @skip_no_kernel
 def test_kernel_accepts_3d_mask():
-    """Kernel should accept 3D mask [batch, q_len, kv_len]."""
+    """3D mask [batch, q_len, kv_len] applies per-query-row gating."""
+    torch.manual_seed(12)
     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")
+    mask[:, 0, 5:] = False  # differs per query row: only the 3D path can apply it
+    mask[:, 1, 6:] = False
 
     out = attn_prefill(q, k, v, mask=mask, is_causal=False)
     assert out.shape == (batch, q_len, n_heads, D)
+    torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05)
 
 
 @skip_no_kernel
 def test_kernel_accepts_4d_mask():
-    """Kernel should accept 4D mask [batch, n_heads, q_len, kv_len]."""
+    """4D mask [batch, 1, q_len, kv_len] broadcasts over heads and gates."""
+    torch.manual_seed(13)
     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)
@@ -49,6 +75,7 @@ def test_kernel_accepts_4d_mask():
 
     out = attn_prefill(q, k, v, mask=mask, is_causal=False)
     assert out.shape == (batch, q_len, n_heads, D)
+    torch.testing.assert_close(out, _reference(q, k, v, mask), atol=0.05, rtol=0.05)
 
 
 @skip_no_kernel

tests/extension/test_linear_dispatch.py

@@ -6,7 +6,6 @@ import torch
 import torch.nn.functional as F
 
 from astrai.extension import is_available, linear
-from astrai.extension.backend import linear as public_linear
 from astrai.extension.dispatch import explain, op_backend, resolve
 
 # The package attribute ``linear`` is the dispatched function; reach the
@@ -36,10 +35,6 @@ def _routes_to_gemv(monkeypatch, x, weight, bias=None) -> bool:
     return linear(x, weight, bias) is sentinel
 
 
-def test_linear_backend_is_public():
-    assert linear is public_linear
-
-
 def test_model_linear_routes_through_backend(monkeypatch):
     sentinel = torch.randn(2, 4)
 

tests/extension/test_swiglu_dispatch.py

@@ -38,10 +38,13 @@ def test_cpu_and_training_calls_fall_back_with_gradients(monkeypatch):
 
 def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
     monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode")
+    x = torch.randn(2, 8)
+    up_weight = torch.randn(4, 8)
+    gate_weight = torch.randn(4, 8)
     with caplog.at_level(logging.WARNING):
-        actual = swiglu(torch.randn(2, 8), torch.randn(4, 8), torch.randn(4, 8))
-    assert actual.shape == (2, 4)
+        actual = swiglu(x, up_weight, gate_weight)
     assert "using auto" in caplog.text
+    torch.testing.assert_close(actual, reference_swiglu(x, up_weight, gate_weight))
 
 
 def test_mlp_routes_through_swiglu_backend(monkeypatch):