[Feature Request] Support `window_size`

[zhuzilin]

Normally we could use other long context methods like deepspeed ulysseus to avoid implementing this.

[conceptofmind]

Do you have any idea of whether Ulysses is plug-and-play with ring attention from this repository? I do understand it works with standard Flash Attention.

[zhuzilin]

Do you have any idea of whether Ulysses is plug-and-play with ring attention from this repository? I do understand it works with standard Flash Attention.

Sorry that I've never tried the origin deepspeed ulysses impl. We have our own write-from-scratch training framework internally, which combined the ulysses and ring attn. In principle, those two algorithms are orthogonal with each other so you could add ring attn to deepspeed with no regard to ulysses.

[conceptofmind]

Hi @zhuzilin,

I appreciate the insight.

One of the deepspeed maintainers stated, "Ulysses is, in principle, attention-type agnostic. Although we haven’t specifically tested Ulysses with Ring Attention, as long as the qkv can be split or sharded along sequence and head dimensions, it should work.".

I will have to test both flash ring attention and deepspeed ulysses.

Thank you,

Enrico

[lwang2070]

Hi, I was wondering whether window_size is currently implemented? It seems that there is no assertion like

assert window_size == (-1, -1)

in your code, and it seems that the window_size parameter is passed into flash_attn api as is

ring-flash-attention/ring_flash_attn/zigzag_ring_flash_attn_varlen.py

Line 129 in 6f1c6f7

block_out, _, _, _, _, block_lse, _, _ = _flash_attn_varlen_forward(**params)

so I'm not sure whether I can just pass in a window_size to the API provided.

Additionally, it is noted in README that

window_size is not supported, because it will be really tricky to implement a varlen version with window_size.

I was wondering if window_size is supported in batch apis?

Thanks a lot in advance:)

Best,
Li

[epwalsh]

I have the same question ☝️