diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py index 20bd6c61e..123898ae9 100644 --- a/lightllm/server/core/objs/sampling_params.py +++ b/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 diff --git a/unit_tests/server/core/objs/test_token_ids_validation.py b/unit_tests/server/core/objs/test_token_ids_validation.py index bf232a209..06a681dbb 100644 --- a/unit_tests/server/core/objs/test_token_ids_validation.py +++ b/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"])