nebius/kvax

A FlashAttention implementation for JAX with support for efficient document mask computation and context parallelism.

2 best points · 0 comments · 1 Hacker News post · 177 GitHub stars · View on GitHub

Hacker News discussions

• Kvax: Fast and easy-to-use flash attention implementation for Jax (2 points, 0 comments, by exizt88, 2025-03-02)