biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[WIP] TPU V1 Support Refactored (#13049)

Browse Source
This commit is contained in:
Alexander Matveev
2025-02-14 03:21:53 -05:00
committed by GitHub
parent b0ccfc565a
commit 45f90bcbba
8 changed files with 1738 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/v1/worker/block_table.py
View File

@@ -61,6 +61,14 @@ class BlockTable:
src, :num_blocks]
self.num_blocks_per_row[tgt] = num_blocks
def swap_row(self, src: int, tgt: int) -> None:
num_blocks_src = self.num_blocks_per_row[src]
num_blocks_tgt = self.num_blocks_per_row[tgt]
self.num_blocks_per_row[src] = num_blocks_tgt
self.num_blocks_per_row[tgt] = num_blocks_src
self.block_table_np[[src, tgt]] = self.block_table_np[[tgt, src]]
def commit(self, num_reqs: int) -> None:
self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
non_blocking=True)