From 4f5891706f5a7b313546cf09281647bb494d929a Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 16:10:10 -0700
Subject: [PATCH 01/14] prompt as list, whole sequence as ndarray

---
 vllm/outputs.py  |   2 +-
 vllm/sequence.py | 113 +++++++++++++++++++++++++++++------------------
 2 files changed, 72 insertions(+), 43 deletions(-)

diff --git a/vllm/outputs.py b/vllm/outputs.py
index 49f526b5f930..a74234db7c1e 100644
--- a/vllm/outputs.py
+++ b/vllm/outputs.py
@@ -126,7 +126,7 @@ def from_seq_group(cls, seq_group: SequenceGroup) -> "RequestOutput":
         outputs = [
             CompletionOutput(seqs.index(seq),
                              seq.get_output_text_to_return(text_buffer_length),
-                             seq.get_output_token_ids(),
+                             seq.get_output_token_ids().tolist(),
                              seq.get_cumulative_logprob(),
                              seq.output_logprobs if include_logprobs else None,
                              SequenceStatus.get_finished_reason(seq.status),
diff --git a/vllm/sequence.py b/vllm/sequence.py
index c618c3692611..47222a07521b 100644
--- a/vllm/sequence.py
+++ b/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.
 
@@ -119,43 +149,47 @@ def __init__(
         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.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.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,15 +220,15 @@ 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
+    def get_output_token_ids(self) -> np.ndarray:
+        return self.tokens[self.num_prompt_tokens:self.num_prompt_tokens +
+                           self.num_output_tokens]
 
     @property
     def stage(self) -> SequenceStage:
@@ -202,8 +236,8 @@ def stage(self) -> SequenceStage:
 
     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().tolist()}, "
                 f"cumulative_logprob={self.cumulative_logprob})")
 
 
@@ -232,10 +266,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 +284,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 +307,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 +335,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]:
@@ -315,8 +344,8 @@ def get_prompt_token_ids(self) -> List[int]:
     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
+    def get_output_token_ids(self) -> np.ndarray:
+        return self.data.get_output_token_ids()
 
     def get_cumulative_logprob(self) -> float:
         return self.data.cumulative_logprob

From e864ed770891506e62d861e4c3d4f6279e8ed09d Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 16:20:01 -0700
Subject: [PATCH 02/14] use separate list

---
 vllm/outputs.py  |  2 +-
 vllm/sequence.py | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/vllm/outputs.py b/vllm/outputs.py
index a74234db7c1e..49f526b5f930 100644
--- a/vllm/outputs.py
+++ b/vllm/outputs.py
@@ -126,7 +126,7 @@ def from_seq_group(cls, seq_group: SequenceGroup) -> "RequestOutput":
         outputs = [
             CompletionOutput(seqs.index(seq),
                              seq.get_output_text_to_return(text_buffer_length),
-                             seq.get_output_token_ids().tolist(),
+                             seq.get_output_token_ids(),
                              seq.get_cumulative_logprob(),
                              seq.output_logprobs if include_logprobs else None,
                              SequenceStatus.get_finished_reason(seq.status),
diff --git a/vllm/sequence.py b/vllm/sequence.py
index 47222a07521b..effc80f6a622 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -156,6 +156,7 @@ def __init__(
         if output_token_ids is None:
             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
@@ -167,6 +168,7 @@ def __init__(
 
     def append_token_id(self, token_id: int, logprob: float) -> None:
         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
 
@@ -226,9 +228,8 @@ def get_last_token_id(self) -> int:
     def get_prompt_token_ids(self) -> List[int]:
         return self.prompt_token_ids_list
 
-    def get_output_token_ids(self) -> np.ndarray:
-        return self.tokens[self.num_prompt_tokens:self.num_prompt_tokens +
-                           self.num_output_tokens]
+    def get_output_token_ids(self) -> List[int]:
+        return self.output_token_ids_list
 
     @property
     def stage(self) -> SequenceStage:
@@ -237,7 +238,7 @@ def stage(self) -> SequenceStage:
     def __repr__(self) -> str:
         return (f"SequenceData("
                 f"prompt_token_ids={self.get_prompt_token_ids()}, "
-                f"output_token_ids={self.get_output_token_ids().tolist()}, "
+                f"output_token_ids={self.get_output_token_ids()}, "
                 f"cumulative_logprob={self.cumulative_logprob})")
 
 
@@ -344,7 +345,7 @@ def get_prompt_token_ids(self) -> List[int]:
     def get_last_token_id(self) -> int:
         return self.data.get_last_token_id()
 
-    def get_output_token_ids(self) -> np.ndarray:
+    def get_output_token_ids(self) -> List[int]:
         return self.data.get_output_token_ids()
 
     def get_cumulative_logprob(self) -> float:

From 713de379ce252f5ec2f32ab3520921f1130a90e9 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 16:35:27 -0700
Subject: [PATCH 03/14] use ndarray for input tokens

---
 vllm/worker/model_runner.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index 181442490a82..26374b125e1c 100644
--- a/vllm/worker/model_runner.py
+++ b/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)

From 66f5154380ebf9a3b8469d8bc1c19e6d653661f7 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 16:51:44 -0700
Subject: [PATCH 04/14] add comments

---
 vllm/sequence.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/vllm/sequence.py b/vllm/sequence.py
index effc80f6a622..261971e92549 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -142,6 +142,16 @@ 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__(

From 567850da4bb5a7a51ef649a0d5afeee08428d44d Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 18:51:09 -0700
Subject: [PATCH 05/14] use max seq len

---
 vllm/sequence.py | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/vllm/sequence.py b/vllm/sequence.py
index 261971e92549..7a2255bb0a64 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -5,6 +5,7 @@
 import math
 import weakref
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
@@ -107,27 +108,20 @@ 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 __init__(self) -> None:
+        self.pool: Dict[int, List[np.ndarray]] = defaultdict(list)
 
-    def alloc_array(self) -> np.ndarray:
-        if self.pool:
-            return self.pool.pop()
-        return np.zeros(self.max_tokens, dtype=np.int64)
+    def alloc_array(self, max_tokens) -> np.ndarray:
+        if max_tokens in self.pool and self.pool[max_tokens]:
+            return self.pool[max_tokens].pop()
+        return np.zeros(max_tokens, dtype=np.int64)
 
     def del_array(self, arr: np.ndarray) -> None:
-        assert arr.size == self.max_tokens
-        self.pool.append(arr)
+        self.pool[len(arr)].append(arr)
 
 
 # for 128k context size
-_SEQUENCE_DATA_POOL = SequenceDataPool(128 * 1024, 32)
+_SEQUENCE_DATA_POOL = SequenceDataPool()
 
 
 class SequenceData:
@@ -158,8 +152,9 @@ def __init__(
         self,
         prompt_token_ids: List[int],
         output_token_ids: Optional[List[int]] = None,
+        max_seq_len: int = 1024,
     ) -> None:
-        self.tokens = _SEQUENCE_DATA_POOL.alloc_array()
+        self.tokens = _SEQUENCE_DATA_POOL.alloc_array(max_seq_len)
         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

From d2defac38653c9a49e2e86d4635e7aa9c6457a72 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 18:57:43 -0700
Subject: [PATCH 06/14] use 16k by default

---
 vllm/sequence.py                    | 7 ++++++-
 vllm/spec_decode/batch_expansion.py | 1 +
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/vllm/sequence.py b/vllm/sequence.py
index 7a2255bb0a64..666b4021090b 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -131,6 +131,10 @@ class SequenceData:
         prompt_token_ids: The token IDs of the prompt.
         output_token_ids: The token IDs of the output. Set to an empty list if
             None.
+        max_seq_len: The maximum sequence length. A buffer of this size is
+            allocated for the tokens. By default, it is set to 16k for test
+            purposes. During inference, it should be set to the maximum
+            sequence length of the model.
 
     Attributes:
         prompt_token_ids: The token IDs of the prompt.
@@ -152,8 +156,9 @@ def __init__(
         self,
         prompt_token_ids: List[int],
         output_token_ids: Optional[List[int]] = None,
-        max_seq_len: int = 1024,
+        max_seq_len: int = 16 * 1024,
     ) -> None:
+        self.max_seq_len = max_seq_len
         self.tokens = _SEQUENCE_DATA_POOL.alloc_array(max_seq_len)
         self.prompt_token_ids_list = prompt_token_ids
         self.num_prompt_tokens = len(prompt_token_ids)
diff --git a/vllm/spec_decode/batch_expansion.py b/vllm/spec_decode/batch_expansion.py
index 40516556344e..08c77dd0811d 100644
--- a/vllm/spec_decode/batch_expansion.py
+++ b/vllm/spec_decode/batch_expansion.py
@@ -283,6 +283,7 @@ def _create_single_target_seq_group_metadata(
             SequenceData(
                 prompt_token_ids=prompt_token_ids,
                 output_token_ids=new_output_token_ids,
+                max_seq_len=seq_data.max_seq_len,
             ),
         }
         # This is a hack. Technically, spec decoding should compute

From 1ebebdbc076c20d9d6b1b76b043e7183ecd0b732 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 19:05:34 -0700
Subject: [PATCH 07/14] pass max seq len from llm engine

---
 vllm/engine/llm_engine.py | 8 ++++++--
 vllm/sequence.py          | 9 ++++++---
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py
index 4b427b1fb2f2..8cadaf3ff567 100644
--- a/vllm/engine/llm_engine.py
+++ b/vllm/engine/llm_engine.py
@@ -464,8 +464,12 @@ def _add_processed_request(
         seq_id = next(self.seq_counter)
         eos_token_id = self._get_eos_token_id(lora_request)
 
-        seq = Sequence(seq_id, processed_inputs, block_size, eos_token_id,
-                       lora_request)
+        seq = Sequence(seq_id,
+                       processed_inputs,
+                       block_size,
+                       eos_token_id,
+                       lora_request,
+                       max_seq_len=self.model_config.max_model_len)
 
         # Create a SequenceGroup based on SamplingParams or PoolingParams
         if isinstance(params, SamplingParams):
diff --git a/vllm/sequence.py b/vllm/sequence.py
index 666b4021090b..b6d690a7e563 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -156,9 +156,9 @@ def __init__(
         self,
         prompt_token_ids: List[int],
         output_token_ids: Optional[List[int]] = None,
-        max_seq_len: int = 16 * 1024,
+        max_seq_len: Optional[int] = None,
     ) -> None:
-        self.max_seq_len = max_seq_len
+        self.max_seq_len = max_seq_len or 16 * 1024
         self.tokens = _SEQUENCE_DATA_POOL.alloc_array(max_seq_len)
         self.prompt_token_ids_list = prompt_token_ids
         self.num_prompt_tokens = len(prompt_token_ids)
@@ -270,6 +270,7 @@ def __init__(
         block_size: int,
         eos_token_id: Optional[int] = None,
         lora_request: Optional[LoRARequest] = None,
+        max_seq_len: Optional[int] = None,
     ) -> None:
         self.seq_id = seq_id
         self.inputs = inputs
@@ -277,7 +278,9 @@ def __init__(
         self.eos_token_id = eos_token_id
         self.lora_request = lora_request
 
-        self.data = SequenceData(self.inputs["prompt_token_ids"])
+        self.max_seq_len = max_seq_len
+        self.data = SequenceData(self.inputs["prompt_token_ids"],
+                                 max_seq_len=max_seq_len)
         self.prompt_token_ids: List[int] = self.inputs["prompt_token_ids"]
         self.prompt: Optional[str] = self.inputs.get("prompt")
         self.output_logprobs: SampleLogprobs = []

From 12f1d5449a57e14a15d733c6f19b999f71fbe71c Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 19:14:49 -0700
Subject: [PATCH 08/14] fix type

---
 vllm/sequence.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vllm/sequence.py b/vllm/sequence.py
index b6d690a7e563..9e7fcefe481c 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -111,10 +111,10 @@ class SequenceDataPool:
     def __init__(self) -> None:
         self.pool: Dict[int, List[np.ndarray]] = defaultdict(list)
 
-    def alloc_array(self, max_tokens) -> np.ndarray:
+    def alloc_array(self, max_tokens: int) -> np.ndarray:
         if max_tokens in self.pool and self.pool[max_tokens]:
             return self.pool[max_tokens].pop()
-        return np.zeros(max_tokens, dtype=np.int64)
+        return np.zeros((max_tokens, ), dtype=np.int64)
 
     def del_array(self, arr: np.ndarray) -> None:
         self.pool[len(arr)].append(arr)
@@ -159,7 +159,7 @@ def __init__(
         max_seq_len: Optional[int] = None,
     ) -> None:
         self.max_seq_len = max_seq_len or 16 * 1024
-        self.tokens = _SEQUENCE_DATA_POOL.alloc_array(max_seq_len)
+        self.tokens = _SEQUENCE_DATA_POOL.alloc_array(self.max_seq_len)
         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

From eead60483c0a69dc7d1e6b33a94dbe71c862912f Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 19:22:58 -0700
Subject: [PATCH 09/14] optimize input_positions

---
 vllm/worker/model_runner.py | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index 26374b125e1c..a0fb31ebb20e 100644
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -314,8 +314,9 @@ def _prepare_model_input_tensors(
 
         If cuda graph is required, this API automatically pads inputs.
         """
+        batch_size = 0
         input_tokens: List[np.ndarray] = []
-        input_positions: List[int] = []
+        input_positions: List[np.ndarray] = []
         slot_mapping: List[int] = []
         lora_index_mapping: List[int] = []
         lora_prompt_mapping: List[int] = []
@@ -475,7 +476,8 @@ def _prepare_model_input_tensors(
                 query_len = sliding_seq_len - sliding_context_len
                 query_lens.append(query_len)
                 input_tokens.append(tokens)
-                input_positions.extend(list(range(context_len, seq_len)))
+                batch_size += len(tokens)
+                input_positions.append(np.arange(context_len, seq_len))
                 lora_id = seq_group_metadata.lora_int_id
 
                 if is_prompt:
@@ -552,8 +554,6 @@ def _prepare_model_input_tensors(
                     slot = block_number * self.block_size + block_offset
                     slot_mapping.append(slot)
 
-        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)
@@ -568,14 +568,11 @@ 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,
-            )
+            zero_pad_array = np.zeros(graph_batch_size - batch_size,
+                                      dtype=np.int64)
+            input_tokens.append(zero_pad_array)
+            input_positions.append(zero_pad_array)
             for _ in range(graph_batch_size - batch_size):
-                input_positions.append(0)
                 slot_mapping.append(_PAD_SLOT_ID)
                 seq_lens.append(1)
                 block_tables.append([])
@@ -615,11 +612,12 @@ def _prepare_model_input_tensors(
                      dtype=seq_start_loc.dtype,
                      out=seq_start_loc[1:])
 
+        input_tokens_array = np.concatenate(input_tokens)
         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)
+            device=self.device, non_blocking=True)
+        input_positions_array = np.concatenate(input_positions)
+        input_positions_tensor = torch.from_numpy(input_positions_array).to(
+            device=self.device, non_blocking=True)
         slot_mapping_tensor = torch.tensor(slot_mapping,
                                            dtype=torch.long,
                                            device=self.device)

From d1db7bd255e89477d65a6298aa2adc57fa495ac1 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 19:37:33 -0700
Subject: [PATCH 10/14] try to fix

---
 vllm/worker/model_runner.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index a0fb31ebb20e..5442b00ec75f 100644
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -476,7 +476,7 @@ def _prepare_model_input_tensors(
                 query_len = sliding_seq_len - sliding_context_len
                 query_lens.append(query_len)
                 input_tokens.append(tokens)
-                batch_size += len(tokens)
+                batch_size += seq_len - context_len
                 input_positions.append(np.arange(context_len, seq_len))
                 lora_id = seq_group_metadata.lora_int_id
 
@@ -614,10 +614,10 @@ def _prepare_model_input_tensors(
 
         input_tokens_array = np.concatenate(input_tokens)
         input_tokens_tensor = torch.from_numpy(input_tokens_array).to(
-            device=self.device, non_blocking=True)
+            device=self.device)
         input_positions_array = np.concatenate(input_positions)
         input_positions_tensor = torch.from_numpy(input_positions_array).to(
-            device=self.device, non_blocking=True)
+            device=self.device)
         slot_mapping_tensor = torch.tensor(slot_mapping,
                                            dtype=torch.long,
                                            device=self.device)

From 0470012b4817b15169731f4b9762a9788e09ab12 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 20:10:07 -0700
Subject: [PATCH 11/14] move pool out

---
 vllm/array_pool.py | 13 +++++++++++++
 vllm/sequence.py   | 26 +++-----------------------
 2 files changed, 16 insertions(+), 23 deletions(-)
 create mode 100644 vllm/array_pool.py

diff --git a/vllm/array_pool.py b/vllm/array_pool.py
new file mode 100644
index 000000000000..aad47792fc18
--- /dev/null
+++ b/vllm/array_pool.py
@@ -0,0 +1,13 @@
+import numpy as np
+from collections import defaultdict
+from typing import Dict, List
+
+_POOL: Dict[int, List[np.ndarray]] = defaultdict(list)
+
+def alloc_array(max_tokens: int) -> np.ndarray:
+    if max_tokens in _POOL and _POOL[max_tokens]:
+        return _POOL[max_tokens].pop()
+    return np.zeros((max_tokens, ), dtype=np.int64)
+
+def del_array(arr: np.ndarray) -> None:
+    _POOL[len(arr)].append(arr)
diff --git a/vllm/sequence.py b/vllm/sequence.py
index 9e7fcefe481c..ceab3cb409c7 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -12,6 +12,7 @@
 import numpy as np
 import torch
 
+from vllm.array_pool import alloc_array, del_array
 from vllm.inputs import LLMInputs
 from vllm.lora.request import LoRARequest
 from vllm.pooling_params import PoolingParams
@@ -104,26 +105,6 @@ class RequestMetrics:
     finished_time: Optional[float] = None
 
 
-class SequenceDataPool:
-    """A pool of numpy array to hold sequence data.
-    """
-
-    def __init__(self) -> None:
-        self.pool: Dict[int, List[np.ndarray]] = defaultdict(list)
-
-    def alloc_array(self, max_tokens: int) -> np.ndarray:
-        if max_tokens in self.pool and self.pool[max_tokens]:
-            return self.pool[max_tokens].pop()
-        return np.zeros((max_tokens, ), dtype=np.int64)
-
-    def del_array(self, arr: np.ndarray) -> None:
-        self.pool[len(arr)].append(arr)
-
-
-# for 128k context size
-_SEQUENCE_DATA_POOL = SequenceDataPool()
-
-
 class SequenceData:
     """Data associated with a sequence.
 
@@ -159,7 +140,7 @@ def __init__(
         max_seq_len: Optional[int] = None,
     ) -> None:
         self.max_seq_len = max_seq_len or 16 * 1024
-        self.tokens = _SEQUENCE_DATA_POOL.alloc_array(self.max_seq_len)
+        self.tokens = alloc_array(self.max_seq_len)
         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
@@ -173,8 +154,7 @@ def __init__(
         # 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)
+        self._finalizer = weakref.finalize(self, del_array, self.tokens)
 
     def append_token_id(self, token_id: int, logprob: float) -> None:
         self.tokens[self.num_prompt_tokens + self.num_output_tokens] = token_id

From 927119046fd06cfcd391599c5b592ea64a6d13ee Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 20:30:04 -0700
Subject: [PATCH 12/14] change blocks to array

---
 vllm/array_pool.py                         |  5 ++-
 vllm/core/block/block_table.py             | 37 +++++++++++++---------
 vllm/core/block/cpu_gpu_block_allocator.py | 12 ++++---
 vllm/core/block/interfaces.py              | 12 ++++---
 vllm/core/block/naive_block.py             | 33 ++++++++++++-------
 vllm/core/block/prefix_caching_block.py    | 22 +++++++------
 vllm/sequence.py                           |  1 -
 vllm/utils.py                              |  5 +++
 8 files changed, 78 insertions(+), 49 deletions(-)

diff --git a/vllm/array_pool.py b/vllm/array_pool.py
index aad47792fc18..b3e9dbed7eb5 100644
--- a/vllm/array_pool.py
+++ b/vllm/array_pool.py
@@ -1,13 +1,16 @@
-import numpy as np
 from collections import defaultdict
 from typing import Dict, List
 
+import numpy as np
+
 _POOL: Dict[int, List[np.ndarray]] = defaultdict(list)
 
+
 def alloc_array(max_tokens: int) -> np.ndarray:
     if max_tokens in _POOL and _POOL[max_tokens]:
         return _POOL[max_tokens].pop()
     return np.zeros((max_tokens, ), dtype=np.int64)
 
+
 def del_array(arr: np.ndarray) -> None:
     _POOL[len(arr)].append(arr)
diff --git a/vllm/core/block/block_table.py b/vllm/core/block/block_table.py
index d705f3d91a07..046b30fb01be 100644
--- a/vllm/core/block/block_table.py
+++ b/vllm/core/block/block_table.py
@@ -1,7 +1,9 @@
 from typing import List, Optional
 
+import numpy as np
+
 from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator
-from vllm.utils import Device, cdiv, chunk_list
+from vllm.utils import Device, cdiv, chunk_array
 
 
 class BlockTable:
@@ -74,7 +76,7 @@ def get_num_required_blocks(token_ids: List[int], block_size: int) -> int:
         return cdiv(len(token_ids), block_size)
 
     def allocate(self,
-                 token_ids: List[int],
+                 token_ids: np.ndarray,
                  device: Device = Device.GPU) -> None:
         """Allocates memory blocks for storing the given sequence of token IDs.
 
@@ -82,19 +84,20 @@ def allocate(self,
         sequence of token IDs.
 
         Args:
-            token_ids (List[int]): The sequence of token IDs to be stored.
+            token_ids (np.ndarray): The sequence of token IDs to be stored.
             device (Device, optional): The device on which the blocks should be
                 allocated. Defaults to Device.GPU.
         """
         assert not self._is_allocated
-        assert token_ids
+        len_token_ids = len(token_ids)
+        assert len_token_ids > 0
         self._blocks = self._allocate_blocks_for_token_ids(prev_block=None,
                                                            token_ids=token_ids,
                                                            device=device)
-        self._num_full_slots = len(token_ids)
+        self._num_full_slots = len_token_ids
 
     def append_token_ids(self,
-                         token_ids: List[int],
+                         token_ids: np.ndarray,
                          num_lookahead_slots: int = 0,
                          num_computed_slots: Optional[int] = None) -> None:
         """Appends a sequence of token IDs to the existing blocks in the
@@ -230,7 +233,8 @@ def physical_block_ids(self) -> List[Optional[int]]:
         assert self._is_allocated
         return [block.block_id for block in self._blocks]
 
-    def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]:
+    def get_unseen_token_ids(self,
+                             sequence_token_ids: np.ndarray) -> np.ndarray:
         """Get the number of "unseen" tokens in the sequence.
 
         Unseen tokens are tokens in the sequence corresponding to this block
@@ -250,10 +254,11 @@ def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]:
         return sequence_token_ids[self.num_full_slots:]
 
     def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
-                                       token_ids: List[int],
+                                       token_ids: np.ndarray,
                                        device: Device) -> List[Block]:
         blocks: List[Block] = []
-        for block_token_ids in chunk_list(token_ids, self._block_size):
+        for i in range(0, len(token_ids), self._block_size):
+            block_token_ids = token_ids[i:i + self._block_size]
             if len(block_token_ids) == self._block_size:
                 # If the block is full, create an immutable block.
                 prev_block = self._allocator.allocate_immutable(
@@ -303,7 +308,7 @@ def num_full_slots(self) -> int:
         return self._num_full_slots
 
     def get_num_blocks_touched_by_append_slots(
-            self, token_ids: List[int], num_lookahead_slots: int) -> int:
+            self, token_ids: np.ndarray, num_lookahead_slots: int) -> int:
         """Determine how many blocks will be "touched" by appending the token
         ids.
 
@@ -311,12 +316,14 @@ def get_num_blocks_touched_by_append_slots(
         continue generation, or if it must be preempted.
         """
 
-        all_token_ids = token_ids + [-1] * num_lookahead_slots
-        token_blocks = self._chunk_token_blocks_for_append(all_token_ids)
-        return len(token_blocks)
+        size = len(token_ids) + num_lookahead_slots
+        first_chunk_size = self._block_size - (self._num_full_slots %
+                                               self._block_size)
+        n_blocks = 1 + len(range(first_chunk_size, size, self._block_size))
+        return n_blocks
 
     def _chunk_token_blocks_for_append(
-            self, token_ids: List[int]) -> List[List[int]]:
+            self, token_ids: np.ndarray) -> List[np.ndarray]:
         """Split the token ids into block-sized chunks so they can be easily
         appended to blocks. The first such "token block" may have less token ids
         than the block size, since the last allocated block may be partially
@@ -324,6 +331,6 @@ def _chunk_token_blocks_for_append(
         """
         first_chunk_size = self._block_size - (self._num_full_slots %
                                                self._block_size)
-        token_blocks = [token_ids[:first_chunk_size]] + chunk_list(
+        token_blocks = [token_ids[:first_chunk_size]] + chunk_array(
             token_ids[first_chunk_size:], self._block_size)
         return token_blocks
diff --git a/vllm/core/block/cpu_gpu_block_allocator.py b/vllm/core/block/cpu_gpu_block_allocator.py
index 255aae9d1731..d4aaa1d24e08 100644
--- a/vllm/core/block/cpu_gpu_block_allocator.py
+++ b/vllm/core/block/cpu_gpu_block_allocator.py
@@ -1,5 +1,7 @@
 from typing import Dict, FrozenSet, List, Optional, Tuple
 
+import numpy as np
+
 from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,
                                         DeviceAwareBlockAllocator)
 from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator
@@ -131,15 +133,15 @@ def allocate_mutable(self, prev_block: Optional[Block],
         return self._allocators[device].allocate_mutable(prev_block)
 
     def allocate_immutable(self, prev_block: Optional[Block],
-                           token_ids: List[int], device: Device) -> Block:
+                           token_ids: np.ndarray, device: Device) -> Block:
         """Allocates a new immutable block with the provided token IDs on the
         specified device.
 
         Args:
             prev_block (Optional[Block]): The previous block in the sequence.
                 Used for prefix hashing.
-            token_ids (List[int]): The list of token IDs to be stored in the new
-                block.
+            token_ids (np.ndarray): The list of token IDs to be stored in the
+                 new block.
             device (Device): The device on which to allocate the new block.
 
         Returns:
@@ -326,7 +328,7 @@ def __init__(self, proxy: Block):
         super().__init__()
         self._proxy = proxy
 
-    def append_token_ids(self, token_ids: List[BlockId]):
+    def append_token_ids(self, token_ids: np.ndarray):
         raise ValueError("null block should not be modified")
 
     @property
@@ -338,7 +340,7 @@ def block_id(self, value: Optional[BlockId]):
         raise ValueError("null block should not be modified")
 
     @property
-    def token_ids(self) -> List[BlockId]:
+    def token_ids(self) -> np.ndarray:
         return self._proxy.token_ids
 
     @property
diff --git a/vllm/core/block/interfaces.py b/vllm/core/block/interfaces.py
index 4b20856a1b42..05b642c2aaf9 100644
--- a/vllm/core/block/interfaces.py
+++ b/vllm/core/block/interfaces.py
@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Dict, FrozenSet, List, Optional, Protocol, Tuple
 
+import numpy as np
+
 from vllm.utils import Device
 
 BlockId = int
@@ -9,7 +11,7 @@
 class Block(ABC):
 
     @abstractmethod
-    def append_token_ids(self, token_ids: List[int]) -> None:
+    def append_token_ids(self, token_ids: np.ndarray) -> None:
         pass
 
     @property
@@ -25,7 +27,7 @@ def block_id(self, value: Optional[int]) -> None:
 
     @property
     @abstractmethod
-    def token_ids(self) -> List[int]:
+    def token_ids(self) -> np.ndarray:
         pass
 
     @property
@@ -70,7 +72,7 @@ class Factory(Protocol):
         def __call__(
             self,
             prev_block: Optional["Block"],
-            token_ids: List[int],
+            token_ids: np.ndarray,
             block_size: int,
             allocator: "BlockAllocator",
             block_id: Optional[int] = None,
@@ -97,7 +99,7 @@ def allocate_mutable(self, prev_block: Optional[Block]) -> Block:
 
     @abstractmethod
     def allocate_immutable(self, prev_block: Optional[Block],
-                           token_ids: List[int]) -> Block:
+                           token_ids: np.ndarray) -> Block:
         pass
 
     @abstractmethod
@@ -180,7 +182,7 @@ def allocate_mutable(self, prev_block: Optional[Block],
 
     @abstractmethod
     def allocate_immutable(self, prev_block: Optional[Block],
-                           token_ids: List[int], device: Device) -> Block:
+                           token_ids: np.ndarray, device: Device) -> Block:
         pass
 
     @abstractmethod
diff --git a/vllm/core/block/naive_block.py b/vllm/core/block/naive_block.py
index 50f27bab3377..48012cb7f0d6 100644
--- a/vllm/core/block/naive_block.py
+++ b/vllm/core/block/naive_block.py
@@ -1,5 +1,9 @@
+import weakref
 from typing import FrozenSet, Iterable, List, Optional, Set, Tuple
 
+import numpy as np
+
+from vllm.array_pool import alloc_array, del_array
 from vllm.core.block.common import (CopyOnWriteTracker, RefCounter,
                                     get_all_blocks_recursively)
 from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device
@@ -51,7 +55,7 @@ def __init__(
 
     def allocate_immutable(self,
                            prev_block: Optional[Block],
-                           token_ids: List[int],
+                           token_ids: np.ndarray,
                            device: Optional[Device] = None) -> Block:
         """Allocates a new immutable block with the given token IDs, linked to
         the previous block.
@@ -87,7 +91,7 @@ def allocate_mutable(self,
         block_id = self._allocate_new_block_id()
         return self._create_block(
             prev_block=prev_block,
-            token_ids=[],
+            token_ids=np.zeros(0, dtype=np.int64),
             block_id=block_id,
             block_size=self._block_size,
             allocator=self,
@@ -286,7 +290,7 @@ class NaiveBlock(Block):
 
     Args:
         prev_block (Block): The previous block in the sequence.
-        token_ids (List[int]): The initial token IDs to be stored in the block.
+        token_ids (np.ndarray): The initial token IDs to be stored in the block.
         block_size (int): The maximum number of token IDs that can be stored in
             the block.
         allocator (BlockAllocator): The block allocator associated with this
@@ -300,12 +304,13 @@ class NaiveBlock(Block):
 
     def __init__(self,
                  prev_block: Optional[Block],
-                 token_ids: List[int],
+                 token_ids: np.ndarray,
                  block_size: int,
                  allocator: BlockAllocator,
                  block_id: Optional[int] = None,
                  _cow_target: Optional[Block] = None):
-        self._token_ids: List[int] = []
+        self._token_ids: np.ndarray = alloc_array(block_size)
+        self._num_tokens = 0
         self._block_size = block_size
         self._prev_block = prev_block
         self._block_id = block_id
@@ -313,8 +318,9 @@ def __init__(self,
         self._cow_target = _cow_target if _cow_target is not None else self
 
         self._append_token_ids_no_cow(token_ids)
+        self._finalizer = weakref.finalize(self, del_array, self._token_ids)
 
-    def append_token_ids(self, token_ids: List[int]) -> None:
+    def append_token_ids(self, token_ids: np.ndarray) -> None:
         """Appends the given token IDs to the block, instructing the allocator
         to perform a copy-on-write if necessary.
 
@@ -327,9 +333,12 @@ def append_token_ids(self, token_ids: List[int]) -> None:
             self._block_id = (self._allocator.cow_block_if_not_appendable(
                 self._cow_target))
 
-    def _append_token_ids_no_cow(self, token_ids: List[int]) -> None:
-        assert self.num_empty_slots >= len(token_ids)
-        self._token_ids.extend(token_ids)
+    def _append_token_ids_no_cow(self, token_ids: np.ndarray) -> None:
+        len_new_tokens = len(token_ids)
+        new_len = self._num_tokens + len_new_tokens
+        assert new_len <= self._block_size
+        self._token_ids[self._num_tokens:new_len] = token_ids
+        self._num_tokens = new_len
 
     @property
     def computed(self) -> bool:
@@ -357,14 +366,14 @@ def block_id(self, value: Optional[int]) -> None:
 
     @property
     def is_full(self) -> bool:
-        return self.num_empty_slots == 0
+        return self._num_tokens == self._block_size
 
     @property
     def num_empty_slots(self) -> int:
-        return self._block_size - len(self._token_ids)
+        return self._block_size - self._num_tokens
 
     @property
-    def token_ids(self) -> List[int]:
+    def token_ids(self) -> np.ndarray:
         return self._token_ids
 
     @property
diff --git a/vllm/core/block/prefix_caching_block.py b/vllm/core/block/prefix_caching_block.py
index 2df7d74e4ff1..ea370c7c46d7 100644
--- a/vllm/core/block/prefix_caching_block.py
+++ b/vllm/core/block/prefix_caching_block.py
@@ -4,6 +4,8 @@
 from os.path import commonprefix
 from typing import Dict, FrozenSet, Iterable, List, Optional, Tuple
 
+import numpy as np
+
 from vllm.core.block.common import (CopyOnWriteTracker,
                                     get_all_blocks_recursively)
 from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device
@@ -76,7 +78,7 @@ def __init__(
     def _create_block(
         self,
         prev_block: Optional[Block],
-        token_ids: List[int],
+        token_ids: np.ndarray,
         block_size: int,
         allocator: BlockAllocator,
         block_id: Optional[int] = None,
@@ -96,7 +98,7 @@ def _create_block(
 
     def allocate_immutable(self,
                            prev_block: Optional[Block],
-                           token_ids: List[int],
+                           token_ids: np.ndarray,
                            device: Optional[Device] = None) -> Block:
         """Allocates an immutable block with the given token IDs, reusing cached
         blocks if possible.
@@ -182,7 +184,7 @@ def allocate_mutable(self,
             # its kvcache
             block = self._create_block(
                 prev_block=prev_block,
-                token_ids=[],
+                token_ids=np.zeros(0, dtype=np.int32),
                 block_size=self._block_size,
                 allocator=self,
                 block_id=block_id,
@@ -504,7 +506,7 @@ class PrefixCachingBlock(Block):
     Args:
         prev_block (Optional[PrefixCachingBlock]): The previous block in the
             sequence.
-        token_ids (List[int]): The initial token IDs to be stored in the block.
+        token_ids (np.ndarray): The initial token IDs to be stored in the block.
         block_size (int): The maximum number of token IDs that can be stored in
             the block.
         prefix_caching_allocator (BlockAllocator): The prefix
@@ -516,7 +518,7 @@ class PrefixCachingBlock(Block):
     def __init__(
         self,
         prev_block: Optional[Block],
-        token_ids: List[int],
+        token_ids: np.ndarray,
         block_size: int,
         prefix_caching_allocator: BlockAllocator,
         block_id: Optional[int] = None,
@@ -560,14 +562,14 @@ def last_accessed(self) -> float:
     def last_accessed(self, last_accessed_ts: float):
         self._last_accessed = last_accessed_ts
 
-    def append_token_ids(self, token_ids: List[int]) -> None:
+    def append_token_ids(self, token_ids: np.ndarray) -> None:
         """Appends the given token IDs to the block and registers the block as
         immutable if the block becomes full.
 
         Internally, the naive block handles CoW.
 
         Args:
-            token_ids (List[int]): The token IDs to be appended to the block.
+            token_ids (np.ndarray): The token IDs to be appended to the block.
         """
         assert token_ids
 
@@ -623,7 +625,7 @@ def block_size(self) -> int:
         return self._block.block_size
 
     @property
-    def token_ids(self) -> List[int]:
+    def token_ids(self) -> np.ndarray:
         return self._block.token_ids
 
     @property
@@ -666,7 +668,7 @@ def content_hash(self) -> Optional[int]:
 
     @staticmethod
     def hash_block_tokens(is_first_block: bool, prev_block_hash: Optional[int],
-                          cur_block_token_ids: List[int]) -> int:
+                          cur_block_token_ids: np.ndarray) -> int:
         """Computes a hash value corresponding to the contents of a block and
         the contents of the preceding block(s). The hash value is used for
         prefix caching.
@@ -678,7 +680,7 @@ def hash_block_tokens(is_first_block: bool, prev_block_hash: Optional[int],
             the sequence.
         - prev_block_hash (Optional[int]): The hash of the previous block. None
             if this is the first block.
-        - cur_block_token_ids (List[int]): A list of token ids in the current
+        - cur_block_token_ids (np.ndarray): A list of token ids in the current
             block. The current block is assumed to be full.
 
         Returns:
diff --git a/vllm/sequence.py b/vllm/sequence.py
index ceab3cb409c7..08c8f4d1f4af 100644
--- a/vllm/sequence.py
+++ b/vllm/sequence.py
@@ -5,7 +5,6 @@
 import math
 import weakref
 from abc import ABC, abstractmethod
-from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
diff --git a/vllm/utils.py b/vllm/utils.py
index 92abdb3fb9b1..c4ccc0c45662 100644
--- a/vllm/utils.py
+++ b/vllm/utils.py
@@ -393,6 +393,11 @@ def chunk_list(lst: List[T], chunk_size: int) -> List[List[T]]:
     return [lst[i:i + chunk_size] for i in range(0, len(lst), chunk_size)]
 
 
+def chunk_array(lst: np.ndarray, chunk_size: int) -> List[np.ndarray]:
+    """Yield successive chunk_size chunks from lst."""
+    return [lst[i:i + chunk_size] for i in range(0, len(lst), chunk_size)]
+
+
 def cdiv(a: int, b: int) -> int:
     """Ceiling division."""
     return -(a // -b)

From 58da18ddc10b78a7892f75a2eecc964746800075 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 20:42:56 -0700
Subject: [PATCH 13/14] fix more types

---
 vllm/core/block/block_table.py             | 32 ++++++++++++----------
 vllm/core/block/cpu_gpu_block_allocator.py |  4 +++
 vllm/core/block/interfaces.py              |  5 ++++
 vllm/core/block/naive_block.py             |  4 +++
 vllm/core/block/prefix_caching_block.py    |  4 +++
 5 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/vllm/core/block/block_table.py b/vllm/core/block/block_table.py
index 046b30fb01be..5825225f5a32 100644
--- a/vllm/core/block/block_table.py
+++ b/vllm/core/block/block_table.py
@@ -54,10 +54,14 @@ def __init__(
         self._max_block_sliding_window = max_block_sliding_window
         # Use helper method instead of directly calculating, as blocks
         # may not be allocated.
-        self._num_full_slots = len(self._get_all_token_ids())
+        _num_full_slots = 0
+        if self._is_allocated:
+            for block in self._blocks:
+                _num_full_slots += block.num_tokens
+        self._num_full_slots = _num_full_slots
 
     @staticmethod
-    def get_num_required_blocks(token_ids: List[int], block_size: int) -> int:
+    def get_num_required_blocks(token_ids: np.ndarray, block_size: int) -> int:
         """Calculates the minimum number of blocks required to store a given
         sequence of token IDs.
 
@@ -65,7 +69,7 @@ def get_num_required_blocks(token_ids: List[int], block_size: int) -> int:
         allocation (e.g. ignoring prefix caching).
 
         Args:
-            token_ids (List[int]): The sequence of token IDs to be stored.
+            token_ids (np.ndarray): The sequence of token IDs to be stored.
             block_size (int): The maximum number of tokens that can be stored in
                 a single block.
 
@@ -113,7 +117,7 @@ def append_token_ids(self,
         separate block.
 
         Args:
-            token_ids (List[int]): The sequence of token IDs to be appended.
+            token_ids (np.ndarray): The sequence of token IDs to be appended.
             num_computed_slots (Optional[int]): The number of KV cache slots
                 that are already filled (computed).
                 When sliding window is enabled, this is used to compute how many
@@ -241,11 +245,11 @@ def get_unseen_token_ids(self,
         table, but are not yet appended to this block table.
 
         Args:
-            sequence_token_ids (List[int]): The list of token ids in the
+            sequence_token_ids (np.ndarray): The list of token ids in the
                 sequence.
 
         Returns:
-            List[int]: The postfix of sequence_token_ids that has not yet been
+            np.ndarray: The postfix of sequence_token_ids that has not yet been
                 appended to the block table.
         """
 
@@ -272,17 +276,15 @@ def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
 
         return blocks
 
-    def _get_all_token_ids(self) -> List[int]:
-        # NOTE: This function is O(seq_len); use sparingly.
-        token_ids: List[int] = []
-
-        if not self._is_allocated:
-            return token_ids
+    def _get_all_token_ids(self) -> np.ndarray:
+        # NOTE: This function is O(seq_len); use only for testing.
+        token_id_arrays: List[np.ndarray] = []
 
-        for block in self._blocks:
-            token_ids.extend(block.token_ids)
+        if self._is_allocated:
+            for block in self._blocks:
+                token_id_arrays.append(block.token_ids)
 
-        return token_ids
+        return np.concatenate(token_id_arrays)
 
     @property
     def _is_allocated(self) -> bool:
diff --git a/vllm/core/block/cpu_gpu_block_allocator.py b/vllm/core/block/cpu_gpu_block_allocator.py
index d4aaa1d24e08..009439708476 100644
--- a/vllm/core/block/cpu_gpu_block_allocator.py
+++ b/vllm/core/block/cpu_gpu_block_allocator.py
@@ -347,6 +347,10 @@ def token_ids(self) -> np.ndarray:
     def num_empty_slots(self) -> BlockId:
         return self._proxy.num_empty_slots
 
+    @property
+    def num_tokens(self) -> BlockId:
+        return self._proxy.num_tokens
+
     @property
     def is_full(self):
         return self._proxy.is_full
diff --git a/vllm/core/block/interfaces.py b/vllm/core/block/interfaces.py
index 05b642c2aaf9..3e25581ee432 100644
--- a/vllm/core/block/interfaces.py
+++ b/vllm/core/block/interfaces.py
@@ -35,6 +35,11 @@ def token_ids(self) -> np.ndarray:
     def num_empty_slots(self) -> int:
         pass
 
+    @property
+    @abstractmethod
+    def num_tokens(self) -> int:
+        pass
+
     @property
     @abstractmethod
     def is_full(self) -> bool:
diff --git a/vllm/core/block/naive_block.py b/vllm/core/block/naive_block.py
index 48012cb7f0d6..6ddf26882940 100644
--- a/vllm/core/block/naive_block.py
+++ b/vllm/core/block/naive_block.py
@@ -372,6 +372,10 @@ def is_full(self) -> bool:
     def num_empty_slots(self) -> int:
         return self._block_size - self._num_tokens
 
+    @property
+    def num_tokens(self) -> int:
+        return self._num_tokens
+
     @property
     def token_ids(self) -> np.ndarray:
         return self._token_ids
diff --git a/vllm/core/block/prefix_caching_block.py b/vllm/core/block/prefix_caching_block.py
index ea370c7c46d7..0cd88b647343 100644
--- a/vllm/core/block/prefix_caching_block.py
+++ b/vllm/core/block/prefix_caching_block.py
@@ -599,6 +599,10 @@ def is_full(self) -> bool:
     def num_empty_slots(self) -> int:
         return self._block.num_empty_slots
 
+    @property
+    def num_tokens(self) -> int:
+        return self._block._num_tokens
+
     @property
     def num_tokens_total(self) -> int:
         """return the total tokens so far.

From c5b29265911e72f288393dc1180cf0f8b1b0cc0e Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@126.com>
Date: Thu, 27 Jun 2024 20:49:13 -0700
Subject: [PATCH 14/14] block table

---
 vllm/core/block/block_table.py | 27 +++++++++------------------
 1 file changed, 9 insertions(+), 18 deletions(-)

diff --git a/vllm/core/block/block_table.py b/vllm/core/block/block_table.py
index 5825225f5a32..5c68c3df186d 100644
--- a/vllm/core/block/block_table.py
+++ b/vllm/core/block/block_table.py
@@ -52,12 +52,10 @@ def __init__(
         self._blocks: List[Block] = _blocks
 
         self._max_block_sliding_window = max_block_sliding_window
-        # Use helper method instead of directly calculating, as blocks
-        # may not be allocated.
+
         _num_full_slots = 0
-        if self._is_allocated:
-            for block in self._blocks:
-                _num_full_slots += block.num_tokens
+        for block in self._blocks:
+            _num_full_slots += block.num_tokens
         self._num_full_slots = _num_full_slots
 
     @staticmethod
@@ -92,7 +90,7 @@ def allocate(self,
             device (Device, optional): The device on which the blocks should be
                 allocated. Defaults to Device.GPU.
         """
-        assert not self._is_allocated
+        assert not self._blocks
         len_token_ids = len(token_ids)
         assert len_token_ids > 0
         self._blocks = self._allocate_blocks_for_token_ids(prev_block=None,
@@ -126,7 +124,7 @@ def append_token_ids(self,
                 Without chunked prefill, it should be the same as
                 _num_full_slots.
         """
-        assert self._is_allocated, "no blocks have been allocated"
+        assert self._blocks, "no blocks have been allocated"
         assert len(self._blocks) > 0
 
         # Drop blocks that are no longer needed due to sliding window
@@ -170,7 +168,7 @@ def ensure_num_empty_slots(self, num_empty_slots: int) -> None:
         # Currently the block table only supports
         # appending tokens to GPU blocks.
         device = Device.GPU
-        assert self._is_allocated
+        assert self._blocks
 
         if self._num_empty_slots >= num_empty_slots:
             return
@@ -197,8 +195,7 @@ def fork(self) -> "BlockTable":
             BlockTable: A new BlockTable instance with a copy of the blocks from
                 the current instance.
         """
-        assert self._is_allocated
-        assert len(self._blocks) > 0
+        assert self._blocks
         forked_blocks = self._allocator.fork(self._blocks[-1])
         return BlockTable(
             block_size=self._block_size,
@@ -215,7 +212,6 @@ def free(self) -> None:
         occupied by each block. After freeing all the blocks, the `_blocks` list
         is set to `None`.
         """
-        assert self._is_allocated
         for block in self._blocks:
             self._allocator.free(block)
         self._blocks = []
@@ -234,7 +230,6 @@ def physical_block_ids(self) -> List[Optional[int]]:
             List[int]: A list of physical block indices for the blocks in the
                 BlockTable.
         """
-        assert self._is_allocated
         return [block.block_id for block in self._blocks]
 
     def get_unseen_token_ids(self,
@@ -280,23 +275,19 @@ def _get_all_token_ids(self) -> np.ndarray:
         # NOTE: This function is O(seq_len); use only for testing.
         token_id_arrays: List[np.ndarray] = []
 
-        if self._is_allocated:
+        if self._blocks:
             for block in self._blocks:
                 token_id_arrays.append(block.token_ids)
 
         return np.concatenate(token_id_arrays)
 
-    @property
-    def _is_allocated(self) -> bool:
-        return len(self._blocks) > 0
-
     @property
     def blocks(self) -> Optional[List[Block]]:
         return self._blocks
 
     @property
     def _num_empty_slots(self) -> int:
-        assert self._is_allocated
+        assert self._blocks
         return len(self._blocks) * self._block_size - self._num_full_slots
 
     @property