[REFACT] [M7 1/3] Extract the sampling module and pin request-local RNG behavior

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 |

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

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
 

tests/test_sampler.py

@@ -0,0 +1,128 @@
+"""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}
+
+
+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)

vllm_rlt/worker/model_runner.py

@@ -12,7 +12,7 @@
 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
 
 
 @dataclass
@@ -79,6 +79,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()
@@ -597,8 +598,9 @@ 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):
-        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

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