cute.interface _flash_attn_fwd fucntion doesn't provide LSE tensor in the return value

[beiw-nv]

For inference, it seems LSE tenor is returned as None here. Is there anyway cute interface can return LSE as in other FA flash_attn_fwd function, for example: here

[tridao]

Right, for inference we don't return LSE since that's typically not needed. Not hard to change, it's just 1 line in pytorch

[beiw-nv]

Can you please point me the line need to change in pytorch?

[tridao]

That's the line you linked to:
https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/cute/interface.py#L118

Just change it to always construct the LSE tensor

[beiw-nv]

I am wondering how multi-gpu run will affect the correctness of _flash_attn_fwd function on B200. I am trying to integrate the latest FA3 cutlass implementation to LongContextAttention here. During 2-gpu testing, I found that gpu 0 always returns the correct out and lse. However, gpu 1 sometime returns mostly 0 values for out and lse depending on the synchronization between two GPUs. The run is launched with torchrun --nproc_per_node=2 so we have 1 proc per gpu.

[tridao]

Maybe we need with torch.cuda.device, sth like this?

flash-attention/flash_attn/ops/triton/rotary.py

Line 158 in 413d07e

with torch.cuda.device(x.device.index):

[beiw-nv]

thanks! let me try that.

[beiw-nv]

Yes, adding

with torch.cuda.device(q.device.index):

tohere fixes the correctness issue in multi-gpu runs completely. We may also need to add it for backward function as well.