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

Simplify flashinfer utilities (#1704)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-10-17 22:54:14 -07:00
committed by GitHub
parent 9e0dac1ad7
commit 6d0fa73ece
8 changed files with 391 additions and 337 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/double_sparsity_backend.py
+2 -2
View File
@@ -134,7 +134,7 @@ class DoubleSparseAttnBackend(AttentionBackend):
)
def init_forward_metadata_capture_cuda_graph(
self, bs: int, req_pool_indices, seq_lens
self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
):
self.forward_metadata = (
self.cuda_graph_start_loc,
@@ -144,7 +144,7 @@ class DoubleSparseAttnBackend(AttentionBackend):
)
def init_forward_metadata_replay_cuda_graph(
self, bs: int, req_pool_indices, seq_lens
self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
):
self.cuda_graph_start_loc.zero_()
self.cuda_graph_start_loc[1:bs] = torch.cumsum(seq_lens[: bs - 1], dim=0)