플래시 어텐션

트랜스포머의 어텐션 계산을 결과는 그대로 두고 GPU 메모리를 덜 오가게 바꾼 알고리즘으로, 모델은 그대로 두고 켜기만 하면 긴 입력을 더 적은 메모리로 더 빠르게 처리한다. 근사 어텐션과 달리 결과가 똑같아서 학습, 파인튜닝, 추론 서버에서 거의 기본처럼 쓰인다.

느린 원인은 계산이 아니라 메모리 왕복

어텐션은 모든 토큰 쌍의 점수를 담은 큰 행렬을 만든다. 일반 구현은 이 행렬을 GPU의 큰 메모리(HBM)에 썼다가 다시 읽는다. 이 왕복이 계산보다 오래 걸리고, 입력이 길어지면 메모리 사용량이 제곱으로 늘어난다.

플래시 어텐션은 입력을 작은 블록으로 쪼개 GPU 칩 안의 빠른 메모리(SRAM)에서 한 번에 처리하고, 큰 행렬은 아예 저장하지 않는다. 학습할 때 필요한 값도 저장하지 않고 다시 계산한다. 계산량은 비슷하지만 메모리 왕복이 줄어서 빨라진다.

Transformers는 옵션 한 줄, 서버는 대개 자동

Hugging Face Transformers에서는 flash-attn 패키지를 설치한 뒤, 모델을 불러올 때 attn_implementation="flash_attention_2"를 넘기면 된다. 이때 모델은 fp16이나 bf16으로 불러와야 한다.

vLLM 같은 추론 서버는 지원하는 GPU에서 이런 최적화 커널을 기본으로 고르는 경우가 많아서 따로 할 일이 거의 없다. PyTorch의 scaled_dot_product_attention도 조건이 맞으면 내부에서 플래시 어텐션 방식의 커널을 쓴다.

켜기 전에 확인할 것

공식 구현은 주로 비교적 최신 세대(Ampere 이후)의 NVIDIA GPU를 지원한다. 오래된 GPU나 CPU에서는 쓸 수 없어서 다른 구현으로 대신해야 한다. 설치할 때 CUDA나 PyTorch 버전이 맞지 않으면 빌드가 실패하는 일이 잦으니, 미리 빌드된 패키지를 쓰는 편이 낫다.

계산량 자체가 줄지는 않는다. 그래서 입력이 짧으면 효과가 작다. 추론할 때 쌓이는 KV 캐시 메모리도 그대로라서, 긴 컨텍스트를 다루려면 다른 최적화도 함께 써야 한다.

버전이 올라가며 더 빨라졌다

FlashAttention-2는 GPU에 작업을 나누는 방식을 개선해서 첫 버전보다 더 빨라졌다. 이후 더 새로운 GPU에 맞춘 버전도 나왔다. 지금 새로 쓴다면 쓰는 라이브러리가 지원하는 최신 구현을 고르면 된다.

긴 입력을 다루는 곳이면 일단 켜 둔다

보통은 직접 고르기보다 라이브러리나 서버가 기본으로 켜 두는 부품으로 쓴다. 메모리가 부족해 입력 길이나 배치 크기를 줄여야 할 때 가장 먼저 확인하는 옵션이기도 하다.