FlashAttention returns all zeros when device is 'cuda:1'

testing code:

import numpy as np
import torch
from flash_attn.flash_attention import FlashAttention
def test(device):
    flash = FlashAttention()
    d_head = 64
    n_heads = 32
    flash.softmax_scale = 1  # / (d_head ** 0.5)
    batch_size = 4
    seq_len = 16
    qkv = torch.ones(batch_size, seq_len, 3, n_heads, d_head, dtype=torch.float16)
    flash = flash.to(device)
    qkv = qkv.to(device)
    out, _ = flash(qkv)
    print(out.shape, torch.abs(out).type(torch.float32).sum())
    return out

if I add the following code and run it,

out = test("cuda:0")

the output is

torch.Size([4, 16, 32, 64]) tensor(131072., device='cuda:0')

if I run this instead

out = test("cuda:1")

the output is

torch.Size([4, 16, 32, 64]) tensor(0., device='cuda:1')

I’ve verified it’s not the hardware issue by adding

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

at the top and it works with cuda:0

[tridao]

I’ve just pushed a commit that fixed this. Let me know if it works on your side.

[geekinglcq]

Great! @geekinglcq Did recompilation fix it or do you need to set torch.cuda.set_device(qkv.device())?

I just recompiled it.