fix: resolve audited training, import, and serving bugs

- shard the Muon Newton-Schulz orthogonalization over the FSDP mesh instead of partial local slices
- import HF checkpoints faithfully: per-head RoPE permutation for q/k projections and qk-norm, qwen3, shared experts, and qk-norm before RoPE (changes numerics for existing use_qk_norm checkpoints)
- make preprocessing and resume self-contained: backfill realigned bucket keys by semantics (masks ones, rest zeros) and snapshot tokenizer files into every checkpoint
- keep RL consistent: sync the offline GRPO old_model each optimizer step and validate online strategies through a public one-off-rollout hook that leaves the replay cache untouched
- fix streaming serving: withhold partial tool-call prefixes with a stream-end flush, stream tool-call arguments from the raw source span, and terminate SSE frames with a blank line
- fix sampling semantics: capture logprobs before top-k/top-p mutate logits in place and detect greedy pipelines polymorphically instead of isinstance bookkeeping

tests/data/test_pipeline_align.py

@@ -0,0 +1,26 @@
+"""Tests for preprocessing pipeline bucket alignment."""
+
+from astrai.preprocessing.pipeline import Pipeline
+
+
+def test_align_bucket_backfills_missing_mask_with_ones():
+    bucket = {
+        "sequence": [[1, 2], [3, 4]],
+        "loss_mask": [[0, 1]],
+        "chosen_mask": [[1]],
+        "position_ids": [[0, 1]],
+    }
+    result = {"sequence": [5, 6, 7]}
+    Pipeline._align_bucket(bucket, result, [5, 6, 7])
+    assert bucket["loss_mask"][-1] == [1, 1, 1]
+    assert bucket["chosen_mask"][-1] == [1, 1, 1]
+    assert bucket["position_ids"][-1] == [0, 0, 0]
+    assert bucket["sequence"] == [[1, 2], [3, 4]]
+
+
+def test_align_bucket_keeps_present_keys():
+    bucket = {"sequence": [[1, 2]], "loss_mask": [[0, 1]]}
+    result = {"sequence": [9], "loss_mask": [1]}
+    Pipeline._align_bucket(bucket, result, [9])
+    assert bucket["loss_mask"] == [[0, 1]]
+    assert bucket["sequence"] == [[1, 2]]

tests/inference/test_sample.py

@@ -3,6 +3,7 @@
 import torch
 
 from astrai.inference.runtime.sample import (
+    BaseSamplingStrategy,
     FrequencyPenaltyStrategy,
     SamplingPipeline,
     TemperatureStrategy,
@@ -295,3 +296,45 @@ def test_greedy_respects_frequency_penalty():
     )
     # Token 0 saw four occurrences: 5 - 2*4 < 4, so the argmax flips.
     assert penalized.tolist() == [1]
+
+
+class _ArgmaxMovingStrategy(BaseSamplingStrategy):
+    """Custom strategy that can move the argmax — must disable greedy."""
+
+    def apply(
+        self, logits, filter_value=-float("inf"), input_ids=None, input_mask=None
+    ):
+        return torch.roll(logits, shifts=1, dims=-1)
+
+
+def test_greedy_detection_is_polymorphic():
+    """Greedy detection asks strategies polymorphically, no isinstance."""
+    base = [TemperatureStrategy(0.0), TopKStrategy(50), TopPStrategy(0.9)]
+    assert SamplingPipeline(list(base)).is_greedy is True
+    assert SamplingPipeline(base + [FrequencyPenaltyStrategy(0.5)]).is_greedy is False
+
+    # A custom argmax-moving strategy disables greedy even though the
+    # pipeline contains a greedy temperature — this is what isinstance
+    # bookkeeping in the old implementation could not see.
+    assert SamplingPipeline(base + [_ArgmaxMovingStrategy()]).is_greedy is False
+
+
+def test_greedy_detection_position_independent():
+    """Greedy temperature anywhere in the pipeline is detected."""
+    pipeline = SamplingPipeline([TopKStrategy(50), TemperatureStrategy(0.0)])
+    assert pipeline.is_greedy is True
+
+
+def test_greedy_detection_composes_across_nested_pipelines():
+    """A nested pipeline participates through the same interface."""
+    inner = SamplingPipeline([TemperatureStrategy(0.0), TopKStrategy(20)])
+    assert inner.is_greedy is True
+    assert SamplingPipeline([TopPStrategy(0.9), inner]).is_greedy is True
+    assert SamplingPipeline([inner, FrequencyPenaltyStrategy(0.5)]).is_greedy is False
+
+
+def test_nongreedy_temperature_is_not_greedy():
+    pipeline = SamplingPipeline(
+        [TemperatureStrategy(0.7), TopKStrategy(0), TopPStrategy(1.0)]
+    )
+    assert pipeline.is_greedy is False

tests/inference/test_tool_parser.py

@@ -565,3 +565,40 @@ def test_parser_uses_token_ids_for_detection():
     parser = TokenIdParser()
     parser.feed("hello", current_token_ids=[1, 999, 3])
     assert parser.has_tool_calls
+
+
+def test_streaming_partial_name_prefix_never_leaks_into_content():
+    parser = SimpleJsonToolParser()
+    parts = ["Hello ", '{"', '{"n', '{"na', '{"name"']
+    emitted = []
+    body = ""
+    for part in parts:
+        body += part
+        for d in parser.feed(body):
+            if "content" in d:
+                emitted.append(d["content"])
+    assert "".join(emitted) == "Hello "
+
+
+def test_finalize_flushes_withheld_plain_json_content():
+    parser = SimpleJsonToolParser()
+    text = 'Answer: {"price": 1}'
+    deltas = parser.feed(text)
+    streamed = "".join(d["content"] for d in deltas if "content" in d)
+    flushed = parser.finalize(text)
+    joined = streamed + "".join(d["content"] for d in flushed if "content" in d)
+    assert joined == text
+    assert not parser.has_tool_calls
+    assert parser.finalize(text) == []
+
+
+def test_streaming_args_concat_matches_parse_complete():
+    parser = SimpleJsonToolParser()
+    # Compact spacing: json.dumps would re-space this and desync the
+    # streamed arguments diff.
+    text = '{"name": "get_weather","arguments": {"city":"Beijing","unit":"c"}}'
+    _, args_chunks = _simulate_streaming(parser, text)
+    streamed = "".join(args_chunks)
+    completed = parser.parse_complete(text)["tool_calls"][0]["function"]["arguments"]
+    assert streamed == completed
+    assert streamed == '"city":"Beijing","unit":"c"'

tests/serialization/test_hf_adapter.py

@@ -15,6 +15,7 @@ from astrai.serialization import (
     looks_like_hf_state_dict,
     save_model,
 )
+from astrai.serialization.hf_adapter import _half_to_interleaved
 from tests.helpers import assert_state_dicts_equal, make_tiny_config
 
 LLAMA_RAW = {
@@ -47,10 +48,38 @@ MOE_RAW = {
 }
 
 
-def to_hf_keys(state_dict):
-    """Rename AstrAI state dict keys to HuggingFace LLaMA-style names."""
+def to_hf_keys(state_dict, head_dim=None):
+    """Rename AstrAI state dict keys to HuggingFace LLaMA-style names.
+
+    When *head_dim* is given, q/k projections and q/k norm weights are
+    also converted from AstrAI interleaved RoPE coordinates to the HF
+    half-split (rotate_half) convention, so the produced state dict is a
+    faithful HF-layout checkpoint.
+    """
     out = {}
     for key, tensor in state_dict.items():
+        if head_dim is not None:
+            name = key.split(".")
+            is_qk_proj = (
+                len(name) >= 4
+                and name[2] == "attention"
+                and name[3] in ("q_proj", "k_proj")
+            )
+            is_qk_norm = (
+                len(name) >= 4
+                and name[2] == "attention"
+                and name[3] in ("q_norm", "k_norm")
+                and name[4] == "weight"
+            )
+            if is_qk_proj or is_qk_norm:
+                inv = torch.argsort(_half_to_interleaved(head_dim))
+                rows = tensor.shape[0]
+                if rows > head_dim:
+                    blocks = torch.arange(rows // head_dim) * head_dim
+                    idx = (blocks[:, None] + inv[None, :]).flatten()
+                else:
+                    idx = inv
+                tensor = tensor.index_select(0, idx)
         if key == "embed_tokens.weight":
             out["model.embed_tokens.weight"] = tensor
         elif key == "norm.weight":
@@ -141,7 +170,9 @@ def test_adapt_config_passthrough():
 def test_convert_hf_weights_dense_roundtrip():
     cfg = make_tiny_config()
     model = AutoRegressiveLM(cfg)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -157,7 +188,10 @@ def test_convert_hf_weights_moe_roundtrip():
     model = AutoRegressiveLM(cfg)
     hf_raw = convert_hf_config(MOE_RAW)
     hf_cfg = ConfigFactory.load(hf_raw)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), hf_cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads),
+        hf_cfg,
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -216,7 +250,9 @@ def test_convert_hf_weights_moe_with_dense_layers_roundtrip():
         decoder_sparse_step=1,
     )
     model = AutoRegressiveLM(cfg)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -230,7 +266,7 @@ def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
         shared_expert_intermediate_size=16,
     )
     model = AutoRegressiveLM(cfg)
-    hf_sd = to_hf_keys(model.state_dict())
+    hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads)
     hf_sd = {
         k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items()
     }
@@ -241,7 +277,9 @@ def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
 def test_convert_hf_weights_gemma_qk_norm_roundtrip():
     cfg = make_tiny_config(use_qk_norm=True)
     model = AutoRegressiveLM(cfg)
-    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    converted = convert_hf_weights(
+        to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg
+    )
     assert_state_dicts_equal(converted, model.state_dict())
 
 
@@ -250,7 +288,9 @@ def test_from_pretrained_hf_directory(tmp_path):
     model = AutoRegressiveLM(cfg).eval()
     save_model(
         config=LLAMA_RAW,
-        state_dict=to_hf_keys(model.state_dict()),
+        state_dict=to_hf_keys(
+            model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
+        ),
         save_directory=str(tmp_path),
     )
     loaded = AutoModel.from_pretrained(tmp_path).eval()
@@ -292,7 +332,9 @@ def test_from_pretrained_weights_format_astrai_rejects_hf(tmp_path):
     model = AutoRegressiveLM(cfg)
     save_model(
         config=LLAMA_RAW,
-        state_dict=to_hf_keys(model.state_dict()),
+        state_dict=to_hf_keys(
+            model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
+        ),
         save_directory=str(tmp_path),
     )
     with pytest.raises(ValueError):
@@ -313,7 +355,7 @@ def test_from_pretrained_invalid_weights_format(tmp_path):
 def test_from_pretrained_hf_directory_sharded(tmp_path):
     cfg = make_tiny_config()
     model = AutoRegressiveLM(cfg).eval()
-    hf_sd = to_hf_keys(model.state_dict())
+    hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads)
     keys = sorted(hf_sd)
     split = len(keys) // 2
     shard_a = {k: hf_sd[k] for k in keys[:split]}
@@ -342,6 +384,128 @@ def test_from_pretrained_hf_directory_sharded(tmp_path):
         )
 
 
+def _half_split_rope(q, theta=10000.0):
+    """HF llama-style rotate_half RoPE on [batch, seq, heads, head_dim]."""
+    b, s, h, d = q.shape
+    inv_freq = theta ** (-torch.arange(0, d, 2, dtype=torch.float64) / d)
+    freqs = torch.outer(torch.arange(s, dtype=torch.float64), inv_freq).float()
+    cos, sin = freqs.cos()[None, :, None, :], freqs.sin()[None, :, None, :]
+    q1, q2 = q[..., : d // 2], q[..., d // 2 :]
+    return torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
+
+
+def _rms_norm_hf(t, weight, eps):
+    t = t.float()
+    t = t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + eps)
+    return weight.float() * t
+
+
+def _hf_reference_attn(
+    x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim, q_norm_w=None, k_norm_w=None, eps=1e-5
+):
+    """Ground-truth HF attention: per-head RMSNorm BEFORE RoPE (half-split)."""
+    import torch.nn.functional as F
+
+    b, s, dim = x.shape
+    q = (x @ Wq.T).reshape(b, s, n_heads, head_dim).float()
+    k = (x @ Wk.T).reshape(b, s, n_kv, head_dim).float()
+    v = (x @ Wv.T).reshape(b, s, n_kv, head_dim).float()
+    if q_norm_w is not None:
+        q = _rms_norm_hf(q, q_norm_w, eps)
+        k = _rms_norm_hf(k, k_norm_w, eps)
+    q, k = _half_split_rope(q), _half_split_rope(k)
+    rep = n_heads // n_kv
+    k = k.repeat_interleave(rep, dim=2).transpose(1, 2)
+    v = v.repeat_interleave(rep, dim=2).transpose(1, 2)
+    out = F.scaled_dot_product_attention(q.transpose(1, 2), k, v, is_causal=True)
+    out = out.transpose(1, 2).reshape(b, s, n_heads * head_dim)
+    return out @ Wo.T
+
+
+def _run_converted_gqa(x, hf_sd, cfg):
+    from astrai.model.components.attention import GQA
+    from astrai.model.components.rope import get_rotary_emb
+
+    attn = GQA(
+        dim=cfg.hidden_size,
+        n_heads=cfg.num_attention_heads,
+        n_kv_heads=cfg.num_key_value_heads,
+        use_qk_norm=cfg.use_qk_norm,
+        norm_eps=cfg.rms_norm_eps,
+        use_gated_attention=False,
+        layer_id=0,
+    ).eval()
+    converted = convert_hf_weights(hf_sd, cfg)
+    local = {
+        k.removeprefix("layers.0.attention."): v
+        for k, v in converted.items()
+        if k.startswith("layers.0.attention.")
+    }
+    attn.load_state_dict(local, strict=True)
+    head_dim = cfg.hidden_size // cfg.num_attention_heads
+    seq = x.shape[1]
+    rot = get_rotary_emb(head_dim, seq)[None, :seq].expand(x.shape[0], seq, -1, -1)
+    with torch.no_grad():
+        return attn(x, rot, is_causal=True)
+
+
+def test_hf_import_rope_permutation_matches_half_split_reference():
+    torch.manual_seed(0)
+    n_heads, n_kv, head_dim = 4, 2, 8
+    dim = n_heads * head_dim
+    Wq = torch.randn(n_heads * head_dim, dim)
+    Wk = torch.randn(n_kv * head_dim, dim)
+    Wv = torch.randn(n_kv * head_dim, dim)
+    Wo = torch.randn(dim, dim)
+    x = torch.randn(2, 16, dim)
+
+    hf_sd = {
+        "model.layers.0.self_attn.q_proj.weight": Wq,
+        "model.layers.0.self_attn.k_proj.weight": Wk,
+        "model.layers.0.self_attn.v_proj.weight": Wv,
+        "model.layers.0.self_attn.o_proj.weight": Wo,
+    }
+    cfg = make_tiny_config(
+        hidden_size=dim, num_attention_heads=n_heads, num_key_value_heads=n_kv
+    )
+    ref = _hf_reference_attn(x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim)
+    out = _run_converted_gqa(x, hf_sd, cfg)
+    torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4)
+
+
+def test_hf_import_qk_norm_matches_norm_before_rope_reference():
+    torch.manual_seed(1)
+    n_heads, n_kv, head_dim = 4, 2, 8
+    dim = n_heads * head_dim
+    Wq = torch.randn(n_heads * head_dim, dim)
+    Wk = torch.randn(n_kv * head_dim, dim)
+    Wv = torch.randn(n_kv * head_dim, dim)
+    Wo = torch.randn(dim, dim)
+    gq = torch.randn(head_dim)
+    gk = torch.randn(head_dim)
+    x = torch.randn(2, 16, dim)
+
+    hf_sd = {
+        "model.layers.0.self_attn.q_proj.weight": Wq,
+        "model.layers.0.self_attn.k_proj.weight": Wk,
+        "model.layers.0.self_attn.v_proj.weight": Wv,
+        "model.layers.0.self_attn.o_proj.weight": Wo,
+        "model.layers.0.self_attn.q_norm.weight": gq,
+        "model.layers.0.self_attn.k_norm.weight": gk,
+    }
+    cfg = make_tiny_config(
+        hidden_size=dim,
+        num_attention_heads=n_heads,
+        num_key_value_heads=n_kv,
+        use_qk_norm=True,
+    )
+    ref = _hf_reference_attn(
+        x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim, q_norm_w=gq, k_norm_w=gk
+    )
+    out = _run_converted_gqa(x, hf_sd, cfg)
+    torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4)
+
+
 def test_from_pretrained_hf_directory_with_moe(tmp_path):
     cfg = make_tiny_config(
         ffn_type="moe",
@@ -354,7 +518,9 @@ def test_from_pretrained_hf_directory_with_moe(tmp_path):
     model = AutoRegressiveLM(cfg).eval()
     save_model(
         config=MOE_RAW,
-        state_dict=to_hf_keys(model.state_dict()),
+        state_dict=to_hf_keys(
+            model.state_dict(), cfg.hidden_size // cfg.num_attention_heads
+        ),
         save_directory=str(tmp_path),
     )
     loaded = AutoModel.from_pretrained(tmp_path).eval()

tests/trainer/test_callbacks.py

@@ -4,7 +4,11 @@ import torch
 
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.serialization import Checkpoint
-from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
+from astrai.trainer.train_callback import (
+    GradientCheckpointingCallback,
+    TrainCallback,
+    _copy_tokenizer_files,
+)
 from astrai.trainer.trainer import Trainer
 from tests.helpers import RandomTokenDataset
 
@@ -174,3 +178,38 @@ def test_checkpoint_captures_completed_optimizer_step(
     assert (
         Path(base_test_env["test_dir"]) / "epoch_0_step_1" / "metric.jsonl"
     ).is_file()
+
+
+def test_checkpoint_snapshots_tokenizer_files(
+    base_test_env, train_config_factory, device, tmp_path
+):
+    """Checkpoints copy tokenizer files from param_path so resume works."""
+    param_dir = tmp_path / "model"
+    param_dir.mkdir()
+    (param_dir / "tokenizer.json").write_text("{}")
+    (param_dir / "tokenizer_config.json").write_text("{}")
+
+    train_config = train_config_factory(
+        model_fn=lambda: base_test_env["model"],
+        dataset=RandomTokenDataset(length=2),
+        test_dir=base_test_env["test_dir"],
+        device=device,
+        batch_per_device=2,
+        ckpt_interval=1,
+    )
+
+    Trainer(train_config).train(param_path=str(param_dir))
+
+    ckpt_dir = Path(base_test_env["test_dir"]) / "epoch_0_step_1"
+    assert (ckpt_dir / "tokenizer.json").is_file()
+    assert (ckpt_dir / "tokenizer_config.json").is_file()
+
+    # Resuming with param_path == checkpoint dir must not raise
+    # (samefile guard).
+    _copy_tokenizer_files(str(ckpt_dir), str(ckpt_dir))
+
+
+def test_copy_tokenizer_files_skips_missing_and_none(tmp_path):
+    _copy_tokenizer_files(None, str(tmp_path))
+    _copy_tokenizer_files(str(tmp_path), str(tmp_path / "out"))
+    assert not (tmp_path / "out").exists() or not any((tmp_path / "out").iterdir())

tests/trainer/test_grpo_strategy.py

@@ -162,3 +162,34 @@ def test_grpo_sync_old_model(grpo_strategy):
         if k in old_sd_after
     )
     assert matches
+
+
+def test_grpo_optimizer_step_syncs_old_model(grpo_strategy):
+    """optimizer_step must refresh old_model after each update."""
+    strategy, device = grpo_strategy
+
+    class _SteppedOptimizer:
+        def step(self):
+            with torch.no_grad():
+                for p in strategy.model.parameters():
+                    p.add_(0.05)
+
+    strategy.optimizer_step(_SteppedOptimizer())
+
+    policy_sd = strategy.model.state_dict()
+    old_sd = strategy.old_model.state_dict()
+    assert all(
+        torch.allclose(policy_sd[k], old_sd[k]) for k in policy_sd if k in old_sd
+    )
+
+
+def test_online_grpo_optimizer_step_skips_sync(grpo_strategy):
+    """old_model=None (online) must not attempt a sync."""
+    strategy, device = grpo_strategy
+    strategy.old_model = None
+
+    class _SteppedOptimizer:
+        def step(self):
+            return None
+
+    strategy.optimizer_step(_SteppedOptimizer())

tests/trainer/test_online_strategy.py

@@ -45,6 +45,7 @@ class _RecordingRunner:
         self._fresh = True
         self.policy_version = result.policy_version
         self.weight_updates = []
+        self.eval_calls = 0
 
     def __call__(self, batch):
         self.calls += 1
@@ -52,6 +53,12 @@ class _RecordingRunner:
         self._fresh = False
         return self.result, fresh
 
+    def evaluate(self, batch):
+        # Mirrors RolloutRunner.evaluate: one-off scoring that never
+        # touches the replay cache or freshness state.
+        self.eval_calls += 1
+        return self.result
+
     def step(self):
         self.step_calls += 1
 
@@ -360,6 +367,26 @@ def test_loss_is_differentiable_dpo(device):
     assert has_grad
 
 
+def test_validate_online_returns_none_without_runner(device):
+    strat = _make_grpo(device)
+    batch = {"input_ids": torch.randint(3, 200, (2, 4), device=device)}
+    assert strat.validate_online(batch) is None
+
+
+def test_validate_online_uses_one_off_rollout_not_replay_cache(device):
+    strat = _make_grpo(device)
+    runner = _RecordingRunner(_make_rollout_result(device=device))
+    strat.set_rollout_runner(runner)
+
+    out = strat.validate_online(
+        {"input_ids": torch.randint(3, 200, (2, 4), device=device)}
+    )
+
+    assert torch.isfinite(out["loss"]).item()
+    assert runner.eval_calls == 1
+    assert runner.calls == 0  # replay cache path untouched
+
+
 def test_ref_model_not_updated_by_backward_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))

tests/trainer/test_rollout.py

@@ -388,6 +388,21 @@ def test_rollout_runner_cache_returns_stale_flag(device):
     assert fresh2 is False
 
 
+def test_rollout_runner_evaluate_leaves_cache_untouched(device):
+    runner, _ = _make_runner(device, rollout_interval=10)
+    batch = _make_instruction_batch()
+    cached, _ = runner(batch)
+
+    eval_batch = _make_instruction_batch(n=1)
+    result = runner.evaluate(eval_batch)
+
+    assert result.rewards.shape == result.responses.shape[:2]
+    replayed, fresh = runner(batch)
+    assert replayed is cached
+    assert fresh is False
+    assert runner._steps_since_rollout == 0
+
+
 def test_rollout_runner_tags_generation_version_and_preserves_cached_behavior(device):
     runner, _ = _make_runner(device, rollout_interval=100)
     batch = _make_instruction_batch(n=1)