File size: 2,366 Bytes
63ca2a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """
Paged Video Latent 内存池与连续批处理引擎 (vLLM Style)
Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119)
"""
import torch
import logging
from typing import Dict, List
logger = logging.getLogger("VerseFlow.PagedLatentPool")
class LatentBlock:
def __init__(self, block_id: int, block_size: int, num_channels: int, height: int, width: int, dtype=torch.float16):
self.block_id = block_id
self.ref_count = 0
self.data = torch.empty((block_size, num_channels, height, width), dtype=dtype, device="cuda")
class PagedVideoLatentMemoryPool:
"""视频扩散模型 Paged Latent 显存池，消除显存碎片"""
def __init__(self, total_blocks: int = 1024, block_size: int = 16, channels: int = 4, height: int = 64, width: int = 64):
self.block_size = block_size
self.free_blocks: List[LatentBlock] = [
LatentBlock(i, block_size, channels, height, width) for i in range(total_blocks)
]
self.allocated_blocks: Dict[int, LatentBlock] = {}
self.block_table: Dict[str, List[int]] = {}
def allocate(self, sequence_id: str, num_frames: int) -> List[int]:
num_blocks_needed = (num_frames + self.block_size - 1) // self.block_size
if len(self.free_blocks) < num_blocks_needed:
raise RuntimeError(f"Paged 显存池开销溢出: 需要 {num_blocks_needed} 页，剩余 {len(self.free_blocks)} 页")
allocated_ids = []
for _ in range(num_blocks_needed):
block = self.free_blocks.pop()
block.ref_count = 1
self.allocated_blocks[block.block_id] = block
allocated_ids.append(block.block_id)
self.block_table[sequence_id] = allocated_ids
logger.info(f"🧠 [PagedLatent] 序列 {sequence_id} 成功分配 {num_blocks_needed} 个 Block 页")
return allocated_ids
def free_sequence(self, sequence_id: str):
if sequence_id not in self.block_table:
return
for block_id in self.block_table[sequence_id]:
block = self.allocated_blocks.pop(block_id, None)
if block:
block.ref_count = 0
self.free_blocks.append(block)
del self.block_table[sequence_id]
logger.info(f"♻️ [PagedLatent] 序列 {sequence_id} 显存页成功回收") |