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

Fix flashinfer (#430)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-05-12 08:18:53 -07:00
committed by GitHub
parent 3fc97f6709
commit 2d580e7a89
3 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/managers/router/model_rpc.py
View File

@@ -113,7 +113,8 @@ class ModelRpcServer:
f"max_prefill_num_token={self.max_prefill_num_token}, "
f"context_len={self.model_config.context_len}, "
)
logger.info(f"server_args: {server_args.print_mode_args()}")
if self.tp_rank == 0:
logger.info(f"server_args: {server_args.print_mode_args()}")
# Init cache
self.tree_cache = RadixCache(disable=server_args.disable_radix_cache)

6
python/sglang/srt/managers/router/model_runner.py
View File

@@ -110,12 +110,12 @@ class InputMetadata:
self.kv_last_page_len = torch.ones(
(self.batch_size,), dtype=torch.int32, device="cuda"
)
req_pool_indices_cpu = self.req_pool_indices.cpu().numpy()
seq_lens_cpu = self.seq_lens.cpu().numpy()
req_pool_indices_cpu = self.req_pool_indices.cpu().tolist()
seq_lens_cpu = self.seq_lens.tolist()
self.kv_indices = torch.cat(
[
self.req_to_token_pool.req_to_token[
req_pool_indices_cpu[i]: seq_lens_cpu[i]
req_pool_indices_cpu[i], : seq_lens_cpu[i]
]
for i in range(self.batch_size)
],