Expose rollout logprobs and full policy publication on shared engines

docs/training_contract.md

@@ -0,0 +1,61 @@
+# Training rollout and full-weight publication
+
+`LLM`, `StageClient`, and `LLMEngine` (the local `StageEngine` alias) expose the same full-update operations. `LLM` also accepts
+an existing `LLMEngine` or `PDEngine`, so framework adapters use the same generation
+and publication calls with either backend.
+
+```python
+llm.start_weight_update(version=1)
+for chunk in converted_hf_weights:
+    llm.update_weights(chunk)  # Iterable[(physical_parameter_name, torch.Tensor)]
+llm.finish_weight_update()
+outputs = llm.generate(prompts, SamplingParams(
+    max_tokens=128, temperature=1.0, seed=42, logprobs=0,
+))
+```
+
+Each output contains aligned selected-token `log_probs`, `exit_depths`, the
+committed `weight_version`, and its effective `sampling_params`. `logprobs=None`
+keeps the ordinary inference path. Reported parameters resolve the model's
+default `max_loops` and cap `top_k` at its vocabulary size. `logprobs=0` requests only the emitted token's
+probability; top-N alternatives are not implemented.
+
+`logprobs_mode="raw"` is the default: FP32 log-softmax of unscaled model logits.
+`"processed"` returns probabilities after positive temperature and top-k/top-p
+filtering. Greedy requests return the raw score in either mode. Speculative
+outputs use the verified **target** policy's score, including correction/bonus
+tokens; proposal or rejection-residual probabilities are never exposed as the
+policy score. Explicit `stop_token_ids` apply even with `ignore_eos=True`.
+`seed=None` resolves an independent seed in `[0, 2**63)` when the request is
+created. The effective seed is returned in `sampling_params`, including for an
+abort before sampling. Reusing those parameters replays the same request on the
+same policy and execution configuration. The default remains `seed=0`.
+
+Publication requires an idle engine. Each chunk is checked for duplicate names,
+unknown names and mismatched shapes before that chunk is copied. `finish_weight_update`
+requires every physical parameter exactly once. Copies preserve parameter addresses
+and captured graphs; a fresh transaction invalidates all cached prefix KV. If a
+transfer fails, generation stays blocked until a fresh full transaction succeeds.
+Starting again discards an unfinished transaction. `pause_generation` stops
+admission while existing requests drain; `resume_generation` requires publication
+to be complete. `abort_request` remains available to cancel requests before an update.
+
+PDEngine waits for all transfers to release, broadcasts CPU-staged weight chunks
+to every P/D worker, and commits a version only after every peer acknowledges it.
+Recovery uses a version newer than **all attempted** PD versions because some
+workers may already have committed. With `terminate_workers_on_failure=False`, fatal workers retain their registered
+memory rather than receive forced process signals; normal `close()` remains
+cooperative. No CUDA tensors change ownership over IPC;
+NIXL's existing GPU KV path and registrations remain intact. Prefix invalidation
+does not replace registered KV storage. PD worker control errors return replies
+without terminating the worker.
+
+Existing exit, async and speculative configuration rules still apply. Self-speculation
+uses a full-depth target. Selected-token scores are supported by synchronous
+self-speculation and ordinary asynchronous decoding; cross-round asynchronous
+self-speculation rejects logprobs requests at admission because its compact result
+buffers contain token IDs only. PD follows the same synchronous and asynchronous
+self-speculation limits on its decode workers. Publication
+adds no such restriction to ordinary early-exit or async rollouts. Trainers must
+replay full-depth prompt prefill and the returned decode depths rather than assume
+the entire sequence used K.

tests/test_pd.py

@@ -484,6 +484,8 @@ def test_engine_yields_while_waiting_for_remote_kv(async_scheduling):
 
     engine = object.__new__(LLMEngine)
     engine.async_speculative = False
+    engine._update_version = None
+    engine._update_failed = False
     engine.speculative_runner = None
     engine.profiling = Profiler("cpu")
     engine.preemption = PreemptionManager(engine)

tests/test_profiling.py

@@ -680,6 +680,13 @@ def start_profile(config, **identity):
         add_request=lambda *a: None,
         abort_request=lambda *a: None,
         stats=lambda: {},
+        get_weight_version=lambda: 0,
+        pause_generation=lambda: None,
+        resume_generation=lambda: None,
+        reset_prefix_cache=lambda: None,
+        start_weight_update=lambda **kwargs: None,
+        update_weights=lambda **kwargs: None,
+        finish_weight_update=lambda: 1,
         start_profile=start_profile,
         stop_profile=lambda: stopped.append(True),
         profile_status=lambda: dict(session_id=session),

tests/test_training_contract.py

@@ -0,0 +1,254 @@
+"""Full-publication recovery and selected-token scores across stage refactoring."""
+
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers import LocalLLM, tiny_ouro_model
+from tests.reference import dense_reference
+from vllm_rlt import (
+    LLM,
+    CacheConfig,
+    ExecutionConfig,
+    ExitConfig,
+    SamplingParams,
+    SpeculativeConfig,
+)
+from vllm_rlt.engine.llm_engine import StageClient
+from vllm_rlt.engine.stage_engine import StageEngine
+from vllm_rlt.pd.config import PDConfig
+from vllm_rlt.pd.engine import PDEngine
+
+
+@pytest.mark.parametrize("asynchronous,speculative", [(False, False), (True, False), (False, True)])
+@pytest.mark.parametrize("mode", ["raw", "processed"])
+@pytest.mark.parametrize("temperature", [0, 0.7])
+def test_selected_scores_match_dense_target_and_keep_tokens(
+    asynchronous, speculative, mode, temperature
+):
+    model = tiny_ouro_model()
+    llm = LocalLLM(
+        model,
+        cache_config=CacheConfig(256, 2),
+        execution_config=ExecutionConfig(async_scheduling=asynchronous),
+        exit_config=ExitConfig("ouro_delayed" if asynchronous else "ouro"),
+        speculative_config=SpeculativeConfig(2) if speculative else None,
+    )
+    params = SamplingParams(
+        max_tokens=7,
+        ignore_eos=True,
+        temperature=temperature,
+        logprobs=0,
+        logprobs_mode=mode,
+    )
+    prompts = [[2, 3], [4, 5, 6]]
+    try:
+        outputs = llm.generate(prompts, params)
+        ordinary = llm.generate(prompts, replace(params, logprobs=None))
+        assert [o.token_ids for o in outputs] == [o.token_ids for o in ordinary]
+        for prompt, output in zip(prompts, outputs):
+            assert len(output.log_probs) == len(output.token_ids) == 7
+            prefix = list(prompt)
+            for token, actual in zip(output.token_ids, output.log_probs):
+                logits = dense_reference(model, prefix, model.config.total_ut_steps)[-1][2][-1]
+                if mode == "processed" and temperature > 0:
+                    logits = logits / temperature
+                expected = logits.float().log_softmax(-1)[token].item()
+                assert actual == pytest.approx(expected, abs=2e-5)
+                prefix.append(token)
+    finally:
+        llm.close()
+
+
+def test_full_publication_invalidates_prefix_and_preserves_parameter_addresses():
+    model = tiny_ouro_model()
+    engine = StageEngine(model, cache_config=CacheConfig(128, 2, enable_prefix_caching=True))
+    llm = LLM(engine)
+    params = SamplingParams(max_tokens=2, ignore_eos=True)
+    try:
+        llm.generate([[2, 3, 4, 5]], params)
+        assert engine.cache_manager.lookup_prefix([2, 3, 4, 5])
+        addresses = {name: value.data_ptr() for name, value in model.named_parameters()}
+        weights = [
+            (name, value.detach().clone() * 0.95) for name, value in model.named_parameters()
+        ]
+        llm.start_weight_update(3)
+        assert not engine.cache_manager.lookup_prefix([2, 3, 4, 5])
+        with pytest.raises(RuntimeError, match="paused"):
+            engine.add_request("blocked", [1], params)
+        llm.update_weights(weights[:1])
+        with pytest.raises(ValueError, match="every physical"):
+            llm.finish_weight_update()
+        llm.update_weights(weights[1:])
+        llm.finish_weight_update()
+        assert llm.get_weight_version() == 3
+        assert addresses == {name: value.data_ptr() for name, value in model.named_parameters()}
+        for (name, value), (_, expected) in zip(model.named_parameters(), weights):
+            torch.testing.assert_close(value, expected)
+        output = llm.generate([[2, 3]], params)[0]
+        assert output.weight_version == 3
+        assert output.sampling_params.max_loops == model.config.total_ut_steps
+    finally:
+        llm.close()
+
+
+@pytest.mark.parametrize("failure", ["duplicate", "shape", "iteration"])
+def test_failed_chunk_blocks_generation_until_new_complete_transaction(failure):
+    model = tiny_ouro_model()
+    engine = StageEngine(model, cache_config=CacheConfig(64, 2))
+    weights = [(name, value.detach().clone()) for name, value in model.named_parameters()]
+
+    def interrupted_chunk():
+        yield weights[0]
+        raise RuntimeError("interrupted source")
+
+    chunk = {
+        "duplicate": [weights[0], weights[0]],
+        "shape": [(weights[0][0], torch.empty(1))],
+        "iteration": interrupted_chunk(),
+    }[failure]
+    try:
+        engine.start_weight_update(1)
+        original = dict(model.named_parameters())[weights[0][0]].detach().clone()
+        with pytest.raises((ValueError, RuntimeError)):
+            engine.update_weights(chunk)
+        torch.testing.assert_close(dict(model.named_parameters())[weights[0][0]], original)
+        for call in (engine.step, engine.resume_generation, engine.finish_weight_update):
+            with pytest.raises(RuntimeError):
+                call()
+        with pytest.raises(ValueError, match="increase"):
+            engine.start_weight_update(1)
+        engine.start_weight_update(2)
+        engine.update_weights(weights)
+        assert engine.finish_weight_update() == 2
+        engine.resume_generation()
+        engine.add_request("recovered", [1], SamplingParams(max_tokens=1))
+        engine.abort_request("recovered")
+    finally:
+        engine.close()
+
+
+def test_publication_waits_for_requests_and_remote_transfer_leases():
+    engine = StageEngine(tiny_ouro_model(), cache_config=CacheConfig(64, 2))
+    try:
+        engine.add_request("live", [1, 2], SamplingParams(max_tokens=1))
+        engine.pause_generation()
+        with pytest.raises(RuntimeError, match="idle"):
+            engine.start_weight_update()
+        engine.abort_request("live")
+        engine.connector = SimpleNamespace(has_pending=True)
+        with pytest.raises(RuntimeError, match="idle"):
+            engine.start_weight_update()
+        engine.connector = None
+        engine.start_weight_update()
+    finally:
+        engine.close()
+
+
+@pytest.mark.parametrize("speculative", [False, True])
+def test_explicit_stop_applies_with_ignore_eos_and_scores_align(speculative):
+    llm = LocalLLM(
+        tiny_ouro_model(),
+        speculative_config=SpeculativeConfig(2) if speculative else None,
+        cache_config=CacheConfig(128, 2),
+    )
+    try:
+        params = SamplingParams(max_tokens=7, ignore_eos=True, logprobs=0)
+        full = llm.generate([[2, 3]], params)[0]
+        stop = full.token_ids[2]
+        stopped = llm.generate([[2, 3]], replace(params, stop_token_ids=(stop,)))[0]
+        count = full.token_ids.index(stop) + 1
+        assert stopped.token_ids == full.token_ids[:count]
+        assert stopped.log_probs == pytest.approx(full.log_probs[:count], abs=2e-5)
+        assert stopped.finish_reason == "stop"
+    finally:
+        llm.close()
+
+
+def test_cpu_stage_clients_publish_to_all_pd_peers_and_recover_control_errors():
+    model = tiny_ouro_model()
+    clients = [
+        StageClient(model, device="cpu", dtype="float32", cache_config=CacheConfig(64, 2))
+        for _ in range(2)
+    ]
+    pd = object.__new__(PDEngine)
+    pd.config = PDConfig()
+    pd.peers = {str(i): SimpleNamespace(client=c) for i, c in enumerate(clients)}
+    pd.requests, pd.transfers = {}, {}
+    pd.closed, pd.failure = False, None
+    pd._weight_version, pd._update_version, pd._update_failed, pd._paused = 0, None, False, False
+    try:
+        clients[0].pause_generation()
+        with pytest.raises(RuntimeError, match="paused"):
+            clients[0].add_request("paused", [1], SamplingParams(max_tokens=1))
+        assert clients[0].process.is_alive() and clients[0].failure is None
+        clients[0].resume_generation()
+        pd.start_weight_update(1)
+        # One rank rejects a duplicate while another has already copied its chunk.
+        first = next(iter(model.named_parameters()))
+        clients[1].update_weights([first])
+        with pytest.raises(RuntimeError, match="Duplicate"):
+            pd.update_weights([first])
+        assert pd.get_weight_version() == 0
+        assert pd._update_failed
+        assert all(c.process.is_alive() and c.failure is None for c in clients)
+        assert all(not c.pending_controls and not c.replies for c in clients)
+        with pytest.raises(RuntimeError):
+            pd.resume_generation()
+        pd.start_weight_update(2)
+        pd.update_weights(model.named_parameters())
+        pd.finish_weight_update()
+        assert pd.get_weight_version() == 2
+        assert [c.get_weight_version() for c in clients] == [2, 2]
+        llm = LLM(clients[0])
+        output = llm.generate([[2, 3]], SamplingParams(max_tokens=2, logprobs=0))[0]
+        assert output.weight_version == 2
+        assert len(output.log_probs) == len(output.token_ids)
+    finally:
+        for client in clients:
+            client.close()
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+def test_mixed_score_admission_and_unspecified_seed_replay(asynchronous):
+    model = tiny_ouro_model()
+    llm = LocalLLM(
+        model,
+        cache_config=CacheConfig(128, 2),
+        execution_config=ExecutionConfig(async_scheduling=asynchronous),
+        exit_config=ExitConfig("ouro_delayed" if asynchronous else "ouro"),
+    )
+    try:
+        scored = SamplingParams(
+            max_tokens=4,
+            temperature=0.8,
+            logprobs=0,
+            seed=None,
+            top_k=model.config.vocab_size + 5,
+            ignore_eos=True,
+        )
+        outputs = llm.generate([[2, 3], [4, 5]], [replace(scored, logprobs=None), scored])
+        assert outputs[0].log_probs is None
+        assert len(outputs[1].log_probs) == 4
+        effective = outputs[1].sampling_params
+        assert effective.top_k == model.config.vocab_size
+        assert 0 <= effective.seed < 2**63
+        replay = llm.generate([[4, 5]], effective)[0]
+        assert replay.token_ids == outputs[1].token_ids
+        assert replay.log_probs == pytest.approx(outputs[1].log_probs, abs=2e-5)
+    finally:
+        llm.close()
+
+
+def test_async_speculative_logprob_rejected_before_any_request_or_transfer():
+    engine = StageEngine(tiny_ouro_model(), cache_config=CacheConfig(64, 2))
+    engine.async_speculative = True  # Admission only: no CUDA speculative runner is constructed.
+    try:
+        with pytest.raises(ValueError, match="logprobs"):
+            engine.add_request("unsupported", [1], SamplingParams(logprobs=0))
+        assert not engine.has_unfinished_requests()
+    finally:
+        engine.async_speculative = False
+        engine.close()

vllm_rlt/core/kv_cache_manager.py

@@ -276,6 +276,12 @@ def lookup_prefix(self, tokens):
             found.append(blocks)
         return tuple(found)
 
+    def reset_prefix_cache(self):
+        self._pending_prefixes.clear()
+        for blocks in self._prefixes.values():
+            self._drop_refs(blocks)
+        self._prefixes.clear()
+
     def allocate(self, request_id: str, max_tokens: int, *, initial_tokens=None, prefix=()) -> bool:
         """Reserve a logical capacity; acquire physical pages for the current frontier."""
         if not isinstance(request_id, str) or not request_id:

vllm_rlt/core/scheduler.py

@@ -1194,7 +1194,9 @@ def _record_token(self, request: Request, token: int, depth: int) -> bool:
         request.generated_token_ids.append(token)
         request.exit_depths.append(depth)
         params = request.sampling_params
-        if token in self.eos_token_ids and not params.ignore_eos:
+        if token in params.stop_token_ids or (
+            token in self.eos_token_ids and not params.ignore_eos
+        ):
             request.finish_reason = FinishReason.STOP
         elif len(request.generated_token_ids) >= params.max_tokens:
             request.finish_reason = FinishReason.LENGTH
@@ -1294,6 +1296,8 @@ def update_from_output(
                         raise RuntimeError("invalid pending output count")
                     request.num_output_placeholders -= 1
                 depth = result.depths[row] if delivered else request.loops_done
+                if request.sampling_params.logprobs is not None:
+                    request.log_probs.append(result.log_probs[row])
                 if self._record_token(request, result.token_ids[row], depth):
                     update.finished.append(request)
                 elif not delivered:
@@ -1317,7 +1321,9 @@ def update_from_output(
                             "speculative GPU frontier disagrees with committed prefix"
                         )
                 emitted = 0
-                for token in spec.token_ids:
+                for index, token in enumerate(spec.token_ids):
+                    if request.sampling_params.logprobs is not None:
+                        request.log_probs.append(spec.log_probs[index])
                     emitted += 1
                     if self._record_token(request, token, self.speculative_config.target_loops):
                         update.finished.append(request)