强烈推荐
https://zhuanlan.zhihu.com/p/680153425
sequnceGroup 存储了相同的prompt对应的不同的sequence, 所以用字典存储
同一个Sequence可能占据多个逻辑Block, 所以在Sequence 中用列表存储
同一个block 要维护tokens_id 列表, 需要添加操作。 还需要判断block 是否还有空位可以放置tokens.
# https://github.com/vllm-project/vllm/tree/v0.2.7/vllm/core/scheduler.py
class Scheduler:
def _schedule(self) -> SchedulerOutputs:
...
if not self.swapped:
...
while self.waiting:
...
seq_group = self.waiting.pop(0)
self._allocate(seq_group)
self.running.append(seq_group)
num_curr_seqs += num_new_seqs
scheduled.append(seq_group)
...
...
return scheduler_outputs
def _allocate(self, seq_group: SequenceGroup) -> None:
# 调用 block manager 的 allocate 方法分配 physical token block
self.block_manager.allocate(seq_group)
# 将 sequence 的状态设置为 RUNNING,即将要被处理
for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
seq.status = SequenceStatus.RUNNING
# https://github.com/vllm-project/vllm/tree/v0.2.7/vllm/core/block_manager.py
class BlockSpaceManager:
def allocate(self, seq_group: SequenceGroup) -> None:
# NOTE: Here we assume that all sequences in the group have the same
# prompt.
seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
# 为请求的 prompt token 分配 physical token block
block_table: BlockTable = []
for logical_idx in range(len(seq.logical_token_blocks)):
if (self.block_sliding_window is not None
and logical_idx >= self.block_sliding_window):
block = block_table[logical_idx % self.block_sliding_window]
else:
block = self.gpu_allocator.allocate()
# 设置 block 的引用数,copy on write 机制会用到
block.ref_count = seq_group.num_seqs()
block_table.append(block)
# Assign the block table for each sequence.
for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
self.block_tables[seq.seq_id] = block_table.copy()
结合上图代码, 对于seq_group 中所有的 seq 分配了相同的block_table( 逻辑块对应的物理块), 因此说明seq_group 中所有的seq 是相同的内容, 即 # NOTE: Here we assume that all sequences in the group have the same prompt.
上面函数只是分配好了逻辑空间与物理block空间,下面函数append_slot将新的token 加入到block 的一个slot 时,可能引发copy on write
scheduler 的_append_slot方法:
# https://github.com/vllm-project/vllm/tree/v0.2.7/vllm/core/scheduler.py
class Scheduler:
def _append_slot(
self,
seq_group: SequenceGroup,
blocks_to_copy: Dict[int, List[int]],
) -> None:
for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
ret = self.block_manager.append_slot(seq)
if ret is not None:
# copy on write 机制
src_block, dst_block = ret
if src_block in blocks_to_copy:
blocks_to_copy[src_block].append(dst_block)
else:
blocks_to_copy[src_block] = [dst_block]
# https://github.com/vllm-project/vllm/tree/v0.2.7/vllm/core/block_manager.py
class BlockSpaceManager:
def append_slot(self, seq: Sequence) -> Optional[Tuple[int, int]]:
"""Allocate a physical slot for a new token."""
logical_blocks = seq.logical_token_blocks
block_table = self.block_tables[seq.seq_id]
if len(block_table) < len(logical_blocks):
if (self.block_sliding_window
and len(block_table) >= self.block_sliding_window):
# re-use a block
block_table.append(block_table[len(block_table) %
self.block_sliding_window])
else:
# sequence 有新的 logical token block
# 所以这里也要分配一个新的 physical token block
block = self.gpu_allocator.allocate()
block_table.append(block)
return None
# We want to append the token to the last physical block.
last_block = block_table[-1]
assert last_block.device == Device.GPU
if last_block.ref_count == 1:
# Not shared with other sequences. Appendable.
return None
else:
# The last block is shared with other sequences.
# Copy on Write: Allocate a new block and copy the tokens.
new_block = self.gpu_allocator.allocate()
block_table[-1] = new_block
self.gpu_allocator.free(last_block)
return last_block.block_number, new_block.block_number