diff --git a/README.md b/README.md index b742543f..7fa75747 100644 --- a/README.md +++ b/README.md @@ -91,7 +91,9 @@ for supported combinations and limitations. [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. diff --git a/benchmarks/cdb_runtime.py b/benchmarks/cdb_runtime.py index 0b75e849..399c05c0 100644 --- a/benchmarks/cdb_runtime.py +++ b/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" diff --git a/docs/async-speculative.md b/docs/async-speculative.md index 543efc5a..316d3337 100644 --- a/docs/async-speculative.md +++ b/docs/async-speculative.md @@ -123,8 +123,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. diff --git a/docs/launching.md b/docs/launching.md index 79ced6bd..9ae862e3 100644 --- a/docs/launching.md +++ b/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 | diff --git a/tests/test_async_speculative_graphs.py b/tests/test_async_speculative_graphs.py index f43b560e..bbd49f9a 100644 --- a/tests/test_async_speculative_graphs.py +++ b/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) diff --git a/tests/test_cdb_runtime.py b/tests/test_cdb_runtime.py index 8b345dcb..ab955c7f 100644 --- a/tests/test_cdb_runtime.py +++ b/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) diff --git a/tests/test_engine.py b/tests/test_engine.py index 1ad00979..b07976ca 100644 --- a/tests/test_engine.py +++ b/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() diff --git a/tests/test_mixed_execution.py b/tests/test_mixed_execution.py index 1eb343f0..ef3eddd0 100644 --- a/tests/test_mixed_execution.py +++ b/tests/test_mixed_execution.py @@ -255,7 +255,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): diff --git a/tests/test_rollout_outputs.py b/tests/test_rollout_outputs.py new file mode 100644 index 00000000..3d206829 --- /dev/null +++ b/tests/test_rollout_outputs.py @@ -0,0 +1,1195 @@ +"""Rollout outputs: per-token logprobs, effective seeds and stop token IDs.""" + +import gc +import math +import random +import time +from collections import defaultdict, deque +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from tests.reference import dense_reference +from tests.test_huginn import tiny_huginn_config +from vllm_rlt import LLM, CacheConfig, ExecutionConfig, ExitConfig, SamplingParams, SchedulerConfig +from vllm_rlt.config import SpeculativeConfig +from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.models import HuginnForCausalLM, OuroForCausalLM +from vllm_rlt.pd.config import PDConfig +from vllm_rlt.pd.engine import PDEngine, Peer +from vllm_rlt.request import Operation, RequestOutput, RequestPhase +from vllm_rlt.serving.protocol import CompletionRequest +from vllm_rlt.worker.async_speculative import SpeculativeRoundTicket +from vllm_rlt.worker.logprobs import token_logprob +from vllm_rlt.worker.model_runner import ReadbackSlot, Submission +from vllm_rlt.worker.output import ModelRunnerOutput, SpeculativeResult +from vllm_rlt.worker.sampler import Sampler +from vllm_rlt.worker.speculative import _CommittedScores + +TRACES = {"a": [4, 1, 3, 2, 4, 1, 2, 3], "b": [4, 2, 1, 4, 3, 2, 1, 4]} # Keyed by request ID. +EXIT_MODES = [ # (async scheduling, exit config, exit_threshold) for every CPU exit policy. + pytest.param(False, ExitConfig("ouro"), 1.0, id="ouro-full-depth"), + pytest.param(False, ExitConfig("ouro"), 0.5, id="ouro-early-exit"), + pytest.param(False, ExitConfig("trace", depths_by_request=TRACES), 1.0, id="trace"), + pytest.param(True, ExitConfig("ouro_delayed"), 0.3, id="async-ouro_delayed"), + pytest.param(True, ExitConfig("random_lookahead"), 0.3, id="async-random_lookahead"), + pytest.param(True, ExitConfig("trace", depths_by_request=TRACES), 1.0, id="async-trace"), +] +ASYNC_EXIT_MODES = [pytest.param(*p.values[1:], id=p.id) for p in EXIT_MODES if p.values[0]] +LOGPROBS_MODES = ["raw_logprobs", "processed_logprobs"] +# How request "a" ends: max_tokens, EOS (honored or ignored), stop IDs, both, abort. +FINISHES = "length eos stop_id ignored_eos eos_with_stop_ids stop_id_at_max eos_at_max abort" + + +def model(seed=123, device="cpu", dtype=torch.float32, **overrides): + torch.manual_seed(seed) + return OuroForCausalLM(tiny_ouro_config(**overrides)).to(device, dtype) + + +def rollout(**overrides): + return SamplingParams(**{"ignore_eos": True, "logprobs": 0, **overrides}) + + +def make_engine(m, *, asynchronous=False, exit_config=None, execution=None, **kwargs): + """Small CPU engine by default; asynchronous engines default to ouro_delayed exits.""" + kwargs.setdefault("cache_config", CacheConfig(128, 2)) + kwargs.setdefault("scheduler_config", SchedulerConfig(max_num_seqs=3, max_num_batched_tokens=3)) + execution = ExecutionConfig(async_scheduling=asynchronous, **(execution or {})) + exit_config = exit_config or ExitConfig("ouro_delayed" if asynchronous else "ouro") + return LLMEngine(m, execution_config=execution, exit_config=exit_config, **kwargs) + + +def check_aligned(output): + n, logprobs = len(output.token_ids), output.logprobs + assert len(output.exit_depths) == n + assert (logprobs is None) == (output.sampling_params.logprobs is None) + assert logprobs is None or len(logprobs) == n + assert all(math.isfinite(value) and value <= 0 for value in logprobs or ()) + + +def generate(engine, prompts, params): + """``LLM.generate`` on an in-process engine, whose runners tests can patch (``LLM`` + runs its engine in a stage worker process).""" + params = params if isinstance(params, list) else [params] * len(prompts) + for i, (prompt, request_params) in enumerate(zip(prompts, params)): + engine.add_request(str(i), prompt, request_params) + outputs = final(stream(engine)) + return [outputs[str(i)] for i in range(len(prompts))] + + +def stream(engine, until=None, streamed=None): + """Step until ``until()`` holds, else drain, checking every output's alignment. A drained + LLMEngine (1000-step limit) holds zero used KV blocks; PD gets a 90 s wall-clock limit.""" + streamed = defaultdict(list) if streamed is None else streamed + done = until or (lambda: not engine.has_unfinished_requests()) + local, deadline, steps = isinstance(engine, LLMEngine), time.monotonic() + 90, 0 + while not done(): + assert engine.has_unfinished_requests() + assert (steps < 1000) if local else (time.monotonic() < deadline) + steps += 1 + for output in engine.step(): + check_aligned(output) + streamed[output.request_id].append(output) + if until is None and local: + assert engine.cache_manager.num_used_blocks == 0 + return streamed + + +def final(streamed): + return {rid: outputs[-1] for rid, outputs in streamed.items()} + + +def assert_same(actual, expected, atol=None): + """Same tokens and exit depths; logprobs equal, or within ``atol`` when given.""" + assert (actual.token_ids, actual.exit_depths) == (expected.token_ids, expected.exit_depths) + logprobs = expected.logprobs + if atol is not None and logprobs is not None: + logprobs = pytest.approx(logprobs, abs=atol) + assert actual.logprobs == logprobs + + +def hold_coda_delivery(monkeypatch): + """Report CODA tickets unfinished, so only the one-output bound delivers them.""" + ready = Submission.ready + monkeypatch.setattr(Submission, "ready", lambda t: t.operation != Operation.CODA and ready(t)) + return ready + + +def lease_cpu_logprob_slots(monkeypatch, engine, count=8): + """Return the logprob readback pool. On CPU, which has none, install ``count`` emulated + slots and monkeypatch ``runner.submit`` to lease one per logprob ticket (CUDA: unchanged).""" + runner, submit = engine.model_runner, engine.model_runner.submit + if runner.device.type == "cuda": + return runner.logprob_slots + runner.logprob_slots = [ReadbackSlot(torch.empty(1)) for _ in range(count)] + + def leasing_submit(*args, **kwargs): + ticket = submit(*args, **kwargs) + if ticket.logprobs is not None: + ticket.logprob_slot = runner._readback_slot(ticket.logprobs, runner.logprob_slots) + return ticket + + monkeypatch.setattr(runner, "submit", leasing_submit) + return runner.logprob_slots + + +def test_valid_rollout_params_are_normalized_and_hashable(): + params = SamplingParams(seed=None, logprobs=0, stop_token_ids=[5, 3, 5]) + assert (params.seed, params.logprobs, params.stop_token_ids) == (None, 0, (5, 3, 5)) + assert {params} == {replace(params, stop_token_ids=(5, 3, 5))} + assert SamplingParams().logprobs is None and SamplingParams(seed=2**63 - 1).seed == 2**63 - 1 + + +def rejected(field, match, *values): + return [pytest.param(field, value, match, id=f"{field}={value!r}") for value in values] + + +@pytest.mark.parametrize( + "field,value,match", + rejected("logprobs", "not supported yet", 1, 5) + + rejected("logprobs", "None or 0", -1, True, 0.0, "0") + + rejected("seed", "seed", True, -1, 1.5, "1") + + [pytest.param("seed", 2**63, "seed", id="seed=2**63")] + + rejected("stop_token_ids", "stop_token_ids", None, 3, "3", [True], [1.0], [-1], {3}), +) +def test_invalid_rollout_params_are_rejected(field, value, match): + with pytest.raises(ValueError, match=match): + SamplingParams(**{field: value}) + + +def test_logprobs_mode_is_validated_and_forwarded_to_stage_workers(monkeypatch): + # LLM and PDEngine run StageEngine in worker processes; capture what they would start. + start = Mock(side_effect=RuntimeError("stop before workers")) + monkeypatch.setattr("vllm_rlt.engine.llm_engine.start_stage_worker", start) + for logprobs_mode in LOGPROBS_MODES: + assert LLMEngine(model(), logprobs_mode=logprobs_mode).model_runner.logprobs_mode == ( + logprobs_mode + ) + for build, source in ((LLM, model), (PDEngine, tiny_ouro_config)): + with pytest.raises(RuntimeError, match="stop before workers"): + build(source(), logprobs_mode=logprobs_mode) + for logprobs_mode in ("raw_logits", "processed_logits", "logprobs", None): + for build, source in ((LLMEngine, model), (LLM, model), (PDEngine, tiny_ouro_config)): + with pytest.raises(ValueError, match="not supported"): + build(source(), logprobs_mode=logprobs_mode) + # Valid modes reach the worker engines (LLM, then PD); invalid ones start no worker. + forwarded = [c.args[3]["engine"]["logprobs_mode"] for c in start.call_args_list] + assert forwarded == [mode for mode in LOGPROBS_MODES for _ in range(2)] + + +def test_admission_rejects_out_of_vocabulary_stop_ids(): + engine = LLMEngine(model()) + for stop in ([64], [3, 1000]): # The tiny vocabulary is IDs 0..63. + with pytest.raises(ValueError, match="stop_token_ids"): + engine.add_request("bad", [1], SamplingParams(stop_token_ids=stop)) + assert not engine.has_unfinished_requests() + engine.add_request("ok", [1], SamplingParams(max_tokens=2, stop_token_ids=[63], logprobs=0)) + assert final(stream(engine))["ok"].finished + + +def test_engine_rejects_missing_logprob_for_requesting_row(): + engine = LLMEngine(model()) + engine.add_request("r", [2], SamplingParams(logprobs=0)) + with pytest.raises(RuntimeError, match="missing logprob"): + engine.scheduler._record_token(engine.scheduler.requests["r"], 3, 4, None) + + +def test_stage_client_round_trips_rollout_outputs(): + """LLM runs StageEngine in a worker process: outputs cross the pipe with their logprobs + and effective params, an abort before the first sample included, and replay exactly.""" + params = rollout(max_tokens=6, temperature=0.9, top_k=8, seed=None) + with LLM(model(), logprobs_mode="processed_logprobs") as llm: + client = llm.engine + client.add_request("early", [2, 3], params) + aborted = client.abort_request("early") + check_aligned(aborted) + assert aborted.finish_reason == "abort" and type(aborted.sampling_params.seed) is int + outputs = llm.generate([[2, 3], [4]], params) + for output in outputs: + check_aligned(output) + assert output.sampling_params == replace(params, seed=output.sampling_params.seed) + assert outputs[0].sampling_params.seed != outputs[1].sampling_params.seed + replay = llm.generate([[2, 3], [4]], [output.sampling_params for output in outputs]) + for actual, expected in zip(replay, outputs): + assert_same(actual, expected) + # A stop ID ends the request in the worker; the stop token keeps its logprob. + stop = outputs[0].token_ids[0] + params = replace(outputs[0].sampling_params, stop_token_ids=[stop]) + (stopped,) = llm.generate([[2, 3]], params) + assert (stopped.finish_reason, stopped.token_ids) == ("stop", [stop]) + assert stopped.logprobs == outputs[0].logprobs[:1] + + +def test_http_request_fields_are_unchanged(): + # tests/test_serving.py covers the other rollout fields. + for value in (0, True): + with pytest.raises(ValueError, match="logprobs"): + CompletionRequest.parse({"model": "m", "prompt": "a", "logprobs": value}, "m") + assert CompletionRequest.parse({"model": "m", "prompt": "a", "seed": 7}, "m").params.seed == 7 + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +@pytest.mark.parametrize("finish", FINISHES.split()) +@pytest.mark.parametrize("asynchronous,exit_config,threshold", EXIT_MODES) +def test_logprobs_are_defined_and_aligned_through_each_finish( + monkeypatch, asynchronous, exit_config, threshold, finish, logprobs_mode +): + rows = defaultdict(list) # CODA logits rows in token order; the baseline runs first. + common = dict(min_loops=1, exit_threshold=threshold) + options = dict(asynchronous=asynchronous, exit_config=exit_config, logprobs_mode=logprobs_mode) + # The unrequested row exercises mixed batches (ignored values). Admitted first, it puts + # "a" behind row 0 in shared batches, so each row must read its own logprob. + unrequested = rollout(max_tokens=6, temperature=0.9, seed=5, logprobs=None, **common) + + def run(params, *, eos=0, abort=False): + engine = make_engine(model(eos_token_id=eos), **options) + runner = engine.model_runner + if params.logprobs is None: # No logprob work when no request asks. + unexpected = AssertionError("logprobs computed without a request") + monkeypatch.setattr(runner, "_sample_with_logprobs", Mock(side_effect=unexpected)) + + def recording(sample): + def record(row, row_params, generator): + if row_params.seed == params.seed: # Request "a"; "b" has seed 5. + rows["a"].append(row.detach().clone()) + return sample(row, row_params, generator) + + return record + + # Processed top-k/top-p rows that request logprobs draw through sample_with_logprob. + for name in ("sample", "sample_with_logprob"): + setattr(runner.sampler, name, recording(getattr(runner.sampler, name))) + engine.add_request("b", [4], unrequested) + engine.add_request("a", [2, 3], replace(params, **common)) + if not abort: + return stream(engine) + request = engine.scheduler.requests["a"] + streamed = stream(engine, lambda: len(request.generated_token_ids) >= 2) + streamed["a"].append(engine.abort_request("a")) + return stream(engine, streamed=streamed) + + params = rollout(max_tokens=6, temperature=1.0, top_k=20, top_p=0.9, seed=7) + baseline = final(run(params))["a"] + tokens = baseline.token_ids + assert baseline.finish_reason == "length" and len(rows["a"]) == len(tokens) == 6 + # Raw scores the token's CODA row; processed, the distribution the sampler draws from. + processed = logprobs_mode == "processed_logprobs" + expected = [ + (Sampler.probabilities(row, params).log() if processed else row.float().log_softmax(-1))[t] + for row, t in zip(rows["a"], tokens) + ] + assert baseline.logprobs == pytest.approx([value.item() for value in expected], abs=1e-6) + # Every policy except fixed full depth must exercise shallow exits here. + shallow = threshold < 1 or exit_config.mode == "trace" + assert any(depth < 4 for depth in baseline.exit_depths) == shallow + # Asking for logprobs changes no token or exit depth. + quiet = final(run(replace(params, logprobs=None)))["a"] + assert (quiet.token_ids, quiet.exit_depths) == (tokens, baseline.exit_depths) + # Stop on the first token that is new at index >= 2, so a stop must cut the stream. + index = next(i for i, t in enumerate(tokens) if i >= 2 and t not in tokens[:i]) + stop, other = tokens[index], next(t for t in range(64) if t not in tokens) + # finish -> (EOS ID, overrides, reason, length); stop IDs, unlike EOS, override ignore_eos. + honor, cut = dict(ignore_eos=False), dict(max_tokens=index + 1) + cases = { + "length": (0, {}, "length", 6), + "eos": (stop, honor, "stop", index + 1), + "stop_id": (0, dict(stop_token_ids=[other, stop]), "stop", index + 1), + "ignored_eos": (stop, {}, "length", 6), + "eos_with_stop_ids": (stop, dict(honor, stop_token_ids=[other]), "stop", index + 1), + "stop_id_at_max": (0, dict(cut, stop_token_ids=[stop]), "stop", index + 1), + "eos_at_max": (stop, dict(honor, **cut), "stop", index + 1), + "abort": (0, {}, "abort", None), # After two delivered tokens. + } + eos, overrides, reason, length = cases[finish] + outputs = run(replace(params, **overrides), eos=eos, abort=finish == "abort")["a"] + n = len(outputs[-1].token_ids) + assert outputs[-1].finish_reason == reason and (2 <= n < 6 if length is None else n == length) + if finish != "abort": # One token per output; a stop token keeps its depth and logprob. + assert [len(o.token_ids) for o in outputs] == list(range(1, n + 1)) + for output in outputs: + # Re-checked after draining: each output must own its lists, not the request's. + check_aligned(output) + k = len(output.token_ids) + assert output.token_ids == tokens[:k] and output.exit_depths == baseline.exit_depths[:k] + assert output.logprobs == baseline.logprobs[:k] + + +# Request ID -> overrides covering each processed-logprob path: one batched raw log_softmax +# (greedy, even with top-k, and T == 1), one per other temperature (0.7 twice, 1.3), the +# sampler's top-k/top-p logits, and a row that asks for no logprobs. +MIXED_ROWS = { + "greedy": dict(temperature=0.0, top_k=3), + "t1": dict(temperature=1.0), + "t07": dict(temperature=0.7), + "t07b": dict(temperature=0.7), + "t13": dict(temperature=1.3), + "top_k": dict(temperature=0.9, top_k=8), + "top_p": dict(temperature=1.0, top_p=0.8), + "top_k_top_p": dict(temperature=0.7, top_k=16, top_p=0.9), + "quiet": dict(temperature=0.8, logprobs=None), +} + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +def test_grouped_processed_logprobs_equal_the_per_row_reference(monkeypatch, asynchronous): + mode, built = "processed_logprobs", [] # built: one entry per processed_logits call. + processed_logits = Sampler.processed_logits + + def counted(logits, params): + built.append(params) + return processed_logits(logits, params) + + monkeypatch.setattr(Sampler, "processed_logits", staticmethod(counted)) + batch_sizes = [] + + def run(logprobs): + scheduler = SchedulerConfig(max_num_seqs=9, max_num_batched_tokens=16) + options = dict(logprobs_mode=mode, scheduler_config=scheduler) + engine = make_engine(model(), asynchronous=asynchronous, **options) + runner = engine.model_runner + batched = runner._sample_with_logprobs + + def checked(logits, params, generators): + before = len(built) + token_ids, values = batched(logits, params, generators) + # Sampling builds each sampled row's processed logits once; logprobs reuse them. + assert len(built) - before == sum(p.temperature > 0 for p in params) + for row, row_params in enumerate(params): # 57d2308 ran token_logprob per row. + if row_params.logprobs is not None: + reference = token_logprob(logits[row], row_params, token_ids[row], mode) + assert torch.equal(values[row], reference) + batch_sizes.append(len(params)) + return token_ids, values + + monkeypatch.setattr(runner, "_sample_with_logprobs", checked) + for i, (rid, overrides) in enumerate(MIXED_ROWS.items()): + params = rollout(max_tokens=6, seed=20 + i, exit_threshold=1.0, **overrides) + engine.add_request(rid, [1 + i], params if logprobs else replace(params, logprobs=None)) + return final(stream(engine)) + + grouped, quiet = run(logprobs=True), run(logprobs=False) + assert max(batch_sizes) == len(MIXED_ROWS) # Every path shared one CODA batch. + for rid, output in quiet.items(): # Logprobs change no token, depth or RNG draw. + assert grouped[rid].token_ids == output.token_ids + assert grouped[rid].exit_depths == output.exit_depths + assert (grouped[rid].logprobs is None) == (rid == "quiet") + + +@pytest.mark.parametrize("layout", ["last_exited", "shared"]) +def test_full_depth_raw_logprobs_vs_teacher_forced_reference(layout): + m = model() + engine = make_engine(m, cache_config=CacheConfig(128, 2, layout)) + prompts = {"a": [2, 3, 5], "s": [7]} + engine.add_request("a", prompts["a"], rollout(max_tokens=6)) + engine.add_request("s", prompts["s"], rollout(max_tokens=6, temperature=0.8, seed=3)) + for rid, output in final(stream(engine)).items(): + assert output.exit_depths == [4] * 6 + tokens = prompts[rid] + output.token_ids + rows = dense_reference(m, tokens, 4)[-1][2][len(prompts[rid]) - 1 : -1] + expected = rows.float().log_softmax(-1)[torch.arange(6), torch.tensor(output.token_ids)] + difference = (torch.tensor(output.logprobs) - expected).abs().max().item() + # Shared KV reads earlier tokens' final-loop KV at every depth: no plain forward matches. + assert difference < 2e-5 if layout == "last_exited" else difference > 1e-2 + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +def test_seed_none_is_resolved_at_admission_and_replays_exactly(asynchronous): + params = rollout(max_tokens=8, temperature=1.0, seed=None) + engine = make_engine(model(), asynchronous=asynchronous) + torch_state, python_state = torch.get_rng_state(), random.getstate() + for rid in ("x", "y", "unsampled"): + engine.add_request(rid, [2, 3], params) + # Seeds come from OS entropy, not from the torch or Python global RNG. + assert torch.equal(torch.get_rng_state(), torch_state) and random.getstate() == python_state + seeds = {rid: r.sampling_params.seed for rid, r in engine.scheduler.requests.items()} + assert all(type(seed) is int and 0 <= seed < 2**63 for seed in seeds.values()) + # Every output, even an abort before the first sample, reports the effective params. + outputs = {"unsampled": engine.abort_request("unsampled")} + assert engine.rng.peek("unsampled") is None # Aborted before its first draw. + sampled = [engine.scheduler.requests[rid] for rid in ("x", "y")] + streamed = stream(engine, lambda: all(r.generated_token_ids for r in sampled)) + # The engine's registry seeded each generator with exactly the seed outputs report. + assert all(engine.rng.peek(r).initial_seed() == seeds[r] for r in ("x", "y")) + outputs |= final(stream(engine, streamed=streamed)) + assert engine.rng.peek("x") is None and engine.rng.peek("y") is None # Dropped at finish. + for rid, output in outputs.items(): + check_aligned(output) + assert output.sampling_params == replace(params, seed=seeds[rid]) + assert len(set(seeds.values())) == 3 and outputs["x"].token_ids != outputs["y"].token_ids + # Replaying the reported params reproduces the rollout exactly. + replay = make_engine(model(), asynchronous=asynchronous) + for rid in ("x", "y"): + replay.add_request(rid, [2, 3], outputs[rid].sampling_params) + for rid, output in final(stream(replay)).items(): + assert_same(output, outputs[rid]) + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +def test_greedy_seed_none_is_resolved_at_admission(asynchronous): + # Greedy Ouro rows never reach the RNG registry, so admission alone resolves their seed + # (Huginn's greedy rows do reach it: see the Huginn test below). + engine = make_engine(model(), asynchronous=asynchronous) + engine.rng.acquire = Mock(side_effect=AssertionError("greedy Ouro rows draw no RNG")) + engine.add_request("g", [2, 3], rollout(max_tokens=3, seed=None)) + seed = engine.scheduler.requests["g"].sampling_params.seed + assert type(seed) is int and 0 <= seed < 2**63 + (output,) = final(stream(engine)).values() + assert output.sampling_params == rollout(max_tokens=3, seed=seed) + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +@pytest.mark.parametrize("temperature", [0.0, 0.8], ids=["greedy", "top_k-top_p"]) +def test_huginn_reports_the_seed_of_its_recurrent_state_generator(temperature, logprobs_mode): + # Huginn draws its random recurrent-state initialization from the request's generator + # at prefill and in every PRELUDE, greedy requests included, so its generator is + # created before any CODA sample. That generator must still carry the reported seed. + torch.manual_seed(42) + m = HuginnForCausalLM(tiny_huginn_config()).eval() + prompts = {"a": [4, 7, 3], "b": [9, 2, 5, 6]} + scheduler = SchedulerConfig(max_num_seqs=3, max_num_batched_tokens=8, prefill_chunk_size=3) + options = dict(cache_config=CacheConfig(128, 2), scheduler_config=scheduler) + + def run(params): + engine = LLMEngine(m, logprobs_mode=logprobs_mode, **options) + acquire, created = engine.rng.acquire, {} + + def recording_acquire(request_id, request_params, device): + generator = acquire(request_id, request_params, device) + created.setdefault(request_id, generator.initial_seed()) + return generator + + engine.rng.acquire = recording_acquire + for rid, prompt in prompts.items(): + engine.add_request(rid, prompt, params[rid]) + return final(stream(engine)), created + + sampling = dict(temperature=temperature, top_k=5, top_p=0.9) if temperature else {} + params = rollout(max_tokens=5, min_loops=3, max_loops=3, seed=None, **sampling) + outputs, created = run(dict.fromkeys(prompts, params)) + assert created == {rid: output.sampling_params.seed for rid, output in outputs.items()} + # Requesting logprobs changes no token, and replaying the reported params is bitwise. + quiet, _ = run({rid: replace(o.sampling_params, logprobs=None) for rid, o in outputs.items()}) + replay, _ = run({rid: output.sampling_params for rid, output in outputs.items()}) + for rid, output in outputs.items(): + assert quiet[rid].token_ids == output.token_ids + assert quiet[rid].exit_depths == output.exit_depths + assert_same(replay[rid], output) + + +PREEMPTED_PHASES = [ + pytest.param(False, RequestPhase.CODA, id="sync-coda"), + pytest.param(False, RequestPhase.PRELUDE, id="sync-prelude"), + pytest.param(False, RequestPhase.RECURRENT, id="sync-recurrent"), + # Asynchronous preemption skips requests with an undelivered output. With a single + # request, one waiting in PRELUDE still has its previous output undelivered, so this + # construction reaches no async PRELUDE victim; the pressure test below does. + pytest.param(True, RequestPhase.CODA, id="async-coda"), + pytest.param(True, RequestPhase.RECURRENT, id="async-recurrent"), +] + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +@pytest.mark.parametrize("asynchronous,phase", PREEMPTED_PHASES) +def test_preemption_at_each_phase_keeps_the_rng_slot(asynchronous, phase, logprobs_mode): + # Top-k/top-p rows: processed mode samples them through _sample_with_logprob. + params = rollout(max_tokens=7, temperature=0.8, top_k=3, top_p=0.9, seed=45, exit_threshold=0) + options = dict(asynchronous=asynchronous, logprobs_mode=logprobs_mode) + scheduler = SchedulerConfig(enable_preemption=True) + options.update(cache_config=CacheConfig(64, 2), scheduler_config=scheduler) + results = [] + for preempt in (False, True): + engine = make_engine(model(), **options) + engine.add_request("a", [1, 2, 3], params) + request = engine.scheduler.requests["a"] + if preempt: + stream( + engine, + lambda: ( + len(request.generated_token_ids) >= 2 + and request.phase == phase + and not request.num_output_placeholders + ), + ) + generator = engine.rng.peek("a") + state = generator.get_state().clone() + engine.add_request("b", [5], SamplingParams(max_tokens=1)) + assert engine.scheduler.preempt(engine.scheduler.requests["b"]) + # Suspension keeps the request's generator, at the same point of its stream. + assert engine.rng.peek("a") is generator + assert torch.equal(generator.get_state(), state) + results.append(final(stream(engine))["a"]) + assert engine.preemption.resumptions == int(preempt) + # A CODA victim resumes into "b"'s CODA batch, so its LM head runs two rows instead of + # one. CPU GEMMs are not batch invariant, so its logprobs match numerically (up to ~5e-7). + # PRELUDE and RECURRENT victims share no batch with "b": theirs stay bitwise. + assert_same(results[1], results[0], atol=1e-6 if phase == RequestPhase.CODA else None) + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +def test_async_prelude_preemption_keeps_the_rng_slot(logprobs_mode): + """Memory pressure preempts asynchronous requests waiting in PRELUDE with no undelivered + output (the construction of tests/test_prefix_growth.py's GPU test for #109). On CPU no + AsyncState exists, so this covers the RNG slot and logprobs, not #109's device token.""" + trace = ExitConfig("trace", depths_by_request={"t": [4, 1, 2, 3, 4, 1, 2, 3]}) + params = rollout(max_tokens=8, min_loops=1, temperature=0.9, top_k=5, top_p=0.95) + scheduler = SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=2, prefill_chunk_size=2, enable_preemption=True + ) + options = dict(asynchronous=True, exit_config=trace, logprobs_mode=logprobs_mode) + results, victims = [], [] + for blocks in (128, 24): + cache = CacheConfig(blocks, 2, incremental_allocation=True) + engine = make_engine(model(72), cache_config=cache, scheduler_config=scheduler, **options) + runner, suspended = engine.model_runner, {} + suspend, resume = runner.suspend, engine.scheduler.resume_callback + + def checked_suspend(request_id, engine=engine, suspend=suspend, suspended=suspended): + request, generator = engine.scheduler.requests[request_id], engine.rng.peek(request_id) + victims.append((request.phase, generator is not None)) + state = None if generator is None else generator.get_state().clone() + suspended[request_id] = generator, state + suspend(request_id) + + def checked_resume(request, engine=engine, resume=resume, suspended=suspended): + if request.request_id in suspended: # Its stream is untouched while suspended. + generator, state = suspended[request.request_id] + assert engine.rng.peek(request.request_id) is generator + assert generator is None or torch.equal(generator.get_state(), state) + return resume(request) + + runner.suspend, engine.scheduler.resume_callback = checked_suspend, checked_resume + for i in range(4): + prompt = [2 * i + 1, 2 * i + 2] + engine.add_request(str(i), prompt, replace(params, seed=10 + i), trace_id="t") + results.append(final(stream(engine))) + assert (engine.preemption.preemptions > 0) == (blocks == 24) + assert engine.preemption.resumptions == engine.preemption.preemptions + assert (RequestPhase.PRELUDE, True) in victims # A PRELUDE victim that had already sampled. + for rid, output in results[0].items(): + # Pressure changes the batch composition, so logprobs match numerically. + assert_same(results[1][rid], output, atol=1e-6) + + +def test_preemption_and_prefix_hits_keep_logprobs(): + params = rollout(max_tokens=7, temperature=0.8, seed=45, exit_threshold=0) + results = [] + for suspend in (False, True): + scheduler = SchedulerConfig(enable_preemption=True) + engine = make_engine(model(), cache_config=CacheConfig(64, 2), scheduler_config=scheduler) + engine.add_request("a", [1, 2, 3], params) + request = engine.scheduler.requests["a"] + stream(engine, lambda: len(request.generated_token_ids) >= 3) + if suspend: + before = list(request.logprobs) + engine.add_request("b", [5], SamplingParams(max_tokens=1)) + assert engine.scheduler.preempt(engine.scheduler.requests["b"]) + assert request.logprobs == before and request.phase == RequestPhase.WAITING + results.append(final(stream(engine))["a"]) + assert engine.preemption.resumptions == int(suspend) + assert_same(results[1], results[0]) + # Memory-pressure preemption matches an unpressured run. + params, results = rollout(max_tokens=8, exit_threshold=0), [] + for blocks, preempt in ((128, False), (24, True)): + cache = CacheConfig(blocks, 2, incremental_allocation=True) + scheduler = SchedulerConfig(max_num_seqs=3, prefill_chunk_size=2, enable_preemption=preempt) + engine = make_engine(model(15), cache_config=cache, scheduler_config=scheduler) + for i in range(3): + engine.add_request(str(i), [1, 2], params) + results.append(final(stream(engine))) + assert (engine.preemption.preemptions > 0) == preempt + for rid, output in results[0].items(): + assert_same(results[1][rid], output, atol=1e-6) + # A prefix-cache hit matches the cold run. + cache = CacheConfig(64, 2, enable_prefix_caching=True) + scheduler = SchedulerConfig(prefill_chunk_size=2) + engine = make_engine(model(15), cache_config=cache, scheduler_config=scheduler) + engine.add_request("cold", [1, 2, 3, 4, 5], replace(params, max_tokens=4)) + cold = final(stream(engine))["cold"] + engine.add_request("warm", [1, 2, 3, 4, 5], replace(params, max_tokens=4)) + engine.step() + assert engine.last_schedule.requests[0].token_start == 4 + assert_same(final(stream(engine))["warm"], cold, atol=1e-6) + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +@pytest.mark.parametrize( + "old,new", [(0, None), (None, 0)], ids=["logprobs-then-none", "none-then-logprobs"] +) +@pytest.mark.parametrize("exit_config,threshold", ASYNC_EXIT_MODES) +def test_lagging_delivery_abort_and_id_reuse_keep_logprobs( + monkeypatch, exit_config, threshold, old, new, logprobs_mode +): + sampling = dict(temperature=0.9, top_p=0.95, seed=7, min_loops=1, exit_threshold=threshold) + params = rollout(max_tokens=6, **sampling) # T != 1 and top-p: processed differs from raw. + options = dict(exit_config=exit_config, logprobs_mode=logprobs_mode) + + def fresh(*, asynchronous): + engine = make_engine(model(), asynchronous=asynchronous, **options) + engine.add_request("a", [5, 6], replace(params, logprobs=new)) + return final(stream(engine))["a"] + + expected = fresh(asynchronous=True) + if exit_config.mode == "trace": # Asynchronous delivery matches the synchronous path. + assert_same(fresh(asynchronous=False), expected, atol=1e-6) + hold_coda_delivery(monkeypatch) + engine = make_engine(model(), asynchronous=True, **options) + slots = lease_cpu_logprob_slots(monkeypatch, engine) + engine.add_request("a", [2, 3], replace(params, logprobs=old)) + engine.add_request("b", [4], replace(params, seed=5)) + request = engine.scheduler.requests["a"] + # Step until the request has delivered a token and has one undelivered CODA output. + stream(engine, lambda: request.generated_token_ids and request.num_output_placeholders == 1) + delivered = len(request.generated_token_ids) + aborted = engine.abort_request("a") + check_aligned(aborted) + # The undelivered token is not part of the abort output. + assert aborted.finish_reason == "abort" and len(aborted.token_ids) == delivered + # Reuse the ID with the opposite setting while the stale ticket (and lease) is pending. + assert engine._pending_coda and (old is None or any(slot.leased for slot in slots)) + engine.add_request("a", [5, 6], replace(params, logprobs=new)) + assert_same(final(stream(engine))["a"], expected, atol=1e-6) + assert not engine._pending_coda and not any(slot.leased for slot in slots) + + +@pytest.mark.parametrize( + "finish,max_tokens,count", + [("round", 9, 8), ("stop_token_ids", 9, 4), ("eos", 9, 5), ("length", 4, 3)], + ids=["length-at-round-end", "stop-id-mid-round", "eos-mid-round", "length-mid-round"], +) +def test_speculative_update_slices_logprobs(monkeypatch, finish, max_tokens, count): + engine = LLMEngine(model(), speculative_config=SpeculativeConfig(3)) + stops = [50] if finish == "stop_token_ids" else [] + params = SamplingParams(max_tokens=max_tokens, ignore_eos=finish != "eos", stop_token_ids=stops) + engine.add_request("r", [2, 3], replace(params, logprobs=0)) + # Each round accepts all but its last draft token; token t has logprob -1 - t / 100. + rounds = ([10, 11], [12, 50, 0, 13], [14, 15], [16, 17]) + results = deque(SpeculativeResult(t, len(t) - 1, [-1 - i / 100 for i in t]) for t in rounds) + runner = engine.speculative_runner + monkeypatch.setattr( + runner, "_collect_round", lambda s: [results.popleft() for _ in s.batch.requests] + ) + result = stream(engine)["r"][-1] + emitted = [10, 11, 12, 50, 0, 13, 14, 15][:count] + reason = "length" if finish in ("round", "length") else "stop" + assert (result.token_ids[1:], result.finish_reason) == (emitted, reason) + assert result.exit_depths == [4] * len(result.token_ids) + assert result.logprobs[1:] == [-1 - t / 100 for t in emitted] + assert math.isfinite(result.logprobs[0]) + assert runner.stats.committed_tokens == len(emitted) + + +SPECULATIVE_SAMPLING = { # Greedy, truncated and unfiltered sampling. + "greedy": {}, + "top_k_top_p": dict(temperature=0.8, top_k=7, top_p=0.8, seed=42), + "unfiltered": dict(temperature=1.0, seed=5), +} + + +def speculative_engine(m, logprobs_mode="raw_logprobs", **kwargs): + return LLMEngine( + m, + speculative_config=SpeculativeConfig(3), + cache_config=CacheConfig(256, 2), + logprobs_mode=logprobs_mode, + **kwargs, + ) + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +def test_speculative_logprobs_match_full_depth_reference(monkeypatch, logprobs_mode): + """Every committed token (accepted, correction, bonus) is scored under the target p on + its verification row, which matches a teacher-forced forward at the model depth + numerically. The mixed batch puts greedy (raw-scored), sampled (processed-scored in + processed mode) and non-requesting rows in the same rounds.""" + m = model() + prompts = [[2, 3], [7, 8, 9], [5], [4, 6]] + runs = {name: [rollout(max_tokens=12, **o)] * 4 for name, o in SPECULATIVE_SAMPLING.items()} + runs["mixed"] = [rollout(max_tokens=12, **o) for o in SPECULATIVE_SAMPLING.values()] + runs["mixed"].append(SamplingParams(max_tokens=12, ignore_eos=True)) + kinds = set() # (sampled, accepted every draft) of each committed round. + mixed_rounds = 0 # Rounds scoring both greedy and sampled requests. + for name, params in runs.items(): + engine = speculative_engine(m, logprobs_mode) + runner = engine.speculative_runner + + def recording(submission, collect=runner._collect_round): + nonlocal mixed_rounds + results = collect(submission) + greedy = set() + for scheduled_request, result in zip(submission.batch.requests, results): + request_params = scheduled_request.request.sampling_params + if request_params.logprobs is None: + assert result.logprobs is None + continue + assert len(result.logprobs) == len(result.token_ids) + sampled = request_params.temperature != 0 + drafts = scheduled_request.token_count - 1 + kinds.add((sampled, result.accepted_count == drafts)) + greedy.add(not sampled) + mixed_rounds += len(greedy) == 2 + return results + + monkeypatch.setattr(runner, "_collect_round", recording) + outputs = generate(engine, prompts, params) + for prompt, output in zip(prompts, outputs): + check_aligned(output) + assert output.exit_depths == [4] * len(output.token_ids) + if output.logprobs is None: + continue + tokens = prompt + output.token_ids + rows = dense_reference(m, tokens, 4)[-1][2][len(prompt) - 1 : -1] + expected = [ + token_logprob(row, output.sampling_params, torch.tensor(token), logprobs_mode) + for row, token in zip(rows, output.token_ids) + ] + assert output.logprobs == pytest.approx(torch.stack(expected).tolist(), abs=2e-5) + if name == "greedy": # Greedy speculation equals the non-speculative engine. + native = generate(LLMEngine(m, logprobs_mode=logprobs_mode), prompts, params) + for actual, expected in zip(outputs, native): + assert_same(actual, expected, atol=2e-5) + # Both the correction (a rejected draft) and the bonus path were scored when sampling. + assert {(True, False), (True, True)} <= kinds + assert mixed_rounds > 0 + + +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +def test_speculative_logprobs_change_no_token_and_replay_exactly(logprobs_mode): + m = model() + prompts = [[2, 3], [7, 8, 9]] + params = rollout(max_tokens=10, **SPECULATIVE_SAMPLING["top_k_top_p"]) + + def run(params): + return generate(speculative_engine(m, logprobs_mode), prompts, params) + + first, second, quiet = run(params), run(params), run(replace(params, logprobs=None)) + for a, b, q in zip(first, second, quiet): + assert_same(a, b) # Same seed: bitwise tokens and logprobs. + # Scoring consumes no RNG, so the tokens do not depend on logprobs. + assert (a.token_ids, a.exit_depths) == (q.token_ids, q.exit_depths) + assert q.logprobs is None + + +def test_speculative_rows_without_logprobs_skip_scoring(monkeypatch): + """Without a requesting row nothing is scored; a mixed round scores only its requesting + request's committed tokens.""" + engine = speculative_engine(model()) + runner, resolve, raw = engine.speculative_runner, _CommittedScores.resolve, [] + collect, rounds = runner._collect_round, [] # Per round: IDs of the scored requests. + + def recording(submission): + rounds.append(None) + results = collect(submission) + indices = rounds[-1] + rounds[-1] = {submission.batch.requests[i].request.request_id for i in indices} + return results + + def resolving(scores, logits): + rounds[-1] = scores.raw_spans.keys() | scores.processed_spans.keys() + return resolve(scores, logits) + + monkeypatch.setattr(runner, "_collect_round", recording) + monkeypatch.setattr(_CommittedScores, "resolve", resolving) + scoring_op = Sampler.raw_logprobs + monkeypatch.setattr( + Sampler, "raw_logprobs", staticmethod(lambda *a: raw.append(1) or scoring_op(*a)) + ) + params = SamplingParams(max_tokens=8, ignore_eos=True) + assert generate(engine, [[2, 3]], params)[0].logprobs is None + assert rounds and not any(rounds) and not raw # No round scored or ran a scoring op. + rounds.clear() + quiet, loud = generate(engine, [[2, 3], [7, 8]], [params, replace(params, logprobs=0)]) + assert quiet.logprobs is None and len(loud.logprobs) == len(loud.token_ids) + assert set().union(*rounds) == {loud.request_id} + + +@pytest.mark.parametrize( + "finish,max_tokens,emitted", + [ + ("round", 9, [10, 50, 0, 13]), + ("stop_token_ids", 9, [10, 50]), + ("eos", 9, [10, 50, 0]), + ("length", 3, [10, 50]), + ], + ids=["continue", "stop-id-mid-round", "eos-mid-round", "length-at-round-end"], +) +def test_async_speculative_collect_uses_shared_finish_rules(finish, max_tokens, emitted): + """Deliver one async speculative round on CPU. The runner needs CUDA, so a stub ticket + stands in for its readback; _collect_speculative itself is the code under test.""" + engine = LLMEngine(model(), speculative_config=SpeculativeConfig(3)) + stops = [50] if finish == "stop_token_ids" else [] + params = SamplingParams(max_tokens=max_tokens, ignore_eos=finish != "eos", stop_token_ids=stops) + engine.add_request("r", [2, 3], params) + request = engine.scheduler.requests["r"] + stream(engine, lambda: request.generated_token_ids) # The bootstrap CODA token. + batch = engine.scheduler.schedule() + [scheduled_request] = batch.requests + # Apply _step_async_speculative's submission bookkeeping, then deliver the round. + engine.scheduler.update_after_submission(batch) + assert request.num_output_placeholders == scheduled_request.token_count + tokens = [10, 50, 0, 13][: scheduled_request.token_count] + frontier = len(request.prompt_token_ids) + len(request.generated_token_ids) - 1 + len(tokens) + result = ModelRunnerOutput( + speculative=[SpeculativeResult(tokens, len(tokens) - 1)], frontiers=[frontier] + ) + ticket = SimpleNamespace( + batch=batch, + states=(SimpleNamespace(delivered=0),), + rounds=(0,), + ready=lambda: True, + collect=lambda: result, + retire=lambda: None, + ) + stats = engine.speculative_runner.stats + engine.speculative_runner = SimpleNamespace( + stats=stats, + readback_waits=0, + discarded_rounds=0, + request_ready=lambda r: True, + release=lambda r: None, + ) + engine.async_speculative, engine._pending_speculative = True, [ticket] + [output] = engine._collect_speculative() + reason = {"round": None, "length": "length"}.get(finish, "stop") + assert (output.token_ids[1:], output.finish_reason) == (emitted, reason) + assert output.exit_depths == [4] * len(output.token_ids) and output.logprobs is None + assert stats.committed_tokens == len(emitted) and request.num_output_placeholders == 0 + if reason is None: + assert list(engine.scheduler.queues[RequestPhase.SPECULATIVE]) == ["r"] + else: + assert not engine.has_unfinished_requests() + assert engine.cache_manager.num_used_blocks == 0 + + +def test_async_speculative_ticket_scores_only_requesting_rows(): + """The async runner needs CUDA, so stub buffers drive SpeculativeRoundTicket.collect: in a + scored round each request that asked gets its committed rows' scores, the others None.""" + + class Done: + def query(self): + return True + + def synchronize(self): + pass + + def collect(logprobs_copied, wants): + # Two requests with 3 and 2 verification rows; 1 and 0 drafts accepted. + buffer = SimpleNamespace( + host_result=torch.tensor([7, 8, 9, 4, 5, 1, 0, 12, 20]), + host_logprobs=torch.tensor([-0.1, -0.2, -0.3, -0.4, -0.5]), + copy_done=Done(), + ) + params = [SamplingParams(logprobs=0 if w else None) for w in wants] + requests = [ + SimpleNamespace(request=SimpleNamespace(sampling_params=p), token_count=n) + for p, n in zip(params, (3, 2)) + ] + batch = SimpleNamespace(requests=requests) + ticket = SpeculativeRoundTicket( + batch, buffer, (), (0, 0), 5, logprobs_copied=logprobs_copied + ) + output = ticket.collect() + return output.speculative, output.frontiers + + results, frontiers = collect(True, (True, False)) + assert [r.token_ids for r in results] == [[7, 8], [4]] and frontiers == [12, 20] + assert results[0].logprobs == pytest.approx([-0.1, -0.2]) and results[1].logprobs is None + results, _ = collect(True, (False, True)) + assert results[0].logprobs is None and results[1].logprobs == pytest.approx([-0.4]) + results, _ = collect(False, (True, True)) + assert [r.logprobs for r in results] == [None, None] + + +def test_pd_coordinator_resolves_seeds_and_mirrors_logprobs(): + engine = object.__new__(PDEngine) # Coordinator only: no worker processes. + engine.closed, engine.failure = False, None + engine.requests, engine.transfers, engine.outputs = {}, {}, deque() + engine.config, engine.exit_config = PDConfig(), ExitConfig("ouro_delayed") + engine.speculative_config, engine._async_speculative = None, False + engine.model_config = tiny_ouro_config() + info, roles = dict(block_size=4, depths=4, num_blocks=1024), ("prefill", "decode") + engine.peers = { + role: Peer(role, role, SimpleNamespace(info=info, send=Mock())) for role in roles + } + with pytest.raises(ValueError, match="stop_token_ids"): + engine.add_request("bad", [1], SamplingParams(stop_token_ids=[64])) + # P and D get these effective params (one resolved seed), which an early abort reports. + for rid in ("unsampled", "r"): + engine.add_request(rid, [1, 2], rollout(max_tokens=8, temperature=1.0, seed=None)) + unsampled = engine.requests["unsampled"].sampling_params + assert engine.abort_request("unsampled").sampling_params == unsampled + request = engine.requests["r"] + assert type(unsampled.seed) is int and type(request.sampling_params.seed) is int + assert request.logprobs == [] + w = next(iter(engine.transfers.values())) + w.p, w.d, w.phase = "prefill", "decode", "decode" + for n in (1, 2): + mirrored = dict(logprobs=[-0.5] * n, sampling_params=request.sampling_params) + output = RequestOutput(w.tid, [1, 2], [5] * n, [4] * n, False, **mirrored) + message = dict(kind="output", request_id=w.tid, output=output) + engine._message(engine.peers["decode"], message) + assert [o.request_id for o in engine.outputs] == ["r", "r"] + assert engine.outputs[-1].logprobs == [-0.5, -0.5] + aborted = engine.abort_request("r") + assert aborted.finish_reason == "abort" and aborted.sampling_params == request.sampling_params + assert (aborted.token_ids, aborted.exit_depths) == ([5, 5], [4, 4]) + assert aborted.logprobs == [-0.5, -0.5] + + +def gpu_engine( + m, *, asynchronous, exit_config=None, seqs=3, logprobs_mode="raw_logprobs", **execution +): + scheduler = SchedulerConfig(max_num_seqs=seqs, max_num_batched_tokens=8, prefill_chunk_size=4) + options = dict(cache_config=CacheConfig(128, 16), scheduler_config=scheduler) + options |= dict(logprobs_mode=logprobs_mode, attention_backend="triton", execution=execution) + exit_config = exit_config or ExitConfig("ouro_delayed") + return make_engine(m, asynchronous=asynchronous, exit_config=exit_config, **options) + + +def gpu_params(rid, **overrides): + values = dict(max_tokens=6 + int(rid), temperature=0.8, top_k=16, seed=40 + int(rid)) + return rollout(**(values | dict(min_loops=2, exit_threshold=0.2) | overrides)) + + +@pytest.mark.gpu +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +def test_gpu_cuda_graph_logprobs_match_eager(asynchronous, logprobs_mode): + # Static eager pads RECURRENT rows to a power of two while graph replay runs len(rows), + # so GEMM reductions may round; FP32 keeps tokens and exits exact and bounds logprobs. + m = model(321, "cuda", torch.float32, head_dim=64) + results = [] + for graphs in (False, True): + # One stream (no prefer_recurrent overlap): batches cannot depend on graph timing. + static = dict(static_buffers=True, pad_to_power_of_two=True, multi_stream=False) + options = dict(asynchronous=asynchronous, logprobs_mode=logprobs_mode, cuda_graphs=graphs) + engine = gpu_engine(m, **options, **static) + for i in range(3): + engine.add_request(str(i), [2 + i, 3, 4] * (i + 1), gpu_params(str(i))) + results.append(final(stream(engine))) + assert not graphs or engine.model_runner.graphs.replays > 0 + eager, graphed = results + assert graphed.keys() == eager.keys() + for rid, output in eager.items(): + assert_same(graphed[rid], output, atol=1e-4) + + +@pytest.mark.gpu +@pytest.mark.parametrize("multi_stream", [False, True], ids=["one-stream", "multi-stream"]) +def test_gpu_async_readback_matches_sync_and_returns_leases(monkeypatch, multi_stream): + # FP32 keeps sampled tokens stable across sync/async batch compositions. + m = model(321, "cuda", torch.float32, head_dim=64) + traces = {str(i): [4, 1, 3, 2, 4, 1, 2, 3, 4, 2, 1, 3] for i in range(4)} + options = dict(exit_config=ExitConfig("trace", depths_by_request=traces), seqs=4) + + def run(*, asynchronous): + engine = gpu_engine(m, asynchronous=asynchronous, multi_stream=multi_stream, **options) + for i in range(4): + params = gpu_params(str(i), max_tokens=12, min_loops=1, exit_threshold=1.0) + engine.add_request(str(i), [2, 3 + i], params) + return engine, final(stream(engine)) + + _, expected = run(asynchronous=False) + # Delivery lags to the one-output bound: logprob leases overlap live exit-score tickets. + hold_coda_delivery(monkeypatch) + engine, actual = run(asynchronous=True) + assert not any(slot.leased for slot in engine.model_runner.logprob_slots) + for rid, output in expected.items(): + assert_same(actual[rid], output, atol=1e-4) + + +@pytest.mark.gpu +def test_gpu_delayed_exit_scores_and_logprobs_do_not_exhaust_readback_pools(monkeypatch): + engine = gpu_engine(model(5, "cuda", torch.bfloat16, head_dim=64), asynchronous=True, seqs=4) + assert len(engine.model_runner.logprob_slots) == 4 + 4 + # Lagging delivery holds a logprob lease per request while delayed exit scores hold theirs. + hold_coda_delivery(monkeypatch) + for round_index in range(3): + for i in range(4): + engine.add_request(f"{round_index}-{i}", [2, 3, i], gpu_params(str(i), max_tokens=16)) + assert len(final(stream(engine))) == 4 + assert not any(slot.leased for slot in engine.model_runner.logprob_slots) + + +@pytest.mark.gpu +def test_gpu_abort_with_pending_coda_returns_lease_and_reuses_id(monkeypatch): + engine = gpu_engine(model(5, "cuda", torch.bfloat16, head_dim=64), asynchronous=True) + slots = engine.model_runner.logprob_slots + ready = hold_coda_delivery(monkeypatch) + for i in range(2): + engine.add_request(str(i), [2, 3, i], gpu_params(str(i), max_tokens=12)) + request = engine.scheduler.requests["0"] + # Step until the request has delivered a token and has one undelivered CODA output. + stream(engine, lambda: request.generated_token_ids and request.num_output_placeholders == 1) + delivered = len(request.generated_token_ids) + aborted = engine.abort_request("0") + check_aligned(aborted) + assert len(aborted.token_ids) == delivered and any(slot.leased for slot in slots) + stale = [t for t in engine._pending_coda if any(i.request is request for i in t.batch.requests)] + assert stale + # release() synchronized the aborted request, so the next step collects its ticket. + monkeypatch.setattr(Submission, "ready", ready) + engine.add_request("0", [2, 3], gpu_params("0", logprobs=None)) + for output in engine.step(): + check_aligned(output) + assert not any(t is s for t in engine._pending_coda for s in stale) + outputs = final(stream(engine)) + assert outputs["0"].logprobs is None and len(outputs["0"].token_ids) == 6 + assert len(outputs["1"].logprobs) == 12 + assert not engine._pending_coda and not any(slot.leased for slot in slots) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_execution_error_returns_undelivered_logprob_leases(monkeypatch, device): + cuda = device == "cuda" + m = model(5, "cuda", torch.bfloat16, head_dim=64) if cuda else model(5) + engine = gpu_engine(m, asynchronous=True, seqs=4) if cuda else make_engine(m, asynchronous=True) + slots = lease_cpu_logprob_slots(monkeypatch, engine) + runner = engine.model_runner + ready, execute = hold_coda_delivery(monkeypatch), runner._execute + held = [] + + def fail(batch, prepared): + result = execute(batch, prepared) + if prepared.batch.operation == Operation.RECURRENT and engine._pending_coda: + held.append(sum(slot.leased for slot in slots)) + raise RuntimeError("injected failure after CODA submission") + return result + + monkeypatch.setattr(runner, "_execute", fail) + for i in range(3): + engine.add_request(str(i), [2, 3, i], gpu_params(str(i), max_tokens=8)) + with pytest.raises(RuntimeError, match="injected"): + stream(engine) + # The failure path drops undelivered tickets without collect(); __del__ returns leases. + assert held and held[0] > 0 and not engine._pending_coda + gc.collect() + assert not any(slot.leased for slot in slots) + monkeypatch.setattr(runner, "_execute", execute) + monkeypatch.setattr(Submission, "ready", ready) + engine.add_request("3", [2, 3], gpu_params("3", max_tokens=4)) + check_aligned(final(stream(engine))["3"]) + + +@pytest.mark.gpu +def test_gpu_async_speculative_stop_token_ids_match_sync(): + """Cross-round async speculation (CUDA, Triton, greedy) stops on stop IDs like sync.""" + m = model(device="cuda") + prompts = {"a": [2, 3, 4, 5], "b": [9], "c": [2, 3, 4, 7]} + + def run(asynchronous, **overrides): + engine = LLMEngine( + m, + attention_backend="triton", + cache_config=CacheConfig(256, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=2 + ), + execution_config=ExecutionConfig(async_scheduling=asynchronous), + speculative_config=SpeculativeConfig(3), + ) + for rid, prompt in prompts.items(): + params = SamplingParams(max_tokens=16, ignore_eos=True, **overrides) + engine.add_request(rid, prompt, params) + outputs = final(stream(engine)) + engine.close() + return outputs + + free = run(False) + # Stop on the latest first occurrence after the bootstrap token, so a round is cut. + first = { + (rid, i): t + for rid, o in free.items() + for i, t in enumerate(o.token_ids) + if i > 0 and t not in o.token_ids[:i] + } + assert first, "greedy outputs repeat their bootstrap token; pick other prompts" + stop = first[max(first, key=lambda key: key[1])] + expected, actual = run(False, stop_token_ids=[stop]), run(True, stop_token_ids=[stop]) + assert any(o.finish_reason == "stop" for o in expected.values()) + for rid, output in expected.items(): + tokens = free[rid].token_ids + cut = tokens[: tokens.index(stop) + 1] if stop in tokens else tokens + assert output.token_ids == cut + assert (actual[rid].token_ids, actual[rid].finish_reason) == ( + output.token_ids, + output.finish_reason, + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("logprobs_mode", LOGPROBS_MODES) +def test_gpu_speculative_logprobs_match_across_graphs_async_and_native(logprobs_mode): + """CUDA/Triton speculation scores like the non-speculative engine, with CUDA Graphs and + with cross-round async speculation (greedy); sampled rounds match with graphs on/off.""" + m = model(device="cuda") + prompts = {"a": [2, 3, 4, 5], "b": [9], "c": [2, 3, 4, 7]} + + def run(speculative=True, asynchronous=False, graphs=False, **overrides): + engine = LLMEngine( + m, + attention_backend="triton", + cache_config=CacheConfig(256, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=2 + ), + execution_config=ExecutionConfig(async_scheduling=asynchronous, cuda_graphs=graphs), + speculative_config=SpeculativeConfig(3) if speculative else None, + logprobs_mode=logprobs_mode, + ) + for rid, prompt in prompts.items(): + engine.add_request(rid, prompt, rollout(max_tokens=16, **overrides)) + outputs = final(stream(engine)) + if speculative and graphs: + assert engine.speculative_runner.coda_graphs.captures > 0 + engine.close() + return outputs + + native = run(speculative=False) + for variant in (run(), run(graphs=True), run(asynchronous=True)): + for rid, expected in native.items(): + assert_same(variant[rid], expected, atol=1e-4) + sampled = SPECULATIVE_SAMPLING["top_k_top_p"] + eager, graphed = run(**sampled), run(graphs=True, **sampled) + for rid, expected in eager.items(): + assert_same(graphed[rid], expected, atol=1e-4) + + +@pytest.mark.gpu +def test_gpu_pd_forwards_logprobs_and_abort_stays_aligned(): + pytest.importorskip("nixl") + if torch.cuda.device_count() < 2: + pytest.skip("requires 2 visible GPUs") + exit_config = ExitConfig("trace", depths_by_request={"t": [4, 2, 3, 1] * 16}) + prefill = SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=5, prefill_chunk_size=3) + transfer = dict(transfer_chunk_bytes=4096, max_inflight_bytes=16384) + caches = dict(prefill_cache_config=CacheConfig(128, 4), decode_cache_config=CacheConfig(128, 4)) + with PDEngine( + tiny_ouro_config(head_dim=64), + pd_config=PDConfig(**transfer, request_timeout=90, startup_timeout=90), # GPUs 0, 1. + prefill_scheduler_config=prefill, + decode_scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=2), + exit_config=exit_config, + execution_config=ExecutionConfig(async_scheduling=True), + attention_backend="triton", + seed=123, + **caches, + ) as e: + m = model(123, "cuda:0", torch.bfloat16, head_dim=64) + options = dict(cache_config=CacheConfig(128, 4), scheduler_config=prefill) + reference = make_engine(m, exit_config=exit_config, attention_backend="triton", **options) + params = rollout(max_tokens=6, min_loops=1, temperature=0.7, top_k=32, seed=947) + for target in (e, reference): + target.add_request("r", [1, 2, 3, 4, 5], params, trace_id="t") + assert_same(final(stream(e))["r"], final(stream(reference))["r"], atol=5e-2) + e.add_request("abort", [1, 2, 3, 4, 5], replace(params, max_tokens=64), trace_id="t") + delivered = defaultdict(list) + stream(e, lambda: len(delivered["abort"]) >= 2, delivered) + output = e.abort_request("abort") + check_aligned(output) + assert output.token_ids == delivered["abort"][-1].token_ids + assert output.logprobs == delivered["abort"][-1].logprobs + stream(e) diff --git a/tests/test_sampler.py b/tests/test_sampler.py index a9595f67..58e8261e 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -1,18 +1,24 @@ -"""Contract tests for the extracted Sampler: greedy, top-k and top-p. +"""Contract tests for sampling: the extracted Sampler, per-token logprobs and seeds. ``Sampler`` only consumes a generator the caller resolved, so these tests create one themselves (mirroring what ``RngRegistry.acquire`` does in production) and keep ``RngRegistry``'s own creation and lifetime contracts separate. The distribution tests state the sampling contract itself, including strict top-k: exactly ``k`` candidates survive by index, so a tie at the k-th logit does not -widen the support. +widen the support. The rest cover ``token_logprob`` (``worker/logprobs.py``) in raw +and processed modes, the batched ``processed_logprobs`` and +``Sampler.sample_with_logprob`` against it, and the registry's rejection of an +unresolved ``seed=None``. """ +import math + import pytest import torch import vllm_rlt.worker.sampler as sampler_module from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.worker.logprobs import processed_logprobs, token_logprob, truncates from vllm_rlt.worker.sampler import RngRegistry, Sampler @@ -250,3 +256,185 @@ def test_rng_registry_drop_removes_only_that_requests_slot(): registry.drop("dropped") assert registry.peek("kept") is kept assert dropped is not kept + + +def test_unresolved_seed_is_rejected_before_any_sampling(): + # The registry is the only generator creator, so it rejects an unresolved seed, + # and without a generator no non-greedy token can be drawn. Engines resolve + # seed=None at admission, greedy requests included (tests/test_rollout_outputs.py). + params = SamplingParams(seed=None, temperature=1.0) + registry = RngRegistry() + with pytest.raises(ValueError, match="seed=None"): + registry.acquire("r", params, torch.device("cpu")) + assert registry.peek("r") is None + with pytest.raises(ValueError, match="requires a generator"): + Sampler().sample(torch.zeros(4), params, None) + + +LOGPROB_LOGITS = torch.tensor([1.5, 0.2, -0.7, 2.1, 0.0, -3.0]) + + +@pytest.mark.parametrize("token", range(len(LOGPROB_LOGITS)), ids=lambda t: f"token={t}") +def test_raw_token_logprob_is_the_model_log_softmax_for_any_token(token): + for dtype in (torch.float32, torch.bfloat16): + logits = LOGPROB_LOGITS.to(dtype) + params = SamplingParams(temperature=0.6, top_k=2, top_p=0.5, seed=1) + value = token_logprob(logits, params, torch.tensor(token), "raw_logprobs") + assert value.dim() == 0 and value.dtype == torch.float32 + expected = torch.log_softmax(logits.float(), -1)[token] + assert torch.equal(value, expected) + + +PROCESSED_PARAMS = [ + SamplingParams(temperature=0.6, seed=1), + SamplingParams(temperature=1.3, top_k=3, seed=1), + SamplingParams(temperature=0.9, top_p=0.7, seed=1), + SamplingParams(temperature=0.8, top_k=4, top_p=0.9, seed=1), +] +PROCESSED_IDS = ["temperature", "temperature-top_k", "temperature-top_p", "temperature-top_k-top_p"] + + +@pytest.mark.parametrize("params", PROCESSED_PARAMS, ids=PROCESSED_IDS) +def test_processed_token_logprob_scores_the_sampling_distribution(params): + probs = Sampler.probabilities(LOGPROB_LOGITS, params) + for token in range(len(LOGPROB_LOGITS)): + value = token_logprob(LOGPROB_LOGITS, params, torch.tensor(token), "processed_logprobs") + if probs[token] == 0: + assert value.item() == -math.inf # Filtered tokens cannot be sampled. + else: + assert value.item() == pytest.approx(probs[token].log().item(), abs=1e-6) + + +def test_processed_token_logprob_edge_cases(): + argmax = torch.tensor(int(LOGPROB_LOGITS.argmax())) + greedy = SamplingParams(temperature=0.0, top_k=1) + # Greedy requests report the unscaled model distribution, as vLLM does. + assert torch.equal( + token_logprob(LOGPROB_LOGITS, greedy, argmax, "processed_logprobs"), + token_logprob(LOGPROB_LOGITS, greedy, argmax, "raw_logprobs"), + ) + # Greedy has no temperature-scaled logits to score, so building them is an error. + for build in (Sampler.processed_logits, Sampler.probabilities): + with pytest.raises(ValueError, match="greedy"): + build(LOGPROB_LOGITS, greedy) + # One token survives top_k=1, or top_p=0.6 after top_k=2: top-p sees the renormalized + # top-2 mass (argmax 0.65), not the full distribution (0.53, which would keep two). + for single in (dict(temperature=0.7, top_k=1), dict(temperature=1.0, top_k=2, top_p=0.6)): + params = SamplingParams(seed=1, **single) + assert token_logprob(LOGPROB_LOGITS, params, argmax, "processed_logprobs").item() == 0.0 + + +def test_processed_token_logprob_scores_the_strict_top_k_support(): + # A tie at the k-th logit: strict top-k keeps exactly k tokens, so the processed + # logprob of the surviving tied token is log(e^4 / (e^5 + e^4)), and the other + # tied token cannot be sampled. + logits = torch.tensor([5.0, 4.0, 4.0, 1.0]) + params = SamplingParams(temperature=1.0, top_k=2, seed=1) + probs = Sampler.probabilities(logits, params) + kept = next(index for index in (1, 2) if probs[index] > 0) + dropped = 3 - kept + value = token_logprob(logits, params, torch.tensor(kept), "processed_logprobs") + assert value.item() == pytest.approx(torch.log_softmax(torch.tensor([5.0, 4.0]), -1)[1].item()) + assert token_logprob(logits, params, torch.tensor(dropped), "processed_logprobs") == -math.inf + + +def test_token_logprob_consumes_no_rng(): + params = SamplingParams(temperature=0.8, top_k=4, top_p=0.9, seed=42) + generator = seeded_generator(params) + token = Sampler().sample(LOGPROB_LOGITS, params, generator) + state, global_state = generator.get_state().clone(), torch.get_rng_state() + for logprobs_mode in ("raw_logprobs", "processed_logprobs"): + token_logprob(LOGPROB_LOGITS, params, token, logprobs_mode) + assert torch.equal(generator.get_state(), state) + assert torch.equal(torch.get_rng_state(), global_state) + + +@pytest.mark.parametrize( + "params", [SamplingParams(top_k=1), *PROCESSED_PARAMS], ids=["greedy", *PROCESSED_IDS] +) +def test_sample_with_logprob_draws_like_sample_and_scores_the_drawn_logits(params): + sampler, plain, scored = Sampler(), seeded_generator(params), seeded_generator(params) + initial = plain.get_state().clone() + for _ in range(4): + token = sampler.sample(LOGPROB_LOGITS, params, plain) + same, logprob = sampler.sample_with_logprob(LOGPROB_LOGITS, params, scored) + assert torch.equal(same, token) + expected = token_logprob(LOGPROB_LOGITS, params, token, "processed_logprobs") + assert torch.equal(logprob, expected) + # Same RNG consumption as sample(): none for greedy rows, which ignore the generator. + assert torch.equal(plain.get_state(), scored.get_state()) + assert torch.equal(scored.get_state(), initial) == (params.temperature == 0) + if params.temperature != 0: # Like sample(), never a silent global-RNG fallback. + with pytest.raises(ValueError, match="requires a generator"): + sampler.sample_with_logprob(LOGPROB_LOGITS, params, None) + + +# One row per processed-logprob path: batched raw (greedy, even with top-k, and T == 1), batched +# per temperature, the sampler's top-k/top-p logits, and (last) a row without logprobs. +BATCH_PARAMS = [ + SamplingParams(temperature=0.0, top_k=3, logprobs=0), + SamplingParams(temperature=1.0, logprobs=0, seed=1), + SamplingParams(temperature=0.7, logprobs=0, seed=2), + SamplingParams(temperature=1.3, logprobs=0, seed=3), + SamplingParams(temperature=0.9, top_k=8, logprobs=0, seed=4), + SamplingParams(temperature=1.0, top_p=0.8, logprobs=0, seed=5), + SamplingParams(temperature=0.7, top_k=16, top_p=0.9, logprobs=0, seed=6), + SamplingParams(temperature=0.8, seed=7), +] + + +ROWS = 48 # Rows laid out by a BATCH_LAYOUTS entry; a row without logprobs follows them. + + +def at(temperature): + return SamplingParams(temperature=temperature, logprobs=0, seed=1) + + +# Layout -> (row -> params, log_softmax passes of processed_logprobs or None: only bounded). +BATCH_LAYOUTS = { + "mixed": (lambda row: BATCH_PARAMS[row % 8], None), + "t1": (lambda row: BATCH_PARAMS[1], 1), # The only group: one pass over the whole batch. + "t0.7": (lambda row: BATCH_PARAMS[2], 1), + "top_k-top_p": (lambda row: BATCH_PARAMS[4 + row % 3], 0), + # Many temperatures: one per row; rows r and ROWS - 1 - r sharing one (nested spans, the + # outer ones scored row by row); two alternating (one pass per overlapping span). + "distinct_t": (lambda row: at(0.5 + 0.01 * row), ROWS), + "nested_t": (lambda row: at(0.5 + 0.01 * min(row, ROWS - 1 - row)), None), + "alternating_t": (lambda row: at((0.7, 1.3)[row % 2]), 2), +} + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("layout", BATCH_LAYOUTS.values(), ids=BATCH_LAYOUTS.keys()) +def test_batched_processed_logprobs_equal_per_row_token_logprob(monkeypatch, layout, dtype): + params = [layout[0](row) for row in range(ROWS)] + [BATCH_PARAMS[-1]] + logits = 4 * torch.randn(len(params), 256, generator=torch.Generator().manual_seed(0)) + logits, sampler, tokens, truncated = logits.to(dtype), Sampler(), [], {} + for row, row_params in enumerate(params): # As ModelRunner._sample_with_logprobs. + # A fresh generator per row, as RngRegistry.acquire creates on a first draw. + generator = None if row_params.temperature == 0 else seeded_generator(row_params) + if row_params.logprobs is not None and truncates(row_params): + token, truncated[row] = sampler.sample_with_logprob(logits[row], row_params, generator) + else: + token = sampler.sample(logits[row], row_params, generator) + tokens.append(token) + token_ids, passes = torch.stack(tokens), [] # passes: rows scored by each log_softmax. + + raw_logprobs = Sampler.raw_logprobs + + def counted(logits, token_ids): + passes.append(len(logits)) + return raw_logprobs(logits, token_ids) + + monkeypatch.setattr(Sampler, "raw_logprobs", staticmethod(counted)) + values = processed_logprobs(logits, params, token_ids, truncated) + monkeypatch.undo() + # However many temperatures mix, at most twice the batched rows are scored (57d2308 scored + # each requesting row once, in a pass of its own). + assert sum(passes) <= 2 * (sum(p.logprobs is not None for p in params) - len(truncated)) + assert layout[1] is None or len(passes) == layout[1] + assert values.shape == (len(params),) and values.dtype == torch.float32 + for row, row_params in enumerate(params): + if row_params.logprobs is not None: + expected = token_logprob(logits[row], row_params, token_ids[row], "processed_logprobs") + assert torch.equal(values[row], expected) diff --git a/tests/test_serving.py b/tests/test_serving.py index 223ab911..f28a6ab0 100644 --- a/tests/test_serving.py +++ b/tests/test_serving.py @@ -14,7 +14,7 @@ from aiohttp.test_utils import TestClient, TestServer from tests.helpers import tiny_ouro_config -from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt import LLM, CacheConfig, SamplingParams, SchedulerConfig from vllm_rlt.engine.llm_engine import LLMEngine from vllm_rlt.models import OuroForCausalLM from vllm_rlt.request import RequestPhase @@ -485,6 +485,27 @@ async def run(): asyncio.run(run()) +def test_offline_stop_token_text_is_omitted_and_http_keeps_its_fields(): + engine, tokenizer = factory() + engine.add_request("direct", tokenizer.encode("abc"), SamplingParams(max_tokens=4)) + outputs = [] + while engine.has_unfinished_requests(): + outputs.extend(engine.step()) + tokens = outputs[-1].token_ids + index = next(i for i, t in enumerate(tokens) if i and t not in tokens[:i]) + stop, expected = tokens[index], tokenizer.decode(tokens[:index]) + # The stop ID stays in token_ids; only its text is omitted, as in vLLM. + llm = LLM(engine.model, tokenizer=tokenizer) + output = llm.generate("abc", SamplingParams(max_tokens=4, stop_token_ids=[stop]))[0] + assert (output.text, output.token_ids) == (expected, tokens[: index + 1]) + output = llm.generate("abc", SamplingParams(max_tokens=4))[0] # Other finishes keep it. + assert (output.text, output.finish_reason) == (tokenizer.decode(tokens), "length") + # HTTP request fields are unchanged until the token-in/token-out API. + for fields in ({"stop_token_ids": [stop]}, {"seed": None}): + with pytest.raises(ValueError): + CompletionRequest.parse(body(**fields), "test-model") + + def test_owner_thread_dynamic_batching_cancellation_and_overload(): async def run(): blocked, release = threading.Event(), threading.Event() diff --git a/tests/test_split_prefill.py b/tests/test_split_prefill.py index 9f5a609f..83ce7e2d 100644 --- a/tests/test_split_prefill.py +++ b/tests/test_split_prefill.py @@ -26,6 +26,7 @@ from vllm_rlt.models import OuroForCausalLM from vllm_rlt.request import Operation, Request, RequestPhase from vllm_rlt.worker.executor import build_execution_batch +from vllm_rlt.worker.model_runner import ExecutionOutput TOTAL_DEPTHS = 4 @@ -494,7 +495,7 @@ def test_cuda_submit_readback_failure_releases_batch_input_buffers(monkeypatch): def execute_with_result(batch, prepared): original_execute(batch, prepared) - return torch.zeros(1, device=runner.device, dtype=torch.float32) + return ExecutionOutput(torch.zeros(1, device=runner.device, dtype=torch.float32)) monkeypatch.setattr(runner, "_execute", execute_with_result) @@ -1600,7 +1601,7 @@ def consume(scheduled, prepared): hidden = prepared.batch_inputs.hidden captured.append(hidden.clone()) runner._store_hidden(batch, hidden + 1, prepared) - return None + return ExecutionOutput(None) runner._execute = consume try: diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index b2442d7b..43d82a7f 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -786,11 +786,24 @@ def needs_exit_signal(self, request: Request) -> bool: maximum = request.sampling_params.max_loops or self.total_ut_steps return self.exit_mode != "trace" and request.loops_done + 1 < maximum - def _record_token(self, request: Request, token: int, depth: int) -> bool: + def _record_token( + self, request: Request, token: int, depth: int, logprob: float | None = None + ) -> bool: + """Record one sampled token; return whether it finished the request. + + EOS stops unless ignore_eos; stop_token_ids stop even with ignore_eos, as + in vLLM. Either stop token stays in the output. Then max_tokens is LENGTH. + """ request.generated_token_ids.append(token) request.exit_depths.append(depth) + if request.logprobs is not None: + if logprob is None: + raise RuntimeError("missing logprob for a request that asked for logprobs") + request.logprobs.append(logprob) 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 @@ -868,7 +881,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 self._record_token(request, result.token_ids[row], depth): + logprob = result.logprobs[row] if result.logprobs is not None else None + if self._record_token(request, result.token_ids[row], depth, logprob): update.finished.append(request) elif not delivered: self.enqueue( @@ -891,9 +905,11 @@ 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): emitted += 1 - if self._record_token(request, token, self.speculative_config.target_loops): + logprob = spec.logprobs[index] if spec.logprobs is not None else None + depth = self.speculative_config.target_loops + if self._record_token(request, token, depth, logprob): update.finished.append(request) break update.committed_tokens += emitted diff --git a/vllm_rlt/engine/llm_engine.py b/vllm_rlt/engine/llm_engine.py index 766c9ade..0f04b95e 100644 --- a/vllm_rlt/engine/llm_engine.py +++ b/vllm_rlt/engine/llm_engine.py @@ -12,6 +12,7 @@ from vllm_rlt.engine.stage_config import StageConfig from vllm_rlt.engine.stage_engine import StageEngine from vllm_rlt.engine.worker import start_stage_worker +from vllm_rlt.worker.logprobs import check_logprobs_mode LLMEngine = StageEngine @@ -39,6 +40,8 @@ def __init__( wait_for_ready=True, **engine_options, ): + if "logprobs_mode" in engine_options: # Fail before starting a worker. + check_logprobs_mode(engine_options["logprobs_mode"]) # An explicitly supplied model is a weight source, never a second runner. # Checkpoint paths load directly inside the worker without a host copy. if isinstance(model, torch.nn.Module): diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 7cf9f42d..a14f9f74 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -11,8 +11,9 @@ from vllm_rlt.engine.stage_config import StageConfig from vllm_rlt.profiling import Profiler from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase -from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.sampling_params import SamplingParams, resolve_seed from vllm_rlt.worker.executor import Executor +from vllm_rlt.worker.logprobs import check_logprobs_mode from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput from vllm_rlt.worker.sampler import RngRegistry @@ -43,7 +44,9 @@ def __init__( speculative_config=None, stage_config=None, connector: StageConnector | None = None, + logprobs_mode="raw_logprobs", ): + check_logprobs_mode(logprobs_mode) self.model = model self.stage_config = stage_config or StageConfig() self.profiling = Profiler(next(model.parameters()).device) @@ -176,6 +179,7 @@ def __init__( execution_config=self.execution_config, scheduler_config=scheduler_config, rng=self.rng, + logprobs_mode=logprobs_mode, ) self.speculative_runner = ( SpeculativeRunner( @@ -186,6 +190,7 @@ def __init__( rng=self.rng, runtime=self.model_runner.runtime, compute_stream=self.model_runner.core_stream, + logprobs_mode=logprobs_mode, ) if speculative_config is not None and not self.async_speculative else None @@ -203,6 +208,7 @@ def __init__( copy_stream=self.model_runner.copy_stream, rng=self.rng, runtime=self.model_runner.runtime, + logprobs_mode=logprobs_mode, ) self.executor = Executor(self.model_runner, speculation=self.speculative_runner) self._pending_speculative = [] @@ -267,7 +273,8 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i """Validate and register a stage-local request from either ingress path.""" if not isinstance(request_id, str) or not request_id: raise ValueError("request_id must be a nonempty string") - params = sampling_params or SamplingParams() + # Every output, including an abort before the first sample, reports this seed. + params = resolve_seed(sampling_params or SamplingParams()) if self.async_speculative and params.temperature != 0: raise ValueError("async speculative decoding currently requires greedy sampling") if request_id in self._retiring_speculative: @@ -277,6 +284,8 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i raise ValueError("prompt must contain at least one token") if any(type(t) is not int or not 0 <= t < config.vocab_size for t in prompt_token_ids): raise ValueError("prompt token IDs must be integers within the model vocabulary") + if any(t >= config.vocab_size for t in params.stop_token_ids): + raise ValueError("stop_token_ids must be within the model vocabulary") max_loops = params.max_loops or config.total_ut_steps if max_loops > config.total_ut_steps or params.min_loops > max_loops: raise ValueError("requested loop bounds exceed the model's supported depth") diff --git a/vllm_rlt/entrypoints/llm.py b/vllm_rlt/entrypoints/llm.py index b3554417..b30b3328 100644 --- a/vllm_rlt/entrypoints/llm.py +++ b/vllm_rlt/entrypoints/llm.py @@ -28,6 +28,7 @@ def __init__( exit_config=None, execution_config=None, speculative_config=None, + logprobs_mode="raw_logprobs", ): self._tokenizer_source = None if isinstance(model, str): @@ -48,6 +49,7 @@ def __init__( exit_config=exit_config, execution_config=execution_config, speculative_config=speculative_config, + logprobs_mode=logprobs_mode, ) self._next_request_id = 0 @@ -84,10 +86,13 @@ def generate(self, prompts, sampling_params=None): for output in self.engine.step(): if output.finished: if self.tokenizer is not None: + # As in vLLM, the stop-terminating ID adds no text. + stop = output.finish_reason == "stop" output = replace( output, text=self.tokenizer.decode( - output.token_ids, skip_special_tokens=True + output.token_ids[:-1] if stop else output.token_ids, + skip_special_tokens=True, ), ) results[output.request_id] = output diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index ce2b3b7d..a43906cf 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -18,7 +18,8 @@ from vllm_rlt.profiling import ProfileConfig from vllm_rlt.profiling_artifacts import _write_json_atomic, profile_timestamp from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase -from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.sampling_params import SamplingParams, resolve_seed +from vllm_rlt.worker.logprobs import check_logprobs_mode from .config import PDConfig from .connector import PDStageConnector @@ -68,7 +69,9 @@ def __init__( revision=None, seed=0, speculative_config=None, + logprobs_mode="raw_logprobs", ): + check_logprobs_mode(logprobs_mode) self.speculative_config = speculative_config if isinstance(model, (str, Path)): model, revision, _ = resolve_model_config(model, revision=revision) @@ -135,6 +138,7 @@ def __init__( speculative_config=speculative_config if role == "decode" else None, execution_config=execution, attention_backend=attention_backend, + logprobs_mode=logprobs_mode, ) self.peers[name] = Peer(name, role, client) deadline = time.monotonic() + self.config.startup_timeout @@ -223,6 +227,8 @@ def _message(self, peer, m): output = replace(m["output"], request_id=w.request.request_id) w.request.generated_token_ids = list(output.token_ids) w.request.exit_depths = list(output.exit_depths) + if output.logprobs is not None: + w.request.logprobs = list(output.logprobs) w.timings.setdefault("first_token", now) self.outputs.append(output) if output.finished: @@ -270,7 +276,8 @@ def add_request(self, request_id, prompt_token_ids, sampling_params=None, *, tra raise ValueError("request ID must be nonempty and unique among active requests") if len(self.transfers) >= self.config.max_pending_requests: raise ValueError("PD pending-request limit reached") - params = sampling_params or SamplingParams() + # P and D receive the same resolved params; only D samples. + params = resolve_seed(sampling_params or SamplingParams()) cfg = self.model_config if self.speculative_config is not None: if ( @@ -286,6 +293,8 @@ def add_request(self, request_id, prompt_token_ids, sampling_params=None, *, tra type(t) is not int or not 0 <= t < cfg.vocab_size for t in prompt_token_ids ): raise ValueError("prompt must contain valid model token IDs") + if any(t >= cfg.vocab_size for t in params.stop_token_ids): + raise ValueError("stop_token_ids must be within the model vocabulary") maximum = params.max_loops or cfg.total_ut_steps if maximum > cfg.total_ut_steps or params.min_loops > maximum: raise ValueError("requested loop bounds exceed model depth") diff --git a/vllm_rlt/request.py b/vllm_rlt/request.py index 537e2bf1..3ae93319 100644 --- a/vllm_rlt/request.py +++ b/vllm_rlt/request.py @@ -42,6 +42,8 @@ class Request: phase: RequestPhase = RequestPhase.WAITING generated_token_ids: list[int] = field(default_factory=list) exit_depths: list[int] = field(default_factory=list) + # Aligned with generated_token_ids when sampling_params.logprobs is set, else None. + logprobs: list[float] | None = None # Scheduling advances before coda's CPU output delivery. num_output_placeholders: int = 0 num_speculative_rounds: int = 0 @@ -66,6 +68,10 @@ class Request: finish_reason: FinishReason | None = None exit_trace: tuple[int, ...] = field(default=(), repr=False) + def __post_init__(self): + if self.logprobs is None and self.sampling_params.logprobs is not None: + self.logprobs = [] + @staticmethod def prefill_boundaries( num_tokens: int, total_depths: int, *, start: int = 0 @@ -199,6 +205,10 @@ class RequestOutput: finished: bool finish_reason: str | None = None text: str = "" + # Log-probability of each token_ids entry, conditional on its exit_depths entry. + logprobs: list[float] | None = None + # Effective parameters, including the engine-chosen seed for seed=None. + sampling_params: SamplingParams | None = None @classmethod def from_request(cls, request: Request): @@ -211,4 +221,6 @@ def from_request(cls, request: Request): finish_reason=request.finish_reason.value if request.finish_reason is not None else None, + logprobs=list(request.logprobs) if request.logprobs is not None else None, + sampling_params=request.sampling_params, ) diff --git a/vllm_rlt/sampling_params.py b/vllm_rlt/sampling_params.py index 30009bde..05691ae4 100644 --- a/vllm_rlt/sampling_params.py +++ b/vllm_rlt/sampling_params.py @@ -1,5 +1,6 @@ import math -from dataclasses import dataclass +import secrets +from dataclasses import dataclass, replace @dataclass(frozen=True) @@ -8,12 +9,17 @@ class SamplingParams: temperature: float = 0.0 top_p: float = 1.0 top_k: int = -1 - seed: int = 0 + # None lets the engine draw a per-request seed; outputs report the seed used. + seed: int | None = 0 min_loops: int = 2 max_loops: int | None = None exit_threshold: float = 1.0 ignore_eos: bool = False priority: int = 0 + # vLLM convention: 0 returns only the sampled token's log-probability. + logprobs: int | None = None + # Stop after emitting any of these IDs, even with ignore_eos (vLLM semantics). + stop_token_ids: tuple[int, ...] = () def __post_init__(self): if type(self.priority) is not int: @@ -34,5 +40,30 @@ def __post_init__(self): raise ValueError("top_k must be -1 or positive") if not 0 <= self.exit_threshold <= 1: raise ValueError("exit_threshold must be in [0, 1]") - if type(self.seed) is not int or not 0 <= self.seed < 2**63: - raise ValueError("seed must be an integer in [0, 2**63)") + if self.seed is not None and (type(self.seed) is not int or not 0 <= self.seed < 2**63): + raise ValueError("seed must be None or an integer in [0, 2**63)") + if self.logprobs is not None: + if type(self.logprobs) is not int or self.logprobs < 0: + raise ValueError("logprobs must be None or 0") + if self.logprobs > 0: + raise ValueError( + "logprobs > 0 (top-k alternatives) is not supported yet; " + "use 0 for the sampled token's log-probability" + ) + if not isinstance(self.stop_token_ids, (list, tuple)) or any( + type(token) is not int or token < 0 for token in self.stop_token_ids + ): + raise ValueError("stop_token_ids must be a list of nonnegative integers") + object.__setattr__(self, "stop_token_ids", tuple(self.stop_token_ids)) + + +def resolve_seed(params: SamplingParams) -> SamplingParams: + """Return ``params`` with ``seed=None`` replaced by a fresh 63-bit seed. + + The seed comes from OS entropy, so torch and Python global RNG streams are + not consumed. Engines call this before creating a Request, so outputs and + every PD worker observe the same effective seed. + """ + if params.seed is not None: + return params + return replace(params, seed=secrets.randbits(63)) diff --git a/vllm_rlt/serving/protocol.py b/vllm_rlt/serving/protocol.py index c08243cd..83d5f9f4 100644 --- a/vllm_rlt/serving/protocol.py +++ b/vllm_rlt/serving/protocol.py @@ -61,7 +61,6 @@ class CompletionRequest: def parse(cls, body, model): if not isinstance(body, dict): raise ValueError("request body must be a JSON object") - params_keys = {field.name for field in fields(SamplingParams)} neutral = { "n": 1, "best_of": 1, @@ -74,6 +73,9 @@ def parse(cls, body, model): "suffix": None, "logit_bias": None, } + # Rollout fields stay offline-only until the token-in/token-out API (#70). + params_keys = {field.name for field in fields(SamplingParams)} - neutral.keys() + params_keys -= {"stop_token_ids"} allowed = ( params_keys | neutral.keys() @@ -119,6 +121,8 @@ def parse(cls, body, model): raise ValueError(f"{key} must be a finite number") if "ignore_eos" in params and type(params["ignore_eos"]) is not bool: raise ValueError("ignore_eos must be a boolean") + if "seed" in params and params["seed"] is None: + raise ValueError("seed must be an integer in [0, 2**63)") trace_id = body.get("trace_id") if "trace_id" in body and (not isinstance(trace_id, str) or not trace_id): raise ValueError("trace_id must be a nonempty string") diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index dd31ebae..cc2886f8 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -10,7 +10,7 @@ from vllm_rlt.request import Operation, Request from vllm_rlt.worker.executor import ForwardBatch from vllm_rlt.worker.output import ModelRunnerOutput, SpeculativeResult -from vllm_rlt.worker.sampler import RngRegistry +from vllm_rlt.worker.sampler import RngRegistry, Sampler from vllm_rlt.worker.speculative import SpeculativeRunner @@ -37,6 +37,11 @@ def __init__(self, runner, scheduler): size = scheduler.max_num_batched_tokens + 2 * scheduler.max_num_seqs self.host_result = torch.empty(size, dtype=torch.int64, pin_memory=True) self.result = torch.empty(size, dtype=torch.int64, device=cache.device) + # FP32 target logprob of each verification row's argmax, copied only for + # rounds with a request that asked for logprobs. + rows = scheduler.max_num_batched_tokens + self.host_logprobs = torch.empty(rows, dtype=torch.float32, pin_memory=True) + self.logprobs = torch.empty(rows, dtype=torch.float32, device=cache.device) self.leased = False self.copy_done: torch.cuda.Event | None = None self.compute_done: torch.cuda.Event | None = None @@ -50,6 +55,7 @@ class SpeculativeRoundTicket: states: tuple[SpeculativeRequestState, ...] rounds: tuple[int, ...] rows: int + logprobs_copied: bool = False cached: ModelRunnerOutput | None = None retired: bool = False @@ -69,13 +75,19 @@ def collect(self) -> ModelRunnerOutput: self.buffer.copy_done.synchronize() b = len(self.batch.requests) values = self.buffer.host_result[: self.rows + 2 * b].tolist() + scores = ( + self.buffer.host_logprobs[: self.rows].tolist() if self.logprobs_copied else None + ) results, start = [], 0 for i, scheduled_request in enumerate(self.batch.requests): accepted = values[self.rows + i] + params = scheduled_request.request.sampling_params + wants = scores is not None and params.logprobs is not None results.append( SpeculativeResult( values[start : start + accepted + 1], accepted, + scores[start : start + accepted + 1] if wants else None, ) ) start += scheduled_request.token_count @@ -117,8 +129,19 @@ def __init__( copy_stream, rng: RngRegistry | None = None, runtime=None, + logprobs_mode="raw_logprobs", ): - super().__init__(model, cache, config, execution, rng=rng, runtime=runtime) + # Async speculation is greedy-only, and greedy rows report raw logprobs in + # both modes, so the mode never changes a score here. + super().__init__( + model, + cache, + config, + execution, + rng=rng, + runtime=runtime, + logprobs_mode=logprobs_mode, + ) self.compute_stream = compute_stream self.copy_stream = copy_stream self.compute_stream.wait_stream(torch.cuda.current_stream(self.device)) @@ -233,7 +256,11 @@ def _draft(self, scheduled_requests, states, tables, allocations): candidates[i].append(proposed[row]) return starts, candidates, hidden_rows - def _verify(self, scheduled_requests, starts, hidden_rows, tables, allocations): + def _verify(self, scheduled_requests, starts, hidden_rows, tables, allocations, *, score=False): + """Return (target, scores): the target argmax of every verification row and, with + ``score``, each row's FP32 raw logprob of that argmax (else None). Committed tokens + are exactly the target argmax of their rows. + """ positions = torch.cat( [ start + torch.arange(scheduled_request.token_count, device=self.device) @@ -249,8 +276,9 @@ def _verify(self, scheduled_requests, starts, hidden_rows, tables, allocations): hidden = torch.stack([h for rows in hidden_rows for h in rows]) for depth in range(self.config.draft_loops, self.config.target_loops): hidden = self._execute_recurrent_step(hidden, positions, row_tables, depth, allocations) - target = self._coda(hidden).argmax(-1) - return target + logits = self._coda(hidden) + target = logits.argmax(-1) + return target, Sampler.raw_logprobs(logits, target) if score else None def _accept_and_update_state( self, scheduled_requests, states, starts, candidates, target, buffer @@ -295,7 +323,15 @@ def submit(self, batch): starts, candidates, hidden_rows = self._draft( scheduled_requests, states, tables, allocations ) - target = self._verify(scheduled_requests, starts, hidden_rows, tables, allocations) + scoring = any( + scheduled_request.request.sampling_params.logprobs is not None + for scheduled_request in scheduled_requests + ) + target, scores = self._verify( + scheduled_requests, starts, hidden_rows, tables, allocations, score=scoring + ) + if scoring: + buffer.logprobs[: len(target)].copy_(scores) count = self._accept_and_update_state( scheduled_requests, states, starts, candidates, target, buffer ) @@ -308,6 +344,8 @@ def submit(self, batch): self.copy_stream.wait_event(buffer.compute_done) n = count + 2 * len(scheduled_requests) buffer.host_result[:n].copy_(buffer.result[:n], non_blocking=True) + if scoring: + buffer.host_logprobs[:count].copy_(buffer.logprobs[:count], non_blocking=True) buffer.copy_done = torch.cuda.Event() buffer.copy_done.record(self.copy_stream) self.stats.rounds += len(scheduled_requests) @@ -315,7 +353,7 @@ def submit(self, batch): scheduled_request.token_count - 1 for scheduled_request in scheduled_requests ) self.stats.verified_rows += count - return SpeculativeRoundTicket(batch, buffer, states, rounds, count) + return SpeculativeRoundTicket(batch, buffer, states, rounds, count, logprobs_copied=scoring) def request_ready(self, request): state = self.states.get(request.request_id) diff --git a/vllm_rlt/worker/logprobs.py b/vllm_rlt/worker/logprobs.py new file mode 100644 index 00000000..90331ed1 --- /dev/null +++ b/vllm_rlt/worker/logprobs.py @@ -0,0 +1,86 @@ +"""Rollout logprobs: the FP32 log-probability of each sampled token. + +Scoring consumes no RNG. Raw logprobs score the model distribution. Processed +logprobs score ``Sampler.processed_logits``, the logits ``Sampler.probabilities`` +normalizes, so they score the distribution the token was actually drawn from. +This module imports ``worker/sampler.py``, never the reverse. +""" + +from collections import defaultdict + +import torch + +from vllm_rlt.worker.sampler import Sampler + +LOGPROBS_MODES = ("raw_logprobs", "processed_logprobs") + + +def check_logprobs_mode(mode): + if mode not in LOGPROBS_MODES: + raise ValueError( + f"logprobs_mode must be one of {LOGPROBS_MODES}, got {mode!r}; " + "raw_logits and processed_logits are not supported" + ) + + +def token_logprob(logits, params, token_id, mode): + """0-dim FP32 log-probability of ``token_id`` (0-dim long tensor) for one logits row. + + raw_logprobs uses the model distribution. processed_logprobs uses the + distribution actually sampled from; greedy rows have no temperature (vLLM). + Consumes no RNG and keeps the result on the logits device. + """ + if mode == "processed_logprobs" and params.temperature != 0: + logits = Sampler.processed_logits(logits, params) + return Sampler.raw_logprobs(logits[None], token_id.view(1))[0] + + +def truncates(params): + """Whether this row is drawn from top-k/top-p truncated logits (see ``processed_logits``). + + Greedy rows never are: ``argmax`` ignores top-k and top-p. + """ + return params.temperature != 0 and (params.top_k > 0 or params.top_p < 1) + + +def scaled_logprobs(logits, token_ids, temperature): + """``raw_logprobs`` of ``logits / temperature``, divided as ``processed_logits`` divides.""" + scaled = logits if temperature == 1 else logits.float() / temperature + return Sampler.raw_logprobs(scaled, token_ids) + + +def processed_logprobs(logits, params, token_ids, truncated): + """FP32 [rows] processed logprobs of ``token_ids``, equal to per-row ``token_logprob``. + + ``logits`` [rows, vocab] are the CODA rows of ``params`` (one per row), at least one of + which requests logprobs. ``truncated`` maps every requesting row that ``truncates`` to + the 0-dim logprob ``Sampler.sample_with_logprob`` computed from the processed logits it + drew from. Other requesting rows are grouped by temperature: greedy and T == 1 rows are + scored by a raw ``log_softmax`` (``x / 1.0 == x``), each other temperature by + ``log_softmax(logits.float() / T)`` with the Python scalar T, as ``processed_logits`` + divides. A group is scored in one pass over its row span ``[first, last]`` when that span + is at most twice its rows, else row by row, so the scored rows never exceed twice the + requesting rows, however many temperatures the batch mixes. The only group, holding at + least half the rows, is scored over the whole batch and returned as is; otherwise one + stack of 0-dim views assembles the rows. No host sync, per-row device write or per-row + device metadata (spans are host-int slices). Rows that did not request logprobs hold + arbitrary values. + """ + groups = defaultdict(list) # Temperature (greedy: 1) -> requesting untruncated rows. + for row, row_params in enumerate(params): + if row_params.logprobs is not None and row not in truncated: + groups[1 if row_params.temperature == 0 else row_params.temperature].append(row) + values = dict(truncated) + for temperature, rows in groups.items(): + if not truncated and len(groups) == 1 and len(params) <= 2 * len(rows): + return scaled_logprobs(logits, token_ids, temperature) # One pass, no assembly. + first, end = rows[0], rows[-1] + 1 + if end - first > 2 * len(rows): # Sparse group: a span pass would mostly score others. + for row in rows: + one = slice(row, row + 1) + values[row] = scaled_logprobs(logits[one], token_ids[one], temperature)[0] + continue + scored = scaled_logprobs(logits[first:end], token_ids[first:end], temperature) + values.update((row, scored[row - first]) for row in rows) + filler = next(iter(values.values())) + return torch.stack([values.get(row, filler) for row in range(len(params))]) diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 58e3414f..92b4bb89 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -18,12 +18,24 @@ ) from vllm_rlt.worker.cuda_graph import CodaGraphs, RecurrentGraphs from vllm_rlt.worker.executor import ForwardBatch, ModelRuntime +from vllm_rlt.worker.logprobs import processed_logprobs, truncates from vllm_rlt.worker.output import ModelRunnerOutput from vllm_rlt.worker.sampler import RngRegistry, Sampler from vllm_rlt.worker.submission_queue import MAX_CONCURRENT_SUBMISSIONS, SubmissionQueue from vllm_rlt.worker.tensors import StagedInputTensor, TensorTransfer, UvaBuffer, pinned_tensor +class ExecutionOutput(NamedTuple): + """One submission's device outputs. + + ``values`` holds RECURRENT exit scores or CODA sampled IDs, else None. + ``logprobs`` holds CODA FP32 logprobs of those IDs when a row asked, else None. + """ + + values: torch.Tensor | None + logprobs: torch.Tensor | None = None + + @dataclass class ReadbackSlot: storage: torch.Tensor @@ -115,10 +127,14 @@ class Submission: positions: tuple[int, ...] = () output_indices: tuple[int, ...] = () operation: Operation | None = None + # CODA: FP32 logprobs of the sampled IDs, read back on the same event. + logprobs: torch.Tensor | None = None + logprob_slot: ReadbackSlot | None = None def __del__(self): - if self.slot is not None: - self.slot.leased = False + for slot in (self.slot, self.logprob_slot): + if slot is not None: + slot.leased = False @property def has_result(self): @@ -133,19 +149,23 @@ def collect(self) -> ModelRunnerOutput: if self.event is not None: self.event.synchronize() values = self.values.tolist() if self.values is not None else None + logprobs = self.logprobs.tolist() if self.logprobs is not None else None is_coda = self.operation == Operation.CODA is_recurrent = self.operation == Operation.RECURRENT self.cached = ModelRunnerOutput( gate_scores=values if is_recurrent else None, token_ids=values if is_coda else None, + logprobs=logprobs if is_coda else None, completion=self.event, depths=self.depths, output_indices=self.output_indices, ) - self.values = None - if self.slot is not None: - self.slot.leased = False - self.slot = None + self.values = self.logprobs = None + for name in ("slot", "logprob_slot"): + slot = getattr(self, name) + if slot is not None: + slot.leased = False + setattr(self, name, None) return self.cached @@ -159,6 +179,7 @@ def __init__( execution_config=None, scheduler_config=None, rng: RngRegistry | None = None, + logprobs_mode="raw_logprobs", ): self.model = model.eval() self.cache_manager = cache_manager @@ -170,6 +191,7 @@ def __init__( # Sampling RNG state lives with the engine, which shares this registry # between the CODA runner and the speculative runner. self.rng = rng if rng is not None else RngRegistry() + self.logprobs_mode = logprobs_mode self.exit_config = exit_config or ExitConfig() self.execution_config = execution_config or ExecutionConfig() scheduler = scheduler_config or SchedulerConfig() @@ -237,6 +259,7 @@ def __init__( ): stream.wait_stream(torch.cuda.current_stream(self.device)) self.readback_slots = [] + self.logprob_slots = [] if self.device.type == "cuda": # cudaHostAlloc during submission can serialize streams. Allocate all # readback storage before work starts, retaining leases until consumed. @@ -247,6 +270,19 @@ def __init__( pinned_tensor(scheduler.max_num_batched_tokens, dtype, self.device) ) ) + # A separate pool keeps exit-score leases from starving CODA logprobs: a + # request whose CODA output is undelivered can already run its next + # token's recurrent loops, so it can hold both kinds of lease at once. + # A CODA ticket keeps its logprob lease until it is delivered or dropped. + # Each pending ticket holds a request with an undelivered output (at most + # one per request), so with the new ticket at most max_num_seqs leases are + # live; a stale ticket of an aborted request is still delivered (its + # results are ignored) or dropped on failure. The +4 is headroom, as for + # the exit-score pool. + for _ in range(scheduler.max_num_seqs + 4): + self.logprob_slots.append( + ReadbackSlot(pinned_tensor(scheduler.max_num_seqs, torch.float32, self.device)) + ) # Retire the oldest consumer BEFORE preparing another batch's inputs. # All rotating prefill pools share this maximum in-flight depth. self.submissions = SubmissionQueue() @@ -1078,12 +1114,12 @@ def _prepare_tokens(self, prepared, size): return tokens @torch.inference_mode() - def _execute(self, scheduled_batch, prepared): + def _execute(self, scheduled_batch, prepared) -> ExecutionOutput: return self._execute_batch(prepared.batch, prepared) - def _execute_batch(self, batch: ExecutionBatch, prepared): + def _execute_batch(self, batch: ExecutionBatch, prepared) -> ExecutionOutput: hidden = None - values = None + values = logprobs = None for entry in prepared.forwards: forward = entry.forward if entry.hidden_rows is not None: @@ -1108,21 +1144,49 @@ def _execute_batch(self, batch: ExecutionBatch, prepared): ) values = selected.float().sigmoid() elif forward.compute_logits: - values = torch.stack( - [ - self.sampler.sample(row, params, generator) - for row, params, generator in zip( - result.logits, prepared.sampling_params, prepared.sampling_generators - ) - ] - ) + params, generators = prepared.sampling_params, prepared.sampling_generators + if any(p.logprobs is not None for p in params): + values, logprobs = self._sample_with_logprobs(result.logits, params, generators) + else: + values = torch.stack( + [ + self.sampler.sample(row, p, generator) + for row, p, generator in zip(result.logits, params, generators) + ] + ) if prepared.routing is not None: prepared.routing.scatter(values, tokens=True) - return values + return ExecutionOutput(values, logprobs) - def _readback_slot(self, result): + def _sample_with_logprobs(self, logits, params, generators): + """Sampled IDs [rows] and FP32 logprobs [rows] of them under ``self.logprobs_mode``. + + Raw mode scores every row in one batched op. Processed mode samples a requesting + top-k/top-p row through ``Sampler.sample_with_logprob``, which scores the processed + logits it drew from, and scores the other requesting rows in temperature batches + (``processed_logprobs``); no row rebuilds its processed logits. All other rows + sample through ``Sampler.sample``. Both draw identically, so tokens and RNG use do + not depend on logprobs. Values of rows that did not request logprobs are arbitrary. + """ + rows = len(params) + if self.logprobs_mode == "raw_logprobs": + token_ids = torch.stack( + [self.sampler.sample(row, p, g) for row, p, g in zip(logits, params, generators)] + ) + return token_ids, Sampler.raw_logprobs(logits[:rows], token_ids) + tokens, truncated = [], {} + for index, (row, p, generator) in enumerate(zip(logits, params, generators)): + if p.logprobs is not None and truncates(p): + token, truncated[index] = self.sampler.sample_with_logprob(row, p, generator) + else: + token = self.sampler.sample(row, p, generator) + tokens.append(token) + token_ids = torch.stack(tokens) + return token_ids, processed_logprobs(logits[:rows], params, token_ids, truncated) + + def _readback_slot(self, result, slots=None): pending = None - for slot in self.readback_slots: + for slot in self.readback_slots if slots is None else slots: if slot.storage.dtype != result.dtype or slot.leased: continue if slot.event is None or slot.event.query(): @@ -1136,6 +1200,13 @@ def _readback_slot(self, result): return pending raise RuntimeError("readback buffers exhausted; retire completed submissions") + @staticmethod + def _copy_to_host(values, slot, stream): + values.record_stream(stream) + host = slot.storage[: values.numel()].view(values.shape) + host.copy_(values, non_blocking=True) + return host + def _synchronize_submission(self, stream, prepared): if self.device.type != "cuda": return @@ -1171,7 +1242,7 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu stream = self.core_stream if batch.loop_count else self.boundary_stream prepared = None routing_finish_attempted = False - slot = None + slot = logprob_slot = None try: if self.submissions.begin() and split_prefill: self.split_prefill_submission_waits += 1 @@ -1191,7 +1262,7 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu self._wait_for_requests(batch, stream) prepared = self.prepare_on_gpu(prepared) try: - result = self._execute(scheduled_batch, prepared) + result, logprobs = self._execute(scheduled_batch, prepared) finally: if prepared.routing is not None: routing_finish_attempted = True @@ -1225,17 +1296,22 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu event = compute_done if result is not None: slot = self._readback_slot(result) + if logprobs is not None: + logprob_slot = self._readback_slot(logprobs, self.logprob_slots) readback_stream = self.readback_stream or compute_stream with torch.cuda.stream(readback_stream): if readback_stream != compute_stream: readback_stream.wait_event(compute_done) - result.record_stream(readback_stream) - host = slot.storage[: result.numel()].view(result.shape) - host.copy_(result, non_blocking=True) + result = self._copy_to_host(result, slot, readback_stream) + if logprobs is not None: + logprobs = self._copy_to_host( + logprobs, logprob_slot, readback_stream + ) event = torch.cuda.Event() event.record(readback_stream) - slot.event = event - result = host + for leased in (slot, logprob_slot): + if leased is not None: + leased.event = event for span in batch.spans: self.events[span.request.request_id] = compute_done self.submissions.complete(compute_done) @@ -1251,6 +1327,8 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu positions=prepared.positions, output_indices=prepared.output_indices, operation=batch.operation, + logprobs=logprobs, + logprob_slot=logprob_slot, ) except Exception: self._synchronize_submission(stream, prepared) @@ -1267,8 +1345,9 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu routing_finish_attempted = True prepared.routing.record_done() finally: - if slot is not None: - slot.leased = False + for leased in (slot, logprob_slot): + if leased is not None: + leased.leased = False raise finally: if prepared is not None and prepared.workspace is not None: diff --git a/vllm_rlt/worker/output.py b/vllm_rlt/worker/output.py index 736f9f42..e9e552b4 100644 --- a/vllm_rlt/worker/output.py +++ b/vllm_rlt/worker/output.py @@ -19,14 +19,27 @@ class ExitSignal: @dataclass(frozen=True) class SpeculativeResult: + """One request's committed tokens from one round. + + ``token_ids[:accepted_count]`` are accepted draft tokens; ``token_ids[accepted_count]`` + is the correction, or the bonus when every draft was accepted. Token ``j`` comes from + verification row ``j``. + """ + token_ids: list[int] accepted_count: int + # Aligned with token_ids when the request asked for logprobs; else None. Each is the + # target's logprob of the committed token on its verification row, never the draft's. + logprobs: list[float] | None = None @dataclass(frozen=True, kw_only=True) class ModelRunnerOutput: gate_scores: list[float] | None = None token_ids: list[int] | None = None + # CODA: FP32 logprob of each token_ids entry; None when no row asked. Rows that did + # not ask hold arbitrary values. + logprobs: list[float] | None = None completion: object = None depths: tuple[int, ...] = () output_indices: tuple[int, ...] = () diff --git a/vllm_rlt/worker/sampler.py b/vllm_rlt/worker/sampler.py index 62bc7846..9d968edb 100644 --- a/vllm_rlt/worker/sampler.py +++ b/vllm_rlt/worker/sampler.py @@ -14,6 +14,10 @@ caller obtained from ``RngRegistry``. The registry is the single owner of RNG creation, lookup and lifetime, and the engine shares one instance between its runners. + +Rollout logprobs (``worker/logprobs.py``) score sampled tokens with +``Sampler.raw_logprobs``, on ``Sampler.processed_logits`` in processed mode. That +module imports this one, never the reverse. """ import torch @@ -25,14 +29,15 @@ class Sampler: """Greedy / top-k / top-p sampling for one logits row.""" @staticmethod - def probabilities(logits: torch.Tensor, params: SamplingParams) -> torch.Tensor: - """Return the temperature/top-k/top-p distribution for one logits row. + def processed_logits(logits: torch.Tensor, params: SamplingParams) -> torch.Tensor: + """Return the FP32 temperature/top-k/top-p logits for one logits row. ``params`` is assumed to be validated by ``SamplingParams.__post_init__``. Greedy decoding has no temperature-scaled distribution, so it is a caller - error rather than a silent argmax. Top-k keeps exactly ``k`` candidates - (or every token when ``k`` reaches the vocabulary); top-p may trim that - support further. + error rather than a silent argmax. Excluded tokens are ``-inf``. Top-k + keeps exactly ``k`` candidates (or every token when ``k`` reaches the + vocabulary); top-p may trim that support further. ``probabilities`` (CODA + and speculative sampling) and processed logprobs share this arithmetic. """ if params.temperature == 0: raise ValueError("greedy decoding has no temperature-scaled distribution") @@ -52,13 +57,30 @@ def probabilities(logits: torch.Tensor, params: SamplingParams) -> torch.Tensor: remove[1:] = remove[:-1].clone() remove[0] = False logits = logits.scatter(0, indices, sorted_logits.masked_fill(remove, -torch.inf)) - return logits.softmax(-1) + return logits + + @staticmethod + def probabilities(logits: torch.Tensor, params: SamplingParams) -> torch.Tensor: + """Return the temperature/top-k/top-p distribution for one logits row. + + It is the softmax of ``processed_logits``, which rejects greedy decoding. + """ + return Sampler.processed_logits(logits, params).softmax(-1) @staticmethod def draw(probs: torch.Tensor, generator: torch.Generator) -> torch.Tensor: """Draw one token from a normalized distribution as a 0-dim long tensor.""" return torch.multinomial(probs, 1, generator=generator).squeeze(0) + @staticmethod + def raw_logprobs(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + """FP32 [rows] log-probabilities of ``token_ids`` [rows] under ``logits`` [rows, vocab]. + + Scoring consumes no RNG and keeps the result on the logits device. + """ + logprobs = torch.log_softmax(logits, -1, dtype=torch.float32) + return logprobs.gather(-1, token_ids[:, None]).squeeze(-1) + def sample( self, logits: torch.Tensor, @@ -80,6 +102,28 @@ def sample( raise ValueError("non-greedy sampling requires a generator") return self.draw(self.probabilities(logits, params), generator) + def sample_with_logprob( + self, + logits: torch.Tensor, + params: SamplingParams, + generator: torch.Generator | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``sample()`` plus the 0-dim FP32 processed logprob of the drawn token. + + The logprob is ``log_softmax`` of the same processed logits the token is + drawn from, so it equals ``token_logprob(..., "processed_logprobs")`` + without rebuilding them; greedy rows report the unscaled distribution. + Tokens, generator requirements and RNG consumption are those of ``sample()``. + """ + if params.temperature == 0: + token = logits.argmax() + return token, self.raw_logprobs(logits[None], token.view(1))[0] + if generator is None: + raise ValueError("non-greedy sampling requires a generator") + processed = self.processed_logits(logits, params) + token = self.draw(processed.softmax(-1), generator) + return token, self.raw_logprobs(processed[None], token.view(1))[0] + class RngRegistry: """Per-request sampling generators for one engine. @@ -105,11 +149,15 @@ def acquire( """Return the request's generator, creating it on ``device`` if needed. The caller passes the device of the distribution it will sample from, so - ``multinomial`` and its generator agree. An existing generator is returned - as-is and never migrated. + ``multinomial`` and its generator agree. A new generator is seeded with + ``params.seed``, which must be engine-resolved, so its ``initial_seed()`` + is the seed the request's outputs report. An existing generator is + returned as-is and never reseeded or migrated. """ generator = self.generators.get(request_id) if generator is None: + if params.seed is None: # Engines resolve seed=None at admission. + raise ValueError("seed=None must be resolved by the engine before sampling") generator = torch.Generator(device=device).manual_seed(params.seed) self.generators[request_id] = generator return generator diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index dc96c6ad..d5c5c9ae 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -5,7 +5,7 @@ """ from contextlib import nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, field import torch @@ -29,6 +29,56 @@ class SpeculativeStats: verified_rows: int = 0 +@dataclass +class _CommittedScores: + """Committed tokens of one round to score, read back together at the end. + + Raw-scored tokens name their verification row in the round's logits; processed-scored + tokens bring the FP32 processed logits their p was normalized from. Spans map a + result index to its tokens' contiguous slice within its group. + """ + + raw_rows: list[int] = field(default_factory=list) + raw_tokens: list[int] = field(default_factory=list) + raw_spans: dict[int, slice] = field(default_factory=dict) + processed_logits: list[torch.Tensor] = field(default_factory=list) + processed_tokens: list[int] = field(default_factory=list) + processed_spans: dict[int, slice] = field(default_factory=dict) + + def add_raw(self, index, first_row, tokens): + start = len(self.raw_tokens) + self.raw_spans[index] = slice(start, start + len(tokens)) + self.raw_rows.extend(range(first_row, first_row + len(tokens))) + self.raw_tokens.extend(tokens) + + def add_processed(self, index, processed_logits, tokens): + start = len(self.processed_tokens) + self.processed_spans[index] = slice(start, start + len(tokens)) + self.processed_logits.extend(processed_logits) + self.processed_tokens.extend(tokens) + + def resolve(self, logits) -> dict[int, list[float]]: + """Map each added result index to its tokens' FP32 logprobs. + + At most two batched scoring ops and one host readback; none if nothing was added. + """ + parts = [] + if self.raw_tokens: + rows = torch.tensor(self.raw_rows, device=logits.device) + tokens = torch.tensor(self.raw_tokens, device=logits.device) + parts.append(Sampler.raw_logprobs(logits[rows], tokens)) + if self.processed_tokens: + tokens = torch.tensor(self.processed_tokens, device=logits.device) + parts.append(Sampler.raw_logprobs(torch.stack(self.processed_logits), tokens)) + if not parts: + return {} + values = torch.cat(parts).tolist() + raw, processed = values[: len(self.raw_tokens)], values[len(self.raw_tokens) :] + scores = {index: raw[span] for index, span in self.raw_spans.items()} + scores.update((index, processed[span]) for index, span in self.processed_spans.items()) + return scores + + def greedy_accept(candidates, targets): """Return (tokens, accepted) for one request from host candidate and target IDs. @@ -107,6 +157,7 @@ def __init__( *, runtime=None, compute_stream=None, + logprobs_mode="raw_logprobs", ): self.model, self.cache, self.config = model, cache, config self.metadata = cache.metadata_builder @@ -121,6 +172,7 @@ def __init__( # The engine injects one registry per engine so both sampling paths share # each request's RNG slot. self.rng = rng if rng is not None else RngRegistry() + self.logprobs_mode = logprobs_mode self.stats = SpeculativeStats() self.graphs = ( RecurrentGraphs(model, cache, execution, False) if execution.cuda_graphs else None @@ -281,37 +333,71 @@ def _collect_round(self, submission): candidates[i].append(candidate) start += len(rows) targets = values[start:] - results, start = [], 0 + committed, scores, start = [], _CommittedScores(), 0 for i, scheduled_request in enumerate(scheduled_requests): rows = logits[start : start + scheduled_request.token_count] - request = scheduled_request.request - if request.sampling_params.temperature == 0: + params = scheduled_request.request.sampling_params + # Greedy rows report raw values in both modes (vLLM). Processed scores of + # sampled rows reuse the processed logits p was normalized from. + score_processed = ( + params.logprobs is not None + and params.temperature != 0 + and self.logprobs_mode == "processed_logprobs" + ) + if params.temperature == 0: tokens, accepted = greedy_accept( candidates[i], targets[start : start + scheduled_request.token_count] ) - start += scheduled_request.token_count - results.append(SpeculativeResult(tokens, accepted)) - continue - start += scheduled_request.token_count - # Sampling keeps sequential rejection; RNG consumption order is unchanged. - tokens, accepted = [], 0 - for offset, candidate in enumerate(candidates[i]): - p = Sampler.probabilities(rows[offset], request.sampling_params) - token, accept = rejection_sample( - candidate, - p, - proposals[i][offset], - self._acquire_generator(request, rows[offset]), - ) - tokens.append(token) - if not accept: - break - accepted += 1 else: - token = Sampler.draw( - Sampler.probabilities(rows[-1], request.sampling_params), - self._acquire_generator(request, rows[-1]), + tokens, accepted, processed = self._sample_round( + scheduled_request.request, + rows, + candidates[i], + proposals[i], + score_processed=score_processed, ) - tokens.append(int(token.item())) - results.append(SpeculativeResult(tokens, accepted)) - return results + # Token j was committed on verification row j. + if score_processed: + scores.add_processed(i, processed, tokens) + elif params.logprobs is not None: + scores.add_raw(i, start, tokens) + start += scheduled_request.token_count + committed.append((tokens, accepted)) + logprobs = scores.resolve(logits) + return [ + SpeculativeResult(tokens, accepted, logprobs.get(i)) + for i, (tokens, accepted) in enumerate(committed) + ] + + def _sample_round(self, request, rows, candidates, proposals, *, score_processed=False): + """Return (tokens, accepted, processed) by sequential rejection over ``rows``. + + RNG consumption order is unchanged. With ``score_processed``, ``processed`` holds + the FP32 processed logits of the row each committed token was drawn on, one per + token; otherwise it is None. A correction is drawn from the residual (p - q)+, + but the round's output is distributed as p, so every token is scored under p. + """ + params = request.sampling_params + processed = [] if score_processed else None + + def target(row): + if processed is None: + return Sampler.probabilities(row, params) + processed.append(Sampler.processed_logits(row, params)) + return processed[-1].softmax(-1) + + tokens, accepted = [], 0 + for offset, candidate in enumerate(candidates): + token, accept = rejection_sample( + candidate, + target(rows[offset]), + proposals[offset], + self._acquire_generator(request, rows[offset]), + ) + tokens.append(token) + if not accept: + return tokens, accepted, processed + accepted += 1 + token = Sampler.draw(target(rows[-1]), self._acquire_generator(request, rows[-1])) + tokens.append(int(token.item())) + return tokens, accepted, processed