58: vLLM 核心模块逐文件：kv_cache.py

class KVCache:
    def __init__(self, 
                 config: KVCacheConfig, 
                 device: torch.device,
                 dtype: torch.dtype):
        """初始化 KVCache"""
        self.config = config
        self.device = device
        self.dtype = dtype
        
        # 初始化块管理器
        self.block_manager = BlockManager(
            num_blocks=config.num_blocks,
            block_size=config.block_size,
            device=device
        )
        
        # 初始化页表
        self.page_table = PageTable()
        
        # 初始化缓存存储
        self.cache_store = CacheStore(
            num_layers=config.num_layers,
            num_heads=config.num_heads,
            head_dim=config.head_dim,
            block_size=config.block_size,
            num_blocks=config.num_blocks,
            device=device,
            dtype=dtype
        )
    
    def allocate(self, seq_len: int) -> KVCacheHandle:
        """为序列分配缓存"""
        # 计算所需块数
        num_blocks = (seq_len + self.config.block_size - 1) // self.config.block_size
        
        # 分配块
        block_ids = self.block_manager.allocate(num_blocks)
        
        # 更新页表
        handle = KVCacheHandle(
            seq_len=seq_len,
            block_ids=block_ids
        )
        
        return handle
    
    def free(self, handle: KVCacheHandle):
        """释放缓存"""
        self.block_manager.free(handle.block_ids)
    
    def get(self, handle: KVCacheHandle, layer_idx: int):
        """获取指定层的 KV 缓存"""
        return self.cache_store.get(handle.block_ids, layer_idx)
    
    def set(self, handle: KVCacheHandle, layer_idx: int, key: torch.Tensor, value: torch.Tensor):
        """设置指定层的 KV 缓存"""
        self.cache_store.set(handle.block_ids, layer_idx, key, value)

class BlockManager:
    def __init__(self, num_blocks: int, block_size: int, device: torch.device):
        """初始化块管理器"""
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.device = device
        
        # 初始化块状态：0 表示空闲，1 表示已分配
        self.block_states = torch.zeros(num_blocks, dtype=torch.bool, device=device)
        
        # 空闲块队列
        self.free_blocks = torch.arange(num_blocks, device=device)
        self.num_free_blocks = num_blocks
    
    def allocate(self, num_blocks: int) -> torch.Tensor:
        """分配指定数量的块"""
        if num_blocks > self.num_free_blocks:
            raise MemoryError(f"Not enough free blocks: requested {num_blocks}, available {self.num_free_blocks}")
        
        # 从空闲队列中取出块
        block_ids = self.free_blocks[:num_blocks]
        self.free_blocks = self.free_blocks[num_blocks:]
        self.num_free_blocks -= num_blocks
        
        # 更新块状态
        self.block_states[block_ids] = True
        
        return block_ids
    
    def free(self, block_ids: torch.Tensor):
        """释放块"""
        # 更新块状态
        self.block_states[block_ids] = False
        
        # 将块添加到空闲队列
        self.free_blocks = torch.cat([self.free_blocks, block_ids])
        self.num_free_blocks += len(block_ids)
    
    def get_num_free_blocks(self) -> int:
        """获取空闲块数量"""
        return self.num_free_blocks

class CacheStore:
    def __init__(self, num_layers: int, num_heads: int, head_dim: int, block_size: int, num_blocks: int, device: torch.device, dtype: torch.dtype):
        """初始化缓存存储"""
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.block_size = block_size
        self.num_blocks = num_blocks
        self.device = device
        self.dtype = dtype
        
        # 初始化 K 缓存：[num_layers, num_heads, num_blocks, block_size, head_dim]
        self.k_cache = torch.empty(
            num_layers, num_heads, num_blocks, block_size, head_dim,
            dtype=dtype, device=device
        )
        
        # 初始化 V 缓存：[num_layers, num_heads, num_blocks, block_size, head_dim]
        self.v_cache = torch.empty(
            num_layers, num_heads, num_blocks, block_size, head_dim,
            dtype=dtype, device=device
        )
    
    def get(self, block_ids: torch.Tensor, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """获取指定块和层的 KV 缓存"""
        # 选择指定层和块的 K 缓存
        k = self.k_cache[layer_idx, :, block_ids]
        
        # 选择指定层和块的 V 缓存
        v = self.v_cache[layer_idx, :, block_ids]
        
        return k, v
    
    def set(self, block_ids: torch.Tensor, layer_idx: int, key: torch.Tensor, value: torch.Tensor):
        """设置指定块和层的 KV 缓存"""
        # 设置 K 缓存
        self.k_cache[layer_idx, :, block_ids] = key
        
        # 设置 V 缓存
        self.v_cache[layer_idx, :, block_ids] = value

class PageTable:
    def __init__(self):
        """初始化页表"""
        self.page_table = {}
    
    def update(self, handle: KVCacheHandle, seq_len: int):
        """更新页表"""
        # 计算所需的页数
        num_pages = (seq_len + handle.block_size - 1) // handle.block_size
        
        # 更新页表项
        self.page_table[handle.id] = {
            'seq_len': seq_len,
            'num_pages': num_pages,
            'pages': handle.block_ids[:num_pages]
        }
    
    def lookup(self, handle_id: int) -> Dict:
        """查找页表项"""
        return self.page_table.get(handle_id, None)

class HybridCacheController:
    def __init__(self, gpu_cache_size: int, dram_cache_size: int, ssd_cache_size: int):
        """初始化混合缓存控制器"""
        self.gpu_cache = GPUCachePool(gpu_cache_size)
        self.dram_cache = DRAMCachePool(dram_cache_size)
        self.ssd_cache = SSDCachePool(ssd_cache_size)
        
        # LRU 策略用于缓存替换
        self.lru = LRUCache()
    
    def get_block(self, block_id: int) -> torch.Tensor:
        """获取块，自动处理缓存层级"""
        if self.gpu_cache.contains(block_id):
            # 更新 LRU
            self.lru.access(block_id)
            return self.gpu_cache.get(block_id)
        
        if self.dram_cache.contains(block_id):
            # 从 DRAM 加载到 GPU
            block = self.dram_cache.get(block_id)
            self._promote_to_gpu(block_id, block)
            return block
        
        if self.ssd_cache.contains(block_id):
            # 从 SSD 加载到 DRAM，再加载到 GPU
            block = self.ssd_cache.get(block_id)
            self.dram_cache.set(block_id, block)
            self._promote_to_gpu(block_id, block)
            return block
        
        raise KeyError(f"Block {block_id} not found in any cache")
    
    def _promote_to_gpu(self, block_id: int, block: torch.Tensor):
        """将块提升到 GPU 缓存"""
        # 如果 GPU 缓存已满，驱逐最久未使用的块
        if self.gpu_cache.is_full():
            evict_block_id = self.lru.evict()
            evict_block = self.gpu_cache.get(evict_block_id)
            self.gpu_cache.remove(evict_block_id)
            self.dram_cache.set(evict_block_id, evict_block)
        
        # 将块添加到 GPU 缓存
        self.gpu_cache.set(block_id, block)
        self.lru.add(block_id)

class DynamicMemoryManager:
    def __init__(self, config: DynamicMemoryConfig):
        """初始化动态内存管理器"""
        self.config = config
        self.memory_monitor = MemoryMonitor()
        self.cache_resizer = CacheResizer()
    
    def adjust_cache_size(self):
        """根据内存使用情况调整缓存大小"""
        # 获取当前内存使用情况
        memory_usage = self.memory_monitor.get_memory_usage()
        
        # 计算目标缓存大小
        if memory_usage > self.config.high_threshold:
            # 内存使用过高，缩小缓存
            target_size = int(self.cache_resizer.current_size * self.config.shrink_ratio)
            self.cache_resizer.resize(target_size)
        elif memory_usage < self.config.low_threshold:
            # 内存使用过低，扩大缓存
            target_size = int(self.cache_resizer.current_size * self.config.grow_ratio)
            self.cache_resizer.resize(target_size)
    
    def on_request_arrival(self):
        """当新请求到达时调用"""
        self.adjust_cache_size()
    
    def on_request_completion(self):
        """当请求完成时调用"""
        self.adjust_cache_size()

class ContextParallelKVCache:
    def __init__(self, num_gpus: int, config: KVCacheConfig):
        """初始化上下文并行 KVCache"""
        self.num_gpus = num_gpus
        self.config = config
        
        # 在每个 GPU 上初始化本地 KVCache
        self.local_caches = [
            KVCache(config, torch.device(f'cuda:{i}'))
            for i in range(num_gpus)
        ]
        
        # 初始化通信控制器
        self.comm_controller = CommunicationController(num_gpus)
    
    def allocate(self, seq_len: int) -> List[KVCacheHandle]:
        """为序列分配分布式缓存"""
        # 计算每个 GPU 负责的序列长度
        per_gpu_seq_len = (seq_len + self.num_gpus - 1) // self.num_gpus
        
        # 在每个 GPU 上分配缓存
        handles = []
        for i in range(self.num_gpus):
            start = i * per_gpu_seq_len
            end = min((i + 1) * per_gpu_seq_len, seq_len)
            local_seq_len = end - start
            
            if local_seq_len > 0:
                handle = self.local_caches[i].allocate(local_seq_len)
                handles.append(handle)
            else:
                handles.append(None)
        
        return handles
    
    def get(self, handles: List[KVCacheHandle], layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """获取分布式 KV 缓存"""
        # 从每个 GPU 获取本地缓存
        local_kvs = []
        for i, handle in enumerate(handles):
            if handle is not None:
                k, v = self.local_caches[i].get(handle, layer_idx)
                local_kvs.append((k, v))
        
        # 合并结果
        if len(local_kvs) == 1:
            return local_kvs[0]
        
        # 跨 GPU 合并
        keys = [kv[0] for kv in local_kvs]
        values = [kv[1] for kv in local_kvs]
        
        # 使用 all-gather 合并
        gathered_keys = self.comm_controller.all_gather(keys)
        gathered_values = self.comm_controller.all_gather(values)
        
        return gathered_keys, gathered_values

from vllm import KVCacheConfig

# 创建 KVCache 配置
kv_cache_config = KVCacheConfig(
    num_blocks=1024,
    block_size=16,
    num_layers=32,
    num_heads=32,
    head_dim=128,
    dtype="float16"
)

from vllm import KVCacheConfig

# 创建混合缓存配置
kv_cache_config = KVCacheConfig(
    num_blocks=1024,
    block_size=16,
    num_layers=32,
    num_heads=32,
    head_dim=128,
    dtype="float16",
    enable_mixed_cache=True,
    dram_cache_size=1024 * 1024 * 1024 * 16,  # 16GB
    ssd_cache_path="/path/to/ssd/cache"
)

from vllm import KVCacheConfig

# 创建上下文并行配置
kv_cache_config = KVCacheConfig(
    num_blocks=1024,
    block_size=16,
    num_layers=32,
    num_heads=32,
    head_dim=128,
    dtype="float16",
    context_parallel_size=4  # 使用4个GPU进行上下文并行
)

from vllm.kv_cache import BlockManager

class CustomBlockManager(BlockManager):
    def __init__(self, num_blocks: int, block_size: int, device: torch.device):
        super().__init__(num_blocks, block_size, device)
        # 初始化自定义策略
        self.custom_policy = "priority_based"
    
    def allocate(self, num_blocks: int, priority: int = 0) -> torch.Tensor:
        """带优先级的块分配"""
        # 自定义分配逻辑
        if self.custom_policy == "priority_based":
            # 优先为高优先级请求分配连续块
            if priority > 0:
                # 尝试分配连续块
                continuous_blocks = self._find_continuous_blocks(num_blocks)
                if continuous_blocks is not None:
                    return continuous_blocks
        
        # 回退到默认分配逻辑
        return super().allocate(num_blocks)
    
    def _find_continuous_blocks(self, num_blocks: int) -> Optional[torch.Tensor]:
        """查找连续的空闲块"""
        # 实现连续块查找逻辑
        # 示例实现省略...
        pass