qqtang/sskj
1
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: remove obsolete scripts/test_block_sparse_attn.py after refactor

Browse Source
This commit is contained in:
Quantong Qiu 2026-07-08 06:01:31 +00:00
parent a5d69dfd9a
commit 07835b9fed
1 changed files with 0 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
scripts/test_block_sparse_attn.py
View File

@ -1,84 +0,0 @@
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!")