minor: some potential bugs (#1044)

This commit is contained in:
Liangsheng Yin
2024-08-11 22:35:44 -07:00
committed by GitHub
parent 14b6493087
commit fb7421db0d
2 changed files with 7 additions and 6 deletions

View File

@@ -140,13 +140,13 @@ class InputMetadata:
if self.forward_mode == ForwardMode.DECODE:
self.extend_seq_lens = self.extend_start_loc = self.extend_no_prefix = None
else:
prefix_lens_cpu = [
extend_lens_cpu = [
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_seq_lens = torch.tensor(extend_lens_cpu, device="cuda")
self.extend_start_loc = torch.zeros_like(self.seq_lens)
self.extend_start_loc[1:] = torch.cumsum(self.extend_seq_lens[:-1], dim=0)
self.extend_no_prefix = all(x == 0 for x in prefix_lens_cpu)
self.extend_no_prefix = all(len(r.prefix_indices) == 0 for r in batch.reqs)
@classmethod
def from_schedule_batch(