A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[BugFix][V1] Fix int32 token index overflow when preparing input ids (#16806)
Yong Hoon Shin committed
Apr 23, 2025 at 12:12 UTC
32d4b669d00d50ad084af128b9c01a5583fb2b7d
2 files changed
+4
-2
vllm/v1/worker/gpu_model_runner.py
+2
-1
@@ -241,10 +241,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
241
device=self.device)
242
243
# OPTIMIZATION: Cache the tensors rather than creating them every step.
244
+ # Keep in int64 to avoid overflow with long context
245
self.arange_np = np.arange(max(self.max_num_reqs + 1,
246
self.max_model_len,
247
self.max_num_tokens),
247
- dtype=np.int32)
248
+ dtype=np.int64)
249
# NOTE(woosuk): These tensors are "stateless", i.e., they are literally
250
# a faster version of creating a new tensor every time. Thus, we should
251
# not make any assumptions about the values in these tensors.
vllm/v1/worker/tpu_model_runner.py
+2
-1
@@ -219,7 +219,8 @@ class TPUModelRunner:
219
220
# Range tensor with values [0 .. self.max_num_tokens - 1].
221
# Used to initialize positions / context_lens / seq_lens
222
- self.arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
222
+ # Keep in int64 to avoid overflow with long context
223
+ self.arange_np = np.arange(self.max_num_tokens, dtype=np.int64)
224
self.num_reqs_paddings = _get_req_paddings(
225
min_req_size=MIN_NUM_SEQS, max_req_size=self.max_num_reqs)
226