torch triangle_attention failed when sequence length large than 3000

[lupengk]

When I pass data with a length over 3000 to triangle attention, torch triangle_attention fails to run, but it can still output the shape.

OUTPUT:

torch.Size([1, 3000, 4, 3000, 32])
Traceback (most recent call last):
  File "test_triattn.py", line 14, in <module>
    print(torch.sum(output))
          ^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

CODE:

import torch
import math
from cuequivariance_torch import triangle_attention


if torch.cuda.is_available():
    q_x=torch.randn([1, 3000,4, 3000, 32],dtype=torch.bfloat16,device=torch.device('cuda'),requires_grad=True)
    kv_x=torch.randn([1, 3000,4, 3000, 32],dtype=torch.bfloat16,device=torch.device('cuda'),requires_grad=True)
    triangle_bias=torch.randn([1, 1, 4, 3000, 3000],dtype=torch.bfloat16,device=torch.device('cuda'),requires_grad=True)
    mask_bias=torch.ones([1, 3000, 1, 1, 3000],dtype=torch.bfloat16,device=torch.device('cuda'))>0
    scale = 1 / math.sqrt(32)
    output = triangle_attention(q_x,kv_x,kv_x,triangle_bias,mask_bias,scale)
    print(output.shape)
    print(torch.sum(output))

[kilianFatras]

Hello,

I am meeting a similar problem but during backward! I suppose these problems are related. Any idea how to solve this in the meantime?

The following example leads to a backward error
RuntimeError: CUDA error at /home/jenkins/agent/workspace/branch-build-pytorch/cuequivariance_ops/cuda/triangle_attention/bwd_fmha.cu:183: an illegal memory access was encountered

import torch
import math
from cuequivariance_torch import triangle_attention

N = 2000
D = 128
if torch.cuda.is_available():
    q_x=torch.randn([1, N, 4, N, D], dtype=torch.bfloat16,device=torch.device('cuda'),requires_grad=True)
    kv_x=torch.randn([1, N, 4, N, D], dtype=torch.bfloat16,device=torch.device('cuda'),requires_grad=True)
    triangle_bias=torch.randn([1, 1, 4, N, N], dtype=torch.bfloat16,device=torch.device('cuda'),requires_grad=True)
    mask_bias=torch.ones([1, N, 1, 1, N], dtype=torch.bfloat16,device=torch.device('cuda'))>0
    scale = 1 / math.sqrt(D)
    output = triangle_attention(q_x,kv_x,kv_x,triangle_bias,mask_bias,scale)
    print(output.shape)

    grad_out = torch.randn_like(output)
    output.backward(grad_out)
    print(q_x.grad.shape)
    print(kv_x.grad.shape)  

[borisfom]

@lupengk : thanks for the report! We are working on this issue and it's going to be fixed in the upcoming 0.6 release.

[mariogeiger]

Closing because it's fixed in 0.6.0