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
This commit is contained in:
2026-07-20 12:52:20 +08:00
parent 754624acf0
commit 95c43368ae
11 changed files with 662 additions and 303 deletions
+11 -2
View File
@@ -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