FlexAttention backend support for sequence packing
Author: vadimkantorov
Pull request: no
Labels: name: Feature request

### Feature request

Ideally, this support compile / fullgraph / cudagraphs (as currently the packing-supporting backend `flash_attention_2` doesn't support fullgraph because of un/pad graph breaks: https://github.com/huggingface/transformers/issues/42950 )

And maybe the inputs should still be padded to a multiple to avoid recompiles (or compile directly with dynamic shapes)

Related:
- https://github.com/huggingface/transformers/issues/27640#issuecomment-2619471784
- https://huggingface.co/blog/packing-with-FA2

### Motivation

More efficient training, not spending cycles on processing padding tokens

### Your contribution

I can try hacking this support, but probably not making a PR at this point

Comments:
for sequence packing we recommend using `continious batching` with `generate_batch` API. Or do you mean for actual training where you want forward pass to handle that? 

The flex attention api is TBH quite cumbersome, but happy to review a PR if its well made!

I meant mainly for training, yes

Hey @vadimkantorov, I think there is a lot of confusion around the usage of packing with a lot of attention flavors, trying to clear a bit up here (also including points in the other issue).

- First, the initial packing with FA relied on position ids only but this is very unfriendly to compile making it only possible with non-fullgraph compile. 

- Second, If you want fullgraph compatibility you will need to prepare the FA kwargs outside of the forward, e.g. via a data collator. You can take look at my [gist](https://gist.github.com/vasqu/c264d14268fbc7f1eb725e09dfb87331) - note the `return_flash_attn_kwargs=True` for the collator, without you fall into the pos ids case only again. I also have this PR #41827 which checked compile compatibility.

  - Side note: FA4 #42435 (when released) will allow to circumvent some of these issues as it no longer uses `max_seqlen_q/k`.

- Third, SDPA + Eager + Flex Attention already support packing within their mask - you simply need to pass the position ids as packed sequences (like the original packed FA version). Keep in mind, not all models pass pos ids for the mask creation - can be multiple reasons but most likely some embeddings always caused issues (e.g. I rewrote parts in bert for this). 

- Fourth, as you may have noted the mask creation would still be within the forward for SDPA + Eager + Flex Attention. Theoretically you could move this outside of the forward by creating it yourself before (use it the same way as in the forward). We ignore already prepared masks in the forward. I mention this specifically for flex attention because the mask API is very slow, so you kinda have to do this outside to not lose too much perf. Edit: Likely also needed for fullgraph compile support, similar issue as in FA (control flow checks)

I can add code examples/pointers if you need them but I hope the gist is clear.

> Third, SDPA + Eager + Flex Attention already support packing within their mask - you simply need to pass the position ids as packed sequences (like the original packed FA version). 

@vasqu, could you please exemplify what you mean by "position ids as packed sequences" - last I checked SDPA ignored pos ids and attended to the whole packed sequence. 

I was passing:
```
position_ids = [0,1,2,3,4,0,1,2,0,1,2,3,4,5]
```
what would be "position ids as packed sequences" if it is not the above?



@stas00 On which model did you try it out? Sadly not all models use the needed mask interface (+ passing position ids to the creation of the mask). Not 100% sure but you still might need `bsz x seq_len`, i.e. in a flattened input `1 x seq_len`, so you could `unsqueeze(0)` or view with `[None, :]` to reshape.

You can see example mask creations at our tests here https://github.com/huggingface/transformers/blob/7f20ad0073ae6d2f6a799c1404448d579496b6c4/tests/utils/test_masking_utils.py#L75

Otherwise, I'll take a proper look tomorrow

> On which model did you try it out?

I worked mainly with llama-3 series of various sizes.

I was always testing with `[1, seqlen]` style - it never worked. I know because it'd take forever to SDPA a large packed seqlen made from very small samples, the attention time would be the same if I were to make one long unpacked sample with the same seqlen. 

The reason I know this is because of https://arxiv.org/abs/2506.13996 where I'd feed `[1,500K]` and higher sizes (the biggest test was `[1,15M]`.

Moreover, the silent ignoring of pos ids - is a very serious (dare I say irresponsible) issue since a user might be unaware that they are training:
(1) with sample cross-contamination 
(2) the training is likely to take many times longer especially at long seqlen because of quadratic compute

this shouldn't be allowed and asserted ideally. Or some other way found to let the user know they have the wrong setup.

For this reason UlyssesSP disallows `sdpa` until it gets fixed. Pretty sure `eager` is in the same boat, but I don't think it even supports pos_ids and doesn't require attn_mask. attn_mask can't be used at very large seqlen, because 4D mask becomes too big to be able to do anything.

> Sadly not all models use the needed mask interface (+ passing position ids to the creation of the mask)

Do you mean in Transformers? There should be no mask created as I mentioned above - the generated 4D attn mask is counter-productive to feasibility/performance because it can't fit into a gpu at 500K seqlen - it should not be created if pos ids are passed - this is intentional.

Here is the section 3.4 from the https://arxiv.org/abs/2506.13996 paper explaining the problem with 4D mask and long seqlen

<img width="664" height="227" alt="Image" src="https://github.com/user-attachments/assets/f021448f-daa8-47d6-9dbc-1dd5877cc964" />

I think there is a misunderstanding in regards to how packing works with sdpa/eager/flex attention: We cannot do the same packing as in flash attention, but rely on the mask to handle avoiding cross contamination (e.g. see the expected test mask https://github.com/huggingface/transformers/blob/7f20ad0073ae6d2f6a799c1404448d579496b6c4/tests/utils/test_masking_utils.py#L40)

> I know because it'd take forever to SDPA a large packed seqlen made from very small samples, the attention time would be the same if I were to make one long unpacked sample with the same seqlen.

Alas, the packed sequence version will use xformers backend while the other might even use flash attention backend. It is kind of expected that we do not gain much perf in this comparison as you have the same number of effective tokens, no? I haven't compared if fa varlen vs base fa would have a big perf gain here but they would also need to process the same number of tokens. The real performance gain is in avoiding padding completely, i.e. you will use all tokens of your input instead of all tokens - all padding tokens. Maybe I'm overseeing something(?)

> this shouldn't be allowed and asserted ideally. Or some other way found to let the user know they have the wrong setup.

It would be indeed nice to possibly add a flag or similar. Let me see what I can add here. I explicitly mentioned that it is not possible with all models because we also have a lot of old models including very exotic ones, encoder-decoder ones etc. - it's hard to support them all. Recent models and especially causal models are usually supporting this (ofc including llama).

> For this reason UlyssesSP disallows sdpa until it gets fixed.

I'm not too familiar with ulysses but we need the lse iirc. At least, most context parallel implementations required it. This will need more care and especially on whether we return the lse in sdpa. It might be possible but not currently as it's implemented in transformers. As I said we hide packing behind the mask atm, not as in flash attention - there is no way around it as sdpa cannot support fa varlen in their backends (although they are working on a separate varlen api on torch main). I know that base CP is also supported in the sdpa interface (likely using the fa backend), see https://docs.pytorch.org/tutorials/unstable/context_parallel.html

> Pretty sure eager is in the same boat, but I don't think it even supports pos_ids and doesn't require attn_mask. attn_mask can't be used at very large seqlen, because 4D mask becomes too big to be able to do anything.

Eager packing works the same as sdpa packing in our implementation, i.e. via the mask. Understandably, the mask will be too big for these high seq lens but maybe I could explain the difference between packing in our implementations. Flex attention will also use their mask so maybe this is more efficient than the naive torch tensor mask as they use their sparse block mask.

> it should not be created if pos ids are passed - this is intentional.

It makes sense and I support this but it is also impossible atm. These backends cannot work without their masks (otherwise we would need to automatically fallback to flash attention and we do not want implicit fallbacks). 

The packing is limited and that is intentional (as we still want to use the base sdpa interface, no fallback). The benefit is moreso to know how many tokens you effectively process / avoid padding in the first place. For real efficiency, flash attention can't be avoided (for now). 


Read a little about ulysses and the best option really would be to just use the integrated fa backend, i.e. force no mask and let the `is_causal` kwarg handle it. 

I see that in this implementation it makes sense to have packing included as you perform attn on a partial subset of the weights (instead of the subset of the sequences so my lse call was wrong). Sadly like I said, no mask in sdpa == no packing. I will work on integrating the varlen api when torch releases it tho, maybe we have a nice middle ground then :)

varlen api as far as I understand, simply uses underlying FA varlen API
(flex is not available for now as backend for sdpa or varlen) - but varlen API is already available in nightly

So in the meanwhile, having some perf benchmarks to prove that Flex is good for packed mode would be awesome :) And a check that it enables fullgraph / cudagraphs.

It also seems that FAv3 should support fullgraph. So maybe also a test for using FAv3 and FAv3's varlen API would be awesome (even before it gets supported as part of PyTorch varlen API) - so maybe lifting up any mask creation (as currently it would fail with max op etc)

In my experience FA w/ varlen makes a huge difference in speed - minutes over hours when packing small sequences into very long sequences. When one wants to train on very long sequences it's very difficult to find a dataset with huge samples (in particular for Instruct), so packing is the only way to train faster.

I'm looking forward to SDPA supporting varlen, since it'd be so much easier to use than overcoming FA2/3 build requirements. Albeit, FA4 is going to make all these headaches go away with JIT cute DSL building.

@stas00 Have you tried Flex support for varlen for this scenario? Does it perform well?

I haven't tried Flex yet. FA2/3 worked too well ;)

As long as it  can support split attn heads across ranks and use pos ids it should work out of the box, just need to add to the list of supported attn implementations. 

One downside is that FA (and thus PyTorch's varlen) doesn't support custom attn_mask :(

> varlen api as far as I understand, simply uses underlying FA varlen API

Still better to be inbuilt in torch directly tho than having to mess with building from source (which happens more oftan than I'd like) 😢 

> So in the meanwhile, having some perf benchmarks to prove that Flex is good for packed mode would be awesome :) And a check that it enables fullgraph / cudagraphs.

You can try it out :p I would consider this a documentation upgrade on perf upgrades (in general our docs need to be updated badly). Not sure about fullgraph but I don't see why it should have issues here. Consider that the mask creation is expensive tho so best to do it outside like I mentioned.

> It also seems that FAv3 should support fullgraph.

Also supporting FA3 here, works the same as FA2! So it only depends on which commit you installed FA3 (or if you use our kernels). No need for torch nightly. Like I said fullgraph needs the data collator preparation, it won't work from the get go.

> In my experience FA w/ varlen makes a huge difference in speed

TIL, good to know. Would be interesting to see how flex fares then indeed.

> I'm looking forward to SDPA supporting varlen, since it'd be so much easier to use than overcoming FA2/3 build requirements. Albeit, FA4 is going to make all these headaches go away with JIT cute DSL building.

Really hoping for FA4 especially removing the annoying max_q/k which causes a lot of device syncs.

I'm in agreement that ideally pytorch/sdpa should be flexible enough to support varlen / perform on par with FA and not require additional packages. 

The cons is that it'd be pytorch version dependent and often fixes can't be used w/o risking the nightly build. Therefore I now simply prebuild various versions of FA2/3 cross-compiled to all archs I need and automate the install from my own gh repo using uv, so it's blazing fast.