fix(core): reject out-of-range raw token ids in token requests - #2347
Open
blasrodri wants to merge 2 commits into
Open
fix(core): reject out-of-range raw token ids in token requests#2347blasrodri wants to merge 2 commits into
blasrodri wants to merge 2 commits into
Conversation
CompletionTokens and EmbeddingTokens requests carry raw token ids straight
from the client. An id >= vocab size indexes past the embedding table; on a
CUDA release build the bounds assert is compiled out, so a single request
like {"input":[4294967294]} to /v1/embeddings triggers an out-of-bounds GPU
read instead of a clean error.
Validate every id against the tokenizer vocab size and return a
ValidationError (HTTP 422) before the tokens reach the model, matching the
existing validation in the same match arm.
Code Metrics Report━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Language Files Lines Code Comments Blanks ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ C Header 23 4454 3116 790 548 CSS 3 282 252 6 24 CUDA 119 23681 19217 1704 2760 Dockerfile 1 35 19 9 7 HTML 2 27 27 0 0 JavaScript 3 577 562 12 3 Jinja2 7 694 656 5 33 JSON 27 15864 15861 0 3 Makefile 1 18 16 0 2 MDX 36 5901 0 4327 1574 Metal Shading Lan| 37 14287 11284 1136 1867 PowerShell 1 657 571 31 55 Python 151 12284 10191 484 1609 Shell 3 1068 849 116 103 Plain Text 53 10687 0 9209 1478 TOML 28 1368 1189 39 140 TypeScript 11 1649 1410 66 173 YAML 3 25 23 2 0 ───────────────────────────────────────────────────────────────────────────────── Jupyter Notebooks 3 122 83 23 16 |- Markdown 1 60 30 22 8 |- Python 1 122 113 1 8 (Total) 304 226 46 32 ───────────────────────────────────────────────────────────────────────────────── Markdown 273 11726 0 8736 2990 |- BASH 23 295 217 46 32 |- Dockerfile 2 5 5 0 0 |- JSON 6 289 289 0 0 |- PowerShell 1 1 1 0 0 |- Python 135 7239 6021 310 908 |- Rust 62 3820 2838 388 594 |- TOML 7 77 65 0 12 (Total) 23452 9436 9480 4536 ───────────────────────────────────────────────────────────────────────────────── Rust 674 300652 267964 5856 26832 |- Markdown 413 9761 452 8126 1183 (Total) 310413 268416 13982 28015 ───────────────────────────────────────────────────────────────────────────────── Svelte 19 1969 1826 51 92 |- CSS 1 4 4 0 0 |- JavaScript 19 921 767 25 129 (Total) 2894 2597 76 221 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Total 1478 430621 345918 41520 43183 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ |
Pull the id-range validation out of handle_request into a pure first_out_of_range_token helper and cover it with unit tests: ids below vocab, the max valid id, an id equal to vocab size (rejected), an out-of-range id reported back, and the empty case.
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.
Problem
CompletionTokensandEmbeddingTokensrequests carry raw token ids straight from the client (e.g.POST /v1/embeddingswith{"input": [[...]]}). These ids were never range-checked before being handed to the model.An id
>= vocab_sizeindexes past the embedding table. On CPU this returns a graceful error, but on a CUDA release build the bounds assertion is compiled out (NDEBUG), so the gather performs an out-of-bounds device memory read. A single unauthenticated request such as{"model":"default","input":[4294967294]}can therefore hit undefined behavior in the engine instead of a clean rejection.Fix
In the shared
CompletionTokens | EmbeddingTokensarm ofEngine::handle_request, validate every id againsttokenizer.get_vocab_size(true)and return aValidationError(HTTP 422) before the tokens reach the model. This mirrors the tokenizer-missing validation already present a few lines above in the same arm.Covers both raw-token entry points (embeddings and completions) since they share the arm.
Tests
The range check is extracted into a pure
first_out_of_range_tokenhelper with unit tests covering: ids below vocab size, the max valid id (vocab_size - 1), an id equal tovocab_size(rejected at the boundary), an out-of-range id reported back to the caller, and the empty-input case.