vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-05 12:09:48 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

fix-cache-lens

Browse source
This commit is contained in:
Atream 2025-04-30 03:37:43 +00:00
parent 8ba7e5d4b8
commit 7adb7281f4
1 changed files with 6 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
ktransformers/server/args.py
View file

@ -1,6 +1,7 @@
import argparse
from ktransformers.server.backend.args import ConfigArgs, default_args
from ktransformers.util.utils import get_free_ports
from transformers import AutoConfig
class ArgumentParser:
def __init__(self, cfg):
@ -138,7 +139,11 @@ class ArgumentParser:
self.cfg.server_port = args.port
self.cfg.user_force_think = args.force_think
args.gpu_memory_size = 4*1024*1024*1024 # TODO: set this to the actual GPU memory size
model_config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=True)
if args.architectures == "Qwen3MoeForCausalLM" or args.architectures == "Qwen2MoeForCausalLM" :
args.gpu_memory_size = args.cache_lens*2*2*model_config.num_hidden_layers*model_config.num_key_value_heads*model_config.head_dim
else:
args.gpu_memory_size = args.cache_lens*2*576*61
self.cfg.gpu_memory_size = args.gpu_memory_size
free_ports = get_free_ports(3, [args.port])
args.sched_port = free_ports[0]