[Question]: when searching the best sparse attention type ,why to caculate the score just pick the 2500 cols?

Describe the issue

I hope this messages find u well!
def stream_llm(vertical_size, slash_size):
q_len = q.shape[2]

    mask = torch.triu(torch.tril(torch.ones(q_len, q_len), 0), -slash_size).to(q)
    mask[:,:vertical_size] = 1
    mask = mask.unsqueeze(0).unsqueeze(1)

    est_attn = torch.tril(mask)
    attn_weights_x = attn_weights * est_attn
    res3 = attn_weights_x[:,:,2500:].sum(-1).mean(-1).squeeze().float().detach().cpu().numpy()
    return res3

why this only pick the 2500: columns? If the head is A-shape, this may lead to mis-classification to V-S type?