Enable overlap scheduler by default for the triton attention backend (#2105)

2024-11-20 02:58:35 -08:00
parent 56a347f7d3
commit 722530fa01
6 changed files with 21 additions and 24 deletions
@@ -53,7 +53,7 @@ class TritonAttnBackend(AttentionBackend):
            start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32)
            start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0)

-            total_num_tokens = torch.sum(forward_batch.seq_lens).item()
+            total_num_tokens = forward_batch.seq_lens_sum
            attn_logits = torch.empty(
                (self.num_head, total_num_tokens),
                dtype=self.reduce_dtype,