drisspg
commented
3 weeks ago The problem is that we dont support wrapping, "create_block_mask" in torch.compile instead BlockMask should be created outside of the region and than ammortized across all the attention Layers. A "hack" to get it to work though is:
import torch
import torch.nn as nn
from torch.nn.attention.flex_attention import flex_attention, create_block_mask, _mask_mod_signature
def generate_sliding_window(window_size: int) -> _mask_mod_signature:
def sliding_window(b, h, q_idx, kv_idx):
del b, h # not used
return torch.abs(q_idx - kv_idx) <= window_size // 2
sliding_window_mask = sliding_window
sliding_window_mask.__name__ = f"sliding_window_{window_size}"
return sliding_window_mask
class Attention(nn.Module):
def __init__(
self,
dim,
heads = 8,
dim_head = 64,
dropout = 0.,
window_size=64,
bias: bool = False,
):
super().__init__()
self.window_size = window_size
self.dim=dim
self.heads = heads
inner_dim = dim_head * heads
self.to_q = nn.Linear(dim, inner_dim, bias = bias)
self.to_kv = nn.Linear(dim, inner_dim * 2, bias = bias)
self.to_out = nn.Linear(inner_dim, dim, bias=bias)
@torch._dynamo.disable()
def create_block_mask(self, l_q, l_k):
sliding_window_mask = generate_sliding_window(window_size=self.window_size)
block_mask = create_block_mask(
sliding_window_mask, B=None, H=None, Q_LEN=l_q, KV_LEN=l_k, _compile=True
)
return block_mask
@torch.compile()
def forward(self, x, context=None):
kv_input = x if context is None else context
q, k, v = self.to_q(x), *self.to_kv(kv_input).chunk(2, dim=-1)
# Reshape q, k, v
b, n, _ = q.shape
q = q.view(b, n, self.heads, -1).transpose(1, 2)
k = k.view(b, n, self.heads, -1).transpose(1, 2)
v = v.view(b, n, self.heads, -1).transpose(1, 2)
######################
# Flex attention block:
l_q = q.shape[2]
l_k = k.shape[2]
print(q.shape, k.shape, v.shape)
block_mask = self.create_block_mask(l_q, l_k)
out = flex_attention(q, k, v, block_mask=block_mask)
# Flex attention block end.
#######################
# Reshape output
out = out.transpose(1, 2).contiguous()
out = out.view(b, n, -1)
out = self.to_out(out)
return out
device= torch.device('cuda')
model = Attention(768, 6, 128, window_size=64)
model.to(device)
test_inp = torch.ones(1,65536,768).to(device)
out = model(test_inp)
fteufel
Chillee
Hi,
Thank you for providing this collection! I'm trying to get local window attention to run. I managed to have a simple example running locally as shown in #15, but I am facing problems now when I try to wrap everything into a module and use it in an actual transformer.
Specifically, I'm facing a triton error that is beyond my understanding:
Code that produces the error:
Any ideas what might be going wrong here?
Torch version
2.5.0.dev20240815+cu118