add .isort.cfg (#378)

This commit is contained in:
Liangsheng Yin
2024-04-22 22:38:09 +08:00
committed by GitHub
parent cf9d8efdd3
commit 9acc6e3504
57 changed files with 125 additions and 139 deletions

View File

@@ -4,6 +4,7 @@ from dataclasses import dataclass
import torch
import torch.distributed as dist
from sglang.srt.managers.router.model_runner import ModelRunner
from sglang.srt.model_config import ModelConfig
@@ -66,9 +67,9 @@ class BenchBatch:
p_idx = prefix_req_idx[i // fork_num].item()
n_idx = self.req_pool_indices[i].item()
req_to_token[n_idx, :prefix_len] = req_to_token[p_idx, :prefix_len]
req_to_token[
n_idx, prefix_len : prefix_len + extend_len
] = self.out_cache_loc[i * extend_len : (i + 1) * extend_len]
req_to_token[n_idx, prefix_len : prefix_len + extend_len] = (
self.out_cache_loc[i * extend_len : (i + 1) * extend_len]
)
def update_decode(self, predict_ids, batch_size):
assert predict_ids.shape[0] == batch_size
@@ -81,9 +82,9 @@ class BenchBatch:
self.out_cache_cont_start,
self.out_cache_cont_end,
) = self.token_to_kv_pool.alloc_contiguous(batch_size)
self.req_to_token_pool.req_to_token[
self.req_pool_indices, self.seq_lens
] = self.out_cache_loc
self.req_to_token_pool.req_to_token[self.req_pool_indices, self.seq_lens] = (
self.out_cache_loc
)
self.seq_lens.add_(1)