54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
import torch
|
|
|
|
def generate_block_sparse_mask_for_function(h, num_blocks, k, device="cuda"):
|
|
"""
|
|
Generate block sparse mask of shape [h, num_blocks, num_blocks].
|
|
|
|
Args:
|
|
h: number of heads
|
|
num_blocks: number of blocks
|
|
k: number of kv blocks each q block attends to
|
|
device: device to create tensors on
|
|
|
|
Returns:
|
|
block_sparse_mask: [h, num_blocks, num_blocks] bool tensor
|
|
"""
|
|
k = min(k, num_blocks)
|
|
scores = torch.rand(h, num_blocks, num_blocks, device=device)
|
|
_, indices = torch.topk(scores, k, dim=-1)
|
|
block_sparse_mask = torch.zeros(h, num_blocks, num_blocks, dtype=torch.bool, device=device)
|
|
|
|
block_sparse_mask = block_sparse_mask.scatter_(2, indices, 1).bool()
|
|
return block_sparse_mask
|
|
|
|
|
|
def create_full_mask_from_block_mask(block_sparse_mask, variable_block_sizes, device="cuda"):
|
|
"""
|
|
Convert block-level sparse mask to full attention mask.
|
|
|
|
Args:
|
|
block_sparse_mask: [h, num_blocks, num_blocks] bool tensor
|
|
variable_block_sizes: [num_blocks] tensor
|
|
device: device to create tensors on
|
|
|
|
Returns:
|
|
full_mask: [h, S, S] bool tensor where S = total sequence length
|
|
"""
|
|
h, num_blocks, _ = block_sparse_mask.shape
|
|
total_seq_len = variable_block_sizes.sum().item()
|
|
cumsum = torch.cat([torch.tensor([0], device=device), variable_block_sizes.cumsum(dim=0)[:-1]])
|
|
|
|
full_mask = torch.zeros(h, total_seq_len, total_seq_len, dtype=torch.bool, device=device)
|
|
|
|
for head in range(h):
|
|
for q_block in range(num_blocks):
|
|
q_start = cumsum[q_block]
|
|
q_end = q_start + variable_block_sizes[q_block]
|
|
|
|
for kv_block in range(num_blocks):
|
|
if block_sparse_mask[head, q_block, kv_block]:
|
|
kv_start = cumsum[kv_block]
|
|
kv_end = kv_start + variable_block_sizes[kv_block]
|
|
full_mask[head, q_start:q_end, kv_start:kv_end] = True
|
|
|
|
return full_mask |