fix: treat near-zero temperature as greedy and bound the seed (#167)

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

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