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/backend/swiglu.py

@@ -1,46 +1,18 @@
 """Inference-only fused SwiGLU selection for dense MLP layers."""
 
-import logging
-import os
-from functools import cache
-
 import torch
 import torch.nn.functional as F
 from torch import Tensor
 
 from astrai.extension.backend.linear import linear
+from astrai.extension.dispatch import env_mode
 from astrai.extension.loader import is_available
 from astrai.extension.ops.swiglu import bf16_swiglu
 
-logger = logging.getLogger(__name__)
-
-# Shape keys are (N, K) for the paired up/gate projections. Automatic entries
-# are populated only after the primitive, MLP chain, and greedy checkpoint
-# gates pass on that architecture.
-_AUTO_SWIGLU_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = {}
-_AUTO_SWIGLU_M = frozenset(
-    m for architecture in _AUTO_SWIGLU_SHAPES.values() for m in architecture
-)
-_VALID_MODES = {"0", "1", "auto"}
-_WARNED_MODES: set[str] = set()
-
-
-def _swiglu_mode() -> str:
-    mode = os.environ.get("ASTRAI_SWIGLU", "auto").strip().lower()
-    if mode in _VALID_MODES:
-        return mode
-    if mode not in _WARNED_MODES:
-        _WARNED_MODES.add(mode)
-        logger.warning(
-            "ASTRAI_SWIGLU=%r is invalid; expected 0, 1, or auto; using auto",
-            mode,
-        )
-    return "auto"
-
 
 def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
     # Keep the existing linear backend in the fallback chain. This preserves
-    # any independently qualified GEMV shapes instead of making the fusion
+    # any independently qualified GEMV batches instead of making the fusion
     # decision suppress linear-level optimizations.
     return linear(x, up_weight) * F.silu(linear(x, gate_weight))
 
@@ -49,11 +21,6 @@ def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
     return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())
 
 
-@cache
-def _device_capability(device_index: int) -> tuple[int, int]:
-    return torch.cuda.get_device_capability(device_index)
-
-
 def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
     return not (
         torch.is_grad_enabled()
@@ -77,33 +44,19 @@ def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
     )
 
 
-def _auto_swiglu_shape(x: Tensor, up_weight: Tensor) -> bool:
-    capability = _device_capability(x.get_device())
-    m = 1 if x.ndim == 1 else x.shape[0]
-    return (up_weight.shape[0], up_weight.shape[1]) in _AUTO_SWIGLU_SHAPES.get(
-        capability, {}
-    ).get(m, ())
-
-
 def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
     """Apply the dense-MLP SwiGLU projection with a safe torch fallback.
 
-    ``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain, ``1`` forces
-    the fused primitive for supported inputs, and ``auto`` uses only
-    architecture/shape bands backed by benchmark and checkpoint evidence.
+    ``ASTRAI_SWIGLU=0`` and ``auto`` keep the unfused linear-backend chain;
+    ``1`` forces the fused primitive for supported inputs. Auto will adopt
+    an M-banded rule mirroring the linear backend once end-to-end evidence
+    qualifies one.
     """
-    mode = _swiglu_mode()
-    if mode == "0" or (mode == "auto" and not _AUTO_SWIGLU_SHAPES):
-        return _unfused_swiglu(x, up_weight, gate_weight)
-    if mode == "auto":
-        m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
-        if m not in _AUTO_SWIGLU_M:
-            return _unfused_swiglu(x, up_weight, gate_weight)
-    if _swiglu_capable(x, up_weight, gate_weight) and (
-        mode == "1" or _auto_swiglu_shape(x, up_weight)
+    if env_mode("ASTRAI_SWIGLU") != "1" or not _swiglu_capable(
+        x, up_weight, gate_weight
     ):
-        return _fused_swiglu(x, up_weight, gate_weight)
-    return _unfused_swiglu(x, up_weight, gate_weight)
+        return _unfused_swiglu(x, up_weight, gate_weight)
+    return _fused_swiglu(x, up_weight, gate_weight)
 
 
 __all__ = ["swiglu"]