refactor: unify rollout onto inference engine KV-cache path

- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass

astrai/inference/core/executor.py

@@ -55,7 +55,23 @@ class Executor:
                 paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
             )
 
-    def execute_decode(self, tasks: List[Task]) -> List[int]:
+    def execute_decode(
+        self, tasks: List[Task], return_logprobs: bool = False
+    ) -> List[int]:
+        """Decode next token for each task.
+
+        Args:
+            return_logprobs: When ``True``, also record (and return)
+                the log-probability of each sampled token under the
+                post-strategy sampling distribution.  The logprob is
+                appended to ``task.output_logprobs`` and the return
+                list becomes ``List[Tuple[int, float]]``.
+
+        Returns:
+            ``List[int]`` of sampled token IDs, or
+            ``List[Tuple[int, float]]`` of ``(token_id, logprob)`` when
+            ``return_logprobs`` is ``True``.
+        """
         if not tasks:
             return []
 
@@ -116,6 +132,23 @@ class Executor:
             )
             logits = outputs["logits"][:, -1, :]
 
+        if return_logprobs:
+            tokens, logprobs = sample(
+                logits,
+                temperature=temperatures,
+                top_k=top_ks,
+                top_p=top_ps,
+                frequency_penalty=freq_penalties,
+                input_ids=padded_ids,
+                input_mask=padded_mask,
+                return_logprobs=True,
+            )
+            tokens_list = tokens.tolist()
+            logprobs_list = logprobs.tolist()
+            for t, lp in zip(tasks, logprobs_list):
+                t.output_logprobs.append(float(lp))
+            return list(zip(tokens_list, logprobs_list))
+
         return sample(
             logits,
             temperature=temperatures,

astrai/inference/core/scheduler.py

@@ -1,8 +1,10 @@
 import logging
 import threading
+import uuid
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch
+from torch import Tensor
 
 from astrai.inference.core.cache import ContiguousCache, KVCache
 from astrai.inference.core.executor import Executor
@@ -194,6 +196,117 @@ class InferenceScheduler:
             self._cache.task_free(task.task_id)
         for task in self._task_mgr.get_waiting_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
+            self._cache.task_free(task.task_id)
         self._task_mgr.clear_queues()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
+
+    def run_batch(
+        self,
+        prompt_ids_list: List[List[int]],
+        *,
+        max_tokens: Optional[int] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 50,
+        frequency_penalty: float = 0.0,
+        rep_window: int = 64,
+        return_logprobs: bool = False,
+    ) -> List[List[int]]:
+        """Synchronous batch generation without the scheduler thread.
+
+        Accepts already-tokenized prompts (no string round-trip) and runs
+        prefill + decode to completion on the calling thread.  Designed for
+        RL rollout, where logprobs of the behaviour policy must be collected
+        alongside generated tokens.
+
+        Args:
+            prompt_ids_list: ``B`` prompts, each a list of token IDs.
+            max_tokens: Maximum tokens to generate per prompt.  ``None``
+                uses ``self.max_seq_len - len(prompt_ids)``.
+            temperature/top_p/top_k/frequency_penalty/rep_window: Sampling
+                parameters (uniform across the batch).
+            return_logprobs: If ``True``, return ``(token_ids, logprobs)``
+                tuples per prompt (logprobs aligned 1-to-1 with token_ids).
+
+        Returns:
+            ``List[List[int]]`` of generated token IDs per prompt, or —
+            when ``return_logprobs`` is ``True`` —
+            ``List[Tuple[List[int], List[float]]]``.
+        """
+        stop_ids = self._task_mgr.tokenizer.stop_ids
+        cache = self._cache
+        seq_cap = self.max_seq_len
+
+        tasks: List[Task] = []
+        for ids in prompt_ids_list:
+            if len(ids) >= seq_cap:
+                tasks.append(None)
+                continue
+            t_max = max_tokens
+            if t_max is None:
+                t_max = seq_cap - len(ids)
+            else:
+                t_max = min(t_max, seq_cap - len(ids))
+            task = Task(
+                task_id=f"batch_{uuid.uuid4().hex[:8]}",
+                prompt_ids=list(ids),
+                max_tokens=t_max,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                frequency_penalty=frequency_penalty,
+                rep_window=rep_window,
+            )
+            if not cache.task_alloc(task.task_id, task.prompt_ids):
+                tasks.append(None)
+                continue
+            task.input_tokens = len(task.prompt_ids)
+            tasks.append(task)
+
+        try:
+            live = [t for t in tasks if t is not None]
+            prefill_groups: Dict[Tuple[int, int], List[Task]] = {}
+            for t in live:
+                key = (len(t.prompt_ids), cache.task_cached(t.task_id))
+                prefill_groups.setdefault(key, []).append(t)
+            for (prompt_len, start_pos), group in prefill_groups.items():
+                self._executor.execute_prefill(group, prompt_len, start_pos)
+
+            while live:
+                valid: List[Task] = []
+                for t in sorted(live, key=lambda x: x.task_id):
+                    if cache.task_extend(t.task_id, t.next_pos):
+                        valid.append(t)
+                    else:
+                        t.status = TaskStatus.ABORTED
+                if not valid:
+                    break
+
+                step_out = self._executor.execute_decode(
+                    valid, return_logprobs=return_logprobs
+                )
+                if return_logprobs:
+                    for t, (ntok, _lp) in zip(valid, step_out):
+                        t.output_ids.append(ntok)
+                        t.output_tokens += 1
+                else:
+                    for t, ntok in zip(valid, step_out):
+                        t.output_ids.append(ntok)
+                        t.output_tokens += 1
+
+                live = [t for t in valid if not t.is_finished(stop_ids)]
+        finally:
+            for t in tasks:
+                if t is not None:
+                    cache.task_free(t.task_id)
+
+        results: List[Any] = []
+        for t in tasks:
+            if t is None:
+                results.append(([], []) if return_logprobs else [])
+            elif return_logprobs:
+                results.append((list(t.output_ids), list(t.output_logprobs)))
+            else:
+                results.append(list(t.output_ids))
+        return results

astrai/inference/core/task.py

@@ -81,6 +81,7 @@ class Task:
 
         self.status = TaskStatus.PENDING
         self.output_ids: List[int] = []
+        self.output_logprobs: List[float] = []
         self.input_tokens: int = 0
         self.output_tokens: int = 0
         self.arrival_time = time.time()

astrai/inference/sample.py

@@ -313,7 +313,8 @@ def sample(
     input_ids: Optional[Tensor] = None,
     input_mask: Optional[Tensor] = None,
     filter_value: float = -float("inf"),
-) -> Tensor:
+    return_logprobs: bool = False,
+):
     """Apply sampling strategies then sample (softmax + multinomial).
 
     Shortcut for ``SamplingPipeline(...).sample(logits)``.
@@ -327,17 +328,39 @@ def sample(
             (0.0 disables, range -2.0~2.0).
         input_ids: Previously generated token IDs ``[batch, seq_len]``.
         input_mask: Boolean mask for ``input_ids`` padding.
+        return_logprobs: If ``True``, also return the log-probability
+            of each sampled token under the (post-strategy) sampling
+            distribution.  Useful for RL rollout: the returned logprob
+            is the behaviour policy's log-prob used in PPO/GRPO
+            importance ratios.
 
     Returns:
-        Sampled token IDs ``[batch]``.
+        Sampled token IDs ``[batch]``, or — when ``return_logprobs`` is
+        ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where
+        ``chosen_logprobs`` has shape ``[batch]``.
     """
     if SamplingPipeline._is_greedy(temperature):
-        return logits.argmax(dim=-1)
+        tokens = logits.argmax(dim=-1)
-    return SamplingPipeline(
+        if not return_logprobs:
+            return tokens
+        log_probs = torch.log_softmax(logits.float(), dim=-1)
+        chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
+        return tokens, chosen
+
+    pipeline = SamplingPipeline(
         [
             TemperatureStrategy(temperature),
             TopKStrategy(top_k),
             TopPStrategy(top_p),
             FrequencyPenaltyStrategy(frequency_penalty),
         ]
-    ).sample(logits, filter_value, input_ids, input_mask)
+    )
+    if not return_logprobs:
+        return pipeline.sample(logits, filter_value, input_ids, input_mask)
+
+    transformed = pipeline.apply(logits, filter_value, input_ids, input_mask)
+    log_probs = torch.log_softmax(transformed.float(), dim=-1)
+    probs = torch.softmax(transformed, dim=-1)
+    tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
+    chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
+    return tokens, chosen

astrai/trainer/rollout.py

@@ -1,26 +1,35 @@
 """Online rollout runner for RL training.
 
 Provides:
-- :class:`RolloutResult` — universal data container for online sampling
+- :class:`RawRollout` — generation output container (no reward yet)
+- :class:`RolloutResult` — a :class:`RawRollout` with rewards attached
 - :class:`BaseRewardModel` — pluggable reward interface
-- :class:`RolloutRunner` — generates + scores batches for any RL strategy
+- :class:`RolloutGenerator` — KV-cache-backed generation of grouped
+  responses + decoding (no reward); delegates the generation loop to
+  :class:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`
+  so rollout and the production inference server share one code path
+- :class:`RolloutRunner` — orchestrates generation + scoring with a
+  step-driven cache; its ``__call__`` returns ``(RolloutResult, is_fresh)``
+  so callers do not need to rely on object identity to detect refreshes.
 """
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.inference.sample import SamplingPipeline
+from astrai.inference.core.scheduler import InferenceScheduler
 
 
-@dataclass
+@dataclass(kw_only=True)
-class RolloutResult:
+class RawRollout:
-    """Universal container produced by :class:`RolloutRunner`.
+    """Generation output before reward scoring.
+
+    Produced by :class:`RolloutGenerator`; consumed by :class:`RolloutRunner`
+    to assemble a :class:`RolloutResult` once rewards are attached.
 
     Fields are designed to cover all common RL algorithms:
     GRPO, PPO, Online DPO, Rejection Sampling, etc.
@@ -35,9 +44,6 @@ class RolloutResult:
     response_mask: Tensor
     """Boolean mask for real (non-pad) response tokens, shape ``[B, G, R_max]``."""
 
-    rewards: Tensor
-    """Reward per response, shape ``[B, G]``."""
-
     logprobs_old: Tensor
     """Per-token log-probs under the behaviour policy, shape ``[B, G, R_max]``."""
 
@@ -48,6 +54,18 @@ class RolloutResult:
     """Decoded response strings, shape ``[B, G]`` (for reward models)."""
 
 
+@dataclass(kw_only=True)
+class RolloutResult(RawRollout):
+    """A :class:`RawRollout` with reward scoring attached.
+
+    Produced by :class:`RolloutRunner` once the :class:`BaseRewardModel`
+    has scored the decoded responses.
+    """
+
+    rewards: Tensor
+    """Reward per response, shape ``[B, G]``."""
+
+
 class BaseRewardModel(ABC):
     """Pluggable reward model interface.
 
@@ -72,115 +90,142 @@ class BaseRewardModel(ABC):
         ...
 
 
-def generate_responses(
+_PAD = 0
-    model: nn.Module,
-    input_ids: Tensor,
-    attention_mask: Tensor,
-    max_new_tokens: int,
-    sampling_pipeline: SamplingPipeline,
-    stop_ids: List[int],
-) -> Dict[str, Tensor]:
-    """Autoregressive generation with log-prob tracking.
 
-    Args:
-        model: Policy model (``forward`` returns ``{"logits": ...}``).
-        input_ids: ``[B, P_len]`` prompt token IDs.
-        attention_mask: ``[B, P_len]`` boolean mask.
-        max_new_tokens: Maximum tokens to generate.
-        sampling_pipeline: Composed sampling strategies.
-        stop_ids: Token IDs that stop generation (eos, etc.).
 
-    Returns:
+class RolloutGenerator:
-        ``dict`` with keys:
+    """Pure generation + decoding for a group of responses per prompt.
-        - ``generated_ids``: ``[B, max_new_tokens]`` (padded to same length)
+
-        - ``generated_mask``: ``[B, max_new_tokens]``
+    Delegates the prefill/decode loop to
-        - ``logprobs``: ``[B, max_new_tokens]`` per-token log-probs
+    :meth:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`,
+    which uses a real KV cache (no O(n²) recompute).  Has no dependency
+    on any reward model; can be reused in isolation for offline
+    generation, qualitative sampling, or eval pipelines.
     """
-    _PAD = 0
-    B, P_len = input_ids.shape
-    device = input_ids.device
-    stop_ids_set = set(stop_ids)
-    done = torch.zeros(B, dtype=torch.bool, device=device)
-    all_ids = input_ids.clone()
-    all_mask = attention_mask.clone()
-    logprob_list: List[Tensor] = []
 
-    for _ in range(max_new_tokens):
+    def __init__(
-        outputs = model(input_ids=all_ids, input_mask=all_mask)
+        self,
-        logits = outputs["logits"][:, -1, :].float()
+        scheduler: InferenceScheduler,
-        log_probs = F.log_softmax(logits, dim=-1)
+        tokenizer,
+        max_tokens: int = 1024,
+        group_size: int = 8,
+        temperature: float = 1.0,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        frequency_penalty: float = 0.0,
+        rep_window: int = 64,
+    ):
+        self.scheduler = scheduler
+        self.tokenizer = tokenizer
+        self.max_tokens = max_tokens
+        self.group_size = group_size
+        self.temperature = temperature
+        self.top_k = top_k
+        self.top_p = top_p
+        self.frequency_penalty = frequency_penalty
+        self.rep_window = rep_window
 
-        logits = sampling_pipeline.apply(logits, input_ids=all_ids, input_mask=all_mask)
+    @torch.no_grad()
-        probs = torch.softmax(logits, dim=-1)
+    def generate(self, batch: Dict[str, Tensor]) -> RawRollout:
-        next_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
+        """Expand prompts by ``group_size`` and generate one response each."""
-
+        prompt_ids = batch["input_ids"] if "input_ids" in batch else batch["prompts"]
-        next_tokens[done] = _PAD
+        prompt_mask = (
-        chosen_logprobs = torch.gather(log_probs, -1, next_tokens.unsqueeze(-1))
+            batch["attention_mask"] if "attention_mask" in batch else (prompt_ids != 0)
-        logprob_list.append(chosen_logprobs)
-
-        all_ids = torch.cat([all_ids, next_tokens.unsqueeze(1)], dim=-1)
-        all_mask = torch.cat([all_mask, (~done).unsqueeze(1)], dim=-1)
-
-        done = done | torch.tensor(
-            [t.item() in stop_ids_set for t in next_tokens],
-            device=device,
         )
-        if done.all():
+        B, _ = prompt_ids.shape
-            break
+        G = self.group_size
 
-    logprobs = torch.cat(logprob_list, dim=-1)
+        prompt_texts: List[str] = []
-    if logprobs.size(1) < max_new_tokens:
+        flat_prompt_ids: List[List[int]] = []
-        pad_len = max_new_tokens - logprobs.size(1)
+        for i in range(B):
-        logprobs = F.pad(logprobs, (0, pad_len), value=0.0)
+            ids = prompt_ids[i, prompt_mask[i]].tolist()
+            text = self.tokenizer.decode(ids, skip_special_tokens=True)
+            for _ in range(G):
+                flat_prompt_ids.append(list(ids))
+            prompt_texts.append(text)
 
-    generated_ids = all_ids[:, P_len:]
+        results = self.scheduler.run_batch(
-    if generated_ids.size(1) < max_new_tokens:
+            flat_prompt_ids,
-        pad_len = max_new_tokens - generated_ids.size(1)
+            max_tokens=self.max_tokens,
-        generated_ids = F.pad(generated_ids, (0, pad_len), value=_PAD)
+            temperature=self.temperature,
+            top_k=self.top_k,
+            top_p=self.top_p,
+            frequency_penalty=self.frequency_penalty,
+            rep_window=self.rep_window,
+            return_logprobs=True,
+        )
 
-    generated_mask = generated_ids != _PAD
+        # Each element is (token_ids, logprobs); pad to max length.
+        max_len = 0
+        for token_ids, _lp in results:
+            max_len = max(max_len, len(token_ids))
+        max_len = max(max_len, 1)
 
-    return {
+        device = prompt_ids.device
-        "generated_ids": generated_ids,
+        responses = torch.full((B, G, max_len), _PAD, dtype=torch.long, device=device)
-        "generated_mask": generated_mask,
+        response_mask = torch.zeros((B, G, max_len), dtype=torch.bool, device=device)
-        "logprobs": logprobs,
+        logprobs_old = torch.zeros((B, G, max_len), dtype=torch.float, device=device)
-    }
+
+        flat_idx = 0
+        response_texts: List[List[str]] = [[] for _ in range(B)]
+        for i in range(B):
+            for g in range(G):
+                token_ids, lps = results[flat_idx]
+                flat_idx += 1
+                n = len(token_ids)
+                if n:
+                    responses[i, g, :n] = torch.tensor(
+                        token_ids, dtype=torch.long, device=device
+                    )
+                    response_mask[i, g, :n] = True
+                    logprobs_old[i, g, :n] = torch.tensor(
+                        lps, dtype=torch.float, device=device
+                    )
+                response_texts[i].append(
+                    self.tokenizer.decode(token_ids, skip_special_tokens=True)
+                )
+
+        return RawRollout(
+            prompts=prompt_ids,
+            responses=responses,
+            response_mask=response_mask,
+            logprobs_old=logprobs_old,
+            prompt_texts=prompt_texts,
+            response_texts=response_texts,
+        )
 
 
 class RolloutRunner:
     """Produces :class:`RolloutResult` from a prompt batch.
 
-    Maintains an internal cache so the same batch prompt can be replayed
+    Composes a :class:`RolloutGenerator` (generation + decoding) with a
-    for multiple gradient steps.  A new rollout is triggered every
+    :class:`BaseRewardModel` (scoring).  Maintains an internal cache so
-    ``rollout_interval`` calls to :meth:`step`.
+    the same batch prompt can be replayed for multiple gradient steps.
+    A new rollout is triggered every ``rollout_interval`` calls to
+    :meth:`step` (or after :meth:`clear_cache`).
+
+    The ``__call__`` contract returns a ``(RolloutResult, is_fresh)``
+    tuple — callers must use the boolean to detect a refreshed rollout
+    rather than relying on object identity.
 
     Usage::
 
-        runner = RolloutRunner(policy, old_policy, tokenizer,
+        generator = RolloutGenerator(policy, tokenizer, pipeline, ...)
-                               reward_model, sampling_pipeline, config)
+        runner = RolloutRunner(generator, reward_model, rollout_interval=512)
-        result = runner(prompt_batch)
+        result, is_fresh = runner(prompt_batch)
+        if is_fresh:
+            ...  # e.g. sync behaviour policy
     """
 
     def __init__(
         self,
-        policy_model: nn.Module,
+        generator: RolloutGenerator,
-        old_model: Optional[nn.Module],
-        tokenizer,
         reward_model: BaseRewardModel,
-        sampling_pipeline: SamplingPipeline,
-        max_tokens: int = 1024,
-        group_size: int = 8,
         rollout_interval: int = 512,
     ):
-        self.policy_model = policy_model
+        self.generator = generator
-        self.old_model = old_model
-        self.tokenizer = tokenizer
         self.reward_model = reward_model
-        self.sampling_pipeline = sampling_pipeline
-        self.max_tokens = max_tokens
-        self.group_size = group_size
         self.rollout_interval = rollout_interval
-        self.stop_ids = getattr(tokenizer, "stop_ids", []) or []
 
         self._cache: Optional[RolloutResult] = None
         self._steps_since_rollout: int = 0
@@ -193,80 +238,28 @@ class RolloutRunner:
         """Force next call to re-run rollout."""
         self._cache = None
 
-    def _tokenize_prompts(self, raw_texts: List[str]) -> Dict[str, Tensor]:
+    def _score(self, raw: RawRollout) -> RolloutResult:
-        ids_list = self.tokenizer.encode(raw_texts, out_ids=True)
+        rewards = self.reward_model.score(raw.prompt_texts, raw.response_texts)
-        B = len(ids_list)
+        device = raw.prompts.device
-        P_max = max(len(ids) for ids in ids_list) if ids_list else 0
-        input_ids = torch.zeros(B, P_max, dtype=torch.long)
-        for i, ids in enumerate(ids_list):
-            input_ids[i, : len(ids)] = torch.tensor(ids[:P_max], dtype=torch.long)
-        attention_mask = input_ids != 0
-        return {"input_ids": input_ids, "attention_mask": attention_mask}
-
-    def _decode(self, token_ids: Tensor, mask: Tensor) -> List[List[str]]:
-        B, G, _ = token_ids.shape
-        texts = []
-        for i in range(B):
-            group_texts = []
-            for g in range(G):
-                ids = token_ids[i, g, mask[i, g]].tolist()
-                group_texts.append(self.tokenizer.decode(ids, skip_special_tokens=True))
-            texts.append(group_texts)
-        return texts
-
-    @torch.no_grad()
-    def _run(self, batch: Dict[str, Tensor]) -> RolloutResult:
-        """Execute the actual generation + reward scoring."""
-        prompt_ids = batch["input_ids"] if "input_ids" in batch else batch["prompts"]
-        prompt_mask = (
-            batch["attention_mask"] if "attention_mask" in batch else (prompt_ids != 0)
-        )
-        B, P_len = prompt_ids.shape
-        G = self.group_size
-        device = prompt_ids.device
-
-        prompt_texts: List[str] = []
-        for i in range(B):
-            ids = prompt_ids[i, prompt_mask[i]].tolist()
-            prompt_texts.append(self.tokenizer.decode(ids, skip_special_tokens=True))
-
-        expanded_ids = prompt_ids.unsqueeze(1).expand(-1, G, -1).reshape(B * G, P_len)
-        expanded_mask = prompt_mask.unsqueeze(1).expand(-1, G, -1).reshape(B * G, P_len)
-
-        gen_out = generate_responses(
-            model=self.policy_model,
-            input_ids=expanded_ids,
-            attention_mask=expanded_mask,
-            max_new_tokens=self.max_tokens,
-            sampling_pipeline=self.sampling_pipeline,
-            stop_ids=self.stop_ids,
-        )
-
-        gen_ids = gen_out["generated_ids"].reshape(B, G, -1)
-        gen_mask = gen_out["generated_mask"].reshape(B, G, -1)
-        gen_logprobs = gen_out["logprobs"].reshape(B, G, -1)
-
-        response_texts = self._decode(gen_ids, gen_mask)
-        reward_tensor = self.reward_model.score(prompt_texts, response_texts)
-        rewards = reward_tensor.to(device=device)
-
         return RolloutResult(
-            prompts=prompt_ids,
+            prompts=raw.prompts,
-            responses=gen_ids,
+            responses=raw.responses,
-            response_mask=gen_mask,
+            response_mask=raw.response_mask,
-            rewards=rewards,
+            rewards=rewards.to(device=device),
-            logprobs_old=gen_logprobs,
+            logprobs_old=raw.logprobs_old,
-            prompt_texts=prompt_texts,
+            prompt_texts=raw.prompt_texts,
-            response_texts=response_texts,
+            response_texts=raw.response_texts,
         )
 
-    def __call__(self, batch: Dict[str, Tensor]) -> RolloutResult:
+    def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]:
-        """Return cached or fresh :class:`RolloutResult`.
+        """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.
         """
         if self._cache is None or self._steps_since_rollout >= self.rollout_interval:
-            self._cache = self._run(batch)
+            raw = self.generator.generate(batch)
+            self._cache = self._score(raw)
             self._steps_since_rollout = 0
-        return self._cache
+            return self._cache, True
+        return self._cache, False

astrai/trainer/strategy.py

@@ -109,7 +109,6 @@ class BaseStrategy(ABC):
         self.executor = kwargs.pop("executor", None)
         self.extra_kwargs = kwargs
         self._rollout_runner = None
-        self._prev_rollout_result = None
 
     @abstractmethod
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
@@ -159,10 +158,9 @@ class BaseStrategy(ABC):
         if self._rollout_runner is None:
             return self.compute_loss(batch)
 
-        result = self._rollout_runner(batch)
+        result, is_fresh = self._rollout_runner(batch)
-        if result is not self._prev_rollout_result:
+        if is_fresh:
             self._on_rollout_refresh()
-            self._prev_rollout_result = result
             if self.executor and self.executor.sync_gradients:
                 self._rollout_runner.step()
 

astrai/trainer/train_context.py

@@ -8,19 +8,14 @@ from torch.utils.data import DataLoader, random_split
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import RDSampler
-from astrai.inference.sample import (
+from astrai.inference.core.scheduler import InferenceScheduler
-    SamplingPipeline,
-    TemperatureStrategy,
-    TopKStrategy,
-    TopPStrategy,
-)
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
 from astrai.serialization import Checkpoint, load_json
 from astrai.tokenize import AutoTokenizer
-from astrai.trainer.rollout import RolloutRunner
+from astrai.trainer.rollout import RolloutGenerator, RolloutRunner
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
 
 
@@ -243,22 +238,27 @@ class TrainContextBuilder:
             tokenizer = AutoTokenizer.from_pretrained(self._param_path)
             reward_model = cfg.reward_model_fn()
 
-            pipeline = SamplingPipeline(
+            scheduler = InferenceScheduler(
-                [
+                model=context.model,
-                    TemperatureStrategy(cfg.rollout_temperature),
+                tokenizer=tokenizer,
-                    TopKStrategy(cfg.rollout_top_k),
+                max_batch_size=strategy_kwargs.get("group_size", 8)
-                    TopPStrategy(cfg.rollout_top_p),
+                * max(1, cfg.batch_size or 1),
-                ]
+                max_seq_len=getattr(context.model.config, "max_len", None),
+                max_prompt_len=getattr(context.model.config, "max_len", 4096),
             )
 
-            runner = RolloutRunner(
+            generator = RolloutGenerator(
-                policy_model=context.model,
+                scheduler=scheduler,
-                old_model=old_model,
                 tokenizer=tokenizer,
-                reward_model=reward_model,
-                sampling_pipeline=pipeline,
                 max_tokens=cfg.rollout_max_tokens,
                 group_size=strategy_kwargs.get("group_size", 8),
+                temperature=cfg.rollout_temperature,
+                top_k=cfg.rollout_top_k,
+                top_p=cfg.rollout_top_p,
+            )
+            runner = RolloutRunner(
+                generator=generator,
+                reward_model=reward_model,
                 rollout_interval=cfg.rollout_interval,
             )
             context.strategy.set_rollout_runner(runner)

tests/inference/test_sample.py

@@ -231,3 +231,54 @@ def test_sample_with_frequency_penalty():
     )
     assert tokens.shape == (1,)
     assert 0 <= tokens[0] < logits.size(-1)
+
+
+def test_sample_return_logprobs_shape():
+    """``return_logprobs=True`` returns ``[batch]`` logprobs aligned to tokens."""
+    logits = torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]])
+    out = sample(logits, temperature=1.0, return_logprobs=True)
+    tokens, logprobs = out
+    assert tokens.shape == (2,)
+    assert logprobs.shape == (2,)
+
+
+def test_sample_return_logprobs_nonpositive():
+    """Probabilities never exceed 1, so logprobs are always ≤ 0."""
+    torch.manual_seed(0)
+    logits = torch.randn(4, 50)
+    _, logprobs = sample(
+        logits, temperature=0.8, top_k=20, top_p=0.9, return_logprobs=True
+    )
+    assert torch.all(logprobs <= 1e-5)
+
+
+def test_sample_return_logprobs_greedy_path():
+    """Greedy decode (temperature 0) also returns logprobs."""
+    logits = torch.tensor([[1.0, 5.0, 2.0]])
+    tokens, logprobs = sample(logits, temperature=0.0, return_logprobs=True)
+    assert tokens[0].item() == 1
+    # log p(token=1) should equal log_softmax(logits)[1]
+    expected = torch.log_softmax(logits.float(), dim=-1)[0, 1]
+    assert torch.allclose(logprobs[0], expected, atol=1e-5)
+
+
+def test_sample_return_logprobs_matches_manual_computation():
+    """Returned logprob equals log_softmax(transformed_logits)[token]."""
+    torch.manual_seed(1)
+    logits = torch.randn(2, 30)
+    tokens, logprobs = sample(logits, temperature=0.7, top_p=0.95, return_logprobs=True)
+    # Recompute with the same pipeline
+    from astrai.inference.sample import (
+        SamplingPipeline,
+        TemperatureStrategy,
+        TopPStrategy,
+    )
+
+    pipeline = SamplingPipeline([TemperatureStrategy(0.7), TopPStrategy(0.95)])
+    transformed = pipeline.apply(logits.clone())
+    expected = torch.gather(
+        torch.log_softmax(transformed.float(), dim=-1),
+        -1,
+        tokens.unsqueeze(-1),
+    ).squeeze(-1)
+    assert torch.allclose(logprobs, expected, atol=1e-5)

tests/inference/test_scheduler.py

@@ -191,3 +191,124 @@ def test_prefill_skips_fully_cached_tasks(mock_model_and_tokenizer):
     task_id = scheduler.add_task("short prompt", stream_callback=lambda t: None)
     scheduler.stop()
     assert task_id.startswith("task_")
+
+
+def _make_real_scheduler(device):
+    """Build a scheduler backed by a tiny real model for run_batch tests."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+    from astrai.model.transformer import AutoRegressiveLM
+
+    class _Tok:
+        stop_ids = [2]
+
+        def encode(self, texts, **_):
+            if isinstance(texts, str):
+                texts = [texts]
+            return [[b for b in t.encode("utf-8")] for t in texts]
+
+        def decode(self, ids, skip_special_tokens=True):
+            return bytes(b for b in ids if b > 2 or not skip_special_tokens).decode(
+                "utf-8", errors="ignore"
+            )
+
+    cfg = AutoRegressiveLMConfig(
+        vocab_size=200,
+        dim=16,
+        n_heads=2,
+        n_kv_heads=1,
+        dim_ffn=32,
+        max_len=64,
+        n_layers=2,
+        norm_eps=1e-5,
+    )
+    model = AutoRegressiveLM(cfg).to(device=device).eval()
+    tokenizer = _Tok()
+    scheduler = InferenceScheduler(
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=8,
+        max_seq_len=64,
+        max_prompt_len=64,
+    )
+    return scheduler, tokenizer, model
+
+
+def test_run_batch_returns_token_sequences():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        prompts = [[10, 20, 30], [5, 6, 7, 8]]
+        results = scheduler.run_batch(prompts, max_tokens=4, temperature=1.0)
+        assert len(results) == 2
+        for ids in results:
+            assert isinstance(ids, list)
+            assert len(ids) <= 4
+            assert all(0 <= i < 200 for i in ids)
+    finally:
+        scheduler.stop()
+
+
+def test_run_batch_return_logprobs_aligned():
+    """return_logprobs=True gives (token_ids, logprobs) tuples with equal len."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        prompts = [[10, 20, 30, 40]]
+        results = scheduler.run_batch(
+            prompts, max_tokens=5, temperature=1.0, return_logprobs=True
+        )
+        assert len(results) == 1
+        token_ids, logprobs = results[0]
+        assert len(token_ids) == len(logprobs)
+        assert all(lp <= 1e-5 for lp in logprobs)  # logprobs ≤ 0
+    finally:
+        scheduler.stop()
+
+
+def test_run_batch_respects_max_tokens():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        prompts = [[10, 20, 30]]
+        results = scheduler.run_batch(prompts, max_tokens=3, temperature=1.0)
+        assert len(results[0]) <= 3
+    finally:
+        scheduler.stop()
+
+
+def test_run_batch_stop_id_terminates():
+    """A token matching stop_ids terminates generation for that prompt."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        prompts = [[10, 20, 30]]
+        results = scheduler.run_batch(prompts, max_tokens=32, temperature=1.0)
+        # If stop token 2 was produced, it is the last token
+        if results[0] and results[0][-1] == 2:
+            # No tokens after stop should exist (since we terminate)
+            assert 2 not in results[0][:-1]
+    finally:
+        scheduler.stop()
+
+
+def test_run_batch_empty_prompts():
+    """Empty prompt list yields empty result list."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        assert scheduler.run_batch([], max_tokens=4) == []
+    finally:
+        scheduler.stop()
+
+
+def test_run_batch_too_long_prompt_skipped():
+    """A prompt longer than max_seq_len yields an empty result slot."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    try:
+        long = list(range(100))  # > max_seq_len=64
+        results = scheduler.run_batch([long, [10, 20]], max_tokens=2)
+        assert results[0] == []
+        assert len(results[1]) <= 2
+    finally:
+        scheduler.stop()

tests/trainer/test_online_strategy.py

@@ -71,22 +71,31 @@ def _make_rollout_result(B=2, G=4, P=6, R=8, device="cpu"):
 
 
 class _RecordingRunner:
-    """Fake RolloutRunner that returns a fixed result and tracks calls."""
+    """Fake RolloutRunner returning a fixed result with freshness tracking.
+
+    Freshness is ``True`` on the first call after construction or after
+    :meth:`swap_result`; ``False`` on subsequent cached calls — mirroring
+    the real ``RolloutRunner`` contract without invoking generation.
+    """
 
     def __init__(self, result):
         self.result = result
         self.calls = 0
         self.step_calls = 0
+        self._fresh = True
 
     def __call__(self, batch):
         self.calls += 1
-        return self.result
+        fresh = self._fresh
+        self._fresh = False
+        return self.result, fresh
 
     def step(self):
         self.step_calls += 1
 
     def swap_result(self, result):
         self.result = result
+        self._fresh = True
 
 
 @pytest.fixture

tests/trainer/test_rollout.py

@@ -1,26 +1,23 @@
 """Unit tests for the online rollout module.
 
-Covers :class:`RolloutResult`, :class:`BaseRewardModel`,
+Covers :class:`RolloutResult` / :class:`RawRollout`, :class:`BaseRewardModel`,
-:func:`generate_responses`, and :class:`RolloutRunner` including
+:class:`RolloutGenerator` (KV-cache-backed via :class:`InferenceScheduler.run_batch`)
-its internal cache and rollout-interval trigger logic.
+and :class:`RolloutRunner` including its internal cache and rollout-interval
+trigger logic.
 """
 
 import pytest
 import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.sample import (
+from astrai.inference.core.scheduler import InferenceScheduler
-    SamplingPipeline,
-    TemperatureStrategy,
-    TopKStrategy,
-    TopPStrategy,
-)
 from astrai.model.transformer import AutoRegressiveLM
 from astrai.trainer.rollout import (
     BaseRewardModel,
+    RawRollout,
+    RolloutGenerator,
     RolloutResult,
     RolloutRunner,
-    generate_responses,
 )
 
 
@@ -57,10 +54,6 @@ class ConstantRewardModel(BaseRewardModel):
         return torch.full((B, G), float(self.value))
 
 
-class _FakeOldModel:
-    """Placeholder old-model; RolloutRunner stores but never calls it."""
-
-
 def _make_config(vocab_size=200, max_len=128):
     return AutoRegressiveLMConfig(
         vocab_size=vocab_size,
@@ -81,9 +74,13 @@ def _make_model(device):
     return m, cfg
 
 
-def _make_pipeline():
+def _make_scheduler(model, tokenizer, max_batch_size=8, max_len=128):
-    return SamplingPipeline(
+    return InferenceScheduler(
-        [TemperatureStrategy(1.0), TopKStrategy(0), TopPStrategy(1.0)]
+        model=model,
+        tokenizer=tokenizer,
+        max_batch_size=max_batch_size,
+        max_seq_len=max_len,
+        max_prompt_len=max_len,
     )
 
 
@@ -93,14 +90,28 @@ def _make_prompt_batch(batch_size=2, prompt_len=6, device="cpu"):
     return {"input_ids": ids, "attention_mask": mask}
 
 
-def test_rollout_result_fields():
+def test_raw_rollout_fields():
+    r = RawRollout(
+        prompts=torch.zeros(2, 4, dtype=torch.long),
+        responses=torch.zeros(2, 3, 5, dtype=torch.long),
+        response_mask=torch.ones(2, 3, 5, dtype=torch.bool),
+        logprobs_old=torch.zeros(2, 3, 5),
+    )
+    assert r.prompts.shape == (2, 4)
+    assert r.responses.shape == (2, 3, 5)
+    assert r.prompt_texts == []
+    assert r.response_texts == []
+
+
+def test_rollout_result_inherits_raw_rollout_fields():
     r = RolloutResult(
         prompts=torch.zeros(2, 4, dtype=torch.long),
         responses=torch.zeros(2, 3, 5, dtype=torch.long),
         response_mask=torch.ones(2, 3, 5, dtype=torch.bool),
-        rewards=torch.zeros(2, 3),
         logprobs_old=torch.zeros(2, 3, 5),
+        rewards=torch.zeros(2, 3),
     )
+    assert r.rewards.shape == (2, 3)
     assert r.prompts.shape == (2, 4)
     assert r.responses.shape == (2, 3, 5)
     assert r.prompt_texts == []
@@ -119,94 +130,96 @@ def test_constant_reward_model_shape():
     assert torch.all(out == 0.5)
 
 
-def test_generate_responses_shapes():
+@pytest.fixture
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+def device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _make_generator(device, **kw):
     model, _ = _make_model(device)
-    pipeline = _make_pipeline()
+    tokenizer = FakeTokenizer()
-    ids = torch.randint(3, 200, (2, 4), device=device)
+    scheduler = _make_scheduler(
-    mask = torch.ones(2, 4, dtype=torch.bool, device=device)
+        model,
-
+        tokenizer,
-    out = generate_responses(
+        max_batch_size=kw.get("max_batch_size", 8),
-        model=model,
+        max_len=kw.get("max_len", 128),
-        input_ids=ids,
-        attention_mask=mask,
-        max_new_tokens=8,
-        sampling_pipeline=pipeline,
-        stop_ids=[],
     )
-    assert out["generated_ids"].shape == (2, 8)
+    generator = RolloutGenerator(
-    assert out["generated_mask"].shape == (2, 8)
+        scheduler=scheduler,
-    assert out["logprobs"].shape == (2, 8)
+        tokenizer=tokenizer,
-
+        max_tokens=kw.get("max_tokens", 8),
-
+        group_size=kw.get("group_size", 2),
-def test_generate_responses_stops_on_stop_id():
+        temperature=kw.get("temperature", 1.0),
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+        top_k=kw.get("top_k", 0),
-    model, _ = _make_model(device)
+        top_p=kw.get("top_p", 1.0),
-    pipeline = _make_pipeline()
-    ids = torch.randint(3, 200, (1, 3), device=device)
-    mask = torch.ones(1, 3, dtype=torch.bool, device=device)
-
-    out = generate_responses(
-        model=model,
-        input_ids=ids,
-        attention_mask=mask,
-        max_new_tokens=16,
-        sampling_pipeline=pipeline,
-        stop_ids=[7],
     )
-    gen = out["generated_ids"][0]
+    return generator, model
-    mask = out["generated_mask"][0]
-    # If a 7 appeared, all tokens after it must be pad (mask False).
-    nonzero_stop = (gen == 7).nonzero()
-    if nonzero_stop.numel():
-        first = nonzero_stop[0].item()
-        assert mask[first + 1 :].sum() == 0
 
 
-def test_generate_responses_logprobs_match_tokens():
+def test_rollout_generator_shapes(device):
-    """logprobs[i] must be the logprob of generated_ids[i]."""
+    gen, _ = _make_generator(device, group_size=3, max_tokens=5)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device)
-    model, _ = _make_model(device)
+    r = gen.generate(batch)
-    pipeline = _make_pipeline()
+    assert r.prompts.shape == (2, 4)
-    ids = torch.randint(3, 200, (1, 2), device=device)
+    assert r.responses.shape == (2, 3, 5)
-    mask = torch.ones(1, 2, dtype=torch.bool, device=device)
+    assert r.response_mask.shape == (2, 3, 5)
+    assert r.logprobs_old.shape == (2, 3, 5)
+    assert len(r.prompt_texts) == 2
+    assert len(r.response_texts) == 2
+    assert len(r.response_texts[0]) == 3
 
-    out = generate_responses(
+
-        model=model,
+def test_rollout_generator_mask_matches_responses(device):
-        input_ids=ids,
+    """Positions beyond a response's length are pad (mask False)."""
-        attention_mask=mask,
+    gen, _ = _make_generator(device, group_size=2, max_tokens=6)
-        max_new_tokens=4,
+    batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device)
-        sampling_pipeline=pipeline,
+    r = gen.generate(batch)
-        stop_ids=[],
+    for i in range(2):
-    )
+        for g in range(2):
-    gen = out["generated_ids"][0]
+            real = r.response_mask[i, g].sum().item()
-    lp = out["logprobs"][0]
+            # Pad positions should be 0
-    for i in range(4):
+            assert r.responses[i, g, real:].sum() == 0
-        if gen[i] == 0 and not out["generated_mask"][0, i]:
+            # logprobs after the real tokens are 0 (padding)
-            continue
+            if real < r.logprobs_old.size(-1):
-        assert lp[i] <= 0.0
+                assert torch.all(r.logprobs_old[i, g, real:] == 0)
+
+
+def test_rollout_generator_logprobs_are_nonpositive(device):
+    """Behaviour-policy logprobs of sampled tokens should be ≤ 0."""
+    gen, _ = _make_generator(device, group_size=2, max_tokens=4)
+    batch = _make_prompt_batch(batch_size=1, prompt_len=3, device=device)
+    r = gen.generate(batch)
+    for i in range(1):
+        for g in range(2):
+            mask = r.response_mask[i, g]
+            lp = r.logprobs_old[i, g][mask]
+            assert torch.all(lp <= 1e-5)
 
 
 def _make_runner(device, **kw):
-    model, _ = _make_model(device)
+    generator, model = _make_generator(
-    rm = ConstantRewardModel(1.0)
+        device,
-    return RolloutRunner(
-        policy_model=model,
-        old_model=_FakeOldModel(),
-        tokenizer=FakeTokenizer(),
-        reward_model=rm,
-        sampling_pipeline=_make_pipeline(),
-        max_tokens=kw.get("max_tokens", 8),
         group_size=kw.get("group_size", 2),
-        rollout_interval=kw.get("rollout_interval", 2),
+        max_tokens=kw.get("max_tokens", 8),
-    ), model
+        max_batch_size=kw.get("max_batch_size", 8),
+        max_len=kw.get("max_len", 128),
+    )
+    rm = ConstantRewardModel(1.0)
+    return (
+        RolloutRunner(
+            generator=generator,
+            reward_model=rm,
+            rollout_interval=kw.get("rollout_interval", 2),
+        ),
+        model,
+    )
 
 
-def test_rollout_runner_shapes():
+def test_rollout_runner_shapes(device):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     runner, _ = _make_runner(device, group_size=3, max_tokens=5)
     batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device)
-    r = runner(batch)
+    r, is_fresh = runner(batch)
+    assert is_fresh
     assert r.prompts.shape == (2, 4)
     assert r.responses.shape == (2, 3, 5)
     assert r.response_mask.shape == (2, 3, 5)
@@ -217,48 +230,52 @@ def test_rollout_runner_shapes():
     assert len(r.response_texts[0]) == 3
 
 
-def test_rollout_runner_cache_returns_same_object():
+def test_rollout_runner_cache_returns_stale_flag(device):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     runner, _ = _make_runner(device, rollout_interval=10)
     batch = _make_prompt_batch(device=device)
-    r1 = runner(batch)
+    r1, fresh1 = runner(batch)
-    r2 = runner(batch)
+    r2, fresh2 = runner(batch)
     assert r1 is r2
+    assert fresh1 is True
+    assert fresh2 is False
 
 
-def test_rollout_runner_step_triggers_new_rollout():
+def test_rollout_runner_step_triggers_new_rollout(device):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     runner, _ = _make_runner(device, rollout_interval=2)
     batch = _make_prompt_batch(device=device)
-    r1 = runner(batch)
+    r1, fresh1 = runner(batch)
+    assert fresh1 is True
     runner.step()
     # interval=2 means trigger when _steps_since_rollout >= 2; 1 step not enough
-    r2 = runner(batch)
+    r2, fresh2 = runner(batch)
-    assert r1 is r2
+    assert r2 is r1
+    assert fresh2 is False
     runner.step()
     # Now _steps_since_rollout == 2 -> re-rollout
-    r3 = runner(batch)
+    r3, fresh3 = runner(batch)
     assert r3 is not r1
+    assert fresh3 is True
 
 
-def test_rollout_runner_clear_cache_forces_rerun():
+def test_rollout_runner_clear_cache_forces_rerun(device):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     runner, _ = _make_runner(device, rollout_interval=100)
     batch = _make_prompt_batch(device=device)
-    r1 = runner(batch)
+    r1, _ = runner(batch)
     runner.clear_cache()
-    r2 = runner(batch)
+    r2, fresh2 = runner(batch)
     assert r2 is not r1
+    assert fresh2 is True
 
 
-def test_rollout_runner_step_resets_counter():
+def test_rollout_runner_step_resets_counter(device):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     runner, _ = _make_runner(device, rollout_interval=1)
     batch = _make_prompt_batch(device=device)
-    r1 = runner(batch)
+    r1, _ = runner(batch)
     runner.step()
-    r2 = runner(batch)
+    r2, fresh2 = runner(batch)
     assert r2 is not r1
+    assert fresh2 is True
     # Counter reset after rollout; second call w/o step should be cached.
-    r3 = runner(batch)
+    r3, fresh3 = runner(batch)
     assert r3 is r2
+    assert fresh3 is False