Skip to content

Support FlashAttention with pad_between_seqs=True #2391

Description

@QiZhangNV

Currently, the dot product attention backend FA3 in TE does not support padding between sequences. As noted here).

I opened a related issue in FlashAttention repo and confirmed that this feature could be supported by using the parameters cu_seqlens_q/k and seqused_q/k.

However, one thing to keep in mind is that we need to manually set the output of the padded positions to zero. I noticed that the fused attention implementation also handles this in a similar way (link).

Given this, it seems like support for padding between sequences could be reasonably added to FA3 in TE.

Metadata

Metadata

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions