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

Fix input_ids && rename to fill_ids (#1021)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-08-10 16:24:12 -07:00
committed by GitHub
parent 54fb1c80c0
commit 43fbb6d919
7 changed files with 28 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -109,7 +109,7 @@ class InputMetadata:
self.positions = torch.tensor(
np.concatenate(
[
np.arange(len(req.prefix_indices), len(req.input_ids))
np.arange(len(req.prefix_indices), len(req.fill_ids))
for req in batch.reqs
],
axis=0,
@@ -124,7 +124,7 @@ class InputMetadata:
[
np.arange(
len(req.prefix_indices) + position_ids_offsets_cpu[i],
len(req.input_ids) + position_ids_offsets_cpu[i],
len(req.fill_ids) + position_ids_offsets_cpu[i],
)
for i, req in enumerate(batch.reqs)
],
@@ -141,7 +141,7 @@ class InputMetadata:
self.extend_seq_lens = self.extend_start_loc = self.extend_no_prefix = None
else:
prefix_lens_cpu = [
len(r.input_ids) - len(r.prefix_indices) for r in batch.reqs
len(r.fill_ids) - len(r.prefix_indices) for r in batch.reqs
]
self.extend_seq_lens = torch.tensor(prefix_lens_cpu, device="cuda")
self.extend_start_loc = torch.zeros_like(self.seq_lens)
@@ -149,7 +149,7 @@ class InputMetadata:
self.extend_no_prefix = all(x == 0 for x in prefix_lens_cpu)
def init_total_num_tokens(self, batch: ScheduleBatch):
self.total_num_tokens = sum(len(req.input_ids) for req in batch.reqs)
self.total_num_tokens = sum(len(req.fill_ids) for req in batch.reqs)
@classmethod
def from_schedule_batch(
@@ -203,7 +203,7 @@ class InputMetadata:
def init_triton_args(self, batch: ScheduleBatch, prefix_lens):
"""Init auxiliary variables for triton attention backend."""
self.triton_max_seq_len = max(len(r.input_ids) for r in batch.reqs)
self.triton_max_seq_len = max(len(r.fill_ids) for r in batch.reqs)
self.triton_prefix_lens = prefix_lens
self.triton_start_loc = torch.zeros_like(self.seq_lens, dtype=torch.int32)
self.triton_start_loc[1:] = torch.cumsum(self.seq_lens[:-1], dim=0)