Commit 58de8cd6 authored by zhuwenwen's avatar zhuwenwen
Browse files

remove attn_masks

parent 99f4d1c2
......@@ -135,8 +135,6 @@ class PagedAttention:
blocksparse_vert_stride: int = 0,
blocksparse_block_size: int = 64,
blocksparse_head_sliding_step: int = 0,
attn_masks: Optional[torch.Tensor] = None,
attn_masks_stride: int = 0
) -> torch.Tensor:
if blocksparse_vert_stride is not None and blocksparse_vert_stride > 1:
# use blocksparse paged attention
......@@ -182,7 +180,7 @@ class PagedAttention:
blocksparse_local_blocks,
blocksparse_vert_stride,
blocksparse_block_size,
blocksparse_head_sliding_step
blocksparse_head_sliding_step,
)
else:
# Run PagedAttention V2.
......@@ -220,7 +218,7 @@ class PagedAttention:
blocksparse_local_blocks,
blocksparse_vert_stride,
blocksparse_block_size,
blocksparse_head_sliding_step
blocksparse_head_sliding_step,
)
return output
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment