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 14, 00:30 UTC
[Core][Bookkeeping] Update cu_num_accepted_tokens for all req_index (#27629)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Oct 30, 2025 at 11:52 UTC
4574d48bab9c4e38b7c0a830eeefc8f0980e8c58
1 file changed
+9
-6
vllm/v1/worker/gpu_model_runner.py
+9
-6
@@ -2323,11 +2323,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2323
sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None
2324
else:
2325
sampled_ids = valid_sampled_token_ids[req_idx]
2326
+
2327
+ num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0
2328
+
2329
+ if cu_num_accepted_tokens is not None:
2330
+ cu_num_accepted_tokens.append(
2331
+ cu_num_accepted_tokens[-1] + num_sampled_ids
2332
+ )
2333
+
2334
if not sampled_ids:
2335
continue
2336
2337
start_idx = self.input_batch.num_tokens_no_spec[req_idx]
2330
- end_idx = start_idx + len(sampled_ids)
2338
+ end_idx = start_idx + num_sampled_ids
2339
assert end_idx <= self.max_model_len, (
2340
"Sampled token IDs exceed the max model length. "
2341
f"Total number of tokens: {end_idx} > max_model_len: "
@@ -2343,11 +2351,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2351
req_state = self.requests[req_id]
2352
req_state.output_token_ids.extend(sampled_ids)
2353
2346
- if cu_num_accepted_tokens is not None:
2347
- cu_num_accepted_tokens.append(
2348
- cu_num_accepted_tokens[-1] + len(sampled_ids)
2349
- )
2350
-
2354
logprobs_lists = (
2355
logprobs_tensors.tolists(cu_num_accepted_tokens)
2356
if not self.use_async_scheduling and logprobs_tensors is not None