| """ |
| 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} 显存页成功回收") |
