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 4, 11:15 UTC
[Perf][Spec Decode] Avoid per-step numpy allocation in prepare_next_t… (#41043)
Signed-off-by: wangluochao902 <wangluochao902@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Luochao Wang committed
Apr 29, 2026 at 14:20 UTC
b58669cb427effa49928b7be5b6e0f4fd707bce5
1 file changed
+5
-9
vllm/v1/spec_decode/llm_base_proposer.py
+5
-9
@@ -883,16 +883,12 @@ class SpecDecodeBaseProposer:
883
is not sampled and comes from `request.get_token_id()` instead. This is denoted
884
the "backup" token id. It also counts rejected tokens via `sampled_token_ids`.
885
"""
886
- # Precompute get_token_id for when there is no valid next token
886
+ # Precompute backup token IDs for discarded requests.
887
num_reqs = gpu_input_batch.num_reqs
888
- seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist()
889
- self.backup_next_token_ids.np[:num_reqs] = np.array(
890
- [
891
- requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])
892
- for i in range(num_reqs)
893
- ],
894
- dtype=np.int32,
895
- )
888
+ for i in range(num_reqs):
889
+ self.backup_next_token_ids.np[i] = requests[
890
+ gpu_input_batch.req_ids[i]
891
+ ].get_token_id(gpu_input_batch.num_tokens_no_spec[i] - 1)
892
self.backup_next_token_ids.copy_to_gpu(num_reqs)
893
backup_tokens_gpu = self.backup_next_token_ids.gpu
894