FastFold icon indicating copy to clipboard operation
FastFold copied to clipboard
verified publisher icon hpcaitech

triton softmax support multi-batch

Open Gy-Lu opened this issue 2 years ago • 2 comments

  1. Support another batch dimension for softmax. In training or batch inference, we may add a batch dimension as the first dimension of some tensors. However, we use the third dimension(tensor.shape[2]) as the head_dim, which would be influenced. In this pr, I modify it to tensor.shape[-3] to solve this problem. CUDA kernel is modified as well.

  2. Enable test_atten_core, this test is skipped by default and never be used.

Gy-Lu avatar Feb 16 '23 06:02 Gy-Lu

I found it would fail when the batch dimension comes to 2.

Gy-Lu avatar Feb 18 '23 07:02 Gy-Lu

Update: It works incorrectly before this commit. For the way getting bias_ptr not supports multi-batch. And I have fixed the triton version. CUDA version may support multi-batch one day :(

Gy-Lu avatar Feb 20 '23 09:02 Gy-Lu