LightweightConv. Incorrect decoding with incremental state

[SolomidHero]

🐛 Bug

I am using LightweightConv layer and it seems to be incorrect when applying with incremental state and step more than 1: https://github.com/pytorch/fairseq/blob/fcca32258c8e8bcc9f9890bf4714fa2f96b6b3e1/fairseq/modules/lightweight_convolution.py#L221

Code and Expected behavior

# LightweightConv1dTBC
conv = LightweightConv(
    1,
    kernel_size=3,
    padding_l=2,
    weight_softmax=True,
    num_heads=1,
    weight_dropout=0.0,
)
long_x = torch.arange(30.0).view(-1, 1, 1)

step = 10
result = []
state = {}
for i in range(0, len(long_x), step):
  with torch.no_grad():
    out = conv(long_x[i:i+step], incremental_state=state).squeeze()
    result.append(out)

result = torch.cat(result)

(conv(long_x).squeeze() == result).all()
# False

Expected True (first few values are True, others False). If step is 1, equation is correct

Environment

• fairseq Version (e.g., 1.0 or main): fairseq==0.10.0
• PyTorch Version (e.g., 1.0): 1.9.0
• OS (e.g., Linux): MacOS Catalina
• How you installed fairseq (pip, source): pip
• CPU only

[SolomidHero]

Hi, any update?