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

Fix memory leak in swapping

Browse Source
This commit is contained in:
Woosuk Kwon 2023-02-14 08:45:59 +00:00
parent 7e5f604e68
commit be1e2163c9
1 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
cacheflow/master/block_manager.py
View File

@ -165,10 +165,11 @@ class BlockSpaceManager:
for cpu_block in block_table: for cpu_block in block_table:
if cpu_block in mapping: if cpu_block in mapping:
new_block_table.append(mapping[cpu_block]) gpu_block = mapping[cpu_block]
continue else:
gpu_block = self.gpu_allocator.allocate() gpu_block = self.gpu_allocator.allocate()
mapping[cpu_block] = gpu_block mapping[cpu_block] = gpu_block
new_block_table.append(gpu_block)
# Free the CPU block swapped in to GPU. # Free the CPU block swapped in to GPU.
self.cpu_allocator.free(cpu_block) self.cpu_allocator.free(cpu_block)
self.block_tables[seq.seq_id] = new_block_table self.block_tables[seq.seq_id] = new_block_table
@ -194,12 +195,11 @@ class BlockSpaceManager:
for gpu_block in block_table: for gpu_block in block_table:
if gpu_block in mapping: if gpu_block in mapping:
new_block_table.append(mapping[gpu_block]) cpu_block = mapping[gpu_block]
continue else:
cpu_block = self.cpu_allocator.allocate() cpu_block = self.cpu_allocator.allocate()
mapping[gpu_block] = cpu_block
new_block_table.append(cpu_block) new_block_table.append(cpu_block)
mapping[gpu_block] = cpu_block
# Free the GPU block swapped out to CPU. # Free the GPU block swapped out to CPU.
self.gpu_allocator.free(gpu_block) self.gpu_allocator.free(gpu_block)
self.block_tables[seq.seq_id] = new_block_table self.block_tables[seq.seq_id] = new_block_table