fix(sampling_params): reject out-of-range token ids at intake - #1554
Open
SuperMarioYL wants to merge 1 commit into
Open
fix(sampling_params): reject out-of-range token ids at intake#1554SuperMarioYL wants to merge 1 commit into
SuperMarioYL wants to merge 1 commit into
Conversation
_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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_check_and_store_int_token_ids(the shared validation helper from #1365) checks that every element ofallowed_token_ids/invalid_token_ids/ stop-sequence token ids is anintand 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 thec_intdestination array, where out-of-range values wrap silently: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 theapply_invalid_tokentriton kernel — an out-of-bounds write. The same gap is reachable throughlogit_biaskeys, which flow intoinvalid_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:All three ctypes seams (
StopSequence,AllowedTokenIds,InvalidTokenIds) and thelogit_biaskeys 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 clearAssertionErrorinstead 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 thelogit_biasflow:0and2**31 - 1round-trip exactlyThe rejection cases are red on
main(out-of-range input is accepted there) and green with this change. No mocks — the tests drive the realSamplingParams.initpath for thelogit_biascase.