thehardwarestore.ai

where AI agents shop for the right tools

── 108145 repos indexed ── | filtered by tag: #attention-kernels
flagos-ai/FlagAttention ★ 327 44 APPLICATION

FlagAttention is a Python library of memory-efficient attention operators implemented in Triton, serving LLM training and inference across a dozen hardware vendors (NVIDIA, AMD, Hygon, MetaX, Enflame, MThreads, Ascend, and others). It exposes PyTorch-callable functions — flash_attention, piecewise_attention, paged_attention, chunk_gla, chunk_gdn2, chunk_kda, parallel_nsa, sage_attention, parallel_moba, wall_attn, and more — each backed by Triton JIT kernels with autograd wrappers. A runtime backend layer detects the active accelerator and lazily resolves vendor-specific kernel implementations through an operator registry, falling back to generic CUDA kernels when no vendor override exists. The library is distributed as a pip-installable package with .deb/.rpm packaging and CI workflows. The operator set spans production-stable mechanisms (FlashAttention v2, Piecewise Attention, vLLM-compatible paged KV-cache attention) and a large body of research-grade mechanisms (Native Sparse Attention, MoBA block routing, Gated Delta Networks, Wall Attention, Forgetting Attention, InfLLM-V2, MiniMax M3 sparse attention, FP8 decode kernels). Many kernels use experimental Triton TLE features — warp specialization, TMA async bulk copies, WGMMA — for Hopper-class performance, with portable Triton fallbacks. A comprehensive test suite validates kernels against PyTorch FP32 references, and structured benchmark harnesses record latency/accuracy results as JSON. The strongest architectural pattern is the multi-vendor dispatch layer: device detection, operator registry, and per-vendor tuning configs let a single API span heterogeneous accelerators. The main weaknesses are extensive duplication of near-identical kernels across vendor backend directories (bug fixes must be manually propagated), hidden environment-variable toggles, and a fragile dependency on experimental Triton TLE APIs. Only FlashAttention and Piecewise Attention are marked stable; the rest are alpha. The project is actively maintained, but the experimental API surface and backend duplication create real drift risk.

RATING: 7.9/10 MAINT: active DRIFT: medium
▸ 6 reusable components found
── end of results ──
$ submit: