diff --git a/docs/async-speculative.md b/docs/async-speculative.md index 52e877ef..da00b9d0 100644 --- a/docs/async-speculative.md +++ b/docs/async-speculative.md @@ -3,11 +3,11 @@ This guide covers cross-round asynchronous block speculation and [pipelined self-speculation](#pipelined-self-speculation). -The cross-round asynchronous block path uses fixed-depth Ouro, `last_exited` KV, greedy sampling, -CUDA execution, and refill scheduling. Triton and compatible paged FlashAttention -backends can consume its GPU-owned context lengths; `auto` selects against this -capability as well as hardware, dtype and block size. Draft depth is configurable, with `d=2, D=4` -as the initial validation configuration. The target depth must match the model. +The cross-round asynchronous block path uses fixed-depth Ouro, `last_exited` KV, greedy or +[random sampling](#random-sampling), CUDA execution, and refill scheduling. Triton and compatible +paged FlashAttention backends can consume its GPU-owned context lengths; `auto` selects against +this capability as well as hardware, dtype and block size. Draft depth is configurable, with +`d=2, D=4` as the initial validation configuration. The target depth must match the model. ## Timeline: synchronous versus asynchronous speculation @@ -158,7 +158,8 @@ if __name__ == "__main__": The CLI combines `--async-scheduling --attention-backend flash_attn_4 --cuda-graphs --speculative-tokens 2 --draft-loops 2 --target-loops 4 --temperature 0`. -For the default block mode, random sampling, adaptive exit and preemption remain rejected. +For the default block mode, adaptive exit and preemption remain rejected; random sampling is +accepted (see [Random sampling](#random-sampling)). PD can enable speculation on its decode stage. Incoming prompts wait for transfer completion and shared scheduler admission before CODA bootstraps local speculative rounds. Cancellation retains GPU resources until the final round completes. @@ -184,6 +185,83 @@ state with the existing Triton routing kernels (`H=2`). Acceptance still compute each request's prefix match and correction independently. Slot allocation, round completion events and deferred retirement retain their existing owners. +## Random sampling + +Async block speculation accepts `temperature`, `top_k`, `top_p` and `seed` in +`SamplingParams`, and greedy (`temperature=0`) and random requests can share one +speculative batch. Pipeline mode and the PD decode stage still reject +`temperature != 0` and require greedy sampling. + +```python +SamplingParams(temperature=0.8, top_k=50, top_p=0.95, seed=1234, + max_loops=4, exit_threshold=1.0, max_tokens=64) +``` + +**Algorithm.** Standard speculative rejection sampling, computed on the device +without host reads during submission. `p` is the target distribution and `q` the +proposal; both are the request's temperature/top-k/top-p distributions, normalized +per row, and every step below uses these normalized values. + +- Drafts are sampled from the actual proposal `q` of the draft pass. +- Draft `x` is accepted if `p(x) >= q(x)`, without consulting a uniform; otherwise + if `u * q(x) < p(x)`, i.e. with probability `min(1, p/q)`. +- At the first rejection the replacement token is sampled from `max(p - q, 0)`. +- If every draft is accepted, the bonus token is sampled from `p` at the next + position. A request with no drafts in a round (effective K=0) samples from `p`. +- Draft and final sampling share one inverse-CDF helper. The threshold is clamped + below the total mass, so a sample always lies in the support of its distribution. + +**Result layout.** Greedy rows keep returning the target argmax prefix. A random +request returns its committed sequence in its result rows: the accepted drafts, +then the correction or bonus token. An accepted draft need not equal the target +argmax, so the committed tokens, the next-round input token on the device and the +CPU-side output all come from this sequence. Rounds containing random requests +append one numerical-error flag per request after the accepted counts and next +positions; greedy-only rounds keep the previous layout and do no extra device work. + +**RNG protocol.** The engine-owned per-request generator (`RngRegistry`, seeded from +`SamplingParams.seed`) is used as follows. + +- Each random request draws one `torch.rand(2K + 1)` per submitted round, with + `K = num_speculative_tokens`: `[0:K]` for draft sampling, `[K:2K]` for acceptance + and `[2K]` for the correction or bonus token. +- Consumption does not depend on the round's draft count or acceptance outcomes, + so no host synchronization is needed. It differs from synchronous speculation, + which draws per token and stops at the first rejection. +- Greedy requests never acquire a generator. +- The first token is sampled by the ordinary CODA path with the same generator. + +**Reproducibility.** A fixed seed reproduces outputs only on the same execution +path: the same mode (eager or CUDA Graph), round shapes and batch plan. Outputs are +not token-identical to synchronous speculation, to ordinary non-speculative +sampling, or between graph and eager execution (padding can change BF16 logits; see +the numerical notes below). + +**Numerical errors.** If a request's probabilities are invalid at its stop point +(the first non-accepted position, or the bonus position), such as a non-finite or +zero-mass row or an empty residual, the whole round fails through the existing +fatal cleanup rather than isolating the request. The check runs after the stale-row +filter and before the scheduler's EOS truncation. Positions after the stop point are +ignored, and results for finished or reused requests are discarded as before. A +failed request's device state is not advanced, so an already queued round reads the +span it had reserved. + +| Combination | Random sampling | +| --- | --- | +| Async block, local stage, eager | Validated on one RTX 5090 (SM120) with the Triton backend | +| Async block, local stage, CUDA Graphs | Validated on one RTX 5090 (SM120) with the Triton backend | +| Pipeline mode (synchronous or asynchronous) | Greedy only | +| PD decode stage with async speculation | Greedy only | +| FA4 | Not yet validated | +| Speculative preemption | Not applicable (preemption is rejected in async block mode) | + +**Validation status.** CPU tests cover probability construction, rejection sampling +against an exact oracle, and host round planning and result decoding. CUDA tests +cover two in-flight rounds, CUDA Graph replay, ragged and `K = 0` rounds, request +lifecycle and numerical-error cleanup, and the public entry point; they pass on one +RTX 5090 (SM120) with the Triton backend. Other GPUs, FA4, multi-GPU and PD +deployments are not validated. This section makes no performance claim. + ## CUDA Graph execution Enable core and LM-head graphs with diff --git a/tests/test_async_speculative.py b/tests/test_async_speculative.py index 4d6d1871..e777e967 100644 --- a/tests/test_async_speculative.py +++ b/tests/test_async_speculative.py @@ -17,6 +17,7 @@ ) from vllm_rlt.core.scheduler import Scheduler from vllm_rlt.engine.llm_engine import LLMEngine +from vllm_rlt.engine.stage_engine import StageEngine from vllm_rlt.request import Request, RequestPhase tiny = tiny_ouro_model @@ -226,8 +227,17 @@ def fail(*args, **kwargs): @pytest.mark.gpu -def test_async_rejects_random_sampling(): - e = make_speculative_engine(tiny_ouro_model("cuda")) +def test_async_pipeline_rejects_random_sampling(): + e = StageEngine( + tiny_ouro_model("cuda"), + attention_backend="triton", + cache_config=CacheConfig(256, 2, enable_prefix_caching=False, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=2 + ), + execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), + speculative_config=SpeculativeConfig(3, mode="pipeline"), + ) with pytest.raises(ValueError, match="greedy"): e.add_request("r", [2], SamplingParams(temperature=0.8)) diff --git a/tests/test_async_speculative_entrypoint.py b/tests/test_async_speculative_entrypoint.py new file mode 100644 index 00000000..fbe85bd6 --- /dev/null +++ b/tests/test_async_speculative_entrypoint.py @@ -0,0 +1,213 @@ +"""Random sampling in asynchronous block speculation, through the public entrypoint. + +Everything here goes through ``vllm_rlt.LLM``: it launches a spawned ``StageWorker`` process +behind a ``StageClient`` and talks to it over the control pipe, as a user would. All tests need +``--run-gpu`` (async speculation requires CUDA). + +The PD decode-stage guard that rejects random sampling under async speculation +(``vllm_rlt/pd/engine.py``) is unchanged by this feature and is not exercised here. + +Observability note. The worker process exposes no speculative statistics: ``StageEngine.stats()`` +(the only stats control call) returns just ``requests`` and ``used_blocks``, and ``RequestOutput`` +carries tokens and exit depths only. The one signal that crosses the process boundary is the +stream of ``RequestOutput`` snapshots. In async speculation the first token comes from the +bootstrap CODA step and every later token is committed by a delivered speculative round, which +emits exactly one snapshot per round. So the snapshot sequence of a request is +``[1, 1 + d1, 1 + d1 + d2, ...]`` with ``1 <= d_i <= K + 1``, and it proves that rounds were +submitted and delivered. A delta above 1 additionally proves that a delivered round accepted +at least one draft, since a plain CODA step commits exactly one token. It cannot show the +round took the random branch: for a request with ``temperature > 0`` a round is always +planned as random, which unit tests cover +(``test_async_speculative_sampling_plan.py``, ``test_async_speculative_lifecycle.py``). +A device-side counter of random rounds would need a new field in ``StageEngine.stats()``. +""" + +import math + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from vllm_rlt import ( + LLM, + CacheConfig, + ExecutionConfig, + ExitConfig, + SamplingParams, + SchedulerConfig, + SpeculativeConfig, +) + +K = 3 +VOCAB = tiny_ouro_config().vocab_size +PROMPT = [2, 3, 4] +MAX_TOKENS = 12 + + +def random_params(**overrides): + values = dict( + max_tokens=MAX_TOKENS, + temperature=1.0, + top_k=16, + top_p=0.95, + seed=11, + ignore_eos=True, + ) + values.update(overrides) + return SamplingParams(**values) + + +def make_llm(mode="block"): + """The public facade over one CUDA worker process, FP32 so greedy rows are reproducible.""" + # The worker initializes weights from LLM's seed (torch.initial_seed()): pin it. + torch.manual_seed(123) + return LLM( + tiny_ouro_config(), + device="cuda", + dtype=torch.float32, + attention_backend="triton", + cache_config=CacheConfig(256, 2, enable_prefix_caching=False, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=2 + ), + exit_config=ExitConfig("ouro"), + execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), + speculative_config=SpeculativeConfig(K, mode=mode), + ) + + +def stream(llm, requests): + """Submit (id, prompt, params) triples and keep every RequestOutput snapshot per request. + + Like ``LLM.generate`` this loops until the client reports no unfinished request; a dead + worker raises from ``step()`` instead of hanging, so no wall-clock bound is involved. + """ + client = llm.engine + for request_id, prompt, params in requests: + client.add_request(request_id, prompt, params) + history = {request_id: [] for request_id, _, _ in requests} + while client.has_unfinished_requests(): + for output in client.step(): + history[output.request_id].append(output) + return history + + +def assert_valid_tokens(token_ids, count): + assert len(token_ids) == count + assert all(type(t) is int and 0 <= t < VOCAB for t in token_ids) + + +def assert_closed_cleanly(client): + assert client.closed and client.failure is None + assert not client.process.is_alive() + assert client.process.exitcode == 0 + assert client.stats() == dict(requests=0, used_blocks=0) + + +# 1. Entry contract ----------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_async_block_engine_accepts_random_sampling(): + params = random_params(max_tokens=4, temperature=0.8, top_k=8, top_p=0.9, seed=7) + with make_llm() as llm: + # stream() would raise ValueError here if the worker rejected the request. + history = stream(llm, [("r", PROMPT, params)])["r"] + final = history[-1] + assert final.finished and final.finish_reason == "length" + assert_valid_tokens(final.token_ids, 4) + assert_closed_cleanly(llm.engine) + + +@pytest.mark.gpu +def test_async_pipeline_engine_still_rejects_random_sampling(): + with make_llm("pipeline") as llm: + client = llm.engine + with pytest.raises(ValueError, match="greedy"): + client.add_request("r", PROMPT, random_params()) + # Rejection leaves no client-side trace and does not take the worker down. + assert client.request_ids() == () + assert not client.has_unfinished_requests() + assert client.process.is_alive() + assert_closed_cleanly(llm.engine) + + +# 2. Public end-to-end --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_random_request_completes_through_speculative_rounds_and_closes(): + # Make draft acceptance near-certain on the random-weight tiny model: at temperature 1.5 + # without top-k/top-p both the draft (2 loops) and target (4 loops) distributions are close + # to uniform over 64 tokens, so sum_x min(p, q) is high for each draft. With 24 tokens + # (at least 6 rounds of up to 3 drafts) the chance of no accepted draft anywhere is negligible. + tokens = 24 + params = random_params(max_tokens=tokens, temperature=1.5, top_k=-1, top_p=1.0) + with make_llm() as llm: + client = llm.engine + history = stream(llm, [("r", PROMPT, params)])["r"] + final = history[-1] + assert final.finished and final.finish_reason == "length" + assert_valid_tokens(final.token_ids, tokens) + assert all(not o.finished for o in history[:-1]) + + # Proof of speculation: one snapshot per delivered round (see the module docstring). + lengths = [len(o.token_ids) for o in history] + deltas = [b - a for a, b in zip([0] + lengths, lengths)] + assert deltas[0] == 1, "the bootstrap CODA step commits exactly the first token" + assert all(0 <= d <= K + 1 for d in deltas[1:]) + assert sum(deltas[1:]) == tokens - 1 + rounds = sum(d > 0 for d in deltas[1:]) + assert rounds >= math.ceil((tokens - 1) / (K + 1)) >= 1 + # Only a delivered speculative round that accepted a draft can commit more than one token. + assert max(deltas[1:]) > 1 + # Every snapshot extends the previous one: rounds only append. + assert all( + a.token_ids == b.token_ids[: len(a.token_ids)] for a, b in zip(history, history[1:]) + ) + + # Secondary evidence that sampling was random rather than silently greedy. + greedy = stream(llm, [("g", PROMPT, SamplingParams(max_tokens=tokens, ignore_eos=True))]) + assert greedy["g"][-1].token_ids != final.token_ids + + # Close while a random request is live, with rounds possibly still in flight. + client.add_request("close-live", PROMPT, random_params(max_tokens=60, seed=5)) + assert_closed_cleanly(llm.engine) + + +# 3. Same-seed replay ----------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_same_seed_replays_identically_through_fresh_llms(): + def run(seed): + with make_llm() as llm: + output = llm.generate([PROMPT], random_params(seed=seed))[0] + assert_valid_tokens(output.token_ids, MAX_TOKENS) + return output.token_ids + + first = run(11) + assert run(11) == first + # A different seed must change the stream, so equality above is not vacuous. + assert run(12) != first + + +# 4. Greedy and random requests together ----------------------------------------------------------- + + +@pytest.mark.gpu +def test_concurrent_greedy_and_random_requests_complete_and_greedy_is_unchanged(): + greedy = SamplingParams(max_tokens=9, ignore_eos=True) + sampled = random_params(max_tokens=13, seed=3) + greedy_prompt, random_prompt = [2, 3, 4, 5], [7, 8] + with make_llm() as llm: + alone = llm.generate([greedy_prompt], greedy)[0] + assert_valid_tokens(alone.token_ids, 9) + mixed = llm.generate([greedy_prompt, random_prompt], [greedy, sampled]) + assert [o.finished for o in mixed] == [True, True] + assert_valid_tokens(mixed[0].token_ids, 9) + assert_valid_tokens(mixed[1].token_ids, 13) + # FP32 tiny model: the argmax of a greedy request does not depend on batch composition + # (the repo's async-vs-sync tests already rely on this across batches). + assert mixed[0].token_ids == alone.token_ids + assert_closed_cleanly(llm.engine) diff --git a/tests/test_async_speculative_lifecycle.py b/tests/test_async_speculative_lifecycle.py index b3d469f7..bba1076b 100644 --- a/tests/test_async_speculative_lifecycle.py +++ b/tests/test_async_speculative_lifecycle.py @@ -393,3 +393,109 @@ def test_ticket_result_format_does_not_depend_on_statistics(mode, with_stats): assert stats.drafted_tokens == stats.verified_rows == int(mode == "pipeline") ticket.retire() assert not buffer.leased + + +def host_ticket(request, counts, host, *, error_flags): + batch = SchedulerOutput( + RequestPhase.SPECULATIVE, + [ScheduledRequest(request, token_start=1, token_count=c) for c in counts], + ) + buffer = SimpleNamespace( + host_result=torch.tensor(host), copy_done=Completion(), leased=True, keepalive=[] + ) + ticket = SpeculativeRoundTicket(batch, buffer, (), (), sum(counts), error_flags=error_flags) + return ticket, buffer + + +def test_random_ticket_decodes_explicit_ranges_and_flags(): + from vllm_rlt.request import Request + + request = Request("r", [1], SamplingParams(max_tokens=8)) + # rows (3 + 2) | accepted | next positions | error flags | stale tail of an older round + host = [7, 8, 9, 5, 6, 1, 0, 11, 22, 0, 1, 999, 999] + ticket, buffer = host_ticket(request, [3, 2], host, error_flags=True) + result = ticket.collect() + assert [r.token_ids for r in result.speculative] == [[7, 8], [5]] + assert [r.accepted_count for r in result.speculative] == [1, 0] + assert [r.numerical_error for r in result.speculative] == [False, True] + # Frontiers are the explicit [rows + b, rows + 2b) range, not the remaining tail. + assert result.frontiers == [11, 22] + assert buffer.copy_done.waits == 1 + + +def test_greedy_ticket_ignores_stale_flag_region_of_a_reused_buffer(): + from vllm_rlt.request import Request + + request = Request("r", [1], SamplingParams(max_tokens=8)) + # A previous random round left nonzero flags where a greedy round's tail now lies. + host = [7, 8, 9, 5, 6, 1, 0, 11, 22, 1, 1, 1] + ticket, _ = host_ticket(request, [3, 2], host, error_flags=False) + result = ticket.collect() + assert [r.numerical_error for r in result.speculative] == [False, False] + assert result.frontiers == [11, 22] + assert [r.token_ids for r in result.speculative] == [[7, 8], [5]] + + +def test_round_buffer_reserves_room_for_error_flags(monkeypatch): + from vllm_rlt.worker.async_speculative import SpeculativeRoundBuffer + + empty = torch.empty + monkeypatch.setattr(torch, "empty", lambda *a, pin_memory=False, **k: empty(*a, **k)) + runner = SimpleNamespace( + cache=SimpleNamespace(block_size=4, max_loops=2, device=torch.device("cpu")), + model=SimpleNamespace(config=SimpleNamespace(max_position_embeddings=16)), + ) + for tokens, seqs, expected in ((64, 4, 76), (8, 4, 24)): + scheduler = SimpleNamespace(max_num_batched_tokens=tokens, max_num_seqs=seqs) + buffer = SpeculativeRoundBuffer(runner, scheduler) + assert buffer.host_result.numel() == buffer.result.numel() == expected + + +def flag_round(ticket, flag): + """Rewrite a prepared one-request ticket into a random round with an error flag.""" + ticket.error_flags = True + ticket.buffer.host_result = torch.cat([ticket.buffer.host_result, torch.tensor([flag])]) + + +@torch.inference_mode() +def test_engine_ignores_flags_of_stale_rows_and_keeps_the_live_request(): + from vllm_rlt.request import Request + + engine, tickets, _ = prepared_engine() + live = engine.scheduler.requests["r"] + # Round 0 belongs to an older request that reused the ID "r" (cancelled and replaced). + stale = Request("r", [2], live.sampling_params) + tickets[0].batch.requests[0] = ScheduledRequest(stale, token_start=1, token_count=3) + flag_round(tickets[0], 1) + # Round 1 is the live request's first round: one accepted draft, frontier 1 + 2 tokens. + tickets[1].rounds = (0,) + tickets[1].buffer.host_result = torch.tensor([13, 14, 1, 3]) + flag_round(tickets[1], 0) + generated = len(live.generated_token_ids) + + assert engine._collect_speculative() == [] # stale: ignored without error + assert engine.speculative_runner.discarded_rounds == 1 + assert len(live.generated_token_ids) == generated + outputs = engine._collect_speculative() + assert [o.request_id for o in outputs] == ["r"] + assert live.generated_token_ids[generated:] == [13, 14] + assert live.finish_reason is None + engine.close() + + +@torch.inference_mode() +def test_engine_raises_on_a_live_flag_before_committing_and_cleans_up(): + engine, tickets, _ = prepared_engine() + runner = engine.speculative_runner + live = engine.scheduler.requests["r"] + generated = list(live.generated_token_ids) + flag_round(tickets[0], 1) + with pytest.raises(RuntimeError, match="invalid speculative sampling distribution"): + engine.step() + # The numeric check precedes _update, so nothing from the failed round was committed. + assert live.generated_token_ids == generated + assert runner.synchronizations == 1 + assert not engine._pending_speculative and not engine._retiring_speculative + assert not engine.scheduler.requests and engine.cache_manager.num_used_blocks == 0 + assert all(not t.buffer.leased and not t.buffer.keepalive for t in tickets) + engine.close() diff --git a/tests/test_async_speculative_random.py b/tests/test_async_speculative_random.py new file mode 100644 index 00000000..3d2aba9d --- /dev/null +++ b/tests/test_async_speculative_random.py @@ -0,0 +1,650 @@ +"""CUDA integration contracts for random sampling in async block speculation (M7 T3'). + +Every test needs a reserved device (``--run-gpu``): the async speculative runner allocates CUDA +streams and events. The tests drive the engine/runner directly (not the public entrypoint) and +observe the runner through thin pass-through wrappers (``RoundRecorder``, ``NanInjector``, +``spy_tickets``) that never change behaviour, except ``NanInjector`` which deliberately corrupts +one draft distribution. + +Reproducibility scope (plan section 3.2): outputs are only compared on the *same* execution path +(same engine configuration, same dtype, same graph mode) and, where noted, with identical request +shapes. Graph and eager outputs are never compared token by token. +""" + +from contextlib import closing +from dataclasses import dataclass + +import pytest +import torch + +import vllm_rlt.worker.async_speculative as async_speculative +from tests.helpers import make_speculative_engine, tiny_ouro_model +from tests.test_async_speculative import drain, make_engine +from vllm_rlt import SamplingParams + +PROMPT = (2, 3, 4) +K = 3 # make_engine's default num_speculative_tokens +TEMPERATURE = 1.5 +DTYPES = [torch.float32, torch.bfloat16] + + +# --------------------------------------------------------------------------- helpers + + +def random_params(seed, max_tokens, temperature=TEMPERATURE, **kwargs): + return SamplingParams( + max_tokens=max_tokens, temperature=temperature, seed=seed, ignore_eos=True, **kwargs + ) + + +def greedy_params(max_tokens): + return SamplingParams(max_tokens=max_tokens, ignore_eos=True) + + +def bootstrap_random(engine, rid, params, prompt=PROMPT): + """Run steps until the bootstrap (ordinary CODA) token of ``rid`` has been delivered. + + Same loop as ``tests.helpers.bootstrap_speculation``, which only builds greedy requests. + """ + engine.add_request(rid, list(prompt), params) + for _ in range(30): + if engine.step(): + request = engine.scheduler.requests[rid] + assert len(request.generated_token_ids) == 1 + return request + pytest.fail("bootstrap did not produce a token") + + +def assert_reclaimed(engine): + """No request, state, leased buffer, slot, block or generator survives a drained engine.""" + runner = engine.speculative_runner + assert not engine.scheduler.requests + assert not engine._pending_speculative and not engine._retiring_speculative + assert engine.cache_manager.num_used_blocks == 0 + assert not runner.states + assert all(not buffer.leased for buffer in runner.buffers) + assert len(runner.free_slots) == len(runner.slot_values) + assert not engine.rng.generators + + +@dataclass +class RoundRecord: + """One submitted block round: host item list plus device tensors cloned at accept time.""" + + items: list # [(request_id, token_count)] in item order + tokens: torch.Tensor # [n, width] input token then drafts (padding zeros) + target: torch.Tensor # [count] target argmax per request-major verify row + random: bool # the round had a random plan + + def view(self, request_id): + """Return (input token + drafts, target argmax rows) of one request in this round.""" + tokens = self.tokens.cpu() + target = self.target.cpu().tolist() + start = 0 + for row, (rid, count) in enumerate(self.items): + if rid == request_id: + return tokens[row, :count].tolist(), target[start : start + count] + start += count + raise KeyError(request_id) + + +class RoundRecorder: + """Record every submitted block round without altering it. + + Wraps ``_round_plan`` (host only) and ``_accept_and_update_state`` (clones its device inputs + on the compute stream; no host read happens during submission). Read ``rounds`` only after + ``runner.synchronize()``. + """ + + def __init__(self, monkeypatch, runner): + self.rounds = [] + self._items = None + round_plan, accept = runner._round_plan, runner._accept_and_update_state + + def recording_round_plan(items, states, buffer): + self._items = [(item.request.request_id, item.token_count) for item in items] + return round_plan(items, states, buffer) + + def recording_accept(plan, states, starts, tokens, target, buffer, **kwargs): + self.rounds.append( + RoundRecord( + self._items, + tokens.detach().clone(), + target.detach().clone(), + plan.random is not None, + ) + ) + return accept(plan, states, starts, tokens, target, buffer, **kwargs) + + monkeypatch.setattr(runner, "_round_plan", recording_round_plan) + monkeypatch.setattr(runner, "_accept_and_update_state", recording_accept) + + def counts(self): + """Per round: {request_id: token_count}.""" + return [dict(record.items) for record in self.rounds] + + +class NanInjector: + """Make the first draft distribution of one speculative round NaN (offset 0 of the draft). + + ``_draft`` computes the draft distribution of every offset with the module-level + ``batched_probabilities``; the verify distribution uses it too, so calls are attributed to + the draft only while a (wrapped) ``_draft`` is executing. Offset 0 always drafts when a round + has more than one row, and an invalid ``q`` at the first draft position is a deterministic + stop-point error (plan 3.3/3.4). ``round_index`` counts ``_draft`` calls from 1. + """ + + def __init__(self, monkeypatch, runner, round_index): + self.round_index = round_index + self.rounds = 0 + self.fired = 0 + self._in_draft = False + self._draft_calls = 0 + draft, probabilities = runner._draft, async_speculative.batched_probabilities + + def wrapped_draft(*args, **kwargs): + self.rounds += 1 + self._in_draft, self._draft_calls = True, 0 + try: + return draft(*args, **kwargs) + finally: + self._in_draft = False + + def wrapped_probabilities(*args, **kwargs): + result = probabilities(*args, **kwargs) + if self._in_draft: + self._draft_calls += 1 + if self.rounds == self.round_index and self._draft_calls == 1: + self.fired += 1 + return torch.full_like(result, float("nan")) + return result + + monkeypatch.setattr(runner, "_draft", wrapped_draft) + monkeypatch.setattr(async_speculative, "batched_probabilities", wrapped_probabilities) + + +def spy_tickets(monkeypatch, runner): + """Collect every speculative ticket the runner returns.""" + tickets = [] + submit = runner.submit + + def spying_submit(batch): + ticket = submit(batch) + tickets.append(ticket) + return ticket + + monkeypatch.setattr(runner, "submit", spying_submit) + return tickets + + +def spy_acquire(monkeypatch, rng): + """Record (request_id, generator) for every ``RngRegistry.acquire`` and pass it through.""" + calls = [] + acquire = rng.acquire + + def spying_acquire(request_id, params, device): + generator = acquire(request_id, params, device) + calls.append((request_id, generator)) + return generator + + monkeypatch.setattr(rng, "acquire", spying_acquire) + return calls + + +def accepted_non_argmax(record, request_id, result): + """Check the CPU result against the recorded GPU round; count accepted non-argmax drafts.""" + row, target = record.view(request_id) + drafts = row[1:] + accepted = result.accepted_count + assert not result.numerical_error + assert 0 <= accepted <= len(drafts) + # Section 3.1: the committed prefix is the accepted drafts, then one correction/bonus token. + assert len(result.token_ids) == accepted + 1 + assert result.token_ids[:accepted] == drafts[:accepted] + return sum(1 for j in range(accepted) if drafts[j] != target[j]) + + +def assert_kv_matches_serial_replay(engine, model, oracle, rid, tokens, dtype): + """Per-depth KV of every position in ``tokens`` equals a serial replay (copied oracle).""" + assert oracle.allocate("serial", len(tokens)) + try: + with torch.inference_mode(): + for position, token in enumerate(tokens): + hidden = model.prelude(torch.tensor([token], device="cuda")) + for depth in range(model.config.total_ut_steps): + hidden, _ = model.recurrent(hidden, ["serial"], [depth], [position], oracle) + allocation = engine.cache_manager._get_allocation(rid) + pos = torch.arange(len(tokens), device="cuda") + atol, rtol = (4e-5, 4e-5) if dtype == torch.float32 else (0.08, 0.04) + for depth in range(model.config.total_ut_steps): + pages = torch.tensor(allocation.block_tables[depth], device="cuda") + blocks, offsets = pages[pos // 2], pos % 2 + for layer in range(model.config.num_hidden_layers): + expected_k, expected_v = oracle.read(layer, "serial", depth, len(tokens)) + torch.testing.assert_close( + engine.cache_manager.key_cache[blocks, layer, offsets], + expected_k, + atol=atol, + rtol=rtol, + ) + torch.testing.assert_close( + engine.cache_manager.value_cache[blocks, layer, offsets], + expected_v, + atol=atol, + rtol=rtol, + ) + finally: + oracle.free("serial") + + +# ----------------------------------------------------------------- 1. key regression (3.1) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("graphs", [False, True]) +def test_accepted_non_argmax_drafts_keep_cpu_sequence_gpu_input_and_kv_consistent( + monkeypatch, dtype, graphs +): + """Plan 3.1: a random round commits accepted drafts + correction/bonus, not target argmax. + + Two rounds are in flight. Per round the recorded GPU draft tokens must equal the CPU-committed + prefix, the second round's GPU input token must be the first round's last committed token, the + slot table must hold the final (token, position), and per-depth KV of every committed + position must match a serial replay of the CPU-committed tokens. Over all seeds at least one + accepted draft must differ from the target argmax, so the test cannot pass vacuously. + Several seeds share one engine (abort + drain between them) to make that overwhelmingly likely. + """ + model = tiny_ouro_model("cuda", dtype) + oracle = make_speculative_engine(model, asynchronous=False).cache_manager + non_argmax = 0 + with closing(make_engine(model, graphs=graphs)) as e: + runner = e.speculative_runner + recorder = RoundRecorder(monkeypatch, runner) + for seed in (31, 32, 33): + rid = f"r{seed}" + request = bootstrap_random(e, rid, random_params(seed, 24, temperature=2.0)) + start = request.position + first = len(recorder.rounds) + e.step() + e.step() + assert len(e._pending_speculative) == 2 + assert runner.submitted_before_collect >= 1 + runner.synchronize() + records = recorder.rounds[first:] + assert len(records) == 2 and all(r.random for r in records) + assert [r.items for r in records] == [[(rid, count)] for count in (K + 1, K + 1)] + results = [ticket.collect().speculative[0] for ticket in e._pending_speculative] + for record, result in zip(records, results): + non_argmax += accepted_non_argmax(record, rid, result) + committed = [token for result in results for token in result.token_ids] + # The next round's GPU input token is the previous round's last committed token. + assert records[1].view(rid)[0][0] == results[0].token_ids[-1] + tokens = request.prompt_token_ids + request.generated_token_ids + committed + assert len(request.generated_token_ids) == 1 + # Slot table: (last committed token, its position); that token has no KV yet. + state = runner.states[rid] + assert int(state.token) == tokens[-1] + assert int(state.position) == start + sum(len(r.token_ids) for r in results) + assert int(state.position) == len(tokens) - 1 + assert_kv_matches_serial_replay(e, model, oracle, rid, tokens[:-1], dtype) + e.abort_request(rid) + drain(e) + assert e.cache_manager.num_used_blocks == 0 + assert_reclaimed(e) + assert non_argmax > 0, "no accepted draft differed from the target argmax: vacuous test" + + +# --------------------------------------------------------------- 2. two rounds in flight + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("graphs", [False, True]) +def test_two_inflight_random_rounds_finish_with_exact_lengths_and_reclaim(dtype, graphs): + with closing(make_engine(tiny_ouro_model("cuda", dtype), graphs=graphs)) as e: + runner = e.speculative_runner + bootstrap_random(e, "r", random_params(1, 20)) + e.step() + e.step() + assert len(e._pending_speculative) == 2 + assert all(ticket.error_flags for ticket in e._pending_speculative) + assert runner.submitted_before_collect == 1 + # A second random request (with top-k/top-p truncation) joins while two rounds fly. + e.add_request("s", list(PROMPT), random_params(2, 11, top_k=16, top_p=0.9)) + outputs = drain(e) + assert runner.submitted_before_collect > 0 + assert {rid: len(o.token_ids) for rid, o in outputs.items()} == {"r": 20, "s": 11} + assert all(o.finished and o.finish_reason == "length" for o in outputs.values()) + vocab = e.model.config.vocab_size + assert all(0 <= t < vocab for o in outputs.values() for t in o.token_ids) + assert_reclaimed(e) + + +# ------------------------------------------------------------------------ 3. mixed batch + + +@pytest.mark.gpu +def test_greedy_request_in_a_mixed_batch_is_unaffected_and_never_acquires_a_generator( + monkeypatch, +): + """A greedy request next to a random one matches its solo run and never touches the RNG. + + FP32 only. The solo run and the mixed runs have different batch shapes, so equality relies on + FP32 row-wise numerics being far from any argmax tie on this model; the second mixed run + (different seed for the random neighbour, same shapes up to acceptance) is the weaker check. + """ + model = tiny_ouro_model("cuda", torch.float32) + + def run(neighbour_seed): + with closing(make_engine(model)) as e: + recorder = RoundRecorder(monkeypatch, e.speculative_runner) + acquired = spy_acquire(monkeypatch, e.rng) + e.add_request("g", list(PROMPT), greedy_params(16)) + if neighbour_seed is not None: + e.add_request("x", list(PROMPT), random_params(neighbour_seed, 16)) + outputs = drain(e) + assert outputs["g"].finish_reason == "length" and len(outputs["g"].token_ids) == 16 + ids = [rid for rid, _ in acquired] + assert "g" not in ids + if neighbour_seed is None: + assert not ids and not any(r.random for r in recorder.rounds) + else: + assert "x" in ids and len(outputs["x"].token_ids) == 16 + # Greedy and random requests really shared speculative batches. + assert any( + r.random and {rid for rid, _ in r.items} == {"g", "x"} for r in recorder.rounds + ) + assert_reclaimed(e) + return {rid: o.token_ids for rid, o in outputs.items()} + + alone = run(None) + mixed = run(3) + other = run(4) + assert mixed["g"] == alone["g"] + assert other["g"] == alone["g"] + assert mixed["x"] != other["x"] # the neighbour really changed + + +# ------------------------------------------------------------------- 4. RNG lifecycle + + +@pytest.mark.gpu +def test_speculative_rounds_consume_the_bootstrap_generator_and_reused_ids_start_fresh( + monkeypatch, +): + with closing(make_engine(tiny_ouro_model("cuda"), k=K)) as e: + bootstrap_random(e, "r", random_params(21, 24)) + generator = e.rng.peek("r") + assert generator is not None # the bootstrap CODA created it + initial = generator.get_state() + acquired = spy_acquire(monkeypatch, e.rng) + # A probe replays exactly one rand((2K + 1,)) per round from the same state. + probe = torch.Generator(device=generator.device) + probe.set_state(initial) + previous = initial + for submitted in (1, 2): + e.step() + assert len(acquired) == submitted + assert acquired[-1][0] == "r" and acquired[-1][1] is generator + torch.rand((2 * K + 1,), generator=probe, device=probe.device, dtype=torch.float32) + state = generator.get_state() + assert not torch.equal(state, previous), "generator did not advance" + assert torch.equal(state, probe.get_state()), "not exactly one draw per round" + previous = state + assert e.rng.peek("r") is generator + torch.cuda.synchronize() + + # Abort drops the entry once the request retires. + assert e.abort_request("r").finish_reason == "abort" + drain(e) + assert e.rng.peek("r") is None and not e.rng.generators + assert_reclaimed(e) + + # A reused ID gets a new generator that starts from the same seed and bootstrap draws. + bootstrap_random(e, "r", random_params(21, 24)) + fresh = e.rng.peek("r") + assert fresh is not None and fresh is not generator + assert torch.equal(fresh.get_state(), initial) + outputs = drain(e) + assert len(outputs["r"].token_ids) == 24 + # Natural completion drops it as well. + assert e.rng.peek("r") is None and not e.rng.generators + assert_reclaimed(e) + + +# ------------------------------------------------------------ 5. same-path fixed-seed replay + + +@pytest.mark.gpu +@pytest.mark.parametrize("graphs", [False, True]) +def test_fixed_seed_single_request_replays_and_a_new_seed_changes_the_output(graphs): + """Single request only, so the batch plan is fixed (reproducibility scope, plan 3.2).""" + model = tiny_ouro_model("cuda", torch.float32) + + def generate(engine, seed): + engine.add_request("r", list(PROMPT), random_params(seed, 32)) + output = drain(engine)["r"] + assert len(output.token_ids) == 32 and output.finish_reason == "length" + return output.token_ids + + with closing(make_engine(model, graphs=graphs)) as first: + expected = generate(first, 7) + assert_reclaimed(first) + with closing(make_engine(model, graphs=graphs)) as second: + assert generate(second, 7) == expected + assert generate(second, 8) != expected + assert_reclaimed(second) + + +# ----------------------------------------------------------------- 6. request isolation + + +@pytest.mark.gpu +def test_changing_one_requests_seed_leaves_its_neighbour_unchanged(): + """Request A is unchanged when only request B's seed changes. + + Assumptions: FP32; ``ignore_eos=True``; two requests of four rows each fit the default + 16-token budget, so both get full rounds in lockstep (the same plan except for B's tail + rounds, whose length depends on B's acceptances). Row numerics are batch-shape independent + to FP32 accuracy, far below any sampling threshold. + """ + model = tiny_ouro_model("cuda", torch.float32) + + def run(seed_b): + with closing(make_engine(model)) as e: + e.add_request("a", list(PROMPT), random_params(11, 21)) + e.add_request("b", list(PROMPT), random_params(seed_b, 17)) + outputs = drain(e) + assert_reclaimed(e) + return {rid: o.token_ids for rid, o in outputs.items()} + + first, second = run(22), run(33) + assert len(first["a"]) == len(second["a"]) == 21 + assert len(first["b"]) == len(second["b"]) == 17 + assert first["a"] == second["a"] + assert first["b"] != second["b"] + + +# --------------------------------------------------------------------- 7. ragged / K = 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype,graphs", [(torch.float32, False), (torch.bfloat16, True)]) +def test_ragged_and_zero_draft_rounds_complete_with_exact_lengths(monkeypatch, dtype, graphs): + """Request lengths 2, 6 and 22 force K = 0 rounds next to full rounds. + + ``a`` (max_tokens=2) has exactly one round with one row (K=0). ``b`` (6) has a full round of + four rows and, scheduled while that is in flight, a one-row round (remaining budget 1). + ``c`` (22) keeps full rounds going, so rounds mixing K=0 and K>0 requests occur. All three + use identical prompts and are admitted together, so they advance in lockstep. + """ + with closing(make_engine(tiny_ouro_model("cuda", dtype), graphs=graphs)) as e: + recorder = RoundRecorder(monkeypatch, e.speculative_runner) + lengths = {"a": 2, "b": 6, "c": 22} + for i, (rid, length) in enumerate(lengths.items()): + e.add_request(rid, list(PROMPT), random_params(10 + i, length)) + outputs = drain(e) + assert {rid: len(o.token_ids) for rid, o in outputs.items()} == lengths + assert all(o.finish_reason == "length" for o in outputs.values()) + counts = recorder.counts() + assert all(record.random for record in recorder.rounds) + assert [c["a"] for c in counts if "a" in c] == [1] + b_counts = [c["b"] for c in counts if "b" in c] + assert b_counts[0] == K + 1 and 1 in b_counts + assert max(max(c.values()) for c in counts) == K + 1 + assert any(min(c.values()) == 1 and max(c.values()) > 1 for c in counts) + assert_reclaimed(e) + + # Every random request of a round has K = 0 (no q buffer): one-row requests only. + seen = len(recorder.rounds) + for rid in ("d", "e"): + e.add_request(rid, list(PROMPT), random_params(40 + ord(rid), 2)) + outputs = drain(e) + assert {rid: len(o.token_ids) for rid, o in outputs.items()} == {"d": 2, "e": 2} + assert recorder.rounds[seen:] + assert all(count == 1 for record in recorder.rounds[seen:] for _, count in record.items) + assert_reclaimed(e) + + +# --------------------------------------------------------------------------- 8. graphs + + +@pytest.mark.gpu +def test_random_and_greedy_requests_replay_graphs_and_produce_valid_tokens(monkeypatch): + """Graph path: replays really happen; outputs are only checked for shape and validity.""" + model = tiny_ouro_model("cuda", torch.bfloat16) + with closing(make_engine(model, graphs=True)) as e: + runner = e.speculative_runner + recorder = RoundRecorder(monkeypatch, runner) + core_before = runner.graphs.replays + coda_before = runner.coda_graphs.replays + e.add_request("x", list(PROMPT), random_params(51, 19)) + e.add_request("y", list(PROMPT), random_params(52, 13, top_k=16, top_p=0.9)) + e.add_request("g", list(PROMPT), greedy_params(10)) + outputs = drain(e) + assert runner.graphs.replays > core_before + assert runner.coda_graphs.replays > coda_before + assert any(record.random for record in recorder.rounds) + assert {rid: len(o.token_ids) for rid, o in outputs.items()} == {"x": 19, "y": 13, "g": 10} + vocab = e.model.config.vocab_size + assert all(0 <= t < vocab for o in outputs.values() for t in o.token_ids) + assert all(o.finish_reason == "length" for o in outputs.values()) + assert_reclaimed(e) + + +# ------------------------------------------------------------------ 9. error handling (3.4) + + +@pytest.mark.gpu +@pytest.mark.parametrize("graphs", [False, True]) +def test_invalid_distribution_in_an_inflight_round_is_fatal_and_the_engine_recovers( + monkeypatch, graphs +): + with closing(make_engine(tiny_ouro_model("cuda"), graphs=graphs)) as e: + runner = e.speculative_runner + injector = NanInjector(monkeypatch, runner, round_index=1) + bootstrap_random(e, "r", random_params(41, 24)) + e.step() + e.step() + # Round 2 was submitted behind the corrupted round 1 (failed rows keep their slot). + assert len(e._pending_speculative) == 2 + assert runner.submitted_before_collect == 1 + with pytest.raises(RuntimeError, match="invalid speculative sampling distribution"): + e.step() + assert injector.fired == 1 + assert not e.has_unfinished_requests() + assert_reclaimed(e) + + e.add_request("r", list(PROMPT), random_params(42, 12)) + output = drain(e)["r"] + assert len(output.token_ids) == 12 and output.finish_reason == "length" + assert injector.fired == 1 + assert_reclaimed(e) + + +@pytest.mark.gpu +def test_stale_error_flag_of_an_aborted_request_is_ignored_when_the_id_is_reused(monkeypatch): + with closing(make_engine(tiny_ouro_model("cuda"))) as e: + runner = e.speculative_runner + injector = NanInjector(monkeypatch, runner, round_index=1) + bootstrap_random(e, "r", random_params(51, 24)) + e.step() + e.step() + assert len(e._pending_speculative) == 2 + runner.synchronize() + assert injector.fired == 1 + # The older ticket does carry the error; it must be dropped with the request. + assert e._pending_speculative[0].collect().speculative[0].numerical_error + e.abort_request("r") + e.add_request("r", list(PROMPT), random_params(52, 14)) + output = drain(e)["r"] # must not raise + assert len(output.token_ids) == 14 and output.finish_reason == "length" + assert runner.discarded_rounds >= 1 + assert_reclaimed(e) + + +@pytest.mark.gpu +def test_greedy_rounds_reusing_random_round_buffers_never_report_errors(monkeypatch): + with closing(make_engine(tiny_ouro_model("cuda"))) as e: + runner = e.speculative_runner + tickets = spy_tickets(monkeypatch, runner) + e.add_request("x", list(PROMPT), random_params(61, 16)) + assert len(drain(e)["x"].token_ids) == 16 + assert tickets and all(ticket.error_flags for ticket in tickets) + assert_reclaimed(e) + + # Fill the tail of both result buffers with "error" flags a buggy reader would see. + runner.synchronize() + torch.cuda.synchronize() + for buffer in runner.buffers: + buffer.host_result.fill_(1) + buffer.result.fill_(1) + torch.cuda.synchronize() + + seen = len(tickets) + e.add_request("g", list(PROMPT), greedy_params(16)) + output = drain(e)["g"] + assert len(output.token_ids) == 16 and output.finish_reason == "length" + greedy = tickets[seen:] + assert greedy and all(not ticket.error_flags for ticket in greedy) + assert all( + ticket.cached is not None + and not any(result.numerical_error for result in ticket.cached.speculative) + for ticket in greedy + ) + assert_reclaimed(e) + + +# ------------------------------------------------------------- 10. greedy path untouched + + +@pytest.mark.gpu +@pytest.mark.parametrize("graphs", [False, True]) +def test_all_greedy_async_batches_never_plan_sample_or_acquire_a_generator(monkeypatch, graphs): + with closing(make_engine(tiny_ouro_model("cuda"), graphs=graphs)) as e: + runner = e.speculative_runner + + def forbidden(*args, **kwargs): + raise AssertionError("an all-greedy speculative round used the random path") + + monkeypatch.setattr(e.rng, "acquire", forbidden) + monkeypatch.setattr(runner, "_uniforms", forbidden) + plans = [] + plan_random_sampling = async_speculative.plan_random_sampling + + def spying_plan(*args, **kwargs): + plans.append(plan_random_sampling(*args, **kwargs)) + return plans[-1] + + monkeypatch.setattr(async_speculative, "plan_random_sampling", spying_plan) + tickets = spy_tickets(monkeypatch, runner) + lengths = {"a": 13, "b": 7, "c": 1} + for i, (rid, length) in enumerate(lengths.items()): + e.add_request(rid, [2 + i, 3, 4], greedy_params(length)) + outputs = drain(e) + assert {rid: len(o.token_ids) for rid, o in outputs.items()} == lengths + assert plans and all(plan is None for plan in plans) + assert tickets and all(not ticket.error_flags for ticket in tickets) + assert e.rng.peek("a") is None + assert_reclaimed(e) diff --git a/tests/test_async_speculative_sampling_plan.py b/tests/test_async_speculative_sampling_plan.py new file mode 100644 index 00000000..002aa439 --- /dev/null +++ b/tests/test_async_speculative_sampling_plan.py @@ -0,0 +1,397 @@ +"""CPU contracts for random sampling in async speculative rounds. + +Covers the pure host plan, the pure tensor merge, and (with stubbed kernels, model and +coda) the runner's draft/verify/accept wiring that otherwise needs CUDA streams. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_rlt import SamplingParams, SpeculativeConfig +from vllm_rlt.core.scheduler import ScheduledRequest +from vllm_rlt.request import Request +from vllm_rlt.worker.async_speculative import ( + AsyncSpeculativeRunner, + merge_random_round, + plan_random_sampling, +) +from vllm_rlt.worker.sampler import RngRegistry +from vllm_rlt.worker.speculative_sampling import ( + AcceptResult, + batched_probabilities, + categorical, + speculative_accept, +) + +VOCAB = 100 + + +def params(temperature=0.0, top_k=-1, top_p=1.0, seed=0): + return SamplingParams(temperature=temperature, top_k=top_k, top_p=top_p, seed=seed) + + +def ints(array): + return np.asarray(array).tolist() + + +def test_mixed_ragged_plan_has_exact_indices_and_parameters(): + plan = plan_random_sampling( + [ + params(0.5, top_k=5, top_p=0.9), + params(), + params(1.0, top_k=VOCAB), + params(2.0, top_k=7), + ], + [3, 1, 2, 3], + VOCAB, + 4, + ) + assert ints(plan.requests) == [0, 2, 3] + assert ints(plan.temperature) == [0.5, 1.0, 2.0] + assert ints(plan.top_k) == [5, VOCAB, 7] + assert ints(plan.top_p) == [0.9, 1.0, 1.0] + assert plan.max_top_k == 7 and plan.apply_top_p + assert plan.width == 3 + # Offset 0 drafts items 0, 2, 3; offset 1 drafts items 0 and 3 (rows 0 and 1 of its block). + assert [(ints(r), ints(k)) for r, k in plan.draft_rows] == [ + ([0, 1, 2], [0, 1, 2]), + ([0, 1], [0, 2]), + ] + # Request starts are [0, 3, 4, 6]; greedy item 1 owns verify row 3. + assert ints(plan.verify_rows) == [0, 1, 2, 4, 5, 6, 7, 8] + assert ints(plan.verify_ranks) == [0, 0, 0, 1, 1, 2, 2, 2] + assert ints(plan.verify_offsets) == [0, 1, 2, 0, 1, 0, 1, 2] + assert ints(plan.verify_temperature) == [0.5] * 3 + [1.0] * 2 + [2.0] * 3 + assert ints(plan.verify_top_k) == [5] * 3 + [VOCAB] * 2 + [7] * 3 + assert ints(plan.verify_top_p) == [0.9] * 3 + [1.0] * 5 + + +def test_greedy_rows_between_random_drafts_keep_block_positions(): + # Item 1 (greedy) drafts at offset 0 too, so random item 2 is row 1 of that block. + plan = plan_random_sampling([params(), params(), params(1.0)], [2, 2, 2], VOCAB, 3) + assert ints(plan.requests) == [2] + assert [(ints(r), ints(k)) for r, k in plan.draft_rows] == [([2], [0])] + assert ints(plan.verify_rows) == [4, 5] + + +def test_all_greedy_has_no_random_plan(): + assert plan_random_sampling([params(), params()], [3, 2], VOCAB, 2) is None + + +def test_random_requests_without_drafts(): + plan = plan_random_sampling([params(1.0), params(), params(0.7)], [1, 3, 1], VOCAB, 2) + assert ints(plan.requests) == [0, 2] and plan.width == 1 + assert plan.draft_rows == () + assert ints(plan.verify_rows) == [0, 4] + assert ints(plan.verify_ranks) == [0, 1] and ints(plan.verify_offsets) == [0, 0] + names = [name for name, _ in plan.uploads()] + assert not any(name.startswith(("r_rows", "r_ranks")) for name in names) + + +def test_truncation_flags_come_from_effective_parameters(): + plan = plan_random_sampling( + [params(1.0), params(1.0, top_k=VOCAB), params(1.0, top_k=VOCAB + 50, top_p=1.0)], + [1, 1, 1], + VOCAB, + 1, + ) + assert plan.max_top_k == 0 and not plan.apply_top_p + plan = plan_random_sampling( + [params(1.0, top_k=VOCAB - 1), params(1.0, top_p=0.5)], [1, 1], 100, 1 + ) + assert plan.max_top_k == VOCAB - 1 and plan.apply_top_p + + +def test_random_request_exceeding_configured_drafts_is_rejected(): + with pytest.raises(ValueError, match="num_speculative_tokens"): + plan_random_sampling([params(1.0)], [4], VOCAB, 2) + # A greedy request is not bound to the uniform budget. + assert plan_random_sampling([params(), params(1.0)], [4, 2], VOCAB, 2) is not None + + +def test_uploads_are_int64_and_floats_round_trip_through_bit_casts(): + plan = plan_random_sampling([params(0.3, top_k=4, top_p=0.7)], [2], VOCAB, 2) + uploads = dict(plan.uploads()) + assert all(v.dtype == np.int64 for v in uploads.values()) + device = torch.from_numpy(np.concatenate(list(uploads.values()))) + names = list(uploads) + start = sum(len(uploads[k]) for k in names[: names.index("r_temperature")]) + decoded = device[start : start + 1].view(torch.float64).float() + assert decoded.tolist() == pytest.approx([0.3]) + assert decoded.dtype == torch.float32 + + +def test_merge_replaces_random_rows_and_keeps_greedy_and_failed_slots(): + # Items: 0 random (3 rows), 1 greedy (2 rows), 2 random (2 rows), failed. + target = torch.tensor([10, 11, 12, 20, 21, 30, 31]) + accepted = torch.tensor([0, 1, 0]) # greedy-computed values + next_tokens = torch.tensor([10, 21, 30]) + starts = torch.tensor([5, 8, 12]) + inputs = torch.tensor([1, 2, 3]) + result = AcceptResult( + # Request 0 accepts drafts 77 (not the argmax 11 of its row) and 66, then final 55. + tokens=torch.tensor([[77, 66, 55], [9, 0, 0]]), + accepted=torch.tensor([2, 0]), + final=torch.tensor([55, 0]), + error=torch.tensor([False, True]), + ) + merged = merge_random_round( + target, + accepted, + next_tokens, + starts, + inputs, + result, + torch.tensor([0, 2]), + torch.tensor([0, 1, 2, 5, 6]), + torch.tensor([0, 0, 0, 1, 1]), + torch.tensor([0, 1, 2, 0, 1]), + ) + assert merged.rows.tolist() == [77, 66, 55, 20, 21, 9, 0] + assert target.tolist() == [10, 11, 12, 20, 21, 30, 31] # inputs are not modified + assert merged.accepted.tolist() == [2, 1, 0] + assert merged.next_tokens.tolist() == [55, 21, 0] + assert merged.next_positions.tolist() == [8, 10, 13] + # The failed request keeps its old (input token, start); the others advance. + assert merged.next_values.tolist() == [[55, 8], [21, 10], [3, 12]] + assert merged.error.tolist() == [0, 0, 1] + assert merged.error.dtype == torch.int64 + + +# --- Runner wiring on CPU: stubbed kernels, model and CODA -------------------------------- + + +class FakeKernel: + def __init__(self, function): + self.function = function + + def __getitem__(self, grid): + return self.function + + +def fake_gather(source, index, out, n, **_): + out.copy_(source[index]) + + +def fake_scatter(values, index, out, **_): + out[index] = values + + +class Harness: + """Drive one round of the real draft/verify/accept methods with a toy Markov model. + + Draft logits are ``E[token] @ W`` and verification logits ``(E[token] @ M) @ W``, so q and + p differ and both depend on the drafted tokens. + """ + + K = 3 + HIDDEN = 8 + + def __init__(self, monkeypatch, seed): + from vllm_rlt.kernels import routing + + monkeypatch.setattr(routing, "gather_kernel", FakeKernel(fake_gather)) + monkeypatch.setattr(routing, "scatter_kernel", FakeKernel(fake_scatter)) + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda self: self) + g = torch.Generator().manual_seed(seed) + self.E = torch.randn(VOCAB, self.HIDDEN, generator=g) + eye = torch.eye(self.HIDDEN) + self.M = eye + 0.3 * torch.randn(self.HIDDEN, self.HIDDEN, generator=g) + self.W = 0.5 * torch.randn(self.HIDDEN, VOCAB, generator=g) + runner = object.__new__(AsyncSpeculativeRunner) + runner.model = SimpleNamespace(config=SimpleNamespace(vocab_size=VOCAB)) + runner.config = SpeculativeConfig( + num_speculative_tokens=self.K, draft_loops=1, target_loops=2 + ) + runner.device = torch.device("cpu") + runner.rng = RngRegistry() + runner.slot_values = torch.zeros((8, 2), dtype=torch.long) + runner.runtime = SimpleNamespace( + forward=lambda batch: SimpleNamespace(hidden=self.E[batch.tokens]) + ) + runner._execute_recurrent_step = lambda hidden, positions, tables, depth, allocations: ( + hidden if depth < 1 else hidden @ self.M + ) + runner._coda = lambda hidden: hidden @ self.W + self.runner = runner + + def run(self, specs): + """specs: (SamplingParams, token_count, input_token, position) per item.""" + runner = self.runner + items = [] + for i, (sampling, count, token, position) in enumerate(specs): + request = Request(f"r{i}", [1], sampling) + items.append(ScheduledRequest(request, token_start=position, token_count=count)) + runner.slot_values[i] = torch.tensor([token, position]) + states = [SimpleNamespace(slot=i, submitted=0) for i in range(len(specs))] + buffer = SimpleNamespace(keepalive=[], result=torch.full((64,), -1, dtype=torch.long)) + plan = runner._round_plan(items, states, buffer) + u = None if plan.random is None else runner._uniforms(items, plan.random) + tables = torch.zeros((len(specs), 1, 1)) + starts, tokens, hidden_blocks, q = runner._draft(plan, tables, None, u) + target, p = runner._verify(plan, starts, hidden_blocks, tables, None) + count = runner._accept_and_update_state( + plan, states, starts, tokens, target, buffer, u=u, q=q, p=p + ) + return SimpleNamespace( + plan=plan, u=u, tokens=tokens, target=target, count=count, buffer=buffer, q=q, p=p + ) + + def q_row(self, token, sampling): + return self.row(self.E[token] @ self.W, sampling) + + def p_row(self, token, sampling): + return self.row(self.E[token] @ self.M @ self.W, sampling) + + @staticmethod + def row(logits, sampling): + k = sampling.top_k + return batched_probabilities( + logits[None], + torch.tensor([sampling.temperature]), + torch.tensor([k]), + torch.tensor([sampling.top_p]), + max_top_k=k if 0 < k < VOCAB else 0, + apply_top_p=sampling.top_p < 1, + )[0] + + +SPECS = [ + # (sampling, token_count, input token, position) + (params(0.8, top_k=20, top_p=0.95, seed=1), 4, 3, 10), + (params(), 3, 4, 20), + (params(1.3, seed=2), 2, 5, 30), + (params(0.6, top_p=0.9, seed=3), 1, 6, 40), + (params(1.0, top_k=5, seed=4), 4, 7, 50), +] + + +@pytest.mark.parametrize("seed", range(8)) +def test_runner_round_matches_per_request_reference(monkeypatch, seed): + h = Harness(monkeypatch, seed) + out = h.run(SPECS) + n, K = len(SPECS), h.K + result = out.buffer.result + accepted_all = result[out.count : out.count + n].tolist() + positions = result[out.count + n : out.count + 2 * n].tolist() + flags = result[out.count + 2 * n : out.count + 3 * n].tolist() + assert flags == [0] * n + slot = h.runner.slot_values + ranks = {i: r for r, i in enumerate(i for i, spec in enumerate(SPECS) if spec[0].temperature)} + start = 0 + for i, (sampling, count, token, position) in enumerate(SPECS): + drafts = count - 1 + row_tokens = out.tokens[i] + assert row_tokens[0].item() == token + if sampling.temperature == 0: + # Greedy: target argmax rows, accepted while the draft equals the argmax. + argmax = [(h.E[row_tokens[j]] @ h.M @ h.W).argmax().item() for j in range(count)] + assert out.target[start : start + count].tolist() == argmax + expected = 0 + while expected < drafts and row_tokens[expected + 1].item() == argmax[expected]: + expected += 1 + committed = argmax[: expected + 1] + else: + gen = torch.Generator().manual_seed(sampling.seed) + u = torch.rand((2 * K + 1,), generator=gen) + assert torch.equal(out.u[ranks[i]], u) + q = ( + torch.stack([h.q_row(row_tokens[j], sampling) for j in range(drafts)]) + if drafts + else None + ) + for j in range(drafts): + drawn, _ = categorical(q[j][None], u[j : j + 1]) + assert drawn.item() == row_tokens[j + 1].item() + p = torch.stack([h.p_row(row_tokens[j], sampling) for j in range(count)]) + res = speculative_accept( + row_tokens[None, 1:count], + torch.tensor([drafts]), + q[None] if drafts else p.new_empty((1, 0, VOCAB)), + p[None], + u[None, K : K + drafts], + u[2 * K : 2 * K + 1], + ) + assert not res.error.item() + expected = res.accepted.item() + committed = res.tokens[0, : expected + 1].tolist() + assert accepted_all[i] == expected + assert result[start : start + expected + 1].tolist() == committed + assert positions[i] == position + expected + 1 + assert slot[i].tolist() == [committed[-1], position + expected + 1] + start += count + assert out.q.shape == (4, 3, VOCAB) and out.p.shape == (4, 4, VOCAB) + + +def test_wiring_exercises_non_argmax_acceptance_and_rejection(monkeypatch): + seen = {"accepted_non_argmax": 0, "rejected": 0, "accepted": 0} + for seed in range(8): + h = Harness(monkeypatch, seed) + out = h.run(SPECS) + n = len(SPECS) + accepted = out.buffer.result[out.count : out.count + n].tolist() + start = 0 + for i, (sampling, count, _, _) in enumerate(SPECS): + if sampling.temperature: + committed = out.buffer.result[start : start + accepted[i]] + argmax = out.target[start : start + accepted[i]] + seen["accepted"] += accepted[i] + seen["rejected"] += accepted[i] < count - 1 + seen["accepted_non_argmax"] += int((committed != argmax).sum()) + start += count + assert all(seen.values()), seen + + +def test_all_greedy_round_builds_no_random_state_and_keeps_greedy_layout(monkeypatch): + h = Harness(monkeypatch, 0) + + def forbidden(*args, **kwargs): + raise AssertionError("greedy rounds must not touch the RNG registry") + + h.runner.rng.acquire = forbidden + out = h.run([(params(), 3, 4, 20), (params(), 2, 5, 30), (params(), 1, 6, 40)]) + assert out.plan.random is None and out.u is None and out.q is None and out.p is None + assert not any(name.startswith(("r_", "v_")) for name in vars(out.plan)) + n, count = 3, out.count + assert count == 6 + result = out.buffer.result + assert result[:count].tolist() == out.target.tolist() + assert result[count + 2 * n].item() == -1 # no flag region is written + assert not h.runner.rng.generators + + +def test_random_draw_consumes_one_fixed_block_per_round_and_only_for_random(monkeypatch): + h = Harness(monkeypatch, 0) + acquired = [] + original = h.runner.rng.acquire + h.runner.rng.acquire = lambda rid, sp, device: ( + acquired.append(rid), + original(rid, sp, device), + )[1] + h.run(SPECS) + assert acquired == ["r0", "r2", "r3", "r4"] + # The generators advanced by exactly one (2K + 1)-draw regardless of draft counts. + for rid, seed in (("r0", 1), ("r3", 3)): + reference = torch.Generator().manual_seed(seed) + torch.rand((2 * h.K + 1,), generator=reference) + expected = torch.rand((1,), generator=reference) + assert torch.rand((1,), generator=h.runner.rng.generators[rid]).equal(expected) + + +def test_numerical_failure_keeps_old_slot_and_flags_only_that_request(monkeypatch): + h = Harness(monkeypatch, 3) + h.E[7] = float("nan") # any row fed token 7 yields NaN logits + specs = [ + (params(1.0, seed=1), 3, 7, 10), # invalid q at offset 0 -> failed + (params(1.0, seed=2), 3, 4, 20), + (params(), 2, 5, 30), + ] + out = h.run(specs) + n, count = 3, out.count + result = out.buffer.result + assert result[count + 2 * n : count + 3 * n].tolist() == [1, 0, 0] + assert h.runner.slot_values[0].tolist() == [7, 10] # not advanced + assert h.runner.slot_values[1, 1].item() > 20 and h.runner.slot_values[2, 1].item() > 30 diff --git a/tests/test_speculative_sampling.py b/tests/test_speculative_sampling.py new file mode 100644 index 00000000..2dc6e4e9 --- /dev/null +++ b/tests/test_speculative_sampling.py @@ -0,0 +1,576 @@ +import pytest +import torch + +from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.worker.sampler import Sampler +from vllm_rlt.worker.speculative_sampling import ( + batched_probabilities, + categorical, + speculative_accept, +) + +DEVICES = ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)] + + +def host_args(top_k, top_p, vocab): + """Derive the host-side arguments the way a caller would from per-row params.""" + effective = [k for k in top_k if 0 < k < vocab] + return { + "max_top_k": max(effective, default=0), + "apply_top_p": any(p < 1 for p in top_p), + } + + +def run(logits, temperature, top_k, top_p): + device = logits.device + vocab = logits.shape[-1] + return batched_probabilities( + logits, + torch.tensor(temperature, dtype=torch.float32, device=device), + torch.tensor(top_k, dtype=torch.int64, device=device), + torch.tensor(top_p, dtype=torch.float32, device=device), + **host_args(top_k, top_p, vocab), + ) + + +def reference(logits, temperature, top_k, top_p): + rows = [ + Sampler.probabilities(logits[i], SamplingParams(temperature=t, top_k=k, top_p=p)) + for i, (t, k, p) in enumerate(zip(temperature, top_k, top_p, strict=True)) + ] + return torch.stack(rows) + + +@pytest.mark.parametrize("device", DEVICES) +def test_matches_row_by_row_sampler(device): + vocab = 17 + generator = torch.Generator().manual_seed(0) + top_ks = [-1, 1, 3, vocab, vocab + 5] + top_ps = [1.0, 0.9, 0.5] + temperatures = [0.5, 0.8, 1.0, 1.5, 2.0] + top_k = [top_ks[i % len(top_ks)] for i in range(30)] + top_p = [top_ps[(i // len(top_ks)) % len(top_ps)] for i in range(30)] + temperature = [temperatures[(i * 3) % len(temperatures)] for i in range(30)] + logits = (torch.randn(30, vocab, generator=generator) * 3).to(device) + + actual = run(logits, temperature, top_k, top_p) + + assert actual.dtype == torch.float32 + expected = reference(logits, temperature, top_k, top_p) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(actual.sum(-1), torch.ones(30, device=device)) + + +@pytest.mark.parametrize("device", DEVICES) +def test_top_k_is_strict_with_ties(device): + logits = torch.tensor([[5.0, 4.0, 4.0, 1.0]], device=device) + + probs = run(logits, [1.0], [2], [1.0])[0].cpu() + + assert (probs > 0).sum().item() == 2 + assert probs[0] > 0 + assert probs[3] == 0 + assert (probs[1] > 0) != (probs[2] > 0) + kept = 1 if probs[1] > 0 else 2 + expected = torch.tensor([5.0, 4.0]).softmax(-1) + torch.testing.assert_close(probs[[0, kept]], expected) + torch.testing.assert_close(probs.sum(), torch.tensor(1.0)) + + +@pytest.mark.parametrize("device", DEVICES) +def test_top_k_all_tied_row_keeps_exactly_k(device): + logits = torch.zeros(1, 6, device=device) + + probs = run(logits, [1.0], [2], [1.0])[0] + + assert (probs > 0).sum().item() == 2 + + +@pytest.mark.parametrize("device", DEVICES) +def test_top_p_one_row_keeps_full_support_next_to_top_p_row(device): + vocab = 64 + generator = torch.Generator().manual_seed(1) + logits = (torch.randn(2, vocab, generator=generator) * 2).to(device) + + probs = run(logits, [1.0, 1.0], [-1, -1], [1.0, 0.5]) + + assert (probs[0] > 0).all() + assert (probs[1] > 0).sum() < vocab + torch.testing.assert_close(probs[0], logits[0].softmax(-1), atol=1e-6, rtol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +def test_non_truncating_rows_keep_full_support_next_to_truncating_rows(device): + vocab = 32 + generator = torch.Generator().manual_seed(2) + logits = (torch.randn(3, vocab, generator=generator) * 2).to(device) + + probs = run(logits, [1.0, 1.0, 1.0], [-1, 3, vocab], [1.0, 1.0, 1.0]) + + assert (probs[0] > 0).all() + assert (probs[1] > 0).sum().item() == 3 + assert (probs[2] > 0).all() + + +@pytest.mark.parametrize("device", DEVICES) +def test_inputs_are_not_modified(device): + generator = torch.Generator().manual_seed(3) + logits = torch.randn(4, 11, generator=generator).to(device) + temperature = torch.tensor([0.7, 1.0, 1.3, 2.0], device=device) + top_k = torch.tensor([-1, 2, 5, 20], device=device) + top_p = torch.tensor([1.0, 0.8, 0.6, 1.0], device=device) + originals = [t.clone() for t in (logits, temperature, top_k, top_p)] + + batched_probabilities(logits, temperature, top_k, top_p, max_top_k=5, apply_top_p=True) + + for tensor, original in zip((logits, temperature, top_k, top_p), originals, strict=True): + assert torch.equal(tensor, original) + + +@pytest.mark.parametrize("device", DEVICES) +def test_bf16_logits_return_float32_matching_float32_reference(device): + generator = torch.Generator().manual_seed(4) + logits = (torch.randn(5, 23, generator=generator) * 2).to(device=device, dtype=torch.bfloat16) + temperature = [0.6, 1.0, 1.2, 0.9, 2.0] + top_k = [-1, 4, 7, 2, -1] + top_p = [1.0, 1.0, 0.7, 0.9, 0.5] + + actual = run(logits, temperature, top_k, top_p) + + assert actual.dtype == torch.float32 + expected = reference(logits.float(), temperature, top_k, top_p) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) + + +def test_host_argument_validation(): + logits = torch.zeros(2, 4) + ones = torch.ones(2) + top_k = torch.full((2,), -1, dtype=torch.int64) + kwargs = {"max_top_k": 0, "apply_top_p": False} + + with pytest.raises(ValueError, match="max_top_k"): + batched_probabilities(logits, ones, top_k, ones, max_top_k=-1, apply_top_p=False) + with pytest.raises(ValueError, match="2-D"): + batched_probabilities(logits[0], ones, top_k, ones, **kwargs) + with pytest.raises(ValueError, match="temperature"): + batched_probabilities(logits, ones[:1], top_k, ones, **kwargs) + with pytest.raises(ValueError, match="top_k"): + batched_probabilities(logits, ones, top_k[:1], ones, **kwargs) + with pytest.raises(ValueError, match="top_p"): + batched_probabilities(logits, ones, top_k, ones[:1], **kwargs) + + +# --------------------------------------------------------------------------- +# categorical +# --------------------------------------------------------------------------- + +BELOW_ONE = torch.nextafter(torch.tensor(1.0), torch.tensor(0.0)).item() + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_midpoint_quadrature_matches_weights(device): + m = 4096 + weights = torch.tensor([3.0, 0.0, 1.0, 4.0, 0.5, 0.0, 2.5], device=device) + u = ((torch.arange(m, dtype=torch.float64) + 0.5) / m).float().to(device) + + index, valid = categorical(weights.expand(m, -1), u) + + assert valid.all() + freq = torch.bincount(index, minlength=weights.numel()).double().cpu() / m + expected = (weights / weights.sum()).double().cpu() + assert (freq - expected).abs().max() <= 1 / m + 1e-9 + assert freq[weights.cpu() == 0].sum() == 0 + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("weight", [1.0, 1 / 41]) +def test_categorical_top_of_range_never_picks_trailing_zeros(device, weight): + row = torch.tensor([weight] * 41 + [0.0, 0.0], dtype=torch.float32, device=device) + mass32 = row.cumsum(-1)[-1].item() + probes = [BELOW_ONE] + for _ in range(3): + probes.append(torch.nextafter(torch.tensor(probes[-1]), torch.tensor(0.0)).item()) + if mass32 < 1: # the FP32 cumulative value is a legal uniform only when below 1 + probes.append(mass32) + u = torch.tensor(probes, dtype=torch.float32, device=device) + + index, valid = categorical(row.expand(len(probes), -1), u) + + assert valid.all() + assert (index < 41).all() + assert index[0] == 40 # the largest u picks the last positive-weight token + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_out_of_contract_uniform_still_lands_on_positive_weight(device): + # u * mass < mass for every u < 1, so the threshold clamp is only a safety net: a + # uniform that reaches 1.0 must not index past the last positive-weight token. + row = torch.tensor([1.0] * 41 + [0.0, 0.0], device=device) + u = torch.tensor([1.0, 1.5], device=device) + + index, valid = categorical(row.expand(2, -1), u) + + assert valid.all() + assert index.tolist() == [40, 40] + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_small_token_after_large_mass_has_its_exact_interval(device): + row = torch.tensor([1.0, 1e-6], dtype=torch.float32, device=device) + upper = torch.tensor(1 - 0.25e-6, dtype=torch.float32) # inside (1 / (1 + 1e-6), 1) + assert upper.item() > 1 / (1 + 1e-6) + u = torch.tensor([0.5, upper.item(), 0.999998], dtype=torch.float32, device=device) + + index, valid = categorical(row.expand(3, -1), u) + + assert valid.all() + assert index.tolist() == [0, 1, 0] # 0.999998 < 1 / (1 + 1e-6) + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_cdf_is_not_quantized_to_float32(device): + # True boundary 1 / (1 + 1e-5) = 0.99999000010 > u, so the exact answer is token 0; an + # FP32 cumulative sum rounds 1 + 1e-5 up and moves this u into token 1. + row = torch.tensor([1.0, 1e-5], dtype=torch.float32, device=device) + u = torch.tensor([0.9999899864196777, 0.99999499], dtype=torch.float32, device=device) + + index, _ = categorical(row.expand(2, -1), u) + + assert index.tolist() == [0, 1] + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_exact_cdf_boundaries_skip_zero_weight_tokens(device): + row = torch.tensor([0.0, 1.0, 0.0, 3.0], device=device) # cumulative [0, 1, 1, 4] + u = torch.tensor([0.0, 0.25, 0.5], device=device) # thresholds 0, 1, 2 (exact) + + index, valid = categorical(row.expand(3, -1), u) + + assert valid.all() + assert index.tolist() == [1, 3, 3] + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_plateau_tail_is_never_chosen(device): + row = torch.tensor([1.0, 1e-20], dtype=torch.float32, device=device) + assert row.double().cumsum(-1)[-1] == 1.0 # the tail increment is rounded away in float64 + u = torch.tensor([0.0, 0.5, 0.999, BELOW_ONE], dtype=torch.float32, device=device) + + index, valid = categorical(row.expand(4, -1), u) + + assert valid.all() + assert (index == 0).all() + + +@pytest.mark.parametrize("device", DEVICES) +def test_categorical_invalid_rows_return_index_zero_and_leave_valid_rows_alone(device): + nan, inf = float("nan"), float("inf") + weights = torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [1.0, nan, 1.0], + [0.0, 0.0, 2.0], + [1.0, inf, 1.0], + ], + device=device, + ) + u = torch.full((5,), 0.7, device=device) + + index, valid = categorical(weights, u) + + assert valid.tolist() == [False, True, False, True, False] + assert index.tolist() == [0, 1, 0, 2, 0] + assert index.dtype == torch.int64 + + +# --------------------------------------------------------------------------- +# speculative_accept +# --------------------------------------------------------------------------- + + +def midpoints(m, device): + return ((torch.arange(m, dtype=torch.float64) + 0.5) / m).float().to(device) + + +def exact_first_position(p0, p1, q, m, device): + """Enumerate every draft x, u_accept on an m-grid and u_final on an m-grid (d=1). + + Returns (first-token marginal, P(accepted == 1), bonus distribution given accepted == 1), + each a float64 vector/scalar weighted by the proposal probability of the draft. + """ + vocab = len(q) + qn = [w / sum(q) for w in q] + xs = [x for x in range(vocab) if qn[x] > 0] + grid = midpoints(m, device) + x_idx = torch.tensor(xs, device=device) + n = len(xs) * m * m + draft = x_idx.repeat_interleave(m * m)[:, None] + ua = grid.repeat_interleave(m).repeat(len(xs))[:, None] + uf = grid.repeat(m * len(xs)) + weight = torch.tensor(qn, dtype=torch.float64, device=device)[draft[:, 0]] + p = torch.tensor([p0, p1], dtype=torch.float32, device=device).expand(n, 2, vocab) + qq = torch.tensor(q, dtype=torch.float32, device=device).expand(n, 1, vocab) + + out = speculative_accept(draft, torch.ones(n, dtype=torch.int64, device=device), qq, p, ua, uf) + + assert not out.error.any() + total = float(m * m) + first = torch.bincount(out.tokens[:, 0], weights=weight, minlength=vocab) / total + accepted = (out.accepted == 1).double() + p_acc = (weight * accepted).sum() / total + bonus = torch.bincount(out.tokens[:, 1], weights=weight * accepted, minlength=vocab) + bonus = bonus / (weight * accepted).sum().clamp_min(1e-30) + return first.cpu(), p_acc.item(), bonus.cpu(), (weight * accepted).sum().item() > 0 + + +def normalized(w): + return torch.tensor([x / sum(w) for x in w], dtype=torch.float64) + + +EXACT_CASES = { + "generic": ([0.5, 4.0, 1.0, 0.0, 2.0], [3.0, 1.0, 0.0, 2.0, 6.0], [2.0, 0.0, 1.0, 3.0, 0.5]), + "disjoint": ([0.0, 0.0, 1.0, 3.0, 2.0], [1.0, 1.0, 2.0, 0.0, 5.0], [1.0, 2.0, 0.0, 0.0, 0.0]), + "equal": ([2.0, 0.0, 1.0, 3.0, 0.5], [4.0, 1.0, 1.0, 1.0, 1.0], [2.0, 0.0, 1.0, 3.0, 0.5]), +} + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("case", list(EXACT_CASES)) +def test_single_draft_matches_target_by_exact_quadrature(device, case): + m = 256 + p0, p1, q = EXACT_CASES[case] + pn0, pn1, qn = normalized(p0), normalized(p1), normalized(q) + + first, p_acc, bonus, has_accepts = exact_first_position(p0, p1, q, m, device) + + assert (first - pn0).abs().max() <= 3 / m + assert abs(p_acc - torch.minimum(pn0, qn).sum().item()) <= 3 / m + if case == "disjoint": + assert p_acc == 0 + if case == "equal": + assert p_acc == pytest.approx(1.0, abs=1e-9) + if has_accepts: + assert (bonus - pn1).abs().max() <= 3 / m + + +@pytest.mark.parametrize("device", DEVICES) +def test_raw_weights_are_normalized_before_comparison(device): + m = 256 + first, *_ = exact_first_position([3.0, 2.0], [1.0, 1.0], [1.0, 3.0], m, device) + + # Normalized: [0.6, 0.4]. Comparing raw weights would give [0.5, 0.5]. + assert (first - torch.tensor([0.6, 0.4], dtype=torch.float64)).abs().max() <= 3 / m + + +@pytest.mark.parametrize("device", DEVICES) +def test_ragged_multi_draft_monte_carlo_matches_target(device): + vocab, w1, reps = 6, 3, 40_000 + counts = [0, 1, 2, 3, 3] + generator = torch.Generator().manual_seed(7) + + def rand_weights(*shape): + w = torch.rand(*shape, generator=generator) + 0.05 + w[torch.rand(*shape, generator=generator) < 0.2] = 0.0 + w[..., 0] += 0.5 # keep every row's mass positive + return w + + q = rand_weights(len(counts), w1, vocab) + p = rand_weights(len(counts), w1 + 1, vocab) + drafts = torch.zeros(len(counts), reps, w1, dtype=torch.int64) + for i, d in enumerate(counts): + for j in range(d): + drafts[i, :, j] = torch.multinomial( + q[i, j], reps, replacement=True, generator=generator + ) + q_in, p_in = q.clone(), p.clone() + for i, d in enumerate(counts): # poison everything past the real drafts + q_in[i, d:] = float("nan") + p_in[i, d + 1 :] = float("nan") + n = len(counts) * reps + + def rep(t): + return t.repeat_interleave(reps, dim=0).to(device) + + u_accept = torch.rand(n, w1, generator=generator).to(device) + u_final = torch.rand(n, generator=generator).to(device) + + out = speculative_accept( + drafts.reshape(n, w1).to(device), + rep(torch.tensor(counts)), + rep(q_in), + rep(p_in), + u_accept, + u_final, + ) + + assert not out.error.any() + tokens, accepted = out.tokens.cpu().view(len(counts), reps, w1 + 1), out.accepted.cpu() + accepted = accepted.view(len(counts), reps) + pn, qn = p / p.sum(-1, keepdim=True), q / q.sum(-1, keepdim=True) + + def assert_distribution(samples, target): + freq = torch.bincount(samples, minlength=vocab).double() / samples.numel() + se = (target.double() * (1 - target.double()) / samples.numel()).sqrt() + assert ((freq - target.double()).abs() <= 5 * se + 1e-6).all() + + for i, d in enumerate(counts): + assert (accepted[i] <= d).all() + assert_distribution(tokens[i, :, 0], pn[i, 0]) # K = 0 rows: a plain draw from p[0] + reach = 1.0 + for k in range(1, d + 1): + reach *= torch.minimum(pn[i, k - 1], qn[i, k - 1]).sum().item() + hit = (accepted[i] >= k).double().mean().item() + assert abs(hit - reach) <= 5 * (reach * (1 - reach) / reps) ** 0.5 + 1e-6 + if d >= 1: # the token after an accepted first draft follows the target at position 1 + assert_distribution(tokens[i, accepted[i] >= 1, 1], pn[i, 1]) + + +def three_state(draft, d, q, p, device, *, u_accept=0.5, u_final=0.5, vocab=4): + w1 = len(draft) + result = speculative_accept( + torch.tensor([draft], dtype=torch.int64, device=device), + torch.tensor([d], dtype=torch.int64, device=device), + torch.tensor([q], dtype=torch.float32, device=device).view(1, w1, vocab), + torch.tensor([p], dtype=torch.float32, device=device).view(1, w1 + 1, vocab), + torch.full((1, w1), u_accept, device=device), + torch.full((1,), u_final, device=device), + ) + return ( + result.tokens[0].tolist(), + result.accepted.item(), + result.final.item(), + result.error.item(), + ) + + +NAN = float("nan") +E0, E1, E2, E3 = ([1.0 if i == j else 0.0 for i in range(4)] for j in range(4)) +BAD = [NAN] * 4 + + +@pytest.mark.parametrize("device", DEVICES) +def test_valid_reject_before_invalid_q_row_delivers_correction(device): + # Position 0: q = e0, p = e2 -> certain reject, residual = e2. Position 2 q is NaN. + tokens, accepted, final, error = three_state( + [0, 1, 0], 3, [E0, E1, BAD], [E2, E1, E0, E0], device + ) + + assert (tokens, accepted, final, error) == ([2, 0, 0, 0], 0, 2, False) + + +@pytest.mark.parametrize("device", DEVICES) +def test_invalid_q_row_after_accepts_is_an_error(device): + tokens, accepted, final, error = three_state( + [0, 1, 0], 3, [E0, E1, BAD], [E0, E1, E0, E0], device + ) + + assert error and accepted == 2 and final == 0 + assert all(0 <= t < 4 for t in tokens) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("bonus", [[0.0] * 4, BAD, [float("inf"), 0.0, 0.0, 0.0]]) +def test_invalid_bonus_row_after_full_accept_is_an_error(device, bonus): + tokens, accepted, final, error = three_state([0, 1], 2, [E0, E1], [E0, E1, bonus], device) + + assert error and accepted == 2 and final == 0 + assert tokens == [0, 1, 0] + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("pad_token", [0, 3, 99, -1]) +def test_nan_in_padding_positions_is_ignored(device, pad_token): + tokens, accepted, final, error = three_state( + [1, pad_token, pad_token], 1, [E1, BAD, BAD], [E1, E2, BAD, BAD], device + ) + + assert (tokens, accepted, final, error) == ([1, 2, 0, 0], 1, 2, False) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("tiny", [1e-38, 1e-30]) +def test_equal_tiny_probabilities_are_accepted_for_any_uniform(device, tiny): + row = [tiny, 1.0, 0.0, 0.0] # u * tiny rounds back to tiny in FP32 when subnormal + + tokens, accepted, final, error = three_state( + [0], 1, [row], [row, E3], device, u_accept=BELOW_ONE + ) + + assert (tokens, accepted, final, error) == ([0, 3], 1, 3, False) + + +@pytest.mark.parametrize("device", DEVICES) +def test_zero_drafts_draw_from_first_target_row(device): + n, vocab = 5, 4 + p = torch.zeros(n, 1, vocab, device=device) + p[torch.arange(n), 0, torch.arange(n) % vocab] = 2.0 + result = speculative_accept( + torch.zeros(n, 0, dtype=torch.int64, device=device), + torch.zeros(n, dtype=torch.int64, device=device), + torch.zeros(n, 0, vocab, device=device), + p, + torch.zeros(n, 0, device=device), + torch.full((n,), 0.3, device=device), + ) + + assert result.accepted.tolist() == [0] * n + assert result.final.tolist() == [i % vocab for i in range(n)] + assert result.tokens.tolist() == [[i % vocab] for i in range(n)] + assert not result.error.any() + + +@pytest.mark.parametrize("device", DEVICES) +def test_output_layout_error_rows_and_inputs_untouched(device): + vocab, w1 = 4, 3 + # row 0: 2 accepts then a reject; row 1: all accepted (bonus e3); row 2: no drafts; + # row 3: invalid q at the first position; row 4: invalid bonus row. + q = torch.tensor( + [ + [E0, E1, E0], + [E2, E2, E2], + [BAD, BAD, BAD], + [BAD, E1, E1], + [E0, E0, E0], + ], + device=device, + ) + p = torch.tensor( + [ + [E0, E1, E3, E3], + [E2, E2, E2, E3], + [E1, BAD, BAD, BAD], + [E0, E1, E1, E1], + [E0, E0, E0, [0.0] * 4], + ], + device=device, + ) + draft_tokens = torch.tensor( + [[0, 1, 0], [2, 2, 2], [0, 0, 0], [0, 1, 1], [0, 0, 0]], device=device + ) + drafts = torch.tensor([3, 3, 0, 3, 3], device=device) + u_accept = torch.full((5, w1), 0.5, device=device) + u_final = torch.full((5,), 0.5, device=device) + originals = [t.clone() for t in (draft_tokens, drafts, q, p, u_accept, u_final)] + + result = speculative_accept(draft_tokens, drafts, q, p, u_accept, u_final) + + assert result.error.tolist() == [False, False, False, True, True] + assert result.accepted.tolist() == [2, 3, 0, 0, 3] + assert result.final.tolist() == [3, 3, 1, 0, 0] + assert result.tokens.tolist() == [ + [0, 1, 3, 0], # drafts[:2], residual at position 2 (q = e0, p = e3), zeros after + [2, 2, 2, 3], # all drafts, then the bonus + [1, 0, 0, 0], # K = 0 + [0, 0, 0, 0], # error row: legal placeholders only + [0, 0, 0, 0], + ] + for out in (result.tokens, result.accepted, result.final): + assert out.dtype == torch.int64 + for out in (result.tokens, result.final): + assert ((out >= 0) & (out < vocab)).all() + assert result.error.dtype == torch.bool + for tensor, original in zip( + (draft_tokens, drafts, q, p, u_accept, u_final), originals, strict=True + ): + torch.testing.assert_close(tensor, original, rtol=0, atol=0, equal_nan=True) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..5ef806a2 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -310,11 +310,9 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i raise ValueError("request_id must be a nonempty string") params = sampling_params or SamplingParams() if ( - self.async_speculative - or (self.speculative_config and self.speculative_config.mode == "pipeline") + self.speculative_config and self.speculative_config.mode == "pipeline" ) and params.temperature != 0: - mode = "async" if self.async_speculative else "pipeline" - raise ValueError(f"{mode} speculative decoding currently requires greedy sampling") + raise ValueError("pipeline speculative decoding currently requires greedy sampling") if request_id in self._retiring_speculative: self._retire_speculative(wait=True, request_id=request_id) config = self.model.config @@ -714,8 +712,8 @@ def _collect_speculative(self): if not ticket.ready(): runner.readback_waits += 1 result = ticket.collect() - for scheduled_request, state, round_id in zip( - ticket.batch.requests, ticket.states, ticket.rounds + for row, (scheduled_request, state, round_id) in enumerate( + zip(ticket.batch.requests, ticket.states, ticket.rounds) ): request = scheduled_request.request if ( @@ -726,6 +724,12 @@ def _collect_speculative(self): continue if state.delivered != round_id: raise RuntimeError("out-of-order speculative round delivery") + # Stale rows were skipped above, so only live requests can fail the round. The + # check precedes _update (and the scheduler's EOS truncation): failure is atomic. + if result.speculative[row].numerical_error: + raise RuntimeError( + f"invalid speculative sampling distribution for request {request.request_id!r}" + ) state.delivered += 1 outputs = self._update(ticket.batch, result) ticket.retire() diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 6b708df2..8b18dfa2 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -1,9 +1,16 @@ -"""Bounded cross-round greedy speculation with device-owned token positions.""" +"""Bounded cross-round speculation with device-owned token positions. + +Block mode supports greedy and random (temperature/top-k/top-p) requests in one batch. Greedy +requests verify by target argmax. Random requests draw drafts from their ``q`` distributions, +verify with speculative rejection sampling against ``p`` (see ``speculative_sampling``) and +return their committed sequence (accepted drafts, then correction or bonus) in the result rows. +A batch without random requests takes the greedy path with no extra device work. +""" import math from dataclasses import dataclass, replace from types import SimpleNamespace -from typing import Literal +from typing import Literal, NamedTuple import numpy as np import torch @@ -25,6 +32,170 @@ SpeculativeRunner, SpeculativeStats, ) +from vllm_rlt.worker.speculative_sampling import ( + AcceptResult, + batched_probabilities, + categorical, + speculative_accept, +) + +# Float plan arrays travel in the int64 upload as float64 bit patterns. +_FLOAT_UPLOADS = frozenset({"r_temperature", "r_top_p", "v_temperature", "v_top_p"}) + + +@dataclass(frozen=True, eq=False) +class RandomSamplingPlan: + """Host plan for the random requests of one block round (numpy only, no device access). + + Random requests are those with ``temperature != 0``; their *rank* is their position among + random requests in item order, and it indexes ``u``/``q``/``p`` rows. ``width`` is the + largest ``token_count`` among random requests (``q`` is ``[r, width - 1, V]``, ``p`` is + ``[r, width, V]``). + """ + + requests: np.ndarray # [r] item index of each random request + temperature: np.ndarray # [r] float64 + top_k: np.ndarray # [r] int64, as given (-1 = full vocabulary) + top_p: np.ndarray # [r] float64 + max_top_k: int # largest k with 0 < k < vocab over random requests, else 0 + apply_top_p: bool # any random request has top_p < 1 + width: int + # Per draft offset o < width - 1: (rows of that offset's drafting block that are random, + # their ranks). Empty arrays when the offset drafts no random request. + draft_rows: tuple[tuple[np.ndarray, np.ndarray], ...] + verify_rows: np.ndarray # flat request-major verify rows of random requests + verify_ranks: np.ndarray # rank owning each such row + verify_offsets: np.ndarray # row index within its request (p / result column) + verify_temperature: np.ndarray # per verify row, float64 + verify_top_k: np.ndarray + verify_top_p: np.ndarray + + def uploads(self): + """(name, int64 array) pairs appended to the round's single index upload.""" + parts = [ + ("r_requests", self.requests), + ("r_temperature", self.temperature), + ("r_top_k", self.top_k), + ("r_top_p", self.top_p), + ("v_rows", self.verify_rows), + ("v_ranks", self.verify_ranks), + ("v_offsets", self.verify_offsets), + ("v_temperature", self.verify_temperature), + ("v_top_k", self.verify_top_k), + ("v_top_p", self.verify_top_p), + ] + for offset, (rows, ranks) in enumerate(self.draft_rows): + if len(rows): + parts += [(f"r_rows{offset}", rows), (f"r_ranks{offset}", ranks)] + return [ + ( + name, + np.asarray(values, dtype=np.float64).view(np.int64) + if name in _FLOAT_UPLOADS + else np.asarray(values, dtype=np.int64), + ) + for name, values in parts + ] + + +def plan_random_sampling(sampling_params, token_counts, vocab_size, num_speculative_tokens): + """Plan the random rows of a block round, or return None when every request is greedy. + + ``sampling_params[i]`` and ``token_counts[i]`` describe item ``i`` (``token_count - 1`` + drafts). Draft ``j`` of a request is proposed at offset ``j`` and checked against its + verify row ``j``; verify rows are request-major. Pure host code: no device access. + """ + counts = np.asarray(token_counts, dtype=np.int64) + is_random = np.array([params.temperature != 0 for params in sampling_params], dtype=bool) + if not is_random.any(): + return None + requests = np.flatnonzero(is_random) + r = len(requests) + rank_of = np.full(len(counts), -1, dtype=np.int64) + rank_of[requests] = np.arange(r) + request_counts = counts[requests] + width = int(request_counts.max()) + if width - 1 > num_speculative_tokens: + raise ValueError("random request has more drafts than num_speculative_tokens") + temperature = np.array([sampling_params[i].temperature for i in requests], dtype=np.float64) + top_k = np.array([sampling_params[i].top_k for i in requests], dtype=np.int64) + top_p = np.array([sampling_params[i].top_p for i in requests], dtype=np.float64) + truncating = top_k[(top_k > 0) & (top_k < vocab_size)] + draft_rows = [] + for offset in range(width - 1): + drafting = np.flatnonzero(counts > offset + 1) + rows = np.flatnonzero(is_random[drafting]) + draft_rows.append((rows, rank_of[drafting[rows]])) + starts = np.cumsum(counts) - counts + verify_rows = np.concatenate([starts[i] + np.arange(counts[i]) for i in requests]).astype( + np.int64 + ) + verify_ranks = np.repeat(np.arange(r), request_counts) + verify_offsets = np.concatenate([np.arange(c) for c in request_counts]).astype(np.int64) + return RandomSamplingPlan( + requests=requests, + temperature=temperature, + top_k=top_k, + top_p=top_p, + max_top_k=int(truncating.max()) if len(truncating) else 0, + apply_top_p=bool((top_p < 1).any()), + width=width, + draft_rows=tuple(draft_rows), + verify_rows=verify_rows, + verify_ranks=verify_ranks, + verify_offsets=verify_offsets, + verify_temperature=np.repeat(temperature, request_counts), + verify_top_k=np.repeat(top_k, request_counts), + verify_top_p=np.repeat(top_p, request_counts), + ) + + +class RandomRoundMerge(NamedTuple): + rows: torch.Tensor # [R] result rows (committed sequences for random requests) + accepted: torch.Tensor # [n] + next_tokens: torch.Tensor # [n] (placeholders for failed requests) + next_positions: torch.Tensor # [n] starts + accepted + 1 (informational for failed rows) + next_values: torch.Tensor # [n, 2] (token, position) to scatter into the slot table + error: torch.Tensor # [n] int64 0/1 flags + + +def merge_random_round( + target, + accepted, + next_tokens, + starts, + input_tokens, + result: AcceptResult, + random_requests, + verify_rows, + verify_ranks, + verify_offsets, +): + """Overlay random-request outcomes on the greedy round tensors; no host reads. + + ``target`` ``[R]`` argmax per request-major verify row, ``accepted``/``next_tokens`` + ``[n]`` the greedy values, ``starts``/``input_tokens`` ``[n]`` the slot values read at the + start of the round. ``result`` is the ``speculative_accept`` outcome of the random requests + ``random_requests`` ``[r]``. Verify row ``start_i + j`` of random request ``rank`` becomes + ``result.tokens[rank, j]`` (accepted drafts, then the final token); greedy rows are + untouched. A failed request keeps its old ``(input_token, start)`` slot values, so a queued + next round reads the span it already reserved; the caller must treat the round as fatal. + """ + rows = target.clone() + rows[verify_rows] = result.tokens[verify_ranks, verify_offsets] + accepted = accepted.clone() + accepted[random_requests] = result.accepted + next_tokens = next_tokens.clone() + next_tokens[random_requests] = result.final + failed = torch.zeros(len(starts), dtype=torch.bool, device=starts.device) + failed[random_requests] = result.error + next_positions = starts + accepted + 1 + next_values = torch.where( + failed[:, None], + torch.stack((input_tokens, starts), 1), + torch.stack((next_tokens, next_positions), 1), + ) + return RandomRoundMerge(rows, accepted, next_tokens, next_positions, next_values, failed.long()) @dataclass @@ -60,9 +231,11 @@ def __init__(self, runner, scheduler): self.host_tables = torch.empty(math.prod(shape), dtype=torch.int32, pin_memory=True) self.tables = torch.empty(math.prod(shape), dtype=torch.int32, device=cache.device) # Result capacity in int64 elements (not bytes), for one submission. - # Block layout: [target token IDs for all verification rows, - # accepted counts per request, next positions per request]. - # Its upper bound is max_num_batched_tokens + 2 * max_num_seqs. + # Block layout: [result token IDs for all verification rows, + # accepted counts per request, next positions per request, + # error flags per request (only when the round has random requests)]. + # Result rows are target argmax IDs for greedy requests and committed sequences for + # random ones. Its upper bound is max_num_batched_tokens + 3 * max_num_seqs. # Pipeline layout: six consecutive values per request: # [prediction, accepted, emitted, next_position, drafted_count, verified_count]. # `emitted` says whether this step produced a target output; prediction is @@ -70,7 +243,7 @@ def __init__(self, runner, scheduler): # Reserve the larger layout in both GPU and pinned host storage; submit() # copies only the used prefix, and collect() decodes the selected layout. size = max( - scheduler.max_num_batched_tokens + 2 * scheduler.max_num_seqs, + scheduler.max_num_batched_tokens + 3 * scheduler.max_num_seqs, 6 * scheduler.max_num_seqs, ) self.host_result = torch.empty(size, dtype=torch.int64, pin_memory=True) @@ -92,6 +265,8 @@ class SpeculativeRoundTicket: stats: SpeculativeStats | None = None cached: ModelRunnerOutput | None = None retired: bool = False + # Block rounds with random requests append one numerical-error flag per request. + error_flags: bool = False @property def has_result(self): @@ -121,7 +296,11 @@ def collect(self) -> ModelRunnerOutput: self.stats.drafted_tokens += sum(row[4] for row in values) self.stats.verified_rows += sum(row[5] for row in values) return self.cached - values = self.buffer.host_result[: self.rows + 2 * b].tolist() + # [rows | accepted (b) | next positions (b) | error flags (b, random rounds only)] + # Only these ranges are read; a reused buffer's tail may hold older rounds' values. + values = self.buffer.host_result[: self.rows + (3 if self.error_flags else 2) * b] + values = values.tolist() + flags = values[self.rows + 2 * b : self.rows + 3 * b] if self.error_flags else None results, start = [], 0 for i, scheduled_request in enumerate(self.batch.requests): accepted = values[self.rows + i] @@ -129,10 +308,13 @@ def collect(self) -> ModelRunnerOutput: SpeculativeResult( values[start : start + accepted + 1], accepted, + numerical_error=bool(flags[i]) if flags is not None else False, ) ) start += scheduled_request.token_count - self.cached = ModelRunnerOutput(speculative=results, frontiers=values[self.rows + b :]) + self.cached = ModelRunnerOutput( + speculative=results, frontiers=values[self.rows + b : self.rows + 2 * b] + ) return self.cached def retire(self, *, discard=False): @@ -330,7 +512,9 @@ def _round_plan(self, items, states, buffer): """Host-built gather/scatter indices for a round, uploaded in one copy. Offset o drafts the requests whose token_count exceeds o, in item order; - verification rows stay request-major in item order, as results expect. + verification rows stay request-major in item order, as results expect. When any + request is random, ``plan.random`` is its host plan and its arrays join the same upload; + for an all-greedy round ``plan.random`` is None and nothing extra is built. """ counts = np.array([item.token_count for item in items], dtype=np.int64) n, width = len(counts), int(counts.max()) @@ -361,13 +545,63 @@ def add(name, values): add("request_rows", request_rows) add("row_offsets", row_offsets) add("hidden_rows", base[row_offsets] + ranks[row_offsets, request_rows]) + random = plan_random_sampling( + [item.request.sampling_params for item in items], + counts, + self.model.config.vocab_size, + self.config.num_speculative_tokens, + ) + if random is not None: + for name, values in random.uploads(): + add(name, values) host = torch.from_numpy(np.concatenate(parts)).pin_memory() values = host.to(self.device, non_blocking=True) buffer.keepalive.extend((host, values)) views = {name: values[start : start + size] for name, (start, size) in layout.items()} - return SimpleNamespace(n=n, width=width, rows=len(request_rows), offsets=offsets, **views) + for name in _FLOAT_UPLOADS & views.keys(): + views[name] = views[name].view(torch.float64).float() + return SimpleNamespace( + n=n, width=width, rows=len(request_rows), offsets=offsets, random=random, **views + ) + + def _uniforms(self, items, random): + """Uniforms ``[r, 2K + 1]`` for the random requests, drawn in item order. + + Columns ``[0:K]`` sample drafts, ``[K:2K]`` decide acceptance and ``[2K]`` samples the + correction or bonus token. Each request consumes exactly one draw per round whatever + its draft count or acceptance outcome; greedy requests never touch the registry. + """ + rows = [] + for i in random.requests: + request = items[i].request + generator = self.rng.acquire(request.request_id, request.sampling_params, self.device) + rows.append( + torch.rand( + (2 * self.config.num_speculative_tokens + 1,), + generator=generator, + device=self.device, + dtype=torch.float32, + ) + ) + return torch.stack(rows) + + def _probabilities(self, plan, logits, temperature, top_k, top_p): + return batched_probabilities( + logits, + temperature, + top_k, + top_p, + max_top_k=plan.random.max_top_k, + apply_top_p=plan.random.apply_top_p, + ) + + def _draft(self, plan, tables, allocations, u=None): + """Propose drafts; returns ``(starts, tokens, hidden_blocks, q)``. - def _draft(self, plan, tables, allocations): + Greedy rows take the CODA argmax. For random rows (``plan.random``) the proposal is + drawn from ``q = batched_probabilities(logits)`` with ``u[rank, offset]`` and ``q`` is + stored in ``q[rank, offset]`` (``q`` is None when no random request drafts). + """ from vllm_rlt.kernels.routing import gather_kernel slot_values = self.slot_values.new_empty((plan.n, 2)) @@ -385,6 +619,14 @@ def _draft(self, plan, tables, allocations): tokens = torch.zeros((plan.n, plan.width), dtype=torch.long, device=self.device) tokens[:, 0] = slot_values[:, 0] hidden_blocks = [] + random = plan.random + q = None + if random is not None and random.width > 1: + q = torch.zeros( + (len(random.requests), random.width - 1, self.model.config.vocab_size), + dtype=torch.float32, + device=self.device, + ) for offset, (active_count, drafting_count) in enumerate(plan.offsets): active = getattr(plan, f"active{offset}", None) if active is None: @@ -402,23 +644,65 @@ def _draft(self, plan, tables, allocations): if drafting_count: rows = getattr(plan, f"draft_rows{offset}", None) drafting = hidden if rows is None else hidden[rows] - proposed = self._coda(drafting).argmax(-1) + # CUDA graphs may return a static buffer that the next _coda replay + # overwrites: everything below is derived from `logits` immediately. + logits = self._coda(drafting) + proposed = logits.argmax(-1) + random_rows = getattr(plan, f"r_rows{offset}", None) + if random_rows is not None: + ranks = getattr(plan, f"r_ranks{offset}") + q_o = self._probabilities( + plan, + logits[random_rows], + plan.r_temperature[ranks], + plan.r_top_k[ranks], + plan.r_top_p[ranks], + ) + # Invalid q rows give placeholder token 0; speculative_accept flags them. + sampled, _ = categorical(q_o, u[ranks, offset]) + proposed[random_rows] = sampled + q[ranks, offset] = q_o tokens[getattr(plan, f"drafting{offset}"), offset + 1] = proposed - return starts, tokens, hidden_blocks + return starts, tokens, hidden_blocks, q def _verify(self, plan, starts, hidden_blocks, tables, allocations): + """Run the target loops; returns ``(target argmax per verify row, p)``. + + ``p`` ``[r, width, V]`` holds the verify distributions of the random requests (zeros + in padding) and is None when every request is greedy. + """ positions = starts[plan.request_rows] + plan.row_offsets row_tables = tables[plan.request_rows] hidden = torch.cat(hidden_blocks)[plan.hidden_rows] for depth in range(self.config.draft_loops, self.config.target_loops): hidden = self._execute_recurrent_step(hidden, positions, row_tables, depth, allocations) - return self._coda(hidden).argmax(-1) + logits = self._coda(hidden) + target = logits.argmax(-1) + random = plan.random + if random is None: + return target, None + p = torch.zeros( + (len(random.requests), random.width, self.model.config.vocab_size), + dtype=torch.float32, + device=self.device, + ) + p[plan.v_ranks, plan.v_offsets] = self._probabilities( + plan, logits[plan.v_rows], plan.v_temperature, plan.v_top_k, plan.v_top_p + ) + return target, p - def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer): + def _accept_and_update_state( + self, plan, states, starts, tokens, target, buffer, *, u=None, q=None, p=None + ): + """Commit the round on the device and fill ``buffer.result``. + + Layout: ``[rows (count) | accepted (n) | next positions (n)]`` plus ``n`` error flags + when ``plan.random`` is set. Greedy requests keep the target-argmax rows; random + requests write their committed sequence (accepted drafts, then correction or bonus). + """ from vllm_rlt.kernels.routing import scatter_kernel count = len(target) - buffer.result[:count].copy_(target) targets = torch.zeros((plan.n, plan.width), dtype=torch.long, device=self.device) targets[plan.request_rows, plan.row_offsets] = target # Candidate j (token column j+1) is checked against target row j; only @@ -429,8 +713,44 @@ def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer) # All target IDs through the first rejection equal the accepted # candidates followed by the correction (or final bonus). next_tokens = targets.gather(1, accepted[:, None]).squeeze(1) - next_positions = starts + accepted + 1 - next_values = torch.stack((next_tokens, next_positions), 1) + rows = target + random = plan.random + if random is None: + next_positions = starts + accepted + 1 + next_values = torch.stack((next_tokens, next_positions), 1) + else: + width = random.width + requests = plan.r_requests + if q is None: + q = p.new_empty((len(random.requests), 0, p.shape[-1])) + k = self.config.num_speculative_tokens + outcome = speculative_accept( + tokens[requests, 1:width], + plan.drafts[requests], + q, + p, + u[:, k : k + width - 1], + u[:, 2 * k], + ) + merged = merge_random_round( + target, + accepted, + next_tokens, + starts, + tokens[:, 0], + outcome, + requests, + plan.v_rows, + plan.v_ranks, + plan.v_offsets, + ) + rows, accepted, next_positions, next_values = ( + merged.rows, + merged.accepted, + merged.next_positions, + merged.next_values, + ) + buffer.result[:count].copy_(rows) scatter_kernel[(plan.n, 1)]( next_values, plan.slots, @@ -442,6 +762,9 @@ def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer) ) buffer.result[count : count + plan.n].copy_(accepted) buffer.result[count + plan.n : count + 2 * plan.n].copy_(next_positions) + if random is not None: + # Written every random round: a reused buffer never exposes older flags. + buffer.result[count + 2 * plan.n : count + 3 * plan.n].copy_(merged.error) for state in states: state.submitted += 1 return count @@ -463,13 +786,18 @@ def submit(self, batch): if any(s.submitted > s.delivered for s in states): self.submitted_before_collect += 1 tables.copy_(host_tables, non_blocking=True) + error_flags = False if self.config.mode == "pipeline": count = self._pipeline_step(scheduled_requests, states, tables, allocations, buffer) else: plan = self._round_plan(scheduled_requests, states, buffer) - starts, tokens, hidden_blocks = self._draft(plan, tables, allocations) - target = self._verify(plan, starts, hidden_blocks, tables, allocations) - count = self._accept_and_update_state(plan, states, starts, tokens, target, buffer) + u = None if plan.random is None else self._uniforms(scheduled_requests, plan.random) + starts, tokens, hidden_blocks, q = self._draft(plan, tables, allocations, u) + target, p = self._verify(plan, starts, hidden_blocks, tables, allocations) + count = self._accept_and_update_state( + plan, states, starts, tokens, target, buffer, u=u, q=q, p=p + ) + error_flags = plan.random is not None buffer.compute_done = torch.cuda.Event() buffer.compute_done.record(self.compute_stream) for state in states: @@ -480,7 +808,7 @@ def submit(self, batch): n = ( 6 * len(scheduled_requests) if self.config.mode == "pipeline" - else count + 2 * len(scheduled_requests) + else count + (3 if error_flags else 2) * len(scheduled_requests) ) buffer.host_result[:n].copy_(buffer.result[:n], non_blocking=True) buffer.copy_done = torch.cuda.Event() @@ -497,6 +825,7 @@ def submit(self, batch): count, mode=self.config.mode, stats=self.stats, + error_flags=error_flags, ) def _pipeline_step(self, spans, states, tables, allocations, buffer): diff --git a/vllm_rlt/worker/output.py b/vllm_rlt/worker/output.py index 19d74f70..e0401f25 100644 --- a/vllm_rlt/worker/output.py +++ b/vllm_rlt/worker/output.py @@ -23,6 +23,8 @@ class SpeculativeResult: accepted_count: int # Exclusive KV extent still needed by retained pipeline tokens. kv_frontier: int | None = None + # Random speculation only: the round's stop point had an invalid distribution. + numerical_error: bool = False @dataclass(frozen=True, kw_only=True) diff --git a/vllm_rlt/worker/speculative_sampling.py b/vllm_rlt/worker/speculative_sampling.py new file mode 100644 index 00000000..061e040f --- /dev/null +++ b/vllm_rlt/worker/speculative_sampling.py @@ -0,0 +1,180 @@ +"""Batched, host-sync-free sampling helpers for asynchronous speculative decoding.""" + +from typing import NamedTuple + +import torch + + +def batched_probabilities( + logits: torch.Tensor, + temperature: torch.Tensor, + top_k: torch.Tensor, + top_p: torch.Tensor, + *, + max_top_k: int, + apply_top_p: bool, +) -> torch.Tensor: + """Return the temperature/top-k/top-p distribution of every row as ``[n, V]`` float32. + + Row ``i`` equals ``Sampler.probabilities(logits[i], SamplingParams(temperature=t_i, + top_k=k_i, top_p=p_i))``: top-k keeps exactly ``min(k, V)`` candidates by index (ties + resolve to an unspecified index), top-p uses the strict ``>`` cumulative test and keeps + the token that crosses the threshold. Rows with ``top_k <= 0`` or ``top_k >= V`` are not + truncated, and rows with ``top_p == 1`` are never trimmed. + + Only non-greedy rows may be passed (``temperature > 0``). No host synchronization + happens here: ``max_top_k`` (the largest effective ``k`` over truncating rows, 0 if no + row truncates) and ``apply_top_p`` (whether any row has ``top_p < 1``) come from the + host-side plan. Inputs are never modified. + """ + if logits.ndim != 2: + raise ValueError("logits must be a 2-D [n, V] tensor") + n, vocab = logits.shape + for name, value in (("temperature", temperature), ("top_k", top_k), ("top_p", top_p)): + if value.shape != (n,): + raise ValueError(f"{name} must have shape [{n}], got {tuple(value.shape)}") + if max_top_k < 0: + raise ValueError("max_top_k must be nonnegative") + + x = logits.float() / temperature.float()[:, None] + if max_top_k > 0: + k_eff = torch.where((top_k > 0) & (top_k < vocab), top_k, vocab) + truncates = k_eff < vocab + values, indices = x.topk(min(max_top_k, vocab), dim=-1) + keep = torch.arange(values.shape[-1], device=x.device) < k_eff[:, None] + masked = torch.full_like(x, -torch.inf).scatter_( + 1, indices, values.masked_fill(~keep, -torch.inf) + ) + x = torch.where(truncates[:, None], masked, x) + if apply_top_p: + p = top_p.float()[:, None] + sorted_x, indices = x.sort(dim=-1, descending=True) + remove = (sorted_x.softmax(-1).cumsum(-1) > p) & (p < 1) + remove[:, 1:] = remove[:, :-1].clone() + remove[:, 0] = False + x = x.scatter(1, indices, sorted_x.masked_fill(remove, -torch.inf)) + return x.softmax(-1) + + +def categorical(weights: torch.Tensor, u: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Inverse-CDF draw from every row of nonnegative, possibly unnormalized ``[m, V]`` weights. + + ``u`` is ``[m]`` uniforms in ``[0, 1)``. Returns ``(index [m] int64, valid [m] bool)``; a + row is valid iff its float64 mass is finite and positive. Invalid rows return the legal + placeholder index 0. Every valid draw lands on a token with a real cumulative increment + (so never on a zero-weight token). No host synchronization happens here. + """ + if weights.ndim != 2 or weights.shape[-1] < 1: + raise ValueError("weights must be a 2-D [m, V] tensor with V >= 1") + if u.shape != weights.shape[:1]: + raise ValueError(f"u must have shape [{weights.shape[0]}], got {tuple(u.shape)}") + vocab = weights.shape[-1] + # Accumulate in float64: each token's selection probability then equals its weight up to + # float64 rounding (an FP32 CDF quantizes small tokens after a large mass). + c = weights.double().cumsum(-1) + mass = c[:, -1] + valid = torch.isfinite(mass) & (mass > 0) + # Scale by the actual mass and clamp the threshold below it (not the index). With + # right=True the chosen i satisfies c[i] > t >= c[i-1], so it lies in an interval with a + # real cumulative increment: never a zero-weight token and never a positive-weight tail + # token whose increment was rounded away (a plateau). t < mass guarantees i <= V - 1. + t = u.double() * mass + t = torch.minimum(t, torch.nextafter(mass, torch.zeros_like(mass))) + index = torch.searchsorted(c, t[:, None], right=True).squeeze(-1) + index = torch.where(valid, index, torch.zeros_like(index)).clamp(0, vocab - 1) + return index, valid + + +class AcceptResult(NamedTuple): + """Device-side speculative verification result; consume only after checking ``error``. + + ``tokens`` ``[n, W1 + 1]``: ``draft_tokens[:s]``, then the final token, then zeros. + ``accepted`` ``[n]``: number of leading accepted drafts ``s``. ``final`` ``[n]``: the + residual/bonus token. ``error`` ``[n]`` bool: the stop point was numerically invalid; + all other fields of such a row are legal placeholder indices (0), not results. + """ + + tokens: torch.Tensor + accepted: torch.Tensor + final: torch.Tensor + error: torch.Tensor + + +def speculative_accept( + draft_tokens: torch.Tensor, + drafts: torch.Tensor, + q: torch.Tensor, + p: torch.Tensor, + u_accept: torch.Tensor, + u_final: torch.Tensor, +) -> AcceptResult: + """Verify sampled drafts with speculative rejection sampling, batched and host-sync-free. + + ``draft_tokens`` ``[n, W1]`` int64 were sampled from ``q`` ``[n, W1, V]``; ``drafts`` ``[n]`` + is the real draft count ``d_i <= W1`` (columns ``j >= d_i`` are padding and are ignored, + even if they hold NaN). ``p`` ``[n, W1 + 1, V]`` are target weights: row ``j`` verifies + draft ``j`` and row ``d_i`` is the bonus row. ``u_accept`` ``[n, W1]`` and ``u_final`` + ``[n]`` are uniforms in ``[0, 1)``. W1 may be 0. + + ``p`` and ``q`` rows are normalized per row here (``w / w.sum()``) and every comparison, + the residual and the bonus use these normalized values; a row is valid iff its raw mass + is finite and positive. Each position ``j < d_i`` is a valid accept, a valid reject, or + invalid (invalid row, ``q(x) == 0``, non-finite value at ``x``, or an out-of-range + token). The stop point ``s`` is the number of leading valid accepts. Then ``final`` is a + residual draw if position ``s`` is a valid reject, a bonus draw from ``p[s]`` if + ``s == d_i``, and an error if position ``s`` is invalid or the final distribution has no + valid mass. Positions after ``s`` never influence the result. Inputs are not modified. + """ + n, w1 = draft_tokens.shape + vocab = p.shape[-1] + if drafts.shape != (n,) or u_final.shape != (n,) or u_accept.shape != (n, w1): + raise ValueError("drafts, u_accept and u_final must match draft_tokens [n, W1]") + if q.shape != (n, w1, vocab) or p.shape != (n, w1 + 1, vocab): + raise ValueError("q must be [n, W1, V] and p must be [n, W1 + 1, V]") + + rows = torch.arange(n, device=p.device) + pos = torch.arange(w1 + 1, device=p.device) + p_mass = p.sum(-1, dtype=torch.float32) # [n, W1 + 1] + p_ok = torch.isfinite(p_mass) & (p_mass > 0) + d = drafts.clamp(0, w1) + error = drafts != d # out-of-range counts cannot be honored; flag instead of guessing + x = draft_tokens.clamp(0, vocab - 1) # legal gather/output indices even for bad drafts + + if w1 > 0: + in_range = (draft_tokens >= 0) & (draft_tokens < vocab) + q_mass = q.sum(-1, dtype=torch.float32) # [n, W1] + q_ok = torch.isfinite(q_mass) & (q_mass > 0) + # Elementwise w / mass gives exactly the entries of the normalized rows. + px = p[:, :w1].gather(2, x[..., None]).squeeze(-1).float() / p_mass[:, :w1] + qx = q.gather(2, x[..., None]).squeeze(-1).float() / q_mass + invalid = ~( + p_ok[:, :w1] & q_ok & in_range & torch.isfinite(px) & torch.isfinite(qx) & (qx > 0) + ) + # p >= q ignores u: u * q can round back to q for tiny q and falsely reject p == q. + accept = ~invalid & ((px >= qx) | (u_accept.float() * qx < px)) + accept &= pos[None, :w1] < d[:, None] + s = accept.long().cumprod(1).sum(1) # leading valid accepts, s <= d + at_end = s == d + s_q = s.clamp(max=w1 - 1) + invalid_at_s = invalid[rows, s_q] & ~at_end + + p_row = p[rows, s].float() / p_mass[rows, s][:, None] + q_row = q[rows, s_q].float() / q_mass[rows, s_q][:, None] + weights = torch.where(at_end[:, None], p_row, (p_row - q_row).clamp_min(0)) + stop_ok = torch.where(at_end, p_ok[rows, s], ~invalid_at_s) + else: + s = torch.zeros_like(d) + weights = p[:, 0].float() / p_mass[:, :1] + stop_ok = p_ok[:, 0] + + final, final_valid = categorical(weights, u_final) + error = error | ~(stop_ok & final_valid) + final = torch.where(error, torch.zeros_like(final), final) + + drafts_padded = torch.nn.functional.pad(x, (0, 1)) + tokens = torch.where( + pos[None] < s[:, None], + drafts_padded, + torch.where(pos[None] == s[:, None], final[:, None], torch.zeros_like(drafts_padded)), + ) + return AcceptResult(tokens, s, final, error)