[FEAT] Rollout outputs: selected-token logprobs, effective seeds and stop token IDs (part of #70)

README.md

@@ -92,7 +92,9 @@ for supported combinations and limitations.
   and [execution scope and constraints](docs/async-speculative.md#cuda-graph-scope-and-ownership).
 - **Offline and online inference.** A Python API, a command-line interface,
   and an OpenAI-compatible completions endpoint with streaming, greedy decoding,
-  and seeded top-k/top-p sampling.
+  and seeded top-k/top-p sampling. RL rollouts can request per-token
+  log-probabilities aligned with exit depths, engine-chosen seeds, and stop
+  token IDs.
 - **Prefill/decode disaggregation.** Separate prefill and decode worker pools
   across GPUs on one host, with NIXL KV transfer and overlap between chunked
   prefill computation and transfer.

benchmarks/cdb_runtime.py

@@ -63,13 +63,18 @@ def benchmark(device="cpu", layout="last_exited", trace=None, seed=123):
         rid: SamplingParams(max_tokens=6, exit_threshold=q, ignore_eos=True)
         for rid, q in zip(prompts, [0.0, 1.0, 0.5])
     }
-    fingerprint = dict(
-        model_config=model.config.to_dict(),
-        seed=seed,
-        dtype=str(dtype),
-        layout=layout,
-        prompts=prompts,
-        sampling={rid: asdict(p) for rid, p in parameters.items()},
+    # JSON-canonical (tuples become lists) to match a --trace-out/--trace-in round trip.
+    fingerprint = json.loads(
+        json.dumps(
+            dict(
+                model_config=model.config.to_dict(),
+                seed=seed,
+                dtype=str(dtype),
+                layout=layout,
+                prompts=prompts,
+                sampling={rid: asdict(p) for rid, p in parameters.items()},
+            )
+        )
     )
     backend = "triton" if device == "cuda" else "torch"
 

docs/async-speculative.md

@@ -244,8 +244,12 @@ use the reserved allocation and stream-ordered writes instead. Only completed
 full prompt pages are eligible for prefix caching.
 
 Every ticket owns its result buffer and page-table snapshot until compute and
-copy events complete. EOS/length stopping is applied in output order on the CPU;
-one already-submitted later round can finish but its outputs are discarded.
+copy events complete. For rounds where a request asked for `logprobs`, each
+verification row's FP32 raw logprob of its argmax is copied with the round's
+results on the same event. Committed tokens are exactly those argmaxes. EOS,
+`stop_token_ids` and length stopping (`Scheduler._record_token`) is applied in
+output order on the CPU; one
+already-submitted later round can finish but its outputs are discarded.
 Cancellation removes scheduling eligibility immediately. KV allocations and
 device state survive until their final GPU user completes. Request object
 identity isolates stale tickets when a request ID is reused.

docs/launching.md

@@ -260,6 +260,70 @@ execution. The first generated token comes from full-depth prefill, reflected
 in `exit_depths`. For dynamic arrivals, stepwise output, and cancellation, see
 the [architecture overview](unified_execution.md).
 
+### Rollout outputs
+
+RL trainers can request, for each generated token, its log-probability under the
+policy that sampled it (RFC [#70](https://github.com/InferMatrixLab/vllm-rlt/issues/70)).
+Names follow vLLM.
+
+| Option | Default | Contract |
+| --- | --- | --- |
+| `SamplingParams.logprobs` | `None` | `0` adds `output.logprobs`, the FP32 log-probability of each sampled token. Positive values (top-k alternatives) raise `ValueError` for now. |
+| `SamplingParams.seed` | `0` | `None` asks the engine for a fresh 63-bit seed from OS entropy at `add_request`, consuming neither the torch nor the Python global RNG. `output.sampling_params` reports the effective parameters, including the seed the request's generator is created with. |
+| `SamplingParams.stop_token_ids` | `()` | Emitting one finishes the request with `stop`, even with `ignore_eos=True`; `finish_reason` does not say which stop ID or EOS matched. The token stays in `token_ids`; `LLM.generate` omits its text. IDs outside the vocabulary are rejected. |
+| `logprobs_mode` of `LLM`, `StageEngine` (`LLMEngine`), `PDEngine` | `"raw_logprobs"` | `"raw_logprobs"` is `log_softmax` of the CODA logits. `"processed_logprobs"` is `log_softmax` of the logits sampling draws from (temperature, then top-k, then top-p); greedy rows report the raw value, as vLLM does. No CLI flag yet. |
+
+Every output of a request with `logprobs=0`, including streamed, `stop`, `length`
+and `abort` outputs, has `len(logprobs) == len(token_ids) == len(exit_depths)`.
+Requesting logprobs consumes no RNG and changes no token or exit depth. Replaying
+the reported `sampling_params` reproduces tokens and logprobs bitwise only with the
+same batch composition, configuration and hardware. A resumed preempted request or
+a prefix-cache hit matches the uninterrupted cold run numerically, not necessarily
+bitwise, because the batch composition (and, for a prefix hit, the prefill
+chunking) changes. Huginn is the exception: it draws its random recurrent-state
+initialization from the request's generator, greedy requests included, so its
+`seed` matters even at `temperature=0`, and the draws depend on how the prompt is
+split into prefill chunks. When the chunking changes, for example under memory
+pressure, a Huginn request's tokens and logprobs can differ, not just numerically.
+
+`logprobs[i]` is conditional on `exit_depths[i]`: the policy is the model truncated
+at that loop depth. A teacher-forced recompute therefore has to replay the depths:
+
+- Prompt positions run at the model depth and the input `token_ids[i]` (position
+  `len(prompt) + i`) runs `exit_depths[i + 1]` loops, so the input depths are
+  `[depth] * len(prompt) + exit_depths[1:]`. The last generated token is never an
+  input.
+- Under `last_exited` KV, a token that exited at depth `d` has its depth-`d` KV
+  copied into the deeper loop planes, and later tokens attend to those copies. An
+  ordinary full-depth forward matches only rollouts whose `exit_depths` all equal
+  the model depth: `exit_threshold=1` and `max_loops` equal to the model depth, in
+  a non-`trace` exit mode (`trace` takes its depths from `depths_by_request`).
+- Under `shared` KV, positions read earlier tokens' final retained KV at every
+  loop depth (see [KV layout examples](kv_layout_computation.md)), so no ordinary
+  forward reproduces the logprobs, even at full depth.
+- Huginn's logprobs also depend on its random recurrent-state initialization, which
+  shares the request's generator with sampling, so no ordinary forward reproduces
+  them either.
+
+Under [speculative decoding](async-speculative.md), every `exit_depths` entry is
+`target_loops`, and each committed token is scored on its target verification row.
+This covers an accepted draft token, the correction after a rejection, and the bonus
+token. The draft distribution q never enters `logprobs`:
+
+- Verification resumes from the hidden state and KV that the draft loops computed
+  with the same weights. The scores therefore match a full-depth forward
+  numerically, not necessarily bitwise. The rollout replays like any other
+  full-depth rollout.
+- A correction is drawn from the residual `(p - q)+`. The round's output is still
+  distributed as p, so its policy logprob is `log p[token]`.
+- For a given seed, the sampled tokens depend on q, because draft draws and
+  rejection tests consume the request's generator. The logprob of a given token
+  does not depend on q. Greedy tokens do not depend on it either.
+- Cross-round async speculation is greedy-only, so it reports raw values in both
+  modes.
+
+The HTTP API does not accept `logprobs`, `stop_token_ids` or `seed: null` yet.
+
 ## Troubleshooting
 
 | Symptom | What to check |

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 `logprobs`, `exit_depths`, the
+committed `weight_version`, and its effective `sampling_params`. `logprobs=None`
+keeps the ordinary inference path. `logprobs=0` requests only the emitted token's
+log-probability; top-N alternatives are not implemented.
+
+`logprobs_mode` is an engine option of `LLM`, `StageClient`, `LLMEngine` and
+`PDEngine`, not a request field. `"raw_logprobs"` is the default: FP32 log-softmax
+of unscaled model logits. `"processed_logprobs"` scores the distribution the token
+was drawn from, 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 engine admits the request. The
+effective seed is returned in `sampling_params`, including for an abort before
+sampling. The default remains `seed=0`. [Rollout outputs](launching.md#rollout-outputs)
+states when replaying those parameters reproduces tokens and logprobs.
+
+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 and
+asynchronous decoding and by synchronous and cross-round asynchronous
+self-speculation; the latter is greedy-only, so it reports raw scores in both
+modes. PD follows the same rules 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_async_speculative_graphs.py

@@ -190,8 +190,8 @@ def test_failed_graph_round_drains_readers_and_releases_buffers(backend, monkeyp
         runner = engine.speculative_runner
         verify = runner._verify
 
-        def fail_after_verify(*args):
-            verify(*args)
+        def fail_after_verify(*args, **kwargs):
+            verify(*args, **kwargs)
             raise RuntimeError("failed after graph submission")
 
         monkeypatch.setattr(runner, "_verify", fail_after_verify)

tests/test_cdb_runtime.py

@@ -1,5 +1,6 @@
 """End-to-end checks of cache semantics, delayed exits and scheduling progress."""
 
+import json
 from dataclasses import replace
 
 import pytest
@@ -455,6 +456,8 @@ def test_synthetic_replay_variants_preserve_outputs():
     assert trace["depths_by_request"]["A"] == [4, 2, 2, 2, 2, 2]
     assert len(results) == 8
     assert all(r["seconds"] > 0 for r in results)
+    # --trace-out/--trace-in round-trip the trace through JSON.
+    assert len(benchmark(trace=json.loads(json.dumps(trace)))[1]) == 8
     with pytest.raises(ValueError, match="does not match"):
         benchmark(layout="shared", trace=trace)
 

tests/test_engine.py

@@ -266,7 +266,7 @@ def unexpected(*args, **kwargs):
         pytest.fail("sampling generators must be resolved before execution")
 
     monkeypatch.setattr(engine.rng, "acquire", unexpected)
-    token = runner._execute(None, prepared)
+    token = runner._execute(None, prepared).values
     assert token.shape == (1,) and token.device == runner.device
     engine.close()
 

tests/test_mixed_execution.py

@@ -316,7 +316,7 @@ def test_prepared_batch_survives_progress_changes_and_contains_no_policy_label()
         scheduled_request.request.generated_token_ids.append(17)
     # The prepared batch owns its addresses and routing; later request progress cannot
     # retarget the submitted operation, even for mixed multi-row state bindings.
-    actual = engine.model_runner._execute(batch, prepared)
+    actual = engine.model_runner._execute(batch, prepared).values
     torch.testing.assert_close(actual, torch.tensor(expected.gate_scores))
     assert not hasattr(prepared.batch, "stage")
     for left, right in zip(prepared.batch.spans, build_execution_batch(reference_batch, 4).spans):

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),