fix: resolve audited dispatch, kernel, and rollout bugs

- re-register the linear family with the operator dispatcher (ASTR_OPS / op_backend / resolve)
- fix bf16 gemv misaligned-address faults and element mispairing for offset weights
- reject misaligned bf16_swiglu inputs with a clear error and fall back in the backend gate
- make the rollout reuse decision, validation, and return atomic under one policy snapshot
- add the documented post-scoring rollout version check
- derive live+1 under the scheduler lock in optimizer_step via apply_weight_update(None, ...)
- reject rollout_max_policy_lag below rollout_interval - 1 at config time
- sync gemv stream-test inputs before switching streams; drop dead loader imports

astrai/extension/backend/linear.py

@@ -6,15 +6,31 @@ selection is keyed on the decode batch size alone (M in [2, 4], where it
 sits at the HBM bandwidth floor and beat the cuBLAS small-M path on every
 measured family). Every training, prefill-sized, out-of-band, or
 unsupported call falls back to PyTorch.
+
+The family stays registered with the shared operator dispatcher, so
+``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
+``explain`` keep working like for attention and rotary. The per-layer
+hot path only consults the dispatcher when one of those selections is
+active, keeping it free of axes dictionaries and record sorting.
 """
 
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
 import torch
 import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.extension.dispatch import env_mode
+from astrai.extension.dispatch import (
+    ImplRecord,
+    Spec,
+    axis,
+    env_mode,
+    env_selection,
+    get_override,
+    register_family,
+    resolve,
+    tensor_axes,
+)
 from astrai.extension.loader import is_available
 from astrai.extension.ops.gemv import bf16_gemv
 
@@ -25,6 +41,10 @@ from astrai.extension.ops.gemv import bf16_gemv
 _AUTO_GEMV_M = frozenset({2, 3, 4})
 
 
+def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
+    return F.linear(x, weight, bias)
+
+
 def _inference_bf16_gemv(
     x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
 ) -> Tensor:
@@ -64,6 +84,108 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
     )
 
 
+def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
+    weight_shape = tuple(weight.shape)
+    m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
+    supported_m = m is not None and 1 <= m <= 8
+    shape_matches = (
+        weight.ndim == 2
+        and x.ndim in (1, 2)
+        and bool(x.shape)
+        and x.shape[-1] == weight_shape[-1]
+    )
+    same_device = x.device == weight.device and (
+        bias is None or bias.device == x.device
+    )
+    bias_supported = bias is None or (
+        bias.ndim == 1
+        and weight.ndim == 2
+        and bias.shape[0] == weight_shape[0]
+        and bias.dtype == torch.bfloat16
+        and bias.is_contiguous()
+    )
+    capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None
+    return tensor_axes(
+        x,
+        mode=env_mode("ASTRAI_GEMV"),
+        m=m,
+        supported_m=supported_m,
+        auto_m=m in _AUTO_GEMV_M,
+        shape_matches=shape_matches,
+        same_device=same_device,
+        weight_dtype=weight.dtype,
+        x_contiguous=x.is_contiguous(),
+        weight_contiguous=weight.is_contiguous(),
+        bias_supported=bias_supported,
+        capability=capability,
+    )
+
+
+_SPEC_CAPABLE = (
+    axis("device_cuda").truthy()
+    & axis("dtype").in_(torch.bfloat16)
+    & axis("weight_dtype").in_(torch.bfloat16)
+    & axis("grad_enabled").eq(False)
+    & axis("supported_m").truthy()
+    & axis("shape_matches").truthy()
+    & axis("same_device").truthy()
+    & axis("x_contiguous").truthy()
+    & axis("weight_contiguous").truthy()
+    & axis("bias_supported").truthy()
+    & Spec.of(
+        lambda ax: ax.get("capability") is not None and ax.get("capability") >= (8, 0),
+        "capability>=sm_80",
+    )
+)
+
+_SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()
+
+
+def _linear_records() -> List[ImplRecord]:
+    mode = env_mode("ASTRAI_GEMV")
+    gemv_priority = 0 if mode == "1" else 100
+    auto_priority = 0 if mode == "auto" else 90
+    torch_priority = 0 if mode == "0" else 50
+    return [
+        ImplRecord(
+            family="linear",
+            name="gemv",
+            obj=_inference_bf16_gemv,
+            spec=_SPEC_CAPABLE,
+            available=lambda: is_available("bf16_gemv"),
+            priority=gemv_priority,
+        ),
+        ImplRecord(
+            family="linear",
+            name="auto_gemv",
+            obj=_inference_bf16_gemv,
+            spec=_SPEC_AUTO,
+            available=lambda: is_available("bf16_gemv"),
+            priority=auto_priority,
+        ),
+        ImplRecord(
+            family="linear",
+            name="torch",
+            obj=_torch_linear,
+            spec=Spec.always(),
+            priority=torch_priority,
+        ),
+    ]
+
+
+def _fallback_record() -> ImplRecord:
+    return ImplRecord(
+        family="linear",
+        name="torch",
+        obj=_torch_linear,
+        spec=Spec.always(),
+        priority=999,
+    )
+
+
+register_family("linear", _axes, _linear_records, _fallback_record)
+
+
 def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
     """Apply a linear projection with safe inference-only GEMV dispatch.
 
@@ -71,12 +193,18 @@ def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
     primitive can safely handle any M in ``{1, ..., 8}``, and ``auto`` (the
     default) uses GEMV for decode batches with M in ``{2, 3, 4}``.
     """
+    # Route through the shared dispatcher whenever a selection is active so
+    # explicit/context/env overrides stay honored; otherwise keep the hot
+    # path free of axes dictionaries and record sorting.
+    if get_override("linear") is not None or env_selection("linear") is not None:
+        return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
+
     mode = env_mode("ASTRAI_GEMV")
     if mode != "0" and _gemv_capable(x, weight, bias):
         m = 1 if x.ndim == 1 else x.shape[0]
         if mode == "1" or m in _AUTO_GEMV_M:
             return _inference_bf16_gemv(x, weight, bias)
-    return F.linear(x, weight, bias)
+    return _torch_linear(x, weight, bias)
 
 
 __all__ = ["linear"]

astrai/extension/backend/swiglu.py

@@ -40,6 +40,11 @@ def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
         or not x.is_contiguous()
         or not up_weight.is_contiguous()
         or not gate_weight.is_contiguous()
+        # The fused kernel reads all streams as uint4; contiguous-but-offset
+        # views are routed to the unfused chain instead of failing.
+        or (x.data_ptr() & 15) != 0
+        or (up_weight.data_ptr() & 15) != 0
+        or (gate_weight.data_ptr() & 15) != 0
         or not is_available("bf16_swiglu")
     )
 

astrai/extension/loader.py

@@ -20,11 +20,8 @@ import glob
 import importlib
 import logging
 import os
-from functools import cache
 from typing import Dict, List
 
-import torch
-
 logger = logging.getLogger(__name__)
 
 _LIB_DIR = os.path.join(os.path.dirname(__file__), "lib")