[Bug] GatedLinearAttention got NaN

[980202006]

Describe the Bug

Forward output nan, GatedLinearAttention

Steps to Reproduce the Bug

        o, recurrent_state = fused_chunk_gla(
            q=q,
            k=k,
            v=v,
            g=gk,
            initial_state=recurrent_state,
            output_final_state=use_cache,
            head_first=False
        )
        if torch.isnan(o).any():
                1==1

head = 1
seq_len=43884
dim=64

Expected Behavior

There is no nan in the forward output

Environment Information

1. Torch:'2.5.1+cu124'
2. Triton:3.0 nightly

[yzhangcs]

@980202006 Hello, could you provide the input tensors resulting in NaNs?

[yzhangcs]

BTW, it's recomended to utilize chunk mode instead of fused_chunk

[yzhangcs]

@980202006 Could you check 458c018 again?