EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix prefill size (#711)

Browse Source
This commit is contained in:
Ying Sheng
2024-07-24 03:41:15 -07:00
committed by GitHub
parent 00e4baa728
commit 4367f4bb8d
2 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/managers/controller/tp_worker.py
View File

@@ -103,6 +103,10 @@ class ModelTpServer:
if server_args.max_running_requests is None if server_args.max_running_requests is None
else server_args.max_running_requests else server_args.max_running_requests
) )
self.max_running_requests = min(
self.max_running_requests,
self.model_runner.req_to_token_pool.size - 1
)
self.int_token_logit_bias = torch.tensor( self.int_token_logit_bias = torch.tensor(
get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size) get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)
) )

1
python/sglang/srt/memory_pool.py
View File

@@ -11,6 +11,7 @@ class ReqToTokenPool:
"""A memory pool that maps a request to its token locations.""" """A memory pool that maps a request to its token locations."""
def __init__(self, size: int, max_context_len: int): def __init__(self, size: int, max_context_len: int):
self.size = size
self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda") self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
self.req_to_token = torch.empty( self.req_to_token = torch.empty(
(size, max_context_len), dtype=torch.int32, device="cuda" (size, max_context_len), dtype=torch.int32, device="cuda"