GPU 기반 텐서 연산을 위한 타일링 주의 스케일드 닷‑프로덕트 어텐션 구현

초록

TiledAttention는 cuTile Python으로 구현된 스케일드 닷‑프로덕트 어텐션(SDPA) 전방 연산으로, 온라인 소프트맥스와 K·V 타일 스트리밍을 이용해 전체 점수 행렬을 메모리에 올리지 않는다. 파이썬 수준에서 타일 크기·스테이징·공유 메모리 레이아웃을 조정할 수 있어 커스터마이징이 용이하고, DGX GB10에서의 재현 가능한 벤치마크 결과는 기존 PyTorch SDPA와 비교해 짧은 시퀀스에서는 경쟁력을, 긴 시퀀스에서는 약간 뒤처짐을 보인다.

상세 분석

본 논문은 최신 트랜스포머 모델에서 핵심 병목인 스케일드 닷‑프로덕트 어텐션(SDPA)의 전방 연산을 고성능 GPU에서 효율적으로 수행하기 위한 새로운 구현체인 TiledAttention를 제안한다. 기존의 고성능 커널은 CUDA /CUTLASS 템플릿이나 Triton과 같은 저수준 DSL에 의존해 최적화가 가능하지만, 코드 수정이 복잡하고 재현성이 떨어지는 단점이 있다. TiledAttention는 이러한 문제를 해결하고자 cuTile Python(또는 TileIR)이라는 고수준 타일 프로그래밍 프레임워크를 활용한다.

핵심 아이디어는 시퀀스 차원(S)을 타일 단위로 나누어 블록‑와이즈로 점수 행렬을 계산하고, 동시에 K와 V를 스트리밍하면서 온라인 소프트맥스 통계를 유지하는 것이다. 구체적으로 각 Cooperative Thread Array(CTA)는 Q 타일을 고정하고 K·V 타일을 순차적으로 로드한다. 로드된 K와 V에 대해 내적을 수행해 부분 점수를 얻고, 마스크를 적용한 뒤 현재까지의 최대값(m)과 정규화 계수(ℓ)를 업데이트한다. 이 과정은 수식 (2)–(4)에 명시된 바와 같이 FP32 누적을 사용해 수치 안정성을 확보한다. 최종적으로 ℓ으로 나누어 정규화된 출력 O를 기록한다.

타일 크기 TM(쿼리 행)와 TN(K·V 열) 그리고 스테이징 단계 수는 파이썬 API를 통해 명시적으로 조정 가능하다. 저자는 TM=64·128, TN=64·128·256 등 여러 조합을 실험했으며, 실험 결과는 타일 크기가 모델의 헤드 차원(D)과 시퀀스 길이(S)에 따라 최적점이 달라짐을 보여준다. 예를 들어 S=1024, D=128인 경우 TM=64, TN=128이 최적이며, 비최적 설정에서는 약 9% 정도 성능 저하가 발생한다.

벤치마크는 NVIDIA DGX GB10(8×A100)에서 수행됐으며, 동일한 워크로드(배치 1, 헤드 8)와 동일한 초기화·스트림·타이밍 방식을 적용해 재현성을 확보했다. 측정 항목은 전방 실행 시간과 토큰당 초당 처리량(tokens/s)이며, median과 p95를 보고한다. 결과는 크게 두 가지 트렌드로 요약된다. 첫째, 짧은 시퀀스(S≤2048)에서는 오버헤드가 지배적이어서 TiledAttention가 PyTorch fused SDPA와 거의 동등하거나 약간 앞서는 경우가 있다. 둘째, 긴 시퀀스(S≥4096)에서는 메모리 대역폭과 공유 메모리 압력이 주요 병목이 되며, 이때는 기존 fused SDPA가 여전히 우위에 있다. 전체 80개의 실험 포인트 중 평균 throughput 비율은 0.632×(median 0.634×)이며, D=128인 경우에만 0.947×에 근접한다.

또한, TiledAttention는 eager attention 및 PyTorch math SDPA와 비교했을 때 각각 14.36×, 28.15×의 평균 속도 향상을 보이며, 이는 연구·프로토타이핑 단계에서 큰 가치를 제공한다. 프로파일링(Nsight Compute/Systems) 결과는 장기 시퀀스에서 공유·전역 메모리 압력이 높아 스테이징 깊이와 공유 메모리 레이아웃이 성능을 좌우함을 확인한다. 반면 짧은 시퀀스에서는 CTA 수와 워프 레벨 병렬도가 제한 요인으로 작용한다.

논문의 한계로는 현재 전방 연산만 구현했으며, 역전파와 KV‑cache, fused epilogue 등 실사용에 필요한 기능은 아직 미구현이다. 또한 A100 계열에 최적화돼 있어 다른 GPU 아키텍처에서는 추가 튜닝이 필요하다. 향후 연구 방향으로는 역전파 구현, 다양한 GPU 아키텍처 포팅, 자동 튜닝 파이프라인 구축 등을 제시한다.

결론적으로 TiledAttention는 파이썬 수준에서 타일 스케줄을 자유롭게 조정할 수 있는 실용적인 연구용 커널을 제공한다. 최고 성능을 목표로 하는 프로덕션 환경에서는 여전히 fused SDPA가 우세하지만, 커스텀 마스크·레이아웃 실험이나 HPC 환경에서 반복적인 커널 튜닝이 필요한 경우 TiledAttention가 높은 생산성을 보장한다.