Value-dim differnet from query-dim/key-dim is not supported

[hahnyuan]

Hello flash_attn Maintainers and Community,

When attempting to execute the provided code snippet utilizing the flash_attn library, a runtime error occurs with the following message:

import torch
import flash_attn

bs=1
seqlen=128
qkdim=128
vdim=256
nheads=8
q=torch.ones(bs,seqlen,nheads,qkdim,dtype=torch.float16).cuda()
k=torch.ones(bs,seqlen,nheads,qkdim,dtype=torch.float16).cuda()
v=torch.ones(bs,seqlen,nheads,vdim,dtype=torch.float16).cuda()

attn=flash_attn.flash_attn_func(q,k,v)

Result: RuntimeError: v must have shape (batch_size, seqlen_k, num_heads_k, head_size_og)

It seems that the error is due to a mismatch between the dimensions of the value tensor v and the query and key tensors q and k, respectively.

Is it possible to extend support to scenarios where the dimensionality of the value tensor v differs from that of the query and key tensors q and k?

[tridao]

No we don't plan to support that.