68 lines
2.6 KiB
Python
68 lines
2.6 KiB
Python
|
|
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 |