einfold가 SQL 25줄에서 FlashAttention 커널을 자동 도출한다

FlashAttention 커널을 직접 작성하지 않고 SQL에서 얻어낼 수 있을까. xqlsystems의 xarray-sql 저장소에 포함된 벤치마크 스크립트가 그 가능성을 데모로 보여준다. 쿼리 옵티마이저 einfold가 약 25줄 분량의 SQL로 표현된 attention을 FlashAttention과 같은 융합 실행 계획으로 바꿔 돌린다.

데모에서 Q, K, V는 토큰·특성·값 세 축을 가진 테이블로 등록된다. attention은 조인 두 번으로 서술된다. 점수는 Q와 K를 특성 축으로 곱해 합산한 뒤 1/√d로 스케일하고, 출력은 exp(점수)로 가중한 V의 평균이다. 같은 모델을 윈도 함수 표기와 GROUP BY 표기 두 가지로 적어 두고, 두 표기가 모두 동일한 융합 파이프라인으로 실행되는지 확인한다. 특성 수 D는 64, 토큰 수는 64와 256 두 규모를 사용한다.

실행 호스트는 DataFusion, einfold를 얹은 DataFusion, DuckDB 세 가지다. 각 호스트의 순전파 결과는 JAX로 같은 계산을 수행한 값과 대조된다. 그래디언트는 SQL 쿼리를 미분하는 ddx가 계산하고 jax.grad(argnums=(0,1,2)) 결과와 비교한다. 손실은 attention 뷰를 읽는 별도 SELECT로 정의되기 때문에, 모델과 손실이 파이썬 코드에서처럼 분리된 채 하나의 SQL 문 안에서 Q, K, V 전체에 대한 역전파가 한 번에 이뤄진다. ddx 백엔드가 없는 호스트는 이 단계를 건너뛴다.

문제의 출발점은 순진한 실행 계획의 비용이다. 일반적인 SQL 엔진은 N×N개의 점수를 전부 물질화한 뒤 softmax를 적용한다. FlashAttention(Dao 등, 2022)은 점수 행렬을 저장하지 않고 키 블록을 쿼리 블록에 흘려보내며 running softmax를 유지해 같은 결과를 얻는다. einfold는 SQL 실행 계획을 배열 알고리즘으로 재작성하는 옵티마이저이고, 여기에 걸린 규칙은 attention을 겨냥한 것이 아니다. 조인에서 점수를 계산하는 softmax 가중 평균이라는 일반 패턴을 융합할 뿐이다.

그 주장을 뒷받침하려고 스크립트는 attention이 아닌 대조군을 함께 돌린다. 리그리딩에 쓰이는 Nadaraya–Watson 추정기, 즉 가우시안 커널(h=0.5)로 거리를 가중한 평균이다. 구조적으로는 attention과 같은 모양이므로, attention을 빠르게 하는 규칙이라면 이 쿼리도 똑같이 빠르게 만들어야 한다. 대조군 역시 einfold의 융합 파이프라인으로 실행되는지가 검사 항목에 들어 있다.

마스크 처리는 SQL 의미론에 맡긴다. 인과 마스크는 조인 조건으로 표현되어, 가려진 (i, j) 쌍은 곱해지는 0이 아니라 아예 존재하지 않는 행이 된다. V에 무한대가 들어가도 허용된 쿼리에만 도달하며, 0·∞가 NaN으로 번지는 경로가 없다. xarray-sql이 NaN을 결측값으로 읽어 SUM이 건너뛰기 때문에, 무한대를 심어 이 동작을 실제로 확인한다.

실무적으로 눈여겨볼 지점은 모델 정의와 실행 계획 선택의 분리다. 사용자는 뷰 하나로 attention을 선언하고, 어떤 호스트에서 어떤 융합 계획이 선택될지는 옵티마이저가 결정한다. 손실과 그래디언트까지 SQL 문법 안에서 표현되므로, 파이썬 프레임워크와 SQL 엔진 사이에서 모델을 두 벌로 유지하던 부담이 줄어들 여지가 있다.

다만 이 스크립트는 einfold가 내건 세 가지 목표 데모 중 첫 번째이며, 규모도 64와 256 토큰 수준이다. 저장소는 einfold 0.1.0rc2 이상과 xarray-sql의 편집 가능 설치를 전제로 하고, 그래디언트 검증은 ddx 백엔드가 붙은 호스트에서만 수행된다. 벤치마크는 `uv run benchmarks/ml/01_attention.py --help`로 직접 실행해 볼 수 있다.