einfold - Автоматический вывод FlashAttention из SQL-запроса
Show HN: Automatic deduction of FlashAttention kernel from ~25 lines of SQL
einfold — это оптимизатор запросов, который превращает обычный SQL в эффективные алгоритмы для массивов. В демо показано, как из ~25 строк SQL, описывающих механизм внимания (attention) в трансформерах, einfold автоматически строит план, эквивалентный FlashAttention, без единого специального правила. Инструмент также умеет дифференцировать SQL-запросы: градиент вычисляется одним запросом и проверяется через jax.grad. Демо включает сравнение с JAX, поддержку масок и контрольный пример с оценкой Nadaraya–Watson. Запуск: uv run benchmarks/ml/01_attention.py.
einfold находит этот план из SQL в обоих вариантах написания, без единого правила, написанного специально для attention: его правило объединяет любое средневзвешенное по softmax, где оценки вычисляются из join, и то же правило объединяет контрольный запрос, который не является attention.