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)