init

vllm_vacc/vllm/v1/worker/block_table.py

@@ -0,0 +1,68 @@
+
+import numpy as np
+import torch
+import os
+
+from vllm.distributed import get_dcp_group
+from vllm.logger import init_logger
+from vllm.utils import cdiv
+
+from vllm_vacc.vllm.model_executor.models.vars import BLOCK_GROUP_SIZE as env_blk_grp_size
+
+logger = init_logger(__name__)
+
+
+class BlockTable:
+
+    def __init__(
+        self,
+        block_size: int,
+        max_num_reqs: int,
+        max_num_blocks_per_req: int,
+        max_num_batched_tokens: int,
+        pin_memory: bool,
+        device: torch.device,
+    ):
+        self.block_size = block_size
+        self.max_num_reqs = max_num_reqs
+        # self.max_num_blocks_per_req = max_num_blocks_per_req
+        max_num_blocks_per_req = (max_num_blocks_per_req + env_blk_grp_size//16 - 1) // (env_blk_grp_size//16) * (env_blk_grp_size//16)
+        self.max_num_blocks_per_req = max_num_blocks_per_req
+        self.max_num_batched_tokens = max_num_batched_tokens
+        self.pin_memory = pin_memory
+        self.device = device
+
+        # self.block_table = torch.zeros(
+        #     (max_num_reqs, max_num_blocks_per_req),
+        #     device=self.device,
+        #     dtype=torch.int32,
+        # )
+        self.block_table = self._make_buffer(max_num_reqs,
+                                        max_num_blocks_per_req,
+                                        dtype=torch.int32)
+        # self.block_table_cpu = torch.zeros(
+        #     (max_num_reqs, max_num_blocks_per_req),
+        #     device="cpu",
+        #     dtype=torch.int32,
+        #     pin_memory=pin_memory,
+        # )
+        # self.block_table_np = self.block_table_cpu.numpy()
+        self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)
+
+        # self.slot_mapping_cpu = torch.zeros(self.max_num_batched_tokens,
+        #                                     dtype=torch.int32,
+        #                                     device="cpu",
+        #                                     pin_memory=self.pin_memory)
+        # self.slot_mapping_np = self.slot_mapping_cpu.numpy()
+        # self.slot_mapping = torch.zeros(self.max_num_batched_tokens,
+        #                                 dtype=torch.int32,
+        #                                 device=self.device)
+        self.slot_mapping = self._make_buffer(self.max_num_batched_tokens,
+                                              dtype=torch.int32)
+        try:
+            self.dcp_world_size = get_dcp_group().world_size
+            self.dcp_rank = get_dcp_group().rank_in_group
+        except AssertionError:
+            # DCP might not be initialized in testing
+            self.dcp_world_size = 1
+            self.dcp_rank = 0