Maintain seq_lens_sum to make more FlashInfer operations non-blocking (#1741)

2024-10-21 01:43:16 -07:00
parent cf470fea32
commit 09603c6dc9
8 changed files with 98 additions and 43 deletions
@@ -144,7 +144,11 @@ class DoubleSparseAttnBackend(AttentionBackend):
        )

    def init_forward_metadata_replay_cuda_graph(
-        self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
    ):
        self.cuda_graph_start_loc.zero_()
        self.cuda_graph_start_loc[1:bs] = torch.cumsum(seq_lens[: bs - 1], dim=0)