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/config/train_config.py

@@ -234,4 +234,15 @@ class TrainConfig(BaseConfig):
                     f"numbers of forward passes and desynchronize the "
                     f"ddp/fsdp collectives, deadlocking NCCL"
                 )
+            if (
+                self.rollout_max_policy_lag is not None
+                and self.rollout_max_policy_lag < self.rollout_interval - 1
+            ):
+                raise ValueError(
+                    f"rollout_max_policy_lag={self.rollout_max_policy_lag} "
+                    f"cannot be below rollout_interval - 1 = "
+                    f"{self.rollout_interval - 1}: the replay cache reuses "
+                    f"rollouts up to that lag, so a tighter bound guarantees "
+                    f"a fatal RolloutVersionError mid-training"
+                )
         return self

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")

astrai/inference/scheduler.py

@@ -176,11 +176,21 @@ class InferenceScheduler:
         return self._commit_weight_version(policy_version)
 
     @_with_weight_lock
-    def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T:
-        """Mutate shared weights and publish their version without generation."""
+    def apply_weight_update(
+        self, policy_version: Optional[int], update: Callable[[], T]
+    ) -> T:
+        """Mutate shared weights and publish their version without generation.
+
+        ``policy_version=None`` derives ``live + 1`` under the same lock, for
+        callers that only need "advance by one" (e.g. ``optimizer.step()``)
+        without a read-compute-write race on the current version.
+        """
         if not callable(update):
             raise TypeError("update must be callable")
-        self._validate_weight_version(policy_version, require_advance=True)
+        if policy_version is None:
+            policy_version = self._policy_version + 1
+        else:
+            self._validate_weight_version(policy_version, require_advance=True)
         self._ensure_weight_update_ready()
 
         result = update()

astrai/trainer/rollout.py

@@ -147,8 +147,15 @@ class RolloutGenerator:
         with self._weight_lock:
             return self.scheduler.update_weights(policy_version)
 
-    def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T:
-        """Apply a shared-model mutation at an atomic generation boundary."""
+    def apply_weight_update(
+        self, policy_version: Optional[int], update: Callable[[], T]
+    ) -> T:
+        """Apply a shared-model mutation at an atomic generation boundary.
+
+        ``policy_version=None`` lets the scheduler derive ``live + 1`` under
+        the policy lock, closing the read-compute-write race for callers
+        that only need to advance by one.
+        """
         with self._weight_lock:
             return self.scheduler.apply_weight_update(policy_version, update)
 
@@ -424,7 +431,9 @@ class RolloutRunner:
         """Publish the shared policy's new version to the rollout backend."""
         return self.generator.update_weights(policy_version)
 
-    def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T:
+    def apply_weight_update(
+        self, policy_version: Optional[int], update: Callable[[], T]
+    ) -> T:
         """Apply a model update and publish its version as one operation."""
         return self.generator.apply_weight_update(policy_version, update)
 
@@ -502,35 +511,43 @@ class RolloutRunner:
         """Return ``(cached or fresh) RolloutResult`` plus an ``is_fresh`` flag.
 
         Triggers a new rollout when ``_steps_since_rollout >= rollout_interval``
-        or when the cache is empty.
+        or when the cache is empty. The reuse decision, its version
+        validation, and the returned object are all captured inside one
+        policy snapshot, so a concurrent commit, refresh, or cache clear
+        can never hand out an object the snapshot has already invalidated.
         """
         cache_key = self._batch_key(batch)
-        if (
-            self._cache is None
-            or cache_key != self._cache_key
-            or self._steps_since_rollout >= self.rollout_interval
-        ):
-            raw = self.generator.generate(batch)
-            self._validate_policy_version(raw)
-            scored = self._score(raw)
 
-            def commit(live_version: int) -> Tuple[RolloutResult, bool]:
-                self._validate_policy_version(scored, live_version=live_version)
-                self._cache = scored
-                self._cache_key = cache_key
-                self._steps_since_rollout = 0
-                return scored, True
-
-            # A weight update cannot land between the final version check and
-            # cache publication. Reward scoring itself intentionally remains
-            # outside the policy lock because it may call an external service.
-            return self.generator.with_policy_snapshot(commit)
-
-        cached = self._cache
-        assert cached is not None
-
-        def reuse(live_version: int) -> Tuple[RolloutResult, bool]:
+        def reuse(live_version: int) -> Optional[Tuple[RolloutResult, bool]]:
+            cached = self._cache
+            if (
+                cached is None
+                or self._cache_key != cache_key
+                or self._steps_since_rollout >= self.rollout_interval
+            ):
+                return None
             self._validate_policy_version(cached, live_version=live_version)
             return cached, False
 
-        return self.generator.with_policy_snapshot(reuse)
+        outcome = self.generator.with_policy_snapshot(reuse)
+        if outcome is not None:
+            return outcome
+
+        raw = self.generator.generate(batch)
+        self._validate_policy_version(raw)
+        scored = self._score(raw)
+        # Post-scoring check: reward scoring may call slow external services;
+        # surface an over-lag policy move before the commit critical section.
+        self._validate_policy_version(scored)
+
+        def commit(live_version: int) -> Tuple[RolloutResult, bool]:
+            self._validate_policy_version(scored, live_version=live_version)
+            self._cache = scored
+            self._cache_key = cache_key
+            self._steps_since_rollout = 0
+            return scored, True
+
+        # A weight update cannot land between the final version check and
+        # cache publication. Reward scoring itself intentionally remains
+        # outside the policy lock because it may call an external service.
+        return self.generator.with_policy_snapshot(commit)

astrai/trainer/strategy.py

@@ -292,8 +292,9 @@ class BaseStrategy(ABC):
         if self._rollout_runner is None:
             return optimizer.step()
 
-        next_version = self.policy_version + 1
-        result = self._rollout_runner.apply_weight_update(next_version, optimizer.step)
+        # None lets the scheduler derive live+1 under the policy lock,
+        # avoiding a read-compute-write race on policy_version.
+        result = self._rollout_runner.apply_weight_update(None, optimizer.step)
         self._rollout_runner.step()
         return result