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

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`,
-:func:`generate_responses`, and :class:`RolloutRunner` including
-its internal cache and rollout-interval trigger logic.
+Covers :class:`RolloutResult` / :class:`RawRollout`, :class:`BaseRewardModel`,
+:class:`RolloutGenerator` (KV-cache-backed via :class:`InferenceScheduler.run_batch`)
+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 (
-    SamplingPipeline,
-    TemperatureStrategy,
-    TopKStrategy,
-    TopPStrategy,
-)
+from astrai.inference.core.scheduler import InferenceScheduler
 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():
-    return SamplingPipeline(
-        [TemperatureStrategy(1.0), TopKStrategy(0), TopPStrategy(1.0)]
+def _make_scheduler(model, tokenizer, max_batch_size=8, max_len=128):
+    return InferenceScheduler(
+        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():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+@pytest.fixture
+def device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _make_generator(device, **kw):
     model, _ = _make_model(device)
-    pipeline = _make_pipeline()
-    ids = torch.randint(3, 200, (2, 4), device=device)
-    mask = torch.ones(2, 4, dtype=torch.bool, device=device)
-
-    out = generate_responses(
-        model=model,
-        input_ids=ids,
-        attention_mask=mask,
-        max_new_tokens=8,
-        sampling_pipeline=pipeline,
-        stop_ids=[],
+    tokenizer = FakeTokenizer()
+    scheduler = _make_scheduler(
+        model,
+        tokenizer,
+        max_batch_size=kw.get("max_batch_size", 8),
+        max_len=kw.get("max_len", 128),
     )
-    assert out["generated_ids"].shape == (2, 8)
-    assert out["generated_mask"].shape == (2, 8)
-    assert out["logprobs"].shape == (2, 8)
-
-
-def test_generate_responses_stops_on_stop_id():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model, _ = _make_model(device)
-    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],
+    generator = RolloutGenerator(
+        scheduler=scheduler,
+        tokenizer=tokenizer,
+        max_tokens=kw.get("max_tokens", 8),
+        group_size=kw.get("group_size", 2),
+        temperature=kw.get("temperature", 1.0),
+        top_k=kw.get("top_k", 0),
+        top_p=kw.get("top_p", 1.0),
     )
-    gen = out["generated_ids"][0]
-    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
+    return generator, model
 
 
-def test_generate_responses_logprobs_match_tokens():
-    """logprobs[i] must be the logprob of generated_ids[i]."""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model, _ = _make_model(device)
-    pipeline = _make_pipeline()
-    ids = torch.randint(3, 200, (1, 2), device=device)
-    mask = torch.ones(1, 2, dtype=torch.bool, device=device)
+def test_rollout_generator_shapes(device):
+    gen, _ = _make_generator(device, group_size=3, max_tokens=5)
+    batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device)
+    r = gen.generate(batch)
+    assert r.prompts.shape == (2, 4)
+    assert r.responses.shape == (2, 3, 5)
+    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,
-        input_ids=ids,
-        attention_mask=mask,
-        max_new_tokens=4,
-        sampling_pipeline=pipeline,
-        stop_ids=[],
-    )
-    gen = out["generated_ids"][0]
-    lp = out["logprobs"][0]
-    for i in range(4):
-        if gen[i] == 0 and not out["generated_mask"][0, i]:
-            continue
-        assert lp[i] <= 0.0
+
+def test_rollout_generator_mask_matches_responses(device):
+    """Positions beyond a response's length are pad (mask False)."""
+    gen, _ = _make_generator(device, group_size=2, max_tokens=6)
+    batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device)
+    r = gen.generate(batch)
+    for i in range(2):
+        for g in range(2):
+            real = r.response_mask[i, g].sum().item()
+            # Pad positions should be 0
+            assert r.responses[i, g, real:].sum() == 0
+            # logprobs after the real tokens are 0 (padding)
+            if real < r.logprobs_old.size(-1):
+                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)
-    rm = ConstantRewardModel(1.0)
-    return RolloutRunner(
-        policy_model=model,
-        old_model=_FakeOldModel(),
-        tokenizer=FakeTokenizer(),
-        reward_model=rm,
-        sampling_pipeline=_make_pipeline(),
-        max_tokens=kw.get("max_tokens", 8),
+    generator, model = _make_generator(
+        device,
         group_size=kw.get("group_size", 2),
-        rollout_interval=kw.get("rollout_interval", 2),
-    ), model
+        max_tokens=kw.get("max_tokens", 8),
+        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():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+def test_rollout_runner_shapes(device):
     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():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+def test_rollout_runner_cache_returns_stale_flag(device):
     runner, _ = _make_runner(device, rollout_interval=10)
     batch = _make_prompt_batch(device=device)
-    r1 = runner(batch)
-    r2 = runner(batch)
+    r1, fresh1 = 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():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+def test_rollout_runner_step_triggers_new_rollout(device):
     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)
-    assert r1 is r2
+    r2, fresh2 = runner(batch)
+    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():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+def test_rollout_runner_clear_cache_forces_rerun(device):
     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():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+def test_rollout_runner_step_resets_counter(device):
     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