如何加快pytorch中的滑动注意力计算？

Question

我正在尝试实现一个使用滑动注意窗口的注意模型。
到目前为止，我通过使用火炬模块 nn.MultiheadAttention 并传递一个注意力掩码来实现这一点，这允许每个位置 n 只关注来自 (n-W, n].

的数据 然而它很慢，我觉得在执行速度方面应该有改进的空间，因为滑动窗口只有实际系列大小的 1/10 左右。
鉴于 T 是我的序列长度和 W 滑动窗口，我觉得应该可以将执行时间从 O(T²) 减少到 O(W*T)。这个（玩具）功能已经做了我想要的，

def custom_matmul(a, b, mask):
    device = a.device
    B, T, _ = a.shape
    ret = torch.zeros((B, T, T)).to(device)
    for row, col in (mask==0).nonzero(as_tuple=False):
        ret[:, row, col] = (a[:,row]*b[:,:,col]).sum(dim=-1)
    return ret

但它仍然比在整个序列 T 上使用 a.matmul(b) 并在之后应用掩码要慢得多。有谁知道更快的实现方式？