[FEAT] Support random sampling in async block speculation

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

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

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)