fix(sampling_params): reject out-of-range token ids at intake

lightllm/server/core/objs/sampling_params.py

@@ -525,11 +525,13 @@ def _check_and_store_int_token_ids(dst, ids: List[int], max_length: int, name: s
 
     The type check runs against the input ``ids`` (not the zero-filled destination buffer), so a non-int
     entry fails fast with a clear message instead of surfacing an opaque ctypes ``TypeError`` at the
-    assignment below.
+    assignment below. The bound check guards the same assignment: ``dst`` stores 32-bit signed ints
+    and ctypes silently wraps out-of-range values (``2**31`` would become ``-2147483648``), and a
+    negative id would later be used as a pointer offset into the logits tensor by the triton kernel.
 
     Args:
         dst: destination ctypes array, declared as ``c_int * max_length``.
-        ids: caller-supplied token ids; every element must be an ``int``.
+        ids: caller-supplied token ids; every element must be an ``int`` in ``[0, 2**31)``.
         max_length: capacity of ``dst``.
         name: field name used in the error messages.
 
@@ -539,5 +541,6 @@ def _check_and_store_int_token_ids(dst, ids: List[int], max_length: int, name: s
     size = len(ids)
     assert size <= max_length, f"Too many {name}: {size} > {max_length}."
     assert all(isinstance(e, int) for e in ids), f"all {name} must be int."
+    assert all(0 <= e < 2 ** 31 for e in ids), f"all {name} must be int in [0, 2**31)."
     dst[:size] = ids[:]
     return size

unit_tests/server/core/objs/test_token_ids_validation.py

@@ -3,6 +3,7 @@
     StopSequence,
     AllowedTokenIds,
     InvalidTokenIds,
+    SamplingParams,
     _check_and_store_int_token_ids,
     STOP_SEQUENCE_MAX_LENGTH,
     ALLOWED_TOKEN_IDS_MAX_LENGTH,
@@ -32,6 +33,22 @@ def test_allowed_token_ids_rejects_too_many():
         allowed_ids.initialize([1] * (ALLOWED_TOKEN_IDS_MAX_LENGTH + 1))
 
 
+@pytest.mark.parametrize("bad_ids", [[-3], [100, -1], [2 ** 31], [2 ** 40]])
+def test_allowed_token_ids_rejects_out_of_range_ids(bad_ids):
+    # Out-of-range ids would wrap silently in the c_int buffer (2**31 -> -2147483648,
+    # 2**40 -> 0), so they must be rejected at intake instead of being stored garbage.
+    allowed_ids = AllowedTokenIds()
+    with pytest.raises(AssertionError, match="allowed token ids"):
+        allowed_ids.initialize(bad_ids)
+
+
+def test_allowed_token_ids_accepts_int32_boundary_ids():
+    allowed_ids = AllowedTokenIds()
+    allowed_ids.initialize([0, 2 ** 31 - 1])
+    assert allowed_ids.size == 2
+    assert allowed_ids.to_list() == [0, 2 ** 31 - 1]
+
+
 def test_invalid_token_ids_accepts_valid_ints():
     invalid_ids = InvalidTokenIds()
     invalid_ids.initialize([4, 5, 6])
@@ -52,12 +69,49 @@ def test_invalid_token_ids_rejects_too_many():
         invalid_ids.initialize([1] * (INVALID_TOKEN_IDS_MAX_LENGTH + 1))
 
 
+def test_invalid_token_ids_rejects_negative():
+    # A negative id survives the GPU-side vocab-size filter and is then used as a
+    # pointer offset into the logits tensor by the apply_invalid_token triton kernel.
+    invalid_ids = InvalidTokenIds()
+    with pytest.raises(AssertionError, match="invalid token ids"):
+        invalid_ids.initialize([-5])
+
+
+@pytest.mark.parametrize("bad_ids", [[2 ** 31], [4, 2 ** 40]])
+def test_invalid_token_ids_rejects_int32_overflow(bad_ids):
+    # An id >= 2**31 would wrap to a different value in the c_int buffer
+    # (2**31 -> -2147483648), silently banning the wrong token.
+    invalid_ids = InvalidTokenIds()
+    with pytest.raises(AssertionError, match="invalid token ids"):
+        invalid_ids.initialize(bad_ids)
+
+
+def test_invalid_token_ids_accepts_int32_boundary_ids():
+    invalid_ids = InvalidTokenIds()
+    invalid_ids.initialize([0, 2 ** 31 - 1])
+    assert invalid_ids.size == 2
+    assert invalid_ids.to_list() == [0, 2 ** 31 - 1]
+
+
 def test_stop_sequence_rejects_non_int():
     seq = StopSequence()
     with pytest.raises(AssertionError):
         seq.initialize([1, "2"])
 
 
+@pytest.mark.parametrize("bad_ids", [[-7], [2 ** 31 + 5]])
+def test_stop_sequence_rejects_out_of_range_ids(bad_ids):
+    seq = StopSequence()
+    with pytest.raises(AssertionError, match="stop token ids"):
+        seq.initialize(bad_ids)
+
+
+def test_stop_sequence_accepts_int32_boundary_ids():
+    seq = StopSequence()
+    seq.initialize([0, 2 ** 31 - 1])
+    assert seq.to_list() == [0, 2 ** 31 - 1]
+
+
 def test_check_and_store_int_token_ids_returns_size_and_writes_buffer():
     import ctypes
 
@@ -75,5 +129,21 @@ def test_check_and_store_int_token_ids_rejects_overflow():
         _check_and_store_int_token_ids(buf, [1, 2, 3], 2, "test ids")
 
 
+@pytest.mark.parametrize("bad_key", ["-3", str(2 ** 31)])
+def test_logit_bias_out_of_range_keys_rejected(bad_key):
+    # logit_bias keys are int()-converted and stored as invalid_token_ids by
+    # SamplingParams.init, so out-of-range keys must fail the same range check
+    # on the real production path, not only when initialize() is called directly.
+    params = SamplingParams()
+    with pytest.raises(AssertionError, match="invalid token ids"):
+        params.init(None, logit_bias={bad_key: 0.5})
+
+
+def test_logit_bias_boundary_keys_round_trip():
+    params = SamplingParams()
+    params.init(None, logit_bias={"0": 1.0, str(2 ** 31 - 1): 0.5})
+    assert params.invalid_token_ids.to_list() == [0, 2 ** 31 - 1]
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])