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 11, 05:00 UTC
[Cleanup] Remove unused ModelRunner V1 `InputBatch.num_tokens` field (#30218)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Dec 18, 2025 at 09:17 UTC
686cbaac643c3412036728dd5e6bc29d6cce1a9f
4 files changed
+12
-36
vllm/v1/worker/gpu_input_batch.py
+9
-19
@@ -128,7 +128,6 @@ class InputBatch:
128
# allocation if max_model_len is big.
129
# Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size)
130
self.req_prompt_embeds: dict[int, torch.Tensor] = {}
131
- self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
131
self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
132
self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
133
self.num_computed_tokens_cpu_tensor = torch.zeros(
@@ -340,9 +339,6 @@ class InputBatch:
339
self.req_prompt_embeds[req_index] = request.prompt_embeds
340
self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids
341
self.is_token_ids[req_index, start_idx:end_idx] = True
343
- # Number of token ids in prompt (token_ids_cpu or prompt_embeds).
344
- # NOTE(woosuk): This may include spec decode tokens.
345
- self.num_tokens[req_index] = request.num_tokens
342
# Number of tokens without spec decode tokens.
343
self.num_tokens_no_spec[req_index] = request.num_tokens
344
@@ -522,10 +518,6 @@ class InputBatch:
518
self.req_id_to_index[old_id_i2],
519
self.req_id_to_index[old_id_i1],
520
)
525
- self.num_tokens[i1], self.num_tokens[i2] = (
526
- self.num_tokens[i2],
527
- self.num_tokens[i1],
528
- )
521
self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = (
522
self.num_tokens_no_spec[i2],
523
self.num_tokens_no_spec[i1],
@@ -661,17 +653,16 @@ class InputBatch:
653
self.req_output_token_ids[last_req_index] = None
654
self.req_id_to_index[req_id] = empty_index
655
664
- if last_req_index != empty_index:
665
- (
666
- self.spec_token_ids[last_req_index],
667
- self.spec_token_ids[empty_index],
668
- ) = (
669
- self.spec_token_ids[empty_index],
670
- self.spec_token_ids[last_req_index],
671
- )
672
- self.spec_token_ids[last_req_index].clear()
656
+ num_tokens = self.num_tokens_no_spec[last_req_index] + len(
657
+ self.spec_token_ids[last_req_index]
658
+ )
659
+
660
+ (self.spec_token_ids[last_req_index], self.spec_token_ids[empty_index]) = (
661
+ self.spec_token_ids[empty_index],
662
+ self.spec_token_ids[last_req_index],
663
+ )
664
+ self.spec_token_ids[last_req_index].clear()
665
674
- num_tokens = self.num_tokens[last_req_index]
666
self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
667
last_req_index, :num_tokens
668
]
@@ -682,7 +673,6 @@ class InputBatch:
673
self.req_prompt_embeds[empty_index] = self.req_prompt_embeds.pop(
674
last_req_index
675
)
685
- self.num_tokens[empty_index] = num_tokens
676
self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
677
last_req_index
678
]
vllm/v1/worker/gpu_model_runner.py
-5
@@ -923,7 +923,6 @@ class GPUModelRunner(
923
self.input_batch.num_prompt_tokens[req_index]
924
+ num_output_tokens
925
)
926
- self.input_batch.num_tokens[req_index] = end_idx
926
self.input_batch.num_tokens_no_spec[req_index] = end_idx
927
928
# Update the block IDs.
@@ -968,7 +967,6 @@ class GPUModelRunner(
967
req_index, start_token_index:end_token_index
968
] = new_token_ids
969
self.input_batch.num_tokens_no_spec[req_index] = end_token_index
971
- self.input_batch.num_tokens[req_index] = end_token_index
970
971
# Add spec_token_ids to token_ids_cpu.
972
spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
@@ -984,8 +982,6 @@ class GPUModelRunner(
982
self.input_batch.token_ids_cpu[
983
req_index, start_index:end_token_index
984
] = spec_token_ids
987
- # NOTE(woosuk): `num_tokens` here may include spec tokens.
988
- self.input_batch.num_tokens[req_index] += num_spec_tokens
985
986
# When speculative decoding is used with structured output,
987
# the scheduler can drop draft tokens that do not
@@ -2702,7 +2698,6 @@ class GPUModelRunner(
2698
self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
2699
self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True
2700
self.input_batch.num_tokens_no_spec[req_idx] = end_idx
2705
- self.input_batch.num_tokens[req_idx] = end_idx
2701
2702
req_id = req_ids[req_idx]
2703
req_state = self.requests[req_id]
vllm/v1/worker/tpu_input_batch.py
+1
-10
@@ -51,7 +51,6 @@ class InputBatch:
51
pin_memory=False,
52
)
53
self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
54
- self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
54
self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
55
self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
56
self.num_computed_tokens_cpu_tensor = torch.zeros(
@@ -200,9 +199,6 @@ class InputBatch:
199
start_idx = num_prompt_tokens
200
end_idx = start_idx + len(request.output_token_ids)
201
self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids
203
- # Number of token ids in token_ids_cpu.
204
- # NOTE(woosuk): This may include spec decode tokens.
205
- self.num_tokens[req_index] = request.num_tokens
202
# Number of tokens without spec decode tokens.
203
self.num_tokens_no_spec[req_index] = request.num_tokens
204
@@ -344,10 +340,6 @@ class InputBatch:
340
self.req_id_to_index[old_id_i2],
341
self.req_id_to_index[old_id_i1],
342
)
347
- self.num_tokens[i1], self.num_tokens[i2] = (
348
- self.num_tokens[i2],
349
- self.num_tokens[i1],
350
- )
343
self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = (
344
self.num_tokens_no_spec[i2],
345
self.num_tokens_no_spec[i1],
@@ -448,11 +440,10 @@ class InputBatch:
440
self.req_output_token_ids[last_req_index] = None
441
self.req_id_to_index[req_id] = empty_index
442
451
- num_tokens = self.num_tokens[last_req_index]
443
+ num_tokens = self.num_tokens_no_spec[last_req_index]
444
self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
445
last_req_index, :num_tokens
446
]
455
- self.num_tokens[empty_index] = num_tokens
447
self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
448
last_req_index
449
]
vllm/v1/worker/tpu_model_runner.py
+2
-2
@@ -1283,7 +1283,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1283
token_id = valid_sampled_token_ids[i][0]
1284
self.input_batch.token_ids_cpu[i, seq_len] = token_id
1285
req_state.output_token_ids.append(token_id)
1286
- self.input_batch.num_tokens[i] += 1
1286
+ self.input_batch.num_tokens_no_spec[i] += 1
1287
1288
else:
1289
valid_mask = selected_token_ids != INVALID_TOKEN_ID
@@ -1291,7 +1291,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1291
valid_sampled_token_ids = [
1292
seq.tolist() for seq in selected_token_ids[valid_mask].split(gen_lens)
1293
]
1294
- self.input_batch.num_tokens[:num_reqs] += gen_lens
1294
+ self.input_batch.num_tokens_no_spec[:num_reqs] += gen_lens
1295
for i, req_state, seq_len in request_seq_lens:
1296
target_slice = slice(seq_len - gen_lens[i] + 1, seq_len + 1)
1297
self.input_batch.token_ids_cpu[i, target_slice] = (