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 13, 10:45 UTC
[Perf] Use np.ndarray instead of list[list[int]] to reduce GC overhead (#28245)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Nov 11, 2025 at 10:28 UTC
4228be7959e98e57d88501bd97aca7ef34ff562e
3 files changed
+17
-10
tests/v1/engine/utils.py
+4
-3
@@ -5,6 +5,7 @@ import random
5
from dataclasses import dataclass
6
from typing import TypeAlias
7
8
+import numpy as np
9
import torch
10
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
11
@@ -369,9 +370,9 @@ class MockEngineCore:
370
self.generated_logprobs_raw[req_idx][token_idx]
371
)
372
logprobs = LogprobsLists(
372
- [logprobs_token_ids_],
373
- [logprobs_],
374
- [sampled_token_ranks_],
373
+ np.array([logprobs_token_ids_]),
374
+ np.array([logprobs_]),
375
+ np.array([sampled_token_ranks_]),
376
)
377
else:
378
logprobs = None
vllm/v1/engine/logprobs.py
+6
-1
@@ -74,7 +74,12 @@ class LogprobsProcessor:
74
75
token_ids_lst, logprobs_lst, ranks_lst, _ = logprobs_lists
76
77
- for rank, logprobs, token_ids in zip(ranks_lst, logprobs_lst, token_ids_lst):
77
+ for rank_np, logprobs_np, token_ids_np in zip(
78
+ ranks_lst, logprobs_lst, token_ids_lst
79
+ ):
80
+ rank = rank_np.tolist()
81
+ logprobs = logprobs_np.tolist()
82
+ token_ids = token_ids_np.tolist()
83
# Detokenize (non-incrementally).
84
decoded_tokens = (
85
NONES
vllm/v1/outputs.py
+7
-6
@@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
5
from dataclasses import dataclass, field
6
from typing import TYPE_CHECKING, NamedTuple
7
8
+import numpy as np
9
import torch
10
11
if TYPE_CHECKING:
@@ -15,11 +16,11 @@ else:
16
17
class LogprobsLists(NamedTuple):
18
# [num_reqs x num_generated_tokens, max_num_logprobs + 1]
18
- logprob_token_ids: list[list[int]]
19
+ logprob_token_ids: np.ndarray
20
# [num_reqs x num_generated_tokens, max_num_logprobs + 1]
20
- logprobs: list[list[float]]
21
+ logprobs: np.ndarray
22
# [num_reqs x num_generated_tokens]
22
- sampled_token_ranks: list[int]
23
+ sampled_token_ranks: np.ndarray
24
# [num_reqs]
25
# Used for slicing the logprobs in cases like speculative
26
# decoding where the number of generated tokens may be
@@ -60,9 +61,9 @@ class LogprobsTensors(NamedTuple):
61
62
def tolists(self, cu_num_generated_tokens: list[int] | None = None):
63
return LogprobsLists(
63
- self.logprob_token_ids.tolist(),
64
- self.logprobs.tolist(),
65
- self.selected_token_ranks.tolist(),
64
+ self.logprob_token_ids.cpu().numpy(),
65
+ self.logprobs.cpu().numpy(),
66
+ self.selected_token_ranks.cpu().numpy(),
67
cu_num_generated_tokens,
68
)
69