로컬 추론 중 Flash Attention은 메모리 트래픽을 어떻게 줄이나요?

Flash Attention은 정확한 어텐션을 타일링하여 중간 점수를 반복해서 GPU HBM에 쓰고 읽는 대신 칩 내에서 처리하도록 함으로써 메모리 트래픽을 줄입니다.

로컬 AI 서버에서는 프롬프트 길이가 늘어날수록 표준 어텐션이 점점 더 큰 중간 작업을 생성하기 때문에 그 효과가 더욱 중요해집니다. Flash Attention은 모델의 지식을 변경하거나 KV 캐시를 제거하지 않습니다. 대신 어텐션 계산을 재구성하여 GPU 메모리 계층에서 가장 느린 단계로 이동하는 데이터의 양을 줄입니다.

표준 어텐션은 큰 중간 텐서를 HBM으로 이동시킵니다

어텐션은 쿼리, 키, 값 텐서에서 시작하여 점수를 계산하고 정규화한 뒤 값과 결합해야 합니다. 단순한 구현에서는 이러한 단계 사이에 큰 중간 행렬을 고대역폭 메모리에 저장합니다.

FlashAttention 원 논문은 어텐션 성능이 연산량뿐만 아니라 HBM과 칩 내 SRAM 사이의 읽기 및 쓰기에 의해서도 제한된다고 설명합니다. 점수 행렬은 토큰 쌍의 수에 따라 커지므로 긴 시퀀스에서는 이러한 트래픽이 더욱 증가합니다.

따라서 홈 GPU는 특히 컨텍스트 길이가 늘어날 때 유용한 행렬 연산보다 데이터 이동에 상당한 시간을 소비할 수 있습니다.

Flash Attention은 문제를 칩 내 작업 집합으로 타일링합니다

Flash Attention은 빠른 칩 내 메모리에 맞는 크기로 쿼리, 키, 값 블록을 나누어 처리합니다. 각 타일은 중간 데이터가 칩 밖으로 나가기 전에 여러 어텐션 단계를 수행합니다.

FlashAttention-3은 이 계열의 핵심 접근 방식으로 중간 전역 메모리 트래픽을 피하는 타일링 전략을 설명합니다. 이 알고리즘은 지원되는 방식에서 여전히 정확한 어텐션을 생성합니다.

타일 크기는 하드웨어 리소스와 커널 구현에 따라 달라집니다. 중요한 메커니즘은 데이터 지역성, 즉 로드한 값을 연산 유닛과 가까이 있을 때 재사용하는 것입니다.

온라인 Softmax는 전체 점수 행렬 저장을 피합니다

일반적인 softmax는 정규화하기 전에 한 행의 모든 점수가 필요해 보이므로, 구현 과정에서 큰 행렬을 저장해야 합니다. 반면 Flash Attention은 타일 전체에서 정규화 통계를 누적하여 유지합니다.

Triton의 융합 어텐션 튜토리얼은 블록 단위 softmax 상태를 사용하는 융합 어텐션을 보여 줍니다. 이를 통해 부분 점수 블록을 사용한 뒤 폐기할 수 있으며, 모든 중간값을 HBM에 다시 쓸 필요가 없습니다.

여기서 메모리 트래픽 감소와 임시 메모리 감소가 맞물립니다. 알고리즘이 계산을 재배치하므로 큰 중간 객체가 전역 메모리에 존재할 필요가 없습니다.

Featured

ZimaCube 2 개인용 클라우드 홈 NAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

커널 융합은 더 많은 어텐션 과정을 하나의 실행 경로에 유지합니다

점수 계산, 스케일링, 마스킹, 정규화, 값 집계는 전체 중간 결과를 전역 메모리를 통해 교환하는 여러 개의 커널을 따로 실행하는 대신 하나의 조정된 커널로 융합할 수 있습니다.

PyTorch의 scaled dot-product attention 인터페이스는 장치, 데이터형, 형태 및 요청된 기능이 지원되는 경우 융합 어텐션 구현을 선택할 수 있습니다.

그 결과 프롬프트 처리 속도가 빨라지고 최대 임시 메모리 사용량이 줄어들 수 있지만, 실제 향상 폭은 런타임이 Flash와 호환되는 커널을 실제로 선택하는지에 따라 달라집니다.

IO 감소는 KV 캐시 제거와 같지 않습니다

자동 회귀 추론 중에는 모델이나 런타임이 별도의 제거, 압축, 윈도잉 또는 오프로딩 방식을 사용하지 않는 한, 이전에 생성된 키와 값을 이후 토큰에서 사용할 수 있도록 계속 유지해야 합니다.

AMD의 ROCm 지침은 Flash Attention이 칩 내 저장 공간과 HBM 사이의 메모리 지역성을 개선한다고 설명합니다. 이 최적화는 활성 상태를 어텐션이 읽고 계산하는 방식을 개선하는 것이지, 과거 KV 상태의 존재 여부를 바꾸는 것이 아닙니다.

이 차이를 이해하면 흔히 발생하는 계획 오류를 피할 수 있습니다. Flash Attention은 작업 메모리와 IO를 줄일 수 있지만, 보존되는 KV 캐시가 계속 증가하기 때문에 긴 컨텍스트를 처리하는 서버의 용량이 여전히 부족해질 수 있습니다.

하드웨어와 형태 지원 여부가 빠른 경로 사용을 결정합니다

Flash 커널은 특수하게 설계됩니다. 헤드 차원, 데이터형, 어텐션 마스크, GPU 아키텍처, 런타임 버전 및 기타 기능에 따라 다른 구현으로 대체될 수 있습니다.

ZimaSpace의 모델 파라미터보다 커지는 어텐션 메모리 관련 글에서는 용량 측면을 다룹니다. Flash Attention은 실행 경로 최적화로서 별도로 측정해야 합니다.

짧은 프롬프트와 긴 프롬프트를 모두 벤치마크하고, 선택된 백엔드를 확인한 뒤 프리필 지연 시간, 최대 할당 메모리, 메모리 대역폭 및 출력 토큰 속도를 기록하세요. ZimaSpace의 로컬 AI 하드웨어 가이드는 커널 최적화로 대체할 수 없는 전반적인 용량 기준을 제공합니다.

NVIDIA의 Transformer Engine은 더 폭넓은 커널 선택 스택의 일부로 최적화된 어텐션 백엔드와 하드웨어에 따른 실행을 다룹니다.