nano-blocks

nano-vllm/nanovllm/engine/block_manager.py · bb823b3 · L26–120 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

26class BlockManager:
27
28    def __init__(self, num_blocks: int, block_size: int):
29        self.block_size = block_size
30        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
31        self.hash_to_block_id: dict[int, int] = dict()
32        self.free_block_ids: deque[int] = deque(range(num_blocks))
33        self.used_block_ids: set[int] = set()
34
35    @classmethod
36    def compute_hash(cls, token_ids: list[int], prefix: int = -1):
37        h = xxhash.xxh64()
38        if prefix != -1:
39            h.update(prefix.to_bytes(8, "little"))
40        h.update(np.array(token_ids).tobytes())
41        return h.intdigest()
42
43    def _allocate_block(self) -> int:
44        block_id = self.free_block_ids.popleft()
45        block = self.blocks[block_id]
46        assert block.ref_count == 0
47        if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
48            del self.hash_to_block_id[block.hash]
49        block.reset()
50        self.used_block_ids.add(block_id)
51        return block_id
52
53    def _deallocate_block(self, block_id: int):
54        assert self.blocks[block_id].ref_count == 0
55        self.used_block_ids.remove(block_id)
56        self.free_block_ids.append(block_id)
57
58    def can_allocate(self, seq: Sequence) -> int:
59        h = -1
60        num_cached_blocks = 0
61        num_new_blocks = seq.num_blocks
62        for i in range(seq.num_blocks - 1):
63            token_ids = seq.block(i)
64            h = self.compute_hash(token_ids, h)
65            block_id = self.hash_to_block_id.get(h, -1)
66            if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
67                break
68            num_cached_blocks += 1
69            if block_id in self.used_block_ids:
70                num_new_blocks -= 1
71        if len(self.free_block_ids) < num_new_blocks:
72            return -1
73        return num_cached_blocks
74
75    def allocate(self, seq: Sequence, num_cached_blocks: int):
76        assert not seq.block_table
77        h = -1
78        for i in range(num_cached_blocks):
79            token_ids = seq.block(i)
80            h = self.compute_hash(token_ids, h)
81            block_id = self.hash_to_block_id[h]
82            block = self.blocks[block_id]
83            if block_id in self.used_block_ids:
84                block.ref_count += 1
85            else:
86                block.ref_count = 1
87                self.free_block_ids.remove(block_id)
88                self.used_block_ids.add(block_id)
89            seq.block_table.append(block_id)
90        for i in range(num_cached_blocks, seq.num_blocks):
91            seq.block_table.append(self._allocate_block())
92        seq.num_cached_tokens = num_cached_blocks * self.block_size
93
94    def deallocate(self, seq: Sequence):
95        for block_id in reversed(seq.block_table):
96            block = self.blocks[block_id]
97            block.ref_count -= 1
98            if block.ref_count == 0:
99                self._deallocate_block(block_id)
100        seq.num_cached_tokens = 0
101        seq.block_table.clear()
102
103    def can_append(self, seq: Sequence) -> bool:
104        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)
105
106    def may_append(self, seq: Sequence):
107        if len(seq) % self.block_size == 1:
108            seq.block_table.append(self._allocate_block())
109
110    def hash_blocks(self, seq: Sequence):
111        start = seq.num_cached_tokens // self.block_size
112        end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size
113        if start == end: return
114        h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1
115        for i in range(start, end):
116            block = self.blocks[seq.block_table[i]]
117            token_ids = seq.block(i)
118            h = self.compute_hash(token_ids, h)
119            block.update(h, token_ids)
120            self.hash_to_block_id[h] = block.block_id