From c11a7eb9b19f73385df6451bedf9ba65ec448825 Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Wed, 23 Sep 2026 15:21:38 +0800
Subject: [PATCH 1/8] refactor: extract Sampler class from
 ModelRunner._sample_tensor

Move the greedy / top-k / top-p arithmetic into worker/sampler.py verbatim.
Sampler holds no per-request state: the caller passes the generator in and
stores the possibly created generator back, so where the RNG lives stays
visible at the call site.

ModelRunner._sample_tensor becomes a thin delegate, preserving the existing
tests/test_async_pipeline.py monkeypatch target and the async CODA path.

No behavior change: sampling arithmetic, dtype promotion, kernel order and
the absence of host synchronization are unchanged.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 vllm_rlt/worker/model_runner.py |  8 +++++-
 vllm_rlt/worker/sampler.py      | 45 +++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 vllm_rlt/worker/sampler.py

diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py
index 9d8ab783..767f02a6 100644
--- a/vllm_rlt/worker/model_runner.py
+++ b/vllm_rlt/worker/model_runner.py
@@ -13,6 +13,7 @@
 from vllm_rlt.worker.buffers import Workspace
 from vllm_rlt.worker.cuda_graph import RecurrentGraphs
 from vllm_rlt.worker.sampling import sample_logits
+from vllm_rlt.worker.sampler import Sampler
 
 
 @dataclass
@@ -79,6 +80,7 @@ def __init__(
         self.cache_manager = cache_manager
         parameter = next(model.parameters())
         self.device = parameter.device
+        self.sampler = Sampler(self.device)
         self.exit_config = exit_config or ExitConfig()
         self.execution_config = execution_config or ExecutionConfig()
         scheduler = scheduler_config or SchedulerConfig()
@@ -601,4 +603,8 @@ def _sample(self, logits: torch.Tensor, request: Request) -> int:
         return int(self._sample_tensor(logits, request).item())
 
     def _sample_tensor(self, logits: torch.Tensor, request: Request):
-        return sample_logits(logits, request)
+        # Thin delegate kept for the async path and for tests that monkeypatch
+        # this method; the algorithm lives in Sampler.
+        token, generator = self.sampler.sample(logits, request.sampling_params, request.generator)
+        request.generator = generator
+        return token
diff --git a/vllm_rlt/worker/sampler.py b/vllm_rlt/worker/sampler.py
new file mode 100644
index 00000000..4bd8f31b
--- /dev/null
+++ b/vllm_rlt/worker/sampler.py
@@ -0,0 +1,45 @@
+"""Greedy, top-k and top-p sampling for a single logits row.
+
+The arithmetic is a verbatim move out of ``ModelRunner._sample_tensor`` so the
+sampling contract can be pinned by unit tests without a runner. This class owns
+no per-request state: the caller stores and passes in the RNG generator.
+"""
+
+import torch
+
+from vllm_rlt.sampling_params import SamplingParams
+
+
+class Sampler:
+    """Greedy / top-k / top-p sampling for one logits row."""
+
+    def __init__(self, device: torch.device):
+        self.device = device
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        params: SamplingParams,
+        generator: torch.Generator | None,
+    ) -> tuple[torch.Tensor, torch.Generator | None]:
+        """Return a 0-dim long device tensor and the possibly created generator.
+
+        ``params`` is assumed to be validated by ``SamplingParams.__post_init__``.
+        No host synchronization happens here: the returned tensor stays on the
+        logits device so the async CODA path can feed it to the next prelude.
+        """
+        if params.temperature == 0:
+            return logits.argmax(), generator
+        logits = logits.float() / params.temperature
+        if params.top_k > 0:
+            threshold = logits.topk(min(params.top_k, logits.numel())).values[-1]
+            logits = logits.masked_fill(logits < threshold, -torch.inf)
+        if params.top_p < 1:
+            sorted_logits, indices = logits.sort(descending=True)
+            remove = sorted_logits.softmax(-1).cumsum(-1) > params.top_p
+            remove[1:] = remove[:-1].clone()
+            remove[0] = False
+            logits = logits.scatter(0, indices, sorted_logits.masked_fill(remove, -torch.inf))
+        if generator is None:
+            generator = torch.Generator(device=self.device).manual_seed(params.seed)
+        return torch.multinomial(logits.softmax(-1), 1, generator=generator).squeeze(0), generator

From 27e8d1a71b73c0fff5c0d240c472934b383e537a Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Wed, 23 Sep 2026 15:28:00 +0800
Subject: [PATCH 2/8] cleanup: remove unused _sample method

ModelRunner._sample had no callers anywhere in the repository: it was a
host-side int wrapper kept next to _sample_tensor. Tests only monkeypatch
_sample_tensor (tests/test_async_pipeline.py:138), and the CODA path calls
_sample_tensor directly (model_runner.py:425).

No behavior change.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 vllm_rlt/worker/model_runner.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py
index 767f02a6..c78a90ca 100644
--- a/vllm_rlt/worker/model_runner.py
+++ b/vllm_rlt/worker/model_runner.py
@@ -599,9 +599,6 @@ def synchronize(self):
             if stream is not None:
                 stream.synchronize()
 
-    def _sample(self, logits: torch.Tensor, request: Request) -> int:
-        return int(self._sample_tensor(logits, request).item())
-
     def _sample_tensor(self, logits: torch.Tensor, request: Request):
         # Thin delegate kept for the async path and for tests that monkeypatch
         # this method; the algorithm lives in Sampler.

From dd71ad425e53cd4dce2320406045459a7d4c9bb7 Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Wed, 23 Sep 2026 15:32:11 +0800
Subject: [PATCH 3/8] test: pin sampling algorithm semantics

Add direct Sampler.sample() contract tests, without a runner:

- greedy returns argmax, a 0-dim long tensor, and never creates a generator
  even when top_k/top_p are set (temperature == 0 short circuits first)
- top_k=1 matches greedy across several logits rows and seeds: the threshold
  is the maximum, so the single candidate cannot depend on the RNG stream
- top_k keeps every token tied at the threshold, so the candidate set can
  exceed k (masked_fill(logits < threshold))
- top_p keeps at least one token when the threshold would empty the set
- top_p=1.0 leaves the full support reachable

These pin current behavior; they do not propose new sampling semantics.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 tests/test_sampler.py | 93 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)
 create mode 100644 tests/test_sampler.py

diff --git a/tests/test_sampler.py b/tests/test_sampler.py
new file mode 100644
index 00000000..a3ac453e
--- /dev/null
+++ b/tests/test_sampler.py
@@ -0,0 +1,93 @@
+"""Contract tests for the extracted Sampler: greedy, top-k and top-p.
+
+These pin the current arithmetic exactly as it behaved inside
+``ModelRunner._sample_tensor``; they are not a specification of what sampling
+should do in the future.
+"""
+
+import pytest
+import torch
+
+from vllm_rlt.sampling_params import SamplingParams
+from vllm_rlt.worker.sampler import Sampler
+
+
+def draw(sampler, logits, params, generator=None, count=1):
+    """Sample ``count`` tokens, threading the generator like the runner does."""
+    tokens = []
+    for _ in range(count):
+        token, generator = sampler.sample(logits, params, generator)
+        tokens.append(token)
+    return tokens, generator
+
+
+def test_greedy_returns_argmax_and_never_creates_a_generator():
+    sampler = Sampler(torch.device("cpu"))
+    logits = torch.tensor([0.1, 2.0, 1.0, -1.0])
+    # top_k/top_p are deliberately set: the temperature == 0 short circuit runs
+    # first, so greedy ignores them and stays RNG-free.
+    params = SamplingParams(temperature=0.0, top_k=1, top_p=0.5)
+    tokens, generator = draw(sampler, logits, params)
+    token = tokens[0]
+    assert token.item() == int(logits.argmax())
+    assert token.dim() == 0
+    assert token.dtype == torch.long
+    assert generator is None
+
+
+GREEDY_CASES = [
+    torch.tensor([0.1, 2.0, 1.0, -1.0]),
+    torch.tensor([-3.0, -3.5, -2.0, -9.0]),
+    torch.tensor([5.0, 4.999, 4.998]),
+]
+
+
+@pytest.mark.parametrize("logits", GREEDY_CASES)
+def test_top_k_one_is_equivalent_to_greedy(logits):
+    sampler = Sampler(torch.device("cpu"))
+    greedy_tokens, _ = draw(sampler, logits, SamplingParams(temperature=0.0))
+    # top_k=1 makes the threshold the maximum, so masked_fill leaves exactly one
+    # candidate and the result cannot depend on the RNG stream. Checking several
+    # logits rows and seeds pins that mechanism, not one coincidental argmax.
+    for seed in (0, 7, 12345):
+        params = SamplingParams(temperature=1.0, top_k=1, seed=seed)
+        top_k_tokens, generator = draw(sampler, logits, params)
+        assert generator is not None
+        assert top_k_tokens[0].item() == greedy_tokens[0].item()
+
+
+def test_top_k_keeps_every_token_tied_at_the_threshold():
+    sampler = Sampler(torch.device("cpu"))
+    # top_k=2 selects a threshold of 3.0, which three tokens share. The mask is
+    # ``logits < threshold``, so the candidate set is larger than k.
+    logits = torch.tensor([3.0, 3.0, 3.0, 0.0])
+    params = SamplingParams(temperature=1.0, top_k=2, seed=11)
+    tokens, _ = draw(sampler, logits, params, count=200)
+    observed = {token.item() for token in tokens}
+    assert observed <= {0, 1, 2}
+    assert len(observed) > 1
+
+
+def test_top_p_keeps_at_least_one_token():
+    sampler = Sampler(torch.device("cpu"))
+    # A top_p far below the head probability would empty the candidate set; the
+    # mask shift keeps the single most likely token instead.
+    logits = torch.linspace(1.0, 0.0, 100)
+    params = SamplingParams(temperature=1.0, top_p=0.005, seed=3)
+    tokens, _ = draw(sampler, logits, params, count=20)
+    assert {token.item() for token in tokens} == {0}
+
+
+def test_top_p_one_does_not_truncate_the_support():
+    sampler = Sampler(torch.device("cpu"))
+    logits = torch.tensor([1.0, 0.9, 0.8, 0.7])
+    truncated = SamplingParams(temperature=1.0, top_p=0.5, seed=13)
+    untruncated = SamplingParams(temperature=1.0, top_p=1.0, seed=13)
+    truncated_tokens, _ = draw(sampler, logits, truncated, count=200)
+    untruncated_tokens, _ = draw(sampler, logits, untruncated, count=200)
+    assert {token.item() for token in truncated_tokens} <= {0, 1}
+    # Index 3 is the least likely token, so it is reachable only while the
+    # support is intact. This asserts reachability over a fixed 200-draw stream
+    # (p(miss) ~ 0.79**200), not the exact draw sequence: the exact sequence
+    # would couple the test to torch's multinomial implementation.
+    assert 3 in {token.item() for token in untruncated_tokens}

From f13a0673fa0d9511ab76237cc30f4cf8ca1f0c27 Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Wed, 23 Sep 2026 16:11:26 +0800
Subject: [PATCH 4/8] test: pin seed reproducibility and generator isolation

Sampler level:
- the same seed replays the same token sequence; different seeds diverge
- later calls advance the generator handed in rather than reseeding it

Engine level:
- a finished request drops its generator, and reusing the request ID with the
  same seed replays the same tokens instead of continuing the old stream
- greedy requests never create a generator at any step

These pin the RNG lifecycle that PR 1 must not change; the runner-side RNG
registry and suspend/release split land in M7 2/3.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 tests/test_engine.py  | 36 ++++++++++++++++++++++++++++++++++++
 tests/test_sampler.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/tests/test_engine.py b/tests/test_engine.py
index 9d60eb12..52e5be47 100644
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -154,6 +154,42 @@ def test_sampled_generation_is_batch_invariant():
     assert serial.token_ids == batched.token_ids
 
 
+def test_same_id_after_finish_starts_from_a_fresh_generator():
+    params = SamplingParams(max_tokens=3, temperature=0.8, seed=17, ignore_eos=True)
+    engine = LLMEngine(tiny_model(), cache_config=CacheConfig(64, 2))
+    engine.add_request("r", [2, 4], params)
+    request = engine.scheduler.requests["r"]
+    for _ in range(20):
+        if request.generator is not None:
+            break
+        engine.step()
+    assert request.generator is not None
+    first, _ = drain(engine)
+    # finish() pops the request from scheduler.requests, but this local name is
+    # still a live reference to the same object, so it can observe the released
+    # RNG: reusing the ID must not continue the previous stream.
+    assert request.generator is None
+    engine.add_request("r", [2, 4], params)
+    assert engine.scheduler.requests["r"] is not request
+    second, _ = drain(engine)
+    assert second["r"].token_ids == first["r"].token_ids
+    assert second["r"].exit_depths == first["r"].exit_depths
+
+
+def test_greedy_requests_never_create_a_generator():
+    engine = LLMEngine(tiny_model(), cache_config=CacheConfig(64, 2))
+    engine.add_request("g", [2, 4], SamplingParams(max_tokens=3, ignore_eos=True))
+    request = engine.scheduler.requests["g"]
+    for _ in range(1000):
+        if not engine.has_unfinished_requests():
+            break
+        engine.step()
+        assert request.generator is None
+    else:
+        pytest.fail("engine failed to finish within bounded scheduler steps")
+    assert request.generator is None
+
+
 def test_execution_failure_releases_affected_requests(monkeypatch):
     engine = LLMEngine(tiny_model())
     engine.add_request("one", [2], SamplingParams(max_tokens=2))
diff --git a/tests/test_sampler.py b/tests/test_sampler.py
index a3ac453e..27c521bf 100644
--- a/tests/test_sampler.py
+++ b/tests/test_sampler.py
@@ -91,3 +91,38 @@ def test_top_p_one_does_not_truncate_the_support():
     # (p(miss) ~ 0.79**200), not the exact draw sequence: the exact sequence
     # would couple the test to torch's multinomial implementation.
     assert 3 in {token.item() for token in untruncated_tokens}
+
+
+def sequence(sampler, logits, params, count=32):
+    tokens, _ = draw(sampler, logits, params, count=count)
+    return [token.item() for token in tokens]
+
+
+def test_same_seed_reproduces_the_same_sequence():
+    sampler = Sampler(torch.device("cpu"))
+    logits = torch.linspace(1.0, -1.0, 64)
+    params = SamplingParams(temperature=0.7, top_k=8, top_p=0.9, seed=42)
+    assert sequence(sampler, logits, params) == sequence(sampler, logits, params)
+
+
+def test_different_seeds_produce_different_sequences():
+    sampler = Sampler(torch.device("cpu"))
+    logits = torch.linspace(1.0, -1.0, 64)
+    base = dict(temperature=0.7, top_k=8, top_p=0.9)
+    assert sequence(sampler, logits, SamplingParams(seed=1, **base)) != sequence(
+        sampler, logits, SamplingParams(seed=2, **base)
+    )
+
+
+def test_generator_is_reused_and_advanced_across_calls():
+    sampler = Sampler(torch.device("cpu"))
+    logits = torch.linspace(1.0, -1.0, 64)
+    params = SamplingParams(temperature=0.7, top_k=8, seed=42)
+    _, generator = draw(sampler, logits, params)
+    assert generator is not None
+    state = generator.get_state().clone()
+    # Later calls must advance the same generator instead of reseeding it, which
+    # is what makes a multi-token request reproducible.
+    _, same_generator = draw(sampler, logits, params, generator=generator)
+    assert same_generator is generator
+    assert not torch.equal(generator.get_state(), state)

From 9709c8b29e451c44c1bbfd01d76ace78599e458b Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Wed, 23 Sep 2026 16:31:35 +0800
Subject: [PATCH 5/8] test: pin preemption RNG continuity

Extend the existing lossless-preemption test with explicit RNG assertions
instead of only comparing tokens:

- before suspension the request owns a generator; preempt() keeps both the
  object and its state, because the RNG survives only via the retained
  Request object and is not part of the KV/hidden snapshot
- after resumption the generator has advanced from that state (no reseed),
  which is what makes the token sequence match the uninterrupted run
- termination then releases the RNG slot

The token and exit-depth equality assertions stay as they were. This pins
the constraint that M7 2/3 must not break when it splits the release path.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 tests/test_prefix_growth.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tests/test_prefix_growth.py b/tests/test_prefix_growth.py
index dbfa9ebf..fc2468e6 100644
--- a/tests/test_prefix_growth.py
+++ b/tests/test_prefix_growth.py
@@ -96,13 +96,26 @@ def test_lossless_preemption_preserves_looped_history():
         while len(e.scheduler.requests["a"].generated_token_ids) < 3:
             e.step()
         if suspend:
+            request = e.scheduler.requests["a"]
+            generator = request.generator
+            assert generator is not None
+            state = generator.get_state().clone()
             e.add_request("b", [5], SamplingParams(max_tokens=1))
             e.scheduler.selected_request_ids.clear()
             assert e.preemption.preempt(e.scheduler.requests["b"])
             assert "a" in e.preemption.snapshots
+            # Suspension shares release() with termination, and the RNG survives
+            # only because the Request object is retained: it is not part of the
+            # KV/hidden snapshot.
+            assert request.generator is generator
+            assert torch.equal(generator.get_state(), state)
         results.append(finish(e)["a"])
         if suspend:
             assert e.preemption.resumptions == 1
+            # Sampling continued from the preserved state instead of reseeding,
+            # and termination then released the RNG slot.
+            assert not torch.equal(generator.get_state(), state)
+            assert request.generator is None
     assert results[0].token_ids == results[1].token_ids
     assert results[0].exit_depths == results[1].exit_depths
 

From a91c256b63fb042579389bb0a59ab29c4c11529e Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Wed, 23 Sep 2026 16:34:25 +0800
Subject: [PATCH 6/8] docs: add sampling module walkthrough

Add docs/sampling_walkthrough.md following the RFC module structure:
baseline and boundary, classes and state, functions and parameters,
line-by-line coverage, evidence-based findings, target design and migration,
and validation.

It records the findings behind this increment, including:

- the FP32 logits promotion before the temperature division is a project
  difference to record under docs/precision-policy.md, not reference behavior
- the per-row CODA sampling loop is a known limitation, not optimized here
- preemption and termination share model_runner.release(), so the RNG
  migration in M7 2/3 must split suspend from release
- the ModelRunner._sample_tensor delegate stays until its monkeypatch target
  in tests/test_async_pipeline.py is updated

GPU validation is marked pending in M7 3/3. Link the note from docs/README.md.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 docs/README.md               |   1 +
 docs/sampling_walkthrough.md | 275 +++++++++++++++++++++++++++++++++++
 2 files changed, 276 insertions(+)
 create mode 100644 docs/sampling_walkthrough.md

diff --git a/docs/README.md b/docs/README.md
index 277c21c8..31940527 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -9,6 +9,7 @@ an environment to receiving the first generated response. The
 | [First-run user guide](launching.md) | Environment setup, model download, server startup, first request, CLI, Python API, and troubleshooting |
 | [Engine design](design.md) | Model stages, scheduling, and KV ownership |
 | [Scheduler walkthrough](scheduler_walkthrough.md) | Responsibilities, admission cases, control flow, and refactoring checklist |
+| [Sampling walkthrough](sampling_walkthrough.md) | Sampling algorithm, RNG lifecycle across preemption and termination, and known limits |
 | [KV layout examples](kv_layout_computation.md) | SHARED and LAST_EXITED semantics and worked attention examples |
 | [Runtime configuration](cdb_runtime.md) | Exit policies, KV layouts, execution options, and CUDA Graphs |
 | [Asynchronous scheduling](https://github.com/hsliuustc0106/vllm-rlt/pull/30) | CPU/GPU pipelining and single-stream or multi-stream execution |
diff --git a/docs/sampling_walkthrough.md b/docs/sampling_walkthrough.md
new file mode 100644
index 00000000..2e59c1cc
--- /dev/null
+++ b/docs/sampling_walkthrough.md
@@ -0,0 +1,275 @@
+# Sampling Module Walkthrough
+
+This note documents the M7 (Sampling) increment of RFC
+[#32](https://github.com/hsliuustc0106/vllm-rlt/issues/32). It covers
+[`worker/sampler.py`](../vllm_rlt/worker/sampler.py) and the sampling call sites in
+[`worker/model_runner.py`](../vllm_rlt/worker/model_runner.py),
+[`request.py`](../vllm_rlt/request.py),
+[`core/scheduler.py`](../vllm_rlt/core/scheduler.py),
+[`engine/preemption.py`](../vllm_rlt/engine/preemption.py),
+[`engine/llm_engine.py`](../vllm_rlt/engine/llm_engine.py) and
+[`pd/worker.py`](../vllm_rlt/pd/worker.py).
+
+M7 is split into three increments: `M7 1/3` extracts the algorithm and pins it with CPU
+contract tests, `M7 2/3` migrates RNG ownership and splits the release path, and
+`M7 3/3` adds GPU validation. This document is written during `M7 1/3`; sections that
+describe later increments say so explicitly.
+
+## 1. Baseline and boundary
+
+### Call chain
+
+```text
+Scheduler.schedule() -> CODA batch
+  -> ModelRunner.execute()                    # sync path
+     -> _execute()
+        -> model.coda(hidden)                 # lm_head logits rows
+        -> _sample_tensor(row, request)       # model_runner.py:425
+        -> result.cpu().tolist()              # the only host readback, in execute()
+  -> ModelRunner.submit()                     # async path
+     -> _execute() samples on the device
+     -> routing.scatter(result, tokens=True)  # sampled IDs stay on the device
+```
+
+The async path requires the sampled value to remain a 0-dim device tensor: it is
+`torch.stack`ed and scattered into the next prelude, so sampling must not call
+`.item()`, `.cpu()`, `synchronize()` or add host work. The sync path's
+`.cpu().tolist()` happens in `execute()`, outside the sampling function.
+
+### In scope
+
+| File | Role |
+| --- | --- |
+| `vllm_rlt/sampling_params.py` | Parameter contract and validation |
+| `vllm_rlt/worker/sampler.py` | Greedy / top-k / top-p algorithm for one logits row |
+| `vllm_rlt/worker/model_runner.py` | Sampling call site and RNG storage (`_sample_tensor`) |
+| `vllm_rlt/request.py` | Current RNG location (`Request.generator`) |
+| `vllm_rlt/core/scheduler.py` | RNG release on termination |
+| `vllm_rlt/engine/preemption.py` | Preemption contract that must preserve the RNG |
+| `tests/test_sampler.py` | Algorithm, seed and generator contracts |
+| `tests/test_engine.py`, `tests/test_prefix_growth.py` | RNG lifecycle across termination and preemption |
+
+### Out of scope
+
+Exit policy and stage transitions (M3), KV allocation and prepared metadata (M4),
+attention backends (M8), the PD connector protocol (M9), new sampling features
+(`min_p`, `repetition_penalty`, `logprobs`, `n > 1`), any change to the sampling
+arithmetic or to user-visible defaults, and the batched rewrite of the per-row CODA
+loop (see "Known limitation" below).
+
+## 2. Classes and state
+
+### `SamplingParams` (frozen dataclass)
+
+Defaults are `max_tokens=16`, `temperature=0.0`, `top_p=1.0`, `top_k=-1`, `seed=0`,
+plus the Ouro extensions `min_loops`, `max_loops`, `exit_threshold`, `ignore_eos` and
+`priority`. `__post_init__` already validates everything sampling depends on:
+`temperature` is finite and nonnegative, `top_p ∈ (0, 1]`, `top_k` is `-1` or a
+positive integer, and `seed ∈ [0, 2**63)`. `Sampler` therefore assumes validated
+input and does not re-check it.
+
+### `Sampler`
+
+`Sampler` owns only the device used to create a new generator. It holds no
+per-request state: `sample()` receives the generator and returns the possibly created
+one, so "where the RNG lives" stays visible at the call site. This is what lets
+`M7 2/3` move RNG storage without touching the algorithm or its tests.
+
+### RNG lifecycle today
+
+| Event | Current behavior | Location |
+| --- | --- | --- |
+| First random sample | Lazily creates `torch.Generator(device=...)` seeded from `params.seed` | `sampler.py:43-44` |
+| Greedy sample | `argmax`, no generator created or advanced | `sampler.py:31-32` |
+| Request finishes | `scheduler.finish` sets `request.generator = None` | `scheduler.py:94` |
+| Request preempted | Not reset; the `Request` object and its generator are retained | `preemption.py:3`, `preemption.py:102` |
+| Request resumed | Not rebuilt; sampling continues on the same generator | `preemption.py:114-144` |
+
+## 3. Functions and parameters
+
+```python
+class Sampler:
+    def __init__(self, device: torch.device): ...
+
+    def sample(
+        self,
+        logits: torch.Tensor,                 # [vocab] 1-dim device tensor, any float dtype; caller guarantees ndim == 1
+        params: SamplingParams,               # validated by __post_init__
+        generator: torch.Generator | None,    # None on the greedy path
+    ) -> tuple[torch.Tensor, torch.Generator | None]:
+        """Return a 0-dim long device tensor and the possibly created generator."""
+```
+
+`ModelRunner` keeps a thin delegate so the async path and existing tests are
+unaffected:
+
+```python
+def _sample_tensor(self, logits: torch.Tensor, request: Request):
+    token, generator = self.sampler.sample(logits, request.sampling_params, request.generator)
+    request.generator = generator
+    return token
+```
+
+## 4. Line-by-line coverage
+
+`Sampler.sample` (`sampler.py:19-45`):
+
+| Lines | Behavior |
+| --- | --- |
+| `31-32` | `temperature == 0` is a greedy short circuit. It runs **before** any RNG work, so greedy requests never create a generator. |
+| `33` | Non-greedy logits are promoted with `float()` and divided by temperature. |
+| `34-36` | `top_k`: the threshold is the k-th largest value and the mask is `logits < threshold`, so every token tied at the threshold survives and the candidate set can exceed `k`. |
+| `37-42` | `top_p`: sort descending, cumulative softmax, then shift the mask by one position and force `remove[0] = False`, so at least one token always survives. |
+| `43-44` | Lazy generator creation with `params.seed`, on `self.device`. |
+| `45` | `torch.multinomial(logits.softmax(-1), 1, generator=...)` returns a 0-dim long tensor on the logits device. |
+
+`ModelRunner` (`model_runner.py`):
+
+| Lines | Behavior |
+| --- | --- |
+| `82` | `self.sampler = Sampler(self.device)` |
+| `425` | CODA calls `_sample_tensor` once per request row and `torch.stack`s the results. |
+| `601-606` | Delegate: reads `request.generator`, calls `Sampler.sample`, stores the result back. |
+
+### Documented project difference
+
+`sampler.py:33` promotes logits to FP32 before dividing by temperature. The pinned
+model file (`modeling_ouro.py` at revision `574fa66`) inherits from `GenerationMixin`
+and delegates sampling to the Transformers generation stack; it contains no
+`multinomial`, `top_k`, `top_p`, `temperature`, or logits-processor code of its own.
+This project does not pin that generation stack, so there is no pinned reference
+sampling path to compare against. The FP32 promotion is therefore a project choice
+awaiting comparison with a real reference sampling path, not a difference from a
+pinned sampling implementation. This is the "record the sampling implementation
+separately" item required by [precision policy](precision-policy.md). `M7 1/3`
+preserves it unchanged.
+
+### Known limitation
+
+`model_runner.py:424-426` samples each request row in a Python loop, so every row pays
+its own `sort`/`topk` over the model vocabulary (`vocab_size = 49152` in `config.py`). vLLM instead batches the multinomial
+across the batch. This is recorded as a known limitation and is deliberately not
+optimized in M7; any batched rewrite must keep the arithmetic order
+`float()/temperature → top_k → top_p → multinomial` and the no-host-sync guarantee.
+
+## 5. Evidence-based findings
+
+**Dead code: `_sample` had no callers (cleanup).** Repository-wide search found
+only the definition, no call site; tests monkeypatch `_sample_tensor`
+(`tests/test_async_pipeline.py:138`). Removed in `M7 1/3`.
+
+**RNG ownership disagrees with the RFC state table.** The RFC assigns "RNG objects
+and their execution/snapshot state" to the worker, while the generator currently lives
+on the scheduler-side `Request` dataclass (`request.py:43`). Behavior is correct today
+because `Request` is a shared handle and preemption/resumption happens to retain it,
+but ownership is implicit: any change that treats `release()` as "termination releases
+everything" can silently reset the RNG.
+
+**Preemption and termination share `model_runner.release()` (the key design
+constraint).** `preemption.py:102` calls `release()` when suspending a victim;
+`llm_engine.py:155` (abort) and `llm_engine.py:261` (finish) call the same method, and
+`pd/worker.py:107` (removal) and `pd/worker.py:305` (prefill compute release) do as
+well. Moving the RNG into the runner and deleting it inside `release()` would reset it
+on preemption, changing the token sequence after restoration. That would violate the
+`preemption.py:3` contract and the RFC acceptance criterion for fixed-seed sampling
+across preemption/restoration. RNG migration must therefore distinguish "termination
+release" from "preemption suspend" (see section 6).
+
+**Sampling had almost no unit coverage (test gap).** Before this increment the
+only coverage was batch invariance (`tests/test_engine.py:147`), a monkeypatched
+EOS path (`tests/test_async_pipeline.py:138`), and an indirect preemption check:
+`tests/test_prefix_growth.py:85` compared token sequences after suspend/resume with
+`temperature=0.8, seed=45`, which would fail if the RNG were reseeded, but did not
+assert generator identity or state directly. Uncovered: greedy never creating a
+generator, `top_k=1` matching greedy, `top_k` tie behavior, `top_p` retaining one
+token, same/different seed reproducibility, direct generator identity and state
+assertions across preemption, and RNG isolation after a request ID is reused.
+`tests/test_sampler.py`, `tests/test_engine.py` and `tests/test_prefix_growth.py`
+now pin these.
+
+**The sampling implementation was not documented separately (documentation gap).**
+`docs/serving.md:55-56` lists only field defaults, and `docs/precision-policy.md:19`
+requires recording the sampling implementation separately. This walkthrough is that
+record.
+
+**Sampling is on the CODA hot path (performance constraint).** The async path
+stacks the sampled results and scatters them with `tokens=True`
+(`model_runner.py:424-428`). Extraction must not add `.item()`, `.cpu()`,
+`synchronize()`, or an extra kernel, and must not reorder the arithmetic.
+
+**The CODA per-row Python loop is a known limitation.** See section 4.
+
+## 6. Target design and migration
+
+`M7 1/3` (this increment) only moves the algorithm into `Sampler` and pins behavior.
+RNG ownership is unchanged: `request.generator` is still read and written by the
+delegate, and `release()` still does not touch it.
+
+`M7 2/3` will, in order:
+
+1. Add a runner-side RNG registry (`dict[str, torch.Generator]`) and have
+   `_sample_tensor` read and write it instead of `request.generator`.
+2. Split the release path: `suspend(request_id)` synchronizes the event and frees the
+   state slot but keeps the RNG (called by preemption); `release(request_id)` does the
+   same and additionally drops the RNG (called by abort, finish and PD removal).
+3. Remove `Request.generator` and the `scheduler.finish` assignment, then delete the
+   temporary adapters.
+
+The target RNG contract:
+
+| Event | Target behavior |
+| --- | --- |
+| First random sample | Create the generator from `params.seed` in the request's RNG slot |
+| Later samples | Reuse and advance the same generator (ordering semantics unchanged) |
+| Greedy request | Never create or advance a generator |
+| Preemption | Keep the generator (or snapshot `get_state()`) and continue from it on resume |
+| Termination (stop/length/abort/PD removal) | Drop the RNG slot so a new request with the same ID starts fresh |
+| Late async result | Must not write into a new request's RNG; the existing object-identity check covers results, and "release on termination" covers the RNG |
+
+### Removal condition for the delegate
+
+`ModelRunner._sample_tensor` exists so that `tests/test_async_pipeline.py:138` keeps
+working. It can be removed once that test patches `Sampler` instead, which is planned
+for `M7 2/3` when the runner-side registry lands.
+
+## 7. Validation
+
+`M7 1/3` CPU evidence (macOS, Python 3.12, PyTorch 2.14.0):
+
+| Suite | Result |
+| --- | --- |
+| `main @ 60af1cf` baseline (`test_engine.py test_prefix_growth.py test_async_pipeline.py`) | 47 passed, 15 skipped |
+| Same three files after this increment (baseline 47 + 2 new engine tests) | 49 passed, 15 skipped |
+| With `test_sampler.py` and `test_serving.py` added | 59 passed, 16 skipped |
+| Full CPU suite (`tests/`) | 211 passed, 115 skipped |
+
+```bash
+# main @ 60af1cf baseline (checkout 60af1cf first)
+OMP_NUM_THREADS=1 python -m pytest -q tests/test_engine.py tests/test_prefix_growth.py tests/test_async_pipeline.py
+# same three files after this increment
+OMP_NUM_THREADS=1 python -m pytest -q tests/test_engine.py tests/test_prefix_growth.py tests/test_async_pipeline.py
+# with new test files
+OMP_NUM_THREADS=1 python -m pytest -q tests/test_sampler.py tests/test_engine.py \
+    tests/test_prefix_growth.py tests/test_async_pipeline.py tests/test_serving.py
+# full CPU suite
+OMP_NUM_THREADS=1 python -m pytest -q tests/
+# lint
+python -m ruff check vllm_rlt/worker/sampler.py vllm_rlt/worker/model_runner.py tests/test_sampler.py
+python -m ruff format --check vllm_rlt/worker/sampler.py vllm_rlt/worker/model_runner.py tests/test_sampler.py
+```
+
+Ruff: all checks passed; 3 files already formatted.
+
+The extraction changes no arithmetic, so no new precision or performance evidence is
+required. Sampling sequences depend on `torch.multinomial`; the numbers above are with
+PyTorch 2.14.0.
+
+`M7 1/3` and `M7 2/3` do **not** include GPU validation. Device-side behavior under
+async scheduling and CUDA streams is covered by `M7 3/3`:
+
+```bash
+python -m pytest -q tests/ -m gpu --run-gpu
+```
+
+Until then, both increments state "GPU validation pending in M7 3/3" in their pull
+request descriptions.

From efaabab54b44dce4522446d2d07d86cf0243e962 Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Thu, 24 Sep 2026 14:06:25 +0800
Subject: [PATCH 7/8] cleanup: drop the stale sample_logits import

ModelRunner._sample_tensor delegates to Sampler after the extraction, so
sample_logits lost its only caller and the import became dead code
(ruff F401/I001 once the branch was rebased onto upstream/main 3314c1b).

Only the import is removed here. worker/sampling.py stays untouched for the
speculative decoding path introduced by PR #44; removing sample_logits() and
collapsing the two modules land in M7 2/3.

Refs #32

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 vllm_rlt/worker/model_runner.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py
index c78a90ca..b0025137 100644
--- a/vllm_rlt/worker/model_runner.py
+++ b/vllm_rlt/worker/model_runner.py
@@ -12,7 +12,6 @@
 from vllm_rlt.request import Request, Stage
 from vllm_rlt.worker.buffers import Workspace
 from vllm_rlt.worker.cuda_graph import RecurrentGraphs
-from vllm_rlt.worker.sampling import sample_logits
 from vllm_rlt.worker.sampler import Sampler
 
 

From de56fc9ecc2d9cbcad6ee698e158e189cc121818 Mon Sep 17 00:00:00 2001
From: ycjcx123 <ycjcx123@gmail.com>
Date: Thu, 24 Sep 2026 14:17:13 +0800
Subject: [PATCH 8/8] docs: update walkthrough for upstream PR #44 coexistence

Rebase onto upstream/main @ 3314c1b introduced worker/sampling.py
(speculative decoding). Update the walkthrough to cover both modules,
their call chains, and the M7 2/3 unification plan.

Signed-off-by: ycjcx123 <ycjcx123@gmail.com>
---
 docs/sampling_walkthrough.md | 129 +++++++++++++++++++++++++++--------
 1 file changed, 99 insertions(+), 30 deletions(-)

diff --git a/docs/sampling_walkthrough.md b/docs/sampling_walkthrough.md
index 2e59c1cc..28e0f1ba 100644
--- a/docs/sampling_walkthrough.md
+++ b/docs/sampling_walkthrough.md
@@ -2,7 +2,9 @@
 
 This note documents the M7 (Sampling) increment of RFC
 [#32](https://github.com/hsliuustc0106/vllm-rlt/issues/32). It covers
-[`worker/sampler.py`](../vllm_rlt/worker/sampler.py) and the sampling call sites in
+[`worker/sampler.py`](../vllm_rlt/worker/sampler.py), the shared helpers in
+[`worker/sampling.py`](../vllm_rlt/worker/sampling.py), the speculative call sites in
+[`worker/speculative.py`](../vllm_rlt/worker/speculative.py) and the sampling call sites in
 [`worker/model_runner.py`](../vllm_rlt/worker/model_runner.py),
 [`request.py`](../vllm_rlt/request.py),
 [`core/scheduler.py`](../vllm_rlt/core/scheduler.py),
@@ -11,12 +13,18 @@ This note documents the M7 (Sampling) increment of RFC
 [`pd/worker.py`](../vllm_rlt/pd/worker.py).
 
 M7 is split into three increments: `M7 1/3` extracts the algorithm and pins it with CPU
-contract tests, `M7 2/3` migrates RNG ownership and splits the release path, and
-`M7 3/3` adds GPU validation. This document is written during `M7 1/3`; sections that
-describe later increments say so explicitly.
+contract tests, `M7 2/3` collapses `worker/sampling.py` into `Sampler`, migrates RNG
+ownership and splits the release path, and `M7 3/3` adds GPU validation. This document is
+written during `M7 1/3`; sections that describe later increments say so explicitly.
 
 ## 1. Baseline and boundary
 
+Pinned baseline: `upstream/main @ 3314c1b`, the merge of PR #44. PR #44 added
+`worker/sampling.py` and routed `ModelRunner._sample_tensor` through
+`sampling.sample_logits()`; this increment replaces that one-line delegate with `Sampler`
+and leaves the speculative decoding path on `sampling.py`. Both modules therefore coexist
+until `M7 2/3` (section 6).
+
 ### Call chain
 
 ```text
@@ -29,6 +37,9 @@ Scheduler.schedule() -> CODA batch
   -> ModelRunner.submit()                     # async path
      -> _execute() samples on the device
      -> routing.scatter(result, tokens=True)  # sampled IDs stay on the device
+  -> SpeculativeRunner.execute()              # speculative path, sync eager only
+     -> sampling.probabilities(logits, ...)   # worker/sampling.py:6
+     -> sampling.draw() / rejection_sample()  # worker/sampling.py:28, :40
 ```
 
 The async path requires the sampled value to remain a 0-dim device tensor: it is
@@ -41,21 +52,25 @@ The async path requires the sampled value to remain a 0-dim device tensor: it is
 | File | Role |
 | --- | --- |
 | `vllm_rlt/sampling_params.py` | Parameter contract and validation |
-| `vllm_rlt/worker/sampler.py` | Greedy / top-k / top-p algorithm for one logits row |
+| `vllm_rlt/worker/sampler.py` | Greedy / top-k / top-p algorithm for one logits row; the M7 boundary |
+| `vllm_rlt/worker/sampling.py` | Distribution helpers and rejection sampling added by PR #44; unchanged here |
+| `vllm_rlt/worker/speculative.py` | Draft/verify loop that calls `sampling.py` directly; unchanged here |
 | `vllm_rlt/worker/model_runner.py` | Sampling call site and RNG storage (`_sample_tensor`) |
 | `vllm_rlt/request.py` | Current RNG location (`Request.generator`) |
 | `vllm_rlt/core/scheduler.py` | RNG release on termination |
 | `vllm_rlt/engine/preemption.py` | Preemption contract that must preserve the RNG |
 | `tests/test_sampler.py` | Algorithm, seed and generator contracts |
 | `tests/test_engine.py`, `tests/test_prefix_growth.py` | RNG lifecycle across termination and preemption |
+| `tests/test_speculative.py` | Speculative-path distribution, replay and rejection-sampling contracts |
 
 ### Out of scope
 
 Exit policy and stage transitions (M3), KV allocation and prepared metadata (M4),
 attention backends (M8), the PD connector protocol (M9), new sampling features
 (`min_p`, `repetition_penalty`, `logprobs`, `n > 1`), any change to the sampling
-arithmetic or to user-visible defaults, and the batched rewrite of the per-row CODA
-loop (see "Known limitation" below).
+arithmetic or to user-visible defaults, the batched rewrite of the per-row CODA loop
+(see "Known limitation" below), and collapsing `worker/sampling.py` into `Sampler`
+(M7 2/3).
 
 ## 2. Classes and state
 
@@ -75,27 +90,44 @@ per-request state: `sample()` receives the generator and returns the possibly cr
 one, so "where the RNG lives" stays visible at the call site. This is what lets
 `M7 2/3` move RNG storage without touching the algorithm or its tests.
 
+### `worker/sampling.py` (unchanged by this increment)
+
+The same arithmetic also exists as module-level helpers, introduced by PR #44:
+`probabilities()` builds the temperature/top-k/top-p distribution, `draw()` performs one
+multinomial draw, and `generator_for()` performs the lazy generator creation for the
+speculative path. `rejection_sample()` is speculative-only and has no counterpart in
+`Sampler`.
+
 ### RNG lifecycle today
 
 | Event | Current behavior | Location |
 | --- | --- | --- |
-| First random sample | Lazily creates `torch.Generator(device=...)` seeded from `params.seed` | `sampler.py:43-44` |
+| First random sample (CODA) | Lazily creates `torch.Generator(device=...)` seeded from `params.seed` | `sampler.py:43-44` |
+| First random sample (speculative) | Lazily creates the same kind of generator on the `Request` | `sampling.py:22-25` |
 | Greedy sample | `argmax`, no generator created or advanced | `sampler.py:31-32` |
-| Request finishes | `scheduler.finish` sets `request.generator = None` | `scheduler.py:94` |
+| Request finishes | `scheduler.finish` sets `request.generator = None` | `scheduler.py:95` |
 | Request preempted | Not reset; the `Request` object and its generator are retained | `preemption.py:3`, `preemption.py:102` |
 | Request resumed | Not rebuilt; sampling continues on the same generator | `preemption.py:114-144` |
 
+Both modules write the same `Request.generator` slot, so the two creation sites must stay
+consistent until `M7 2/3` moves the RNG behind the runner (section 6). Speculative
+decoding requires synchronous eager execution, refill scheduling and no preemption
+(`llm_engine.py:37-50`), so only the CODA path exercises the preemption rows above.
+
 ## 3. Functions and parameters
 
 ```python
 class Sampler:
     def __init__(self, device: torch.device): ...
 
+    # logits: [vocab] 1-dim device tensor, any float dtype; caller guarantees ndim == 1
+    # params: validated by SamplingParams.__post_init__
+    # generator: None on the greedy path
     def sample(
         self,
-        logits: torch.Tensor,                 # [vocab] 1-dim device tensor, any float dtype; caller guarantees ndim == 1
-        params: SamplingParams,               # validated by __post_init__
-        generator: torch.Generator | None,    # None on the greedy path
+        logits: torch.Tensor,
+        params: SamplingParams,
+        generator: torch.Generator | None,
     ) -> tuple[torch.Tensor, torch.Generator | None]:
         """Return a 0-dim long device tensor and the possibly created generator."""
 ```
@@ -131,6 +163,16 @@ def _sample_tensor(self, logits: torch.Tensor, request: Request):
 | `425` | CODA calls `_sample_tensor` once per request row and `torch.stack`s the results. |
 | `601-606` | Delegate: reads `request.generator`, calls `Sampler.sample`, stores the result back. |
 
+`worker/sampling.py` (unchanged; kept for the speculative path):
+
+| Lines | Behavior |
+| --- | --- |
+| `6-19` | `probabilities`: the same temperature/top-k/top-p statements as `Sampler.sample`, returning the normalized distribution. |
+| `22-25` | `generator_for`: lazily creates `request.generator` from `params.seed`. |
+| `28-29` | `draw`: one `torch.multinomial` draw followed by `squeeze(0)`. |
+| `32-37` | `sample_logits`: greedy short circuit or `draw`. This was `_sample_tensor`'s body before this increment, and it has no caller now. |
+| `40-56` | `rejection_sample`: exact speculative rejection with residual resampling. |
+
 ### Documented project difference
 
 `sampler.py:33` promotes logits to FP32 before dividing by temperature. The pinned
@@ -142,15 +184,17 @@ sampling path to compare against. The FP32 promotion is therefore a project choi
 awaiting comparison with a real reference sampling path, not a difference from a
 pinned sampling implementation. This is the "record the sampling implementation
 separately" item required by [precision policy](precision-policy.md). `M7 1/3`
-preserves it unchanged.
+preserves it unchanged; the speculative path repeats the same promotion in
+`sampling.py:9`.
 
 ### Known limitation
 
 `model_runner.py:424-426` samples each request row in a Python loop, so every row pays
-its own `sort`/`topk` over the model vocabulary (`vocab_size = 49152` in `config.py`). vLLM instead batches the multinomial
-across the batch. This is recorded as a known limitation and is deliberately not
-optimized in M7; any batched rewrite must keep the arithmetic order
-`float()/temperature → top_k → top_p → multinomial` and the no-host-sync guarantee.
+its own `sort`/`topk` over the model vocabulary (`vocab_size = 49152` in `config.py`).
+vLLM instead batches the multinomial across the batch. This is recorded as a known
+limitation and is deliberately not optimized in M7; any batched rewrite must keep the
+arithmetic order `float()/temperature → top_k → top_p → multinomial` and the
+no-host-sync guarantee.
 
 ## 5. Evidence-based findings
 
@@ -158,16 +202,26 @@ optimized in M7; any batched rewrite must keep the arithmetic order
 only the definition, no call site; tests monkeypatch `_sample_tensor`
 (`tests/test_async_pipeline.py:138`). Removed in `M7 1/3`.
 
+**Sampling exists in two modules after PR #44 (structural finding).**
+`worker/sampling.py` arrived with the speculative decoding change, and its
+`probabilities()`/`draw()` carry the same statements this increment moved into
+`Sampler.sample`; `_sample_tensor` was a one-line delegate to `sampling.sample_logits()`
+before it. This increment deliberately keeps both: `Sampler` becomes the M7 boundary for
+the CODA path and is pinned by the new tests, while `sampling.py` keeps serving the
+speculative path untouched. The cost is one duplicated distribution construction for one
+increment, and `sample_logits()` loses its only caller, so its import is dropped here and
+the function itself is removed in `M7 2/3` (section 6).
+
 **RNG ownership disagrees with the RFC state table.** The RFC assigns "RNG objects
 and their execution/snapshot state" to the worker, while the generator currently lives
-on the scheduler-side `Request` dataclass (`request.py:43`). Behavior is correct today
+on the scheduler-side `Request` dataclass (`request.py:44`). Behavior is correct today
 because `Request` is a shared handle and preemption/resumption happens to retain it,
 but ownership is implicit: any change that treats `release()` as "termination releases
 everything" can silently reset the RNG.
 
 **Preemption and termination share `model_runner.release()` (the key design
 constraint).** `preemption.py:102` calls `release()` when suspending a victim;
-`llm_engine.py:155` (abort) and `llm_engine.py:261` (finish) call the same method, and
+`llm_engine.py:191` (abort) and `llm_engine.py:332` (finish) call the same method, and
 `pd/worker.py:107` (removal) and `pd/worker.py:305` (prefill compute release) do as
 well. Moving the RNG into the runner and deleting it inside `release()` would reset it
 on preemption, changing the token sequence after restoration. That would violate the
@@ -203,18 +257,28 @@ stacks the sampled results and scatters them with `tokens=True`
 
 `M7 1/3` (this increment) only moves the algorithm into `Sampler` and pins behavior.
 RNG ownership is unchanged: `request.generator` is still read and written by the
-delegate, and `release()` still does not touch it.
+`_sample_tensor` delegate and by `sampling.generator_for` on the speculative path, and
+`release()` still does not touch it.
 
 `M7 2/3` will, in order:
 
-1. Add a runner-side RNG registry (`dict[str, torch.Generator]`) and have
-   `_sample_tensor` read and write it instead of `request.generator`.
-2. Split the release path: `suspend(request_id)` synchronizes the event and frees the
+1. Collapse `worker/sampling.py` into `Sampler`: move `probabilities()`, `draw()` and
+   `generator_for()` onto the class, retarget the `speculative.py` call sites, and delete
+   `sample_logits()`.
+2. Add a runner-side RNG registry (`dict[str, torch.Generator]`) and have both sampling
+   paths read and write it instead of `request.generator`.
+3. Split the release path: `suspend(request_id)` synchronizes the event and frees the
    state slot but keeps the RNG (called by preemption); `release(request_id)` does the
    same and additionally drops the RNG (called by abort, finish and PD removal).
-3. Remove `Request.generator` and the `scheduler.finish` assignment, then delete the
+4. Remove `Request.generator` and the `scheduler.finish` assignment, then delete the
    temporary adapters.
 
+The unification must keep each path's RNG call order: the CODA path draws one
+`multinomial` per token, while the speculative path additionally draws draft candidates,
+`torch.rand` acceptance uniforms and residual resamples. The two paths therefore produce
+different sequences from the same seed today, and only their distributions agree; `M7 2/3`
+must preserve that rather than align the streams.
+
 The target RNG contract:
 
 | Event | Target behavior |
@@ -222,6 +286,7 @@ The target RNG contract:
 | First random sample | Create the generator from `params.seed` in the request's RNG slot |
 | Later samples | Reuse and advance the same generator (ordering semantics unchanged) |
 | Greedy request | Never create or advance a generator |
+| Speculative request | Same slot; draft draws, acceptance uniforms and residual draws all come from it |
 | Preemption | Keep the generator (or snapshot `get_state()`) and continue from it on resume |
 | Termination (stop/length/abort/PD removal) | Drop the RNG slot so a new request with the same ID starts fresh |
 | Late async result | Must not write into a new request's RNG; the existing object-identity check covers results, and "release on termination" covers the RNG |
@@ -230,7 +295,7 @@ The target RNG contract:
 
 `ModelRunner._sample_tensor` exists so that `tests/test_async_pipeline.py:138` keeps
 working. It can be removed once that test patches `Sampler` instead, which is planned
-for `M7 2/3` when the runner-side registry lands.
+for `M7 2/3` when the runner-side registry lands and both sampling paths read it.
 
 ## 7. Validation
 
@@ -238,13 +303,14 @@ for `M7 2/3` when the runner-side registry lands.
 
 | Suite | Result |
 | --- | --- |
-| `main @ 60af1cf` baseline (`test_engine.py test_prefix_growth.py test_async_pipeline.py`) | 47 passed, 15 skipped |
+| `upstream/main @ 3314c1b` baseline (`test_engine.py test_prefix_growth.py test_async_pipeline.py`) | 47 passed, 15 skipped |
 | Same three files after this increment (baseline 47 + 2 new engine tests) | 49 passed, 15 skipped |
 | With `test_sampler.py` and `test_serving.py` added | 59 passed, 16 skipped |
-| Full CPU suite (`tests/`) | 211 passed, 115 skipped |
+| `upstream/main @ 3314c1b` full CPU suite | 240 passed, 121 skipped |
+| Full CPU suite (`tests/`) | 252 passed, 121 skipped |
 
 ```bash
-# main @ 60af1cf baseline (checkout 60af1cf first)
+# upstream/main @ 3314c1b baseline (checkout 3314c1b first)
 OMP_NUM_THREADS=1 python -m pytest -q tests/test_engine.py tests/test_prefix_growth.py tests/test_async_pipeline.py
 # same three files after this increment
 OMP_NUM_THREADS=1 python -m pytest -q tests/test_engine.py tests/test_prefix_growth.py tests/test_async_pipeline.py
@@ -254,11 +320,14 @@ OMP_NUM_THREADS=1 python -m pytest -q tests/test_sampler.py tests/test_engine.py
 # full CPU suite
 OMP_NUM_THREADS=1 python -m pytest -q tests/
 # lint
-python -m ruff check vllm_rlt/worker/sampler.py vllm_rlt/worker/model_runner.py tests/test_sampler.py
+python -m ruff check .
 python -m ruff format --check vllm_rlt/worker/sampler.py vllm_rlt/worker/model_runner.py tests/test_sampler.py
 ```
 
-Ruff: all checks passed; 3 files already formatted.
+The full suite grows by 12 tests over the baseline: 10 in `tests/test_sampler.py` and the
+2 new engine tests. Ruff: `ruff check .` passes and the touched Python files are already
+formatted. Three pre-existing markdown files are not `ruff format` clean on the baseline
+either, so the documentation check stays scoped to the touched Python files.
 
 The extraction changes no arithmetic, so no new precision or performance evidence is
 required. Sampling sequences depend on `torch.multinomial`; the numbers above are with