Post #115

Соседняя команда проделала большую работу и сегодня выпустила нашу версию FlashAttention для JAX, которая поддерживает Context Parallelism (CP) и Document Masks. При работе с агентами часто встречаются траектории огромной длины вплоть до 128к+ токенов, поэтому CP становится необходимым. При этом, если дополнять каждую последовательность до max_seq_len, ресурсы будут использоваться неэффективно из-за бесполезных pad токенов, т.к. траектории все-таки могут быть и короткими. Отсюда хочется уметь паковать несколько траекторий в один ряд в батче и разделять их масками, чтобы при вычислении attention скоров они не влияли друг на друга. Все это реализовано в нашей либе, за что ребятам большое спасибо ☺️ В посте есть подробное описание, как это работает, а также сравнения с cudnn FA и FlexAttention в различных сценариях, думаю, всем причастным должно быть интересно https://x.com/hr0nix/status/1895096056570093966