[BugFix][V1] Fix int32 token index overflow when preparing input ids

[sarckk]

Fixes #16414

The issue can be reproduced when setting max-model-len to a large value (e.g. 3600000) and sending requests without limiting the max number of requests (e.g. not restricting --max-num-seqs, sending requests without --max-concurrency) and/or not limiting max. number of tokens to be processed in a single iteration (--max-num-batched-tokens).

The error is IndexError: index out of range in self which happens in _prepare_inputs:

        torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),
                           0,
                           torch.from_numpy(token_indices),
                           out=self.input_ids_cpu[:total_num_scheduled_tokens])

When it fails, the last several items of token_indices contains indices that are negative, due to int32 overflow. token_indices is given by:

        token_indices = (positions_np +
                         req_indices * self.input_batch.token_ids_cpu.shape[1])

Here, req_indices is a np.int32 array, e.g. np.array([0,1,2,...598]), and self.input_batch.token_ids_cpu.shape[1] is the max model len, e.g. 3,600,000. As an example, multiplying 598 * 3,600,000 gives us 2,152,800,000 which is greater np.int32 max value of 2,147,483,647 so it overflows and gives us a negative number for token index.

To avoid overflow, we can cast to int64 before the multiplication. the max token index can be given by max_req_index * max_model_len - 1.

To verify, I checked that below completes successfully on H200 (previously it would fail with index error after ~1000 requests). Without this fix, the max no. of concurrent requests we can serve with 3.6M context length would be 596.

VLLM_DISABLE_COMPILE_CACHE=1 vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct --disable-log-requests --gpu-memory-utilization=0.95 --max-model-len=3600000 -tp 8

python3 benchmarks/benchmark_serving.py --port=8000 --dataset-name=sharegpt --dataset-path=ShareGPT_V3_unfiltered_cleaned_split.json --tokenizer=meta-llama/Llama-4-Scout-17B-16E-Instruct --request-rate=64 --backend=vllm --num-prompts=19200 --model=meta-llama/Llama-4-Scout-17B-16E-Instruct --sharegpt-output-len=1024

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[luccafong]

shall we just create req_indices as token_indices_dtype instead of upcasting here?

[houseroad]

req_indices may be used in the following logic.

Another idea is always having int64 index for such use cases. But probably require a bigger change / refactor, thoughts, @WoosukKwon ? Shall we just land this fix, and think about a more thorough check in the following PR?

[NickLucche]

I think we can create self.arange_np with dtype self.token_indices_dtype here. Increase in host memory is negligible even at 10M and we don't pay the cast at each scheduler step.
Computation is in i64 anyways as self.positions_cpu is already at i64, so that wouldn't be the only instance.

[sarckk]

makes sense, maybe we can just keep self.arange_np dtype in int64 similar to self.positions_cpu?

[houseroad]

Btw, @sarckk , could you lint the code?

[NickLucche]

Thanks for fixing this!

[NickLucche]

I think we can create self.arange_np with dtype self.token_indices_dtype here. Increase in host memory is negligible even at 10M and we don't pay the cast at each scheduler step.
Computation is in i64 anyways as self.positions_cpu is already at i64, so that wouldn't be the only instance.

[sarckk]

Btw, @sarckk , could you lint the code?

ran lint with pre-commit

[WoosukKwon]

LGTM! Thanks for the PR!

[houseroad]

Looks good. If it becomes a perf bottleneck, we can apply the optimization later.

[houseroad]

Rerunning the CI and see how it goes.

[NickLucche]

Looks clean, thanks!

[NickLucche]

PS I think there's some commits missing a sign-off.