chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

[minor] simplify the TokenToKVPoolAllocator (#7414)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-06-22 12:37:18 +08:00
committed by GitHub
parent b7a2df0a44
commit 05c9bc8956
14 changed files with 165 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/tp_worker.py
+3 -2
View File
@@ -35,7 +35,8 @@ from sglang.srt.managers.io_struct import (
UpdateWeightsFromTensorReqInput,
)
from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_args_dict
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import ServerArgs
@@ -57,7 +58,7 @@ class TpModelWorker:
nccl_port: int,
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[TokenToKVPoolAllocator] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
):
# Parse args
self.tp_size = server_args.tp_size