85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
import torch
|
|
import block_sparse_attn
|
|
from block_sparse_attn import block_sparse_attn_func, token_streaming_attn_func, block_streaming_attn_func
|
|
|
|
print("block_sparse_attn version:", block_sparse_attn.__version__)
|
|
print("Available funcs:", [k for k in dir(block_sparse_attn) if not k.startswith("_")])
|
|
|
|
device = "cuda:0"
|
|
batch_size = 1
|
|
nheads = 4
|
|
head_dim = 64
|
|
seqlen = 256
|
|
|
|
total_len = batch_size * seqlen
|
|
q = torch.randn(total_len, nheads, head_dim, device=device, dtype=torch.float16)
|
|
k = torch.randn(total_len, nheads, head_dim, device=device, dtype=torch.float16)
|
|
v = torch.randn(total_len, nheads, head_dim, device=device, dtype=torch.float16)
|
|
|
|
cu_seqlens_q = torch.arange(0, total_len + 1, seqlen, device=device, dtype=torch.int32)
|
|
cu_seqlens_k = cu_seqlens_q.clone()
|
|
|
|
# 1. dense attention (all heads mask_type=0)
|
|
head_mask_type_dense = torch.zeros(nheads, device=device, dtype=torch.int32)
|
|
streaming_info = torch.tensor([1, 3] * nheads, device=device, dtype=torch.int32)
|
|
out_dense = block_sparse_attn_func(
|
|
q, k, v,
|
|
cu_seqlens_q, cu_seqlens_k,
|
|
head_mask_type_dense,
|
|
streaming_info,
|
|
None,
|
|
seqlen, seqlen,
|
|
0.0,
|
|
is_causal=False,
|
|
)
|
|
print("dense output shape:", out_dense.shape)
|
|
assert out_dense.shape == q.shape
|
|
|
|
# 2. block-sparse attention (2 dense heads + 2 blocksparse heads)
|
|
num_blocksparse_heads = 2
|
|
nrow = seqlen // 128
|
|
ncol = seqlen // 128
|
|
base_blockmask = torch.zeros(batch_size, num_blocksparse_heads, nrow, ncol, device=device, dtype=torch.bool)
|
|
base_blockmask[:, :, :, :ncol//2] = True # attend to first half
|
|
|
|
head_mask_type_sparse = torch.tensor([0, 0, 1, 1], device=device, dtype=torch.int32)
|
|
out_sparse = block_sparse_attn_func(
|
|
q, k, v,
|
|
cu_seqlens_q, cu_seqlens_k,
|
|
head_mask_type_sparse,
|
|
streaming_info,
|
|
base_blockmask,
|
|
seqlen, seqlen,
|
|
0.0,
|
|
is_causal=False,
|
|
)
|
|
print("block-sparse output shape:", out_sparse.shape)
|
|
assert out_sparse.shape == q.shape
|
|
|
|
# 3. token-level streaming attention
|
|
head_mask_type_stream = torch.full((nheads,), -1, device=device, dtype=torch.int32)
|
|
out_stream = token_streaming_attn_func(
|
|
q, k, v,
|
|
cu_seqlens_q, cu_seqlens_k,
|
|
head_mask_type_stream,
|
|
streaming_info,
|
|
seqlen, seqlen,
|
|
)
|
|
print("token-streaming output shape:", out_stream.shape)
|
|
assert out_stream.shape == q.shape
|
|
|
|
# 4. block-level streaming attention
|
|
out_block_stream = block_streaming_attn_func(
|
|
q, k, v,
|
|
cu_seqlens_q, cu_seqlens_k,
|
|
head_mask_type_stream,
|
|
streaming_info,
|
|
seqlen, seqlen,
|
|
0.0,
|
|
is_causal=True,
|
|
)
|
|
print("block-streaming output shape:", out_block_stream.shape)
|
|
assert out_block_stream.shape == q.shape
|
|
|
|
print("All smoke tests passed!")
|