38. GPU 显存碎片问题如何解决：vLLM的碎片管理技术深度解析

class PagedKVCache:
    def __init__(self, block_size: int, total_blocks: int):
        self.block_size = block_size
        self.total_blocks = total_blocks
        
        # 预分配显存
        self.device_memory = torch.empty(
            (total_blocks, block_size), 
            dtype=torch.float16, 
            device="cuda"
        )
        
        # 块管理
        self.free_blocks = set(range(total_blocks))
        self.used_blocks = set()
        
        # 页表：虚拟页到物理块的映射
        self.page_tables = defaultdict(dict)  # request_id -> {page_id: block_id}
    
    def allocate(self, request_id: int, num_pages: int) -> Dict[int, int]:
        """为请求分配虚拟页"""
        if len(self.free_blocks) < num_pages:
            # 没有足够的空闲块，触发驱逐
            self.evict(num_pages - len(self.free_blocks))
        
        # 分配块
        allocated_blocks = random.sample(self.free_blocks, num_pages)
        for block_id in allocated_blocks:
            self.free_blocks.remove(block_id)
            self.used_blocks.add(block_id)
        
        # 建立虚拟页到物理块的映射
        page_map = {}
        for i, block_id in enumerate(allocated_blocks):
            page_map[i] = block_id
        
        self.page_tables[request_id] = page_map
        return page_map
    
    def free(self, request_id: int):
        """释放请求的虚拟页"""
        if request_id not in self.page_tables:
            return
        
        # 释放所有块
        page_map = self.page_tables.pop(request_id)
        for block_id in page_map.values():
            self.used_blocks.remove(block_id)
            self.free_blocks.add(block_id)
    
    def evict(self, num_blocks: int):
        """驱逐最久未使用的块"""
        # 简化实现，实际会使用LRU等策略
        evict_blocks = random.sample(self.used_blocks, num_blocks)
        for block_id in evict_blocks:
            self.used_blocks.remove(block_id)
            self.free_blocks.add(block_id)
            
            # 查找并更新使用该块的页表
            for request_id, page_map in self.page_tables.items():
                for page_id, bid in page_map.items():
                    if bid == block_id:
                        del page_map[page_id]
                        break

class FragmentationDetector:
    def __init__(self, memory_manager):
        self.memory_manager = memory_manager
    
    def calculate_fragmentation_rate(self) -> float:
        """计算显存碎片率"""
        # 获取显存使用情况
        mem_stats = self.memory_manager.get_memory_stats()
        
        # 计算总显存和已用显存
        total_memory = mem_stats['total_memory']
        used_memory = mem_stats['used_memory']
        free_memory = total_memory - used_memory
        
        if free_memory == 0:
            return 0.0
        
        # 获取空闲块列表
        free_blocks = mem_stats['free_blocks']
        
        # 计算最大连续空闲块大小
        max_contiguous_free = 0
        current_contiguous = 0
        
        # 按地址排序空闲块
        sorted_free_blocks = sorted(free_blocks, key=lambda x: x['address'])
        
        for i, block in enumerate(sorted_free_blocks):
            if i == 0:
                current_contiguous = block['size']
            else:
                prev_block = sorted_free_blocks[i-1]
                if block['address'] == prev_block['address'] + prev_block['size']:
                    # 连续块，合并大小
                    current_contiguous += block['size']
                else:
                    # 非连续块，更新最大连续大小
                    max_contiguous_free = max(max_contiguous_free, current_contiguous)
                    current_contiguous = block['size']
        
        # 最后一次更新
        max_contiguous_free = max(max_contiguous_free, current_contiguous)
        
        # 碎片率 = 1 - (最大连续空闲块大小 / 总空闲内存)
        fragmentation_rate = 1.0 - (max_contiguous_free / free_memory)
        
        return fragmentation_rate
    
    def detect_fragmentation(self) -> Dict:
        """检测显存碎片情况"""
        fragmentation_rate = self.calculate_fragmentation_rate()
        
        # 碎片等级划分
        if fragmentation_rate < 0.2:
            level = "low"
        elif fragmentation_rate < 0.5:
            level = "medium"
        else:
            level = "high"
        
        return {
            "fragmentation_rate": fragmentation_rate,
            "level": level,
            "timestamp": time.time()
        }

class CuMemMonitor:
    def __init__(self):
        self.is_monitoring = False
        self.memory_events = []
    
    def start_monitoring(self):
        """启动cuMem监控"""
        if self.is_monitoring:
            return
        
        # 初始化cuMem监控
        ctypes.CDLL('libcudart.so').cuMemSetAccessCallback(self._memory_callback)
        self.is_monitoring = True
        
    def stop_monitoring(self):
        """停止cuMem监控"""
        if not self.is_monitoring:
            return
        
        # 停止cuMem监控
        ctypes.CDLL('libcudart.so').cuMemSetAccessCallback(None)
        self.is_monitoring = False
    
    def _memory_callback(self, event_type, ptr, size, user_data):
        """cuMem事件回调函数"""
        event = {
            "type": event_type,
            "address": ptr,
            "size": size,
            "timestamp": time.time()
        }
        self.memory_events.append(event)
    
    def get_memory_stats(self) -> Dict:
        """获取内存统计信息"""
        # 简化实现，实际会调用cuMemGetInfo等API
        import torch
        
        # 获取当前GPU设备
        device = torch.cuda.current_device()
        
        # 获取总显存和剩余显存
        total_memory = torch.cuda.get_device_properties(device).total_memory
        reserved_memory = torch.cuda.memory_reserved(device)
        allocated_memory = torch.cuda.memory_allocated(device)
        free_memory = total_memory - allocated_memory
        
        return {
            "total_memory": total_memory,
            "reserved_memory": reserved_memory,
            "allocated_memory": allocated_memory,
            "free_memory": free_memory,
            "free_blocks": self._get_free_blocks(),
            "timestamp": time.time()
        }
    
    def _get_free_blocks(self) -> List[Dict]:
        """获取空闲块信息"""
        # 简化实现，实际会通过cuMemGetAddressRange等API获取
        return [
            {"address": 0x10000000, "size": 16 * 1024 * 1024},
            {"address": 0x20000000, "size": 8 * 1024 * 1024},
            {"address": 0x30000000, "size": 32 * 1024 * 1024}
        ]

class FragmentationDefragmenter:
    def __init__(self, memory_manager):
        self.memory_manager = memory_manager
        self.stats = {
            "defragmentation_count": 0,
            "total_time": 0.0,
            "total_blocks_moved": 0
        }
    
    def defragment(self) -> Dict:
        """执行碎片整理"""
        start_time = time.time()
        
        # 获取当前显存使用情况
        mem_stats = self.memory_manager.get_memory_stats()
        free_blocks = mem_stats['free_blocks']
        
        # 按地址排序空闲块
        sorted_free_blocks = sorted(free_blocks, key=lambda x: x['address'])
        
        # 计算需要整理的空闲块
        # 简化实现，实际会更复杂
        blocks_to_defrag = []
        for block in sorted_free_blocks:
            if block['size'] < 16 * 1024 * 1024:  # 小于16MB的块需要整理
                blocks_to_defrag.append(block)
        
        if not blocks_to_defrag:
            return {
                "status": "no_defragmentation_needed",
                "message": "No small free blocks found"
            }
        
        # 执行碎片整理
        # 1. 收集所有使用中块的数据
        # 2. 释放所有使用中块
        # 3. 重新分配连续块
        # 4. 将数据写回新分配的块
        
        # 简化实现，实际会使用更高效的算法
        blocks_moved = len(blocks_to_defrag)
        
        end_time = time.time()
        duration = end_time - start_time
        
        # 更新统计信息
        self.stats["defragmentation_count"] += 1
        self.stats["total_time"] += duration
        self.stats["total_blocks_moved"] += blocks_moved
        
        return {
            "status": "success",
            "blocks_moved": blocks_moved,
            "duration": duration,
            "fragmentation_rate_before": self.memory_manager.fragmentation_detector.calculate_fragmentation_rate(),
            "fragmentation_rate_after": self.memory_manager.fragmentation_detector.calculate_fragmentation_rate(),
            "stats": self.stats
        }
    
    def get_defragmentation_stats(self) -> Dict:
        """获取碎片整理统计信息"""
        return self.stats

Fragmentation Rate = 1 - (Max Contiguous Free Memory / Total Free Memory)