zijie-tian/nano-vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

support CUDA_VISIBLE_DEVICES

Browse Source
This commit is contained in:
GeeeekExplorer
2025-06-12 23:14:01 +08:00
parent ec3c60d96f
commit 98a1551a7d
2 changed files with 13 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
nanovllm/engine/model_runner.py
View File

@@ -35,7 +35,7 @@ class ModelRunner:
total, used, _ = get_gpu_memory()
free = total * gpu_memory_utilization - used
block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * hf_config.num_key_value_heads * hf_config.head_dim * hf_config.torch_dtype.itemsize
config.num_kvcache_blocks = int(free * 1e6) // block_bytes
config.num_kvcache_blocks = int(free) // block_bytes
self.kv_cache = torch.zeros(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, hf_config.num_key_value_heads, hf_config.head_dim)
layer_id = 0
for module in self.model.modules():