drisspg
commented
1 month ago Ah, yeah, we need to better document this. There are two things going on here:
block_mask = create_block_mask(
sliding_window_mask, B=None, H=None, Q_LEN=SEQ_LEN, KV_LEN=SEQ_LEN, _compile=False
)You need to call this with _compile=True. We essentially map your block mask over a full Q_LEN x KV_LEN matrix in order to produce the block mask. Without compile, we need to materialize this full thing, and it can cause OOMs on long sequences.
As well, you need to run flex_attention = torch.compile(flex_attention). Without compile, flex falls back to a non-fused eager implementation that is great for debugging, but it is much slower and materializes the full scores matrix.
mishooax
foreverpiano
fteufel
Hi - I'm trying to implement sliding window attention with flex-attention, as described in the snippet below, inspired by the sliding window attn example in the attention-gym. Note I use a rather long sequence length (40320), but a small sliding window (512).
Unfortunately, this OOMs:
Looks like flex-attention is trying to create a
SEQ_LEN x SEQ_LENmatrix at theQ * K^Tstep. What am I doing wrong here? Flash attention'sflash_attn_funccan handle this sequence len just fine, as long as i pass in the correctwindow_size. Thank you!