perf: dispatch linear gemv by decode batch size and unify extension style

- replace the per-shape auto tables in the linear backend with an M-banded rule (M in [2,4] on compute capability 8.0+) that measured at the HBM bandwidth floor across every family, and fold the capability check into the capable guard
- drop the unreachable swiglu auto shape-table machinery so both backends share one env-mode ladder via the new dispatch.env_mode helper
- add __all__ across extension modules, name the rotary registration records, and unify typing to the typing-module style
- rewrite test_linear_dispatch.py around behavioral routing assertions and document the M-banded policy in the developer docs
- Benchmark: L20 SM89, Python dispatch overhead 2.9us to 1.5us, auto now covers every projection shape at M in [2,4].

astrai/extension/ops/attention.py

@@ -190,3 +190,12 @@ def attn_paged_prefill(
         mask,
         causal_offset=causal_offset,
     )
+
+
+__all__ = [
+    "TensorLayout",
+    "attn_decode",
+    "attn_paged_decode",
+    "attn_paged_prefill",
+    "attn_prefill",
+]

astrai/extension/ops/fp8.py

@@ -114,3 +114,6 @@ def mm_fp8(
     BF16; FP8 output is a separate quantize operation.
     """
     return get_module("fp8_ops").mm_fp8(a, b, scale, trans_a, trans_b, bias)
+
+
+__all__ = ["mm_fp8", "quantize", "quantize_dual"]

astrai/extension/ops/gemv.py

@@ -21,3 +21,6 @@ def bf16_gemv(
     fallback or model-level dispatch.
     """
     return get_module("bf16_gemv").bf16_gemv(x, weight, bias)
+
+
+__all__ = ["bf16_gemv"]

astrai/extension/ops/rotary.py

@@ -29,3 +29,6 @@ def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
     if not freqs_cis.is_contiguous():
         freqs_cis = freqs_cis.contiguous()
     return mod.rotary_emb(x, freqs_cis)
+
+
+__all__ = ["rotary_emb"]