测试MultiHeadAttention将bias设成False出现报错

版本信息:
MegEngine版本:1.13.1
python版本:3.8

运行如下代码可以正常运行:
import megengine as mge
import numpy as np
import megengine.functional as F

embed_dim, num_heads = 768, 12
seq_length, batch_size = 1024, 2

query = F.ones((batch_size, seq_length, embed_dim)) * 1.0

multihead_attn = mge.module.MultiHeadAttention(embed_dim, num_heads, bias=True)

attn_output = multihead_attn(query, query, query)[0]

print(attn_output)

但是如果把bias设置成False,就会有如下报错:

请问是怎么回事?

用的是什么显卡呀,查了下 cudnnMultiHeadAttnForward(),报错是因为 Insufficient amount of shared memory to launch a GPU kernel,换成小shape可以运行。

改小shape确实可以跑了,谢谢。那设置bias为False时候,就不能跑大shape了吗?

这个和显存大小有关,如果 bias 为 True 就需要额外的显存去存 tensor,显存够大的话是可以的