Skip to content

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

Open
SuperMarioYL wants to merge 1 commit into
ModelTC:mainfrom
SuperMarioYL:fix/reject-out-of-range-token-ids
Open

fix(sampling_params): reject out-of-range token ids at intake#1554
SuperMarioYL wants to merge 1 commit into
ModelTC:mainfrom
SuperMarioYL:fix/reject-out-of-range-token-ids

Conversation

@SuperMarioYL

Copy link
Copy Markdown
Contributor

Summary

_check_and_store_int_token_ids (the shared validation helper from #1365) checks that every element of allowed_token_ids / invalid_token_ids / stop-sequence token ids is an int and fits the capacity — but not that it is in the valid 32-bit range. A negative id or an id ≥ 2**31 is accepted and copied into the c_int destination array, where out-of-range values wrap silently:

>>> p = InvalidTokenIds(); p.initialize([-5, 2**31]); p.to_list()
[-5, -2147483648]          # 2**31 wraps to INT32_MIN
>>> a = AllowedTokenIds(); a.initialize([-3, 100, 2**40]); a.to_list()
[-3, 100, 0]               # 2**40 wraps to 0

The negative id then survives the GPU-side filter [e for e in invalid_token_ids if e < vocab_size] (infer_batch.py) and is used as a logits pointer offset in the apply_invalid_token triton kernel — an out-of-bounds write. The same gap is reachable through logit_bias keys, which flow into invalid_token_ids (logit_bias={"-3": 0.5} is accepted today).

#1457 fixed exactly this (non-negative + 32-bit range checks) but was closed when the more thorough #1365 landed — and #1365 only carried the element-type half. This PR completes the dropped range-check half, in the same helper and the same assert style as #1365. Refs #1457.

Change

One assertion in _check_and_store_int_token_ids, after the existing type check:

assert all(0 <= e < 2 ** 31 for e in ids), f"all {name} must be int in [0, 2**31)."

All three ctypes seams (StopSequence, AllowedTokenIds, InvalidTokenIds) and the logit_bias keys flow route through this helper, so one check covers every intake path. Valid boundary values (0, 2**31 - 1) still round-trip exactly; the only observable change is that out-of-range input now fails fast with a clear AssertionError instead of wrapping silently.

Tests

unit_tests/server/core/objs/test_token_ids_validation.py (the file added by #1365; separate from the in-flight #1350 test changes) gains, for each of the three structures and the logit_bias flow:

  • negative ids are rejected
  • ids ≥ 2**31 are rejected (no silent wrap)
  • boundary values 0 and 2**31 - 1 round-trip exactly

The rejection cases are red on main (out-of-range input is accepted there) and green with this change. No mocks — the tests drive the real SamplingParams.init path for the logit_bias case.

_check_and_store_int_token_ids validated only element types, so negative
ids and ids >= 2**31 were accepted and stored into c_int destination
arrays where ctypes wraps them silently (2**31 -> -2147483648). A
negative id then survives the GPU-side vocab-size filter and reaches the
apply_invalid_token triton kernel as a pointer offset in front of the
logits tensor.

Add one bound assertion (0 <= id < 2**31) to the shared helper so every
intake seam (allowed_token_ids, invalid_token_ids / logit_bias keys,
stop_sequences) rejects out-of-range ids with a field-named message,
while boundary values 0 and 2**31-1 keep working.

Completes the range-check half of ModelTC#1457, which was closed when the
type-check half landed as ModelTC#1365.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant