attn_num_mem_kv > 0 and attn_one_kv_head = True error

[pfeatherstone]

If you set both attn_num_mem_kv and attn_one_kv_head, you get an error.

Repro:

lm = TransformerWrapper (
    num_tokens          = 32,
    max_seq_len         = 0,
    num_memory_tokens   = 20,
    attn_layers = Decoder (
        dim             = 512,
        depth           = 4,
        heads           = 4,
        rotary_pos_emb  = True,
        attn_flash      = True,
        attn_onnxable   = True,
        attn_num_mem_kv  = 20,
        attn_one_kv_head = True
    )
)

x = torch.randint(0, 32, (8, 1024))
logits = lm(x)
print(x.shape, x.dtype)
print(logits.shape, logits.dtype)

[lucidrains]

@pfeatherstone ah nice catch

corrected

[pfeatherstone]

The first fix is:

https://github.com/lucidrains/x-transformers/blob/029ec31957a86dbfaa4a9865f0f23403ef31e4aa/x_transformers/x_transformers.py#L789-L792

self.num_mem_kv = num_mem_kv
        if num_mem_kv > 0:
            self.mem_k = nn.Parameter(torch.randn(kv_heads, num_mem_kv, dim_head))
            self.mem_v = nn.Parameter(torch.randn(kv_heads, num_mem_kv, dim_head))

[pfeatherstone]

Not so fast

[lucidrains]

@pfeatherstone there's another issue?

[pfeatherstone]

it fails at https://github.com/lucidrains/x-transformers/blob/029ec31957a86dbfaa4a9865f0f23403ef31e4aa/x_transformers/attend.py#L135

[lucidrains]

weird, let me try your example, runs for me

[lucidrains]

oh i see, it is an issue with flash attention

[lucidrains]

ok, fixed for real, thanks for reporting