[Fix] Validate user-supplied input_ids against the vocabulary

[v-code01]

Part of #31597.

Motivation

_validate_input_ids_in_vocab exists but has zero callers, so user-supplied input_ids (native /generate, and OpenAI token-id prompts) are never validated. And the method itself only checks id >= vocab_size — it misses id < 0. Consequences:

• an id >= vocab_size indexes the model's embedding (nn.Embedding(vocab, hidden)) out of bounds. On CUDA that is a device-side assert that poisons the context and takes down the whole server from a single request (e.g. POST /generate {"input_ids":[1000000000],"sampling_params":{"max_new_tokens":1}}).
• a negative id wraps to the end of the embedding table (silently wrong output).

This is the same untrusted-input class as logit_bias / token_ids_logprob / stop_token_ids, which are already validated.

Modifications

Wire _validate_input_ids_in_vocab into _validate_one_request, guarded on obj.input_ids is not None so it only scans user-provided token ids (tokenizer output is always valid, so text requests pay nothing). Reject both out-of-range and negative ids, and add an empty-input guard.

Tests

test_tokenizer_manager_input_ids_validation.py: ids >= vocab_size, == vocab_size, and negative all raise (single and batch); in-range and empty pass. The negative case used to pass under the old >= vocab_size-only check.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.

---

CI States

Latest PR Test (Base): ❌ Run #29813122122
Latest PR Test (Extra): ❌ Run #29813121808

[v-code01]

cc @hnyls2002 @merrymercy (managers owners). Part of #31597 — this wires up the existing-but-dead _validate_input_ids_in_vocab and adds the missing < 0 guard, so user token-id prompts can no longer index the embedding out of bounds. Companion to the already-open #31561 / #31567.

[gemini-code-assist]

Code Review

This pull request adds validation for user-supplied input_ids to prevent out-of-bounds token IDs (both negative and exceeding vocabulary size) from causing CUDA device-side asserts or silent corruption, along with a new unit test suite. The reviewer suggests extending this validation to embedding requests (EmbeddingReqInput) as well, rather than limiting it to generation requests, to prevent similar server crashes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The vocabulary validation for user-supplied input_ids is currently restricted to GenerateReqInput requests. However, EmbeddingReqInput requests can also accept user-supplied input_ids (e.g., via the embedding endpoints). If an out-of-bounds token ID is supplied to an embedding request, it will also trigger a CUDA device-side assert or index out-of-bounds error, crashing the server or causing silent corruption.

Moving this validation outside of the isinstance(obj, GenerateReqInput) block ensures that both generation and embedding requests are properly validated.

Suggested change

	# Validate generation-specific fields
	if isinstance(obj, GenerateReqInput):
	# Only user-supplied token ids can be out of range; ids produced by
	# the tokenizer are always valid, so skip the scan for text requests.
	if obj.input_ids is not None:
	self._validate_input_ids_in_vocab(
	input_ids, self.model_config.vocab_size
	)
	# Only user-supplied token ids can be out of range; ids produced by
	# the tokenizer are always valid, so skip the scan for text requests.
	if obj.input_ids is not None:
	self._validate_input_ids_in_vocab(
	input_ids, self.model_config.vocab_size
	)
	# Validate generation-specific fields
	if isinstance(obj, GenerateReqInput):

[guanxingithub]

Independent confirmation of this one, in case it helps it move — we hit the same crash in
QA fuzzing, arrived at the same root cause, and then found this PR already open.

Reproduced. Cosmos-Reason1-7B (vocab 152064), single GPU, TP=1, on a GB300 host. One
POST /v1/completions with {"prompt": [200000], "max_tokens": 16} kills the server
process; every later request including /health is refused until restart.

An in-vocab prompt was sent first and served 200 with the process alive, so the result is
specific to the out-of-vocabulary id rather than to the request shape:

control in-vocab [100,200]   -> HTTP 200, server alive, /health 200
test    oov      [200000]    -> connection dropped, server DEAD, /health refused

Server side, the assert this PR prevents:

ATen/native/cuda/Indexing.cu:1528: indexSelectSmallIndex:
  Assertion `srcIndex < srcSelectDimSize` failed.     (3584 occurrences)
-> torch.AcceleratorError: CUDA error: device-side assert triggered
-> Fatal Python error: Aborted   -> scheduler gone

This fix shape closes it. We applied a functionally equivalent patch — wire
_validate_input_ids_in_vocab into _validate_one_request, reject both bounds — and re-ran:
out-of-range id returns HTTP 400 with the server alive and zero asserts, while in-vocab
and a multimodal image request both still return 200. The server then drained cleanly on
SIGTERM, i.e. still serving rather than merely answering /health.

A regression lane, proven both ways on one host in one job, patch as the only difference:

arm	result
unfixed	fails — all four out-of-range/negative/batched probes kill the server; 3584 asserts
fixed	passes — all four return 400 with the server alive; 0 asserts; in-vocab control still 200

The probes include a negative id and a batched prompt with one bad id, both of which pass
on the fixed build. Worth noting because the negative case is silent rather than fatal — it
wraps to the end of the embedding table — so it is easy to miss without a test, which is
exactly why the one in this PR is worth having.

One question, not a review blocker. The call sits inside
if isinstance(obj, GenerateReqInput):. EmbeddingReqInput also accepts user-supplied token
ids (io_struct.py: "one can either specify text or input_ids"), so the embedding path
looks like it stays unvalidated. We have not tested an embedding request with an
out-of-range id — reproduction went through /v1/completions only — so this is a question
rather than a finding. If embeddings reach the same nn.Embedding lookup, moving the call
above the isinstance branch would cover both with no change in behaviour for text requests.

Happy to help with anything that would get this landed.

[guanxingithub]

Follow-up with a concrete downstream data point, in case it helps prioritize this.

We now ship this same fix in NVIDIA's sglang fork (our 26.08 release branch and main), verified end-to-end on hardware both ways: OOV / negative / batched-with-one-bad-id prompts all return HTTP 400 with the server still serving, in-vocab and multimodal controls unaffected, and zero device-side asserts in the server log.

When this PR merges we would like to drop our fork patch entirely and converge on upstream. The one difference currently preventing that is the placement question from my earlier comment: with the call inside if isinstance(obj, GenerateReqInput):, EmbeddingReqInput.input_ids reaches the same embedding lookup unvalidated, so we would have to keep a residual delta. Moving the call just above that branch (after the input-length check) covers both request types — that is the placement we run in production, and the multimodal path is unaffected because validation runs before the mm processor rewrites placeholder ids.

Happy to send that as a suggestion on the diff if useful, or share our verification setup. Not pushing on scope otherwise — the PR as it stands already closes the crash we hit.