[core][optimization] use a pool of numpy ndarray to hold seq data

vllm/sequence.py

@@ -1,11 +1,14 @@
 """Sequence and its related classes."""
 import copy
 import enum
+import hashlib
 import math
+import weakref
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 
 from vllm.inputs import LLMInputs
@@ -100,6 +103,33 @@ class RequestMetrics:
     finished_time: Optional[float] = None
 
 
+class SequenceDataPool:
+    """A pool of numpy array to hold sequence data.
+    """
+
+    def __init__(self, max_tokens: int, initial_pool_size: int) -> None:
+        self.max_tokens = max_tokens
+        self.pool: List[np.ndarray] = []
+        if initial_pool_size > 0:
+            self.pool = [
+                np.zeros(max_tokens, dtype=np.int64)
+                for _ in range(initial_pool_size)
+            ]
+
+    def alloc_array(self) -> np.ndarray:
+        if self.pool:
+            return self.pool.pop()
+        return np.zeros(self.max_tokens, dtype=np.int64)
+
+    def del_array(self, arr: np.ndarray) -> None:
+        assert arr.size == self.max_tokens
+        self.pool.append(arr)
+
+
+# for 128k context size
+_SEQUENCE_DATA_POOL = SequenceDataPool(128 * 1024, 32)
+
+
 class SequenceData:
     """Data associated with a sequence.
 
@@ -112,50 +142,66 @@ class SequenceData:
         prompt_token_ids: The token IDs of the prompt.
         output_token_ids: The token IDs of the output.
         cumulative_logprob: The cumulative log probability of the output.
+        tokens: array of all the tokens (prompt + output)
+
+    NOTE: special care must be taken regarding data copy and returning `list`
+        or `np.ndarray` from this class. `get_prompt_token_ids` and 
+        `get_output_token_ids` return `list`. They are used to construct
+        request output that will be returned to the user. They need to be
+        Python lists. And they are actually quite cheap, because the function
+        returns a reference to the list that is already stored in the object.
+        `get_token_ids` returns a view of `np.ndarray`. It avoids data copy,
+        and also allows array operations to be performed on the data.
     """
 
     def __init__(
         self,
         prompt_token_ids: List[int],
         output_token_ids: Optional[List[int]] = None,
     ) -> None:
+        self.tokens = _SEQUENCE_DATA_POOL.alloc_array()
+        self.prompt_token_ids_list = prompt_token_ids
+        self.num_prompt_tokens = len(prompt_token_ids)
+        self.tokens[:self.num_prompt_tokens] = prompt_token_ids
         if output_token_ids is None:
             output_token_ids = []
-
-        self.prompt_token_ids = prompt_token_ids
-        self._prompt_token_ids_tuple = tuple(prompt_token_ids)
-        self.output_token_ids = output_token_ids
+        self.num_output_tokens = len(output_token_ids)
+        self.output_token_ids_list = output_token_ids
+        self.tokens[self.num_prompt_tokens:self.num_prompt_tokens +
+                    self.num_output_tokens] = output_token_ids
         self.cumulative_logprob = 0.0
         # The number of tokens that are computed (that run against the model).
         self._num_computed_tokens = 0
         self._stage: SequenceStage = SequenceStage.PREFILL
+        self._finalizer = weakref.finalize(self, _SEQUENCE_DATA_POOL.del_array,
+                                           self.tokens)
 
     def append_token_id(self, token_id: int, logprob: float) -> None:
-        self.output_token_ids.append(token_id)
+        self.tokens[self.num_prompt_tokens + self.num_output_tokens] = token_id
+        self.output_token_ids_list.append(token_id)
+        self.num_output_tokens += 1
         self.cumulative_logprob += logprob
 
     def get_len(self) -> int:
-        return len(self.output_token_ids) + len(self.prompt_token_ids)
+        return self.num_prompt_tokens + self.num_output_tokens
 
     def get_prompt_len(self) -> int:
-        return len(self.prompt_token_ids)
+        return self.num_prompt_tokens
 
     def get_output_len(self) -> int:
-        return len(self.output_token_ids)
+        return self.num_output_tokens
 
-    def get_token_ids(self) -> List[int]:
-        return self.prompt_token_ids + self.output_token_ids
+    def get_token_ids(self) -> np.ndarray:
+        return self.tokens[:self.num_prompt_tokens + self.num_output_tokens]
 
-    def get_prefix_token_ids(
-            self, num_tokens: int
-    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
+    def hash_prefix_token_ids(self, num_tokens: int) -> bytes:
         """Get prefix tokens, and make the return value hashable"""
-        prompt_length = len(self.prompt_token_ids)
-        if num_tokens > prompt_length:
-            return (self._prompt_token_ids_tuple,
-                    tuple(self.output_token_ids[:num_tokens - prompt_length]))
-        else:
-            return (self._prompt_token_ids_tuple[:num_tokens], None)
+        data = self.tokens[:num_tokens]
+        # get a memory view of the underlying data
+        buffer = memoryview(data)  # type: ignore
+        # hash the memory view
+        hash_value = hashlib.sha256(buffer).digest()
+        return hash_value
 
     def get_num_computed_tokens(self) -> int:
         """Return the number of prefill tokens that are already computed."""
@@ -186,24 +232,23 @@ def get_num_uncomputed_tokens(self) -> int:
         return self.get_len() - self.get_num_computed_tokens()
 
     def get_last_token_id(self) -> int:
-        if not self.output_token_ids:
-            return self.prompt_token_ids[-1]
-        return self.output_token_ids[-1]
+        return int(self.tokens[self.num_prompt_tokens +
+                               self.num_output_tokens - 1])
 
     def get_prompt_token_ids(self) -> List[int]:
-        return self.prompt_token_ids
+        return self.prompt_token_ids_list
 
     def get_output_token_ids(self) -> List[int]:
-        return self.output_token_ids
+        return self.output_token_ids_list
 
     @property
     def stage(self) -> SequenceStage:
         return self._stage
 
     def __repr__(self) -> str:
         return (f"SequenceData("
-                f"prompt_token_ids={self.prompt_token_ids}, "
-                f"output_token_ids={self.output_token_ids}, "
+                f"prompt_token_ids={self.get_prompt_token_ids()}, "
+                f"output_token_ids={self.get_output_token_ids()}, "
                 f"cumulative_logprob={self.cumulative_logprob})")
 
 
@@ -232,10 +277,13 @@ def __init__(
         self.eos_token_id = eos_token_id
         self.lora_request = lora_request
 
-        self.data = SequenceData(self.prompt_token_ids)
+        self.data = SequenceData(self.inputs["prompt_token_ids"])
+        self.prompt_token_ids: List[int] = self.inputs["prompt_token_ids"]
+        self.prompt: Optional[str] = self.inputs.get("prompt")
         self.output_logprobs: SampleLogprobs = []
         self.output_text = ""
 
+        # Initialize the logical token blocks with the prompt token ids.
         self.status = SequenceStatus.WAITING
         self.stop_reason: Union[int, str, None] = None
 
@@ -247,15 +295,7 @@ def __init__(
 
     @property
     def n_blocks(self) -> int:
-        return math.ceil(self.get_len() / self.block_size)
-
-    @property
-    def prompt(self) -> Optional[str]:
-        return self.inputs.get("prompt")
-
-    @property
-    def prompt_token_ids(self) -> List[int]:
-        return self.inputs["prompt_token_ids"]
+        return math.ceil(self.data.get_len() / self.block_size)
 
     @property
     def multi_modal_data(self) -> Optional["MultiModalData"]:
@@ -278,8 +318,8 @@ def hash_of_block(self, logical_idx: int) -> int:
         # TODO: The current hashing function is O(L^2). We should optimize
         # this in the future.
         num_tokens = self.num_hashed_tokens_of_block(logical_idx)
-        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
-        return hash((hashed_tokens, self.lora_int_id))
+        tokens_hash = self.data.hash_prefix_token_ids(num_tokens)
+        return hash((tokens_hash, self.lora_int_id))
 
     def num_hashed_tokens_of_block(self, logical_idx: int):
         return logical_idx * self.block_size + self.block_size
@@ -306,7 +346,7 @@ def get_prompt_len(self) -> int:
     def get_output_len(self) -> int:
         return self.data.get_output_len()
 
-    def get_token_ids(self) -> List[int]:
+    def get_token_ids(self) -> np.ndarray:
         return self.data.get_token_ids()
 
     def get_prompt_token_ids(self) -> List[int]:
@@ -316,7 +356,7 @@ def get_last_token_id(self) -> int:
         return self.data.get_last_token_id()
 
     def get_output_token_ids(self) -> List[int]:
-        return self.data.output_token_ids
+        return self.data.get_output_token_ids()
 
     def get_cumulative_logprob(self) -> float:
         return self.data.cumulative_logprob

vllm/worker/model_runner.py

@@ -314,7 +314,7 @@ def _prepare_model_input_tensors(
 
         If cuda graph is required, this API automatically pads inputs.
         """
-        input_tokens: List[int] = []
+        input_tokens: List[np.ndarray] = []
         input_positions: List[int] = []
         slot_mapping: List[int] = []
         lora_index_mapping: List[int] = []
@@ -390,9 +390,7 @@ def _prepare_model_input_tensors(
                 if is_prompt:
                     tokens = seq_data.get_token_ids()[context_len:seq_len]
                 else:
-                    # Optimization. get_token_ids requires the entire copy of
-                    # tokens.
-                    tokens = [seq_data.get_last_token_id()]
+                    tokens = seq_data.get_token_ids()[-1:]
 
                 # Prefix cache was hit.
                 # Prefix is not supported with sliding_window
@@ -476,7 +474,7 @@ def _prepare_model_input_tensors(
                 context_lens.append(sliding_context_len)
                 query_len = sliding_seq_len - sliding_context_len
                 query_lens.append(query_len)
-                input_tokens.extend(tokens)
+                input_tokens.append(tokens)
                 input_positions.extend(list(range(context_len, seq_len)))
                 lora_id = seq_group_metadata.lora_int_id
 
@@ -554,7 +552,8 @@ def _prepare_model_input_tensors(
                     slot = block_number * self.block_size + block_offset
                     slot_mapping.append(slot)
 
-        batch_size = len(input_tokens)
+        input_tokens_array = np.concatenate(input_tokens)
+        batch_size = len(input_tokens_array)
         max_query_len = max(query_lens)
         max_prefill_seq_len = max(prefill_seq_lens, default=0)
         max_decode_seq_len = max(decode_seq_lens, default=0)
@@ -569,8 +568,13 @@ def _prepare_model_input_tensors(
         if use_captured_graph:
             graph_batch_size = _get_graph_batch_size(batch_size)
             assert graph_batch_size >= batch_size
+            input_tokens_array = np.pad(
+                input_tokens_array,
+                (0, graph_batch_size - batch_size),
+                mode="constant",
+                constant_values=0,
+            )
             for _ in range(graph_batch_size - batch_size):
-                input_tokens.append(0)
                 input_positions.append(0)
                 slot_mapping.append(_PAD_SLOT_ID)
                 seq_lens.append(1)
@@ -611,9 +615,8 @@ def _prepare_model_input_tensors(
                      dtype=seq_start_loc.dtype,
                      out=seq_start_loc[1:])
 
-        input_tokens_tensor = torch.tensor(input_tokens,
-                                           dtype=torch.long,
-                                           device=self.device)
+        input_tokens_tensor = torch.from_numpy(input_tokens_array).to(
+            device=self.device)
         input_positions_tensor = torch.tensor(input_positions,
                                               dtype=torch.long,
                                               device=self.device)