nano-cache-allocation

nano-vllm/nanovllm/engine/model_runner.py · bb823b3 · L103–122 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

103    def allocate_kv_cache(self):
104        config = self.config
105        hf_config = config.hf_config
106        free, total = torch.cuda.mem_get_info()
107        used = total - free
108        peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
109        current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
110        num_kv_heads = hf_config.num_key_value_heads // self.world_size
111        head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)
112        block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.dtype.itemsize
113        config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes
114        assert config.num_kvcache_blocks > 0
115        self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim)
116        layer_id = 0
117        for module in self.model.modules():
118            if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
119                module.k_cache = self.kv_cache[0, layer_id]
120                module.v_cache = self.kv_cache[1, layer_id]
121                layer_id += 1
122