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

tests/trainer/test_online_e2e.py

@@ -144,3 +144,38 @@ def test_online_rollout_end_to_end(
     checkpoint = Checkpoint.load(checkpoint_dir)
     assert checkpoint.meta["policy_version"] == 2
     assert len(created_reference_models) == 1
+
+
+def _minimal_online_config(**overrides):
+    """A TrainConfig for online GRPO that only needs field overrides."""
+    defaults = dict(
+        strategy="online_grpo",
+        model_fn=lambda: torch.nn.Linear(2, 2),
+        dataset=InstructionDataset(),
+        optimizer_fn=lambda m: torch.optim.SGD(m.parameters(), lr=0.0),
+        scheduler_fn=lambda o: SchedulerFactory.create(
+            "cosine", o, warmup_steps=1, lr_decay_steps=4, min_rate=0.05
+        ),
+        reward_model_fn=LengthRewardModel,
+    )
+    defaults.update(overrides)
+    return TrainConfig(**defaults)
+
+
+def test_online_config_rejects_contradictory_policy_lag():
+    """rollout_max_policy_lag below rollout_interval - 1 guarantees a fatal
+    RolloutVersionError mid-training; it must fail at config time instead."""
+    with pytest.raises(ValueError, match="rollout_max_policy_lag=0"):
+        _minimal_online_config(rollout_interval=3, rollout_max_policy_lag=0)
+
+    # lag == interval - 1 (including the derived default) stays valid.
+    config = _minimal_online_config(rollout_interval=3, rollout_max_policy_lag=2)
+    assert config.rollout_max_policy_lag == 2
+    config = _minimal_online_config(rollout_interval=3)
+    assert config.rollout_max_policy_lag is None
+
+    # Offline strategies never consult the rollout window.
+    config = _minimal_online_config(
+        strategy="sft", rollout_interval=3, rollout_max_policy_lag=0
+    )
+    assert config.rollout_max_policy_lag == 0

tests/trainer/test_online_strategy.py

@@ -62,6 +62,9 @@ class _RecordingRunner:
 
     def apply_weight_update(self, policy_version, update):
         result = update()
+        if policy_version is None:
+            # Mirror the scheduler: None derives live+1 under the lock.
+            policy_version = self.policy_version + 1
         self.update_weights(policy_version)
         return result
 

tests/trainer/test_rollout.py

@@ -459,7 +459,7 @@ def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device):
         nonlocal validation_calls
         validation_calls += 1
         original_validate(result, live_version=live_version)
-        if validation_calls == 2:
+        if validation_calls == 3:
             final_validation_started.set()
             assert allow_final_validation_to_finish.wait(timeout=5)
 
@@ -503,6 +503,58 @@ def test_rollout_runner_derives_default_policy_lag_from_interval(device):
     assert runner.max_policy_lag == 3
 
 
+def _interleave_before_snapshot(runner, callback):
+    """Wrap ``with_policy_snapshot`` so ``callback`` runs just before a
+    named snapshot callback enters the generator/scheduler locks."""
+    original_snapshot = runner.generator.with_policy_snapshot
+
+    def wrapper(inspect):
+        if inspect.__name__ == "reuse":
+            callback()
+        return original_snapshot(inspect)
+
+    runner.generator.with_policy_snapshot = wrapper
+
+
+def test_rollout_runner_reuse_reads_cache_inside_the_snapshot(device):
+    """The reuse decision must observe the cache under the policy snapshot
+    (regression: the cache was read outside the lock, so a concurrent
+    commit between the read and the lock silently handed the trainer a
+    stale rollout — a lost update)."""
+    import dataclasses
+
+    runner, _ = _make_runner(device, rollout_interval=100)
+    batch = _make_instruction_batch(n=1)
+    first, _ = runner(batch)
+    assert first.policy_version == 0
+
+    def concurrent_refresh():
+        runner.update_weights(1)
+        runner._cache = dataclasses.replace(first, policy_version=1)
+        runner._steps_since_rollout = 0
+
+    _interleave_before_snapshot(runner, concurrent_refresh)
+    result, fresh = runner(batch)
+    assert fresh is False
+    assert result is not first
+    assert result.policy_version == 1
+
+
+def test_rollout_runner_recovers_when_cache_cleared_before_reuse_snapshot(device):
+    """A cache clear between the reuse decision and the snapshot must
+    trigger a fresh rollout instead of an assertion failure (regression:
+    ``assert cached is not None`` fired because the object was captured
+    outside the lock)."""
+    runner, _ = _make_runner(device, rollout_interval=100)
+    batch = _make_instruction_batch(n=1)
+    first, _ = runner(batch)
+
+    _interleave_before_snapshot(runner, runner.clear_cache)
+    result, fresh = runner(batch)
+    assert fresh is True
+    assert result is not first
+
+
 @pytest.mark.parametrize(
     ("kwargs", "message"),
     [