From aa23aaa98b2cc0daa1f65bcda0e7f845a7c35970 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:56:42 +0900 Subject: [PATCH] fix: treat near-zero temperature as greedy and bound the seed (#167) temperature=1e-40 turned logits into inf/NaN and seeds >= ~9.2e12 overflowed manual_seed(seed * 1_000_003 + p). Both passed SamplingParams validation, so the error surfaced inside the batcher's step() and poisoned the service. SamplingParams now treats temperature < 1e-5 as greedy and accepts seeds in [0, 2**32); the HTTP parser already maps the ValueError to a 400. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/accel/pd/contracts.py | 9 +++++---- opt_prime/tests/test_pd_sampling.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/opt_prime/opt_prime/accel/pd/contracts.py b/opt_prime/opt_prime/accel/pd/contracts.py index 6ac3c26c..6ddd1c01 100644 --- a/opt_prime/opt_prime/accel/pd/contracts.py +++ b/opt_prime/opt_prime/accel/pd/contracts.py @@ -41,7 +41,8 @@ class SamplingParams: """Token selection. The default is greedy; engines without `supports_sampling` accept only that. `seed` makes sampling reproducible across processes (each position draws from its own seeded stream), so a prefill and a decode stage on different hosts continue one sequence. Generation ends early at any of - `stop_token_ids` (the stop token is kept in the output).""" + `stop_token_ids` (the stop token is kept in the output). A temperature below 1e-5 is greedy (dividing logits by + it overflows), and the seed is bounded so a position's stream seed fits torch's 64-bit `manual_seed`.""" temperature: float = 0.0 top_p: float = 1.0 top_k: int = 0 @@ -55,14 +56,14 @@ def __post_init__(self): raise ValueError("top_p must be in (0, 1]") if type(self.top_k) is not int or self.top_k < 0: raise ValueError("top_k must be a nonnegative integer") - if self.seed is not None and (type(self.seed) is not int or self.seed < 0): - raise ValueError("seed must be a nonnegative integer") + if self.seed is not None and (type(self.seed) is not int or not 0 <= self.seed < 2 ** 32): + raise ValueError("seed must be an integer in [0, 2**32)") if not isinstance(self.stop_token_ids, tuple) or any(type(t) is not int or t < 0 for t in self.stop_token_ids): raise ValueError("stop token IDs must be a tuple of nonnegative integers") @property def greedy(self): - return self.temperature == 0 or self.top_k == 1 + return self.temperature < 1e-5 or self.top_k == 1 GREEDY = SamplingParams() diff --git a/opt_prime/tests/test_pd_sampling.py b/opt_prime/tests/test_pd_sampling.py index d92eb8d1..615e55c5 100644 --- a/opt_prime/tests/test_pd_sampling.py +++ b/opt_prime/tests/test_pd_sampling.py @@ -89,3 +89,20 @@ def test_service_finish_reason_eos_stop_string_and_greedy_only_engines(): assert err.value.status == 400 with pytest.raises(StageError): engine.execute("prefill", "u", dict(body(1), top_p=0)) + + +def test_extreme_sampling_values_cannot_poison_a_batch(): + # #167: a near-zero temperature is greedy (logits / 1e-40 overflow to NaN), a seed that would overflow + # manual_seed(seed * 1_000_003 + p) is rejected up front (HTTP 400), so a normal neighbour completes + from opt_prime.accel.pd.batching import ContinuousBatcher + model, _ = _runtime() + b = ContinuousBatcher(model, slots=2, max_context=64) + tiny = b.submit(PDRequest("tiny", PROMPT, 6, SamplingParams(temperature=1e-40, seed=1))) + good = b.submit(PDRequest("good", PROMPT, 6, SamplingParams(temperature=1.3, seed=5))) + b.run_until_idle() + assert tiny.result() == _hf_greedy(model, 6) and len(good.result()) == 6 + with pytest.raises(ValueError): + SamplingParams(temperature=1.0, seed=2 ** 62) + with pytest.raises(StageError) as err: + service().execute("prefill", "big", dict(body(1), seed=2 ** 62)) + assert err.value.status == 400