FlashAttention 핵심 해독을 위한 펜앤페이퍼 수학

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

이 게시물은 FlashAttention 구현의 핵심 메모리·연산 흐름을 필기 다이어그램으로 정리한 자료이다. 전체 Q/K/V 텐서는 대용량의 글로벌 HBM에 보관하고, 처리에 필요한 활성 서브블록만 고속의 로컬 SRAM으로 청크 단위로 옮겨 연산을 수행하는 구조를 중심으로 구성되어 있다. 다이어그램에는 청크 로드 루프, 연산 엔진에서의 Q·K 내적과 온라인 소프트맥스 트래킹, 결과 누적 및 쓰기 복귀 루프가 명시되어 있어 메모리 이동과 연산 중간 상태 관리 방식이 드러난다.

로컬 SRAM으로의 블록 타일링은 입력 청크를 로드한 뒤 Compute Engine에서 s_ij = Q_i @ K_j^T를 계산하고, 온라인으로 소프트맥스 정규화 계수를 추적하며 V의 가중합을 누적하는 처리 파이프라인을 따른다. 이 방식은 전체 토큰 쌍을 한꺼번에 메모리에 올리지 않음으로써 HBM 대역폭과 로컬 메모리 한계를 조율하고, 스트리밍으로 블록을 순차 처리해 메모리 사용량을 제어한다. 다이어그램의 요소들은 구현 관점에서 어느 단계에서 상태를 유지해야 하는지와 데이터 이동 비용의 트레이드오프를 명확하게 제시한다.

이 접근 방식은 메모리 병목을 완화해 대규모 컨텍스트에서 Attention 연산의 실효 처리량을 높일 수 있다는 실무적 시사점을 제공한다. 다만 이미지 자체가 손글씨 다이어그램이므로 세부 파라미터나 성능 수치가 명시되어 있지 않아 바로 재현 가능한 구현 가이드는 포함되지 않는다. 따라서 이 자료는 설계 아이디어와 데이터 흐름의 개념 증거로 유용하며, 실제 구현을 위해서는 추가적인 수치와 코드화 작업이 필요하다.

섹션별 상세

대규모 어텐션 연산에서 발생하는 메모리 병목 문제를 해결하기 위해 전체 Q/K/V 텐서를 느린 대용량 HBM에 유지하고 필요한 서브블록만 빠른 SRAM으로 옮겨 처리하는 전략이 제시되어 있다. 구체적으로 입력 토큰에서 생성된 Q_i와 K_j, V_j를 블록 단위로 HBM에서 로드한 뒤 로컬 SRAM에 적재하는 로드-청크 루프가 운영된다. 다이어그램은 이러한 데이터 이동이 어떻게 연산 타이밍과 메모리 사용량을 절감하는지 구조적으로 나타내고 있으며, 이 접근은 메모리 대역폭 제약을 완화해 대규모 시퀀스 처리에서 실효 처리량을 개선할 수 있다는 의미를 가진다.

손글씨 다이어그램은 글로벌 HBM과 로컬 SRAM 간의 데이터 흐름, 블록 타일링, Compute Engine의 연산 흐름 및 온라인 소프트맥스 추적을 도식화하고 있다. — Diagram이미지는 HBM에 저장된 전체 Q/K/V에서 필요한 서브블록을 로컬 SRAM으로 스트리밍하는 로드 루프와, SRAM에서 실행되는 s_ij 계산 및 온라인 소프트맥스 트래킹, 마지막으로 결과를 HBM으로 쓰기 복귀하는 흐름을 시각적으로 연결한다. 이 구조적 다이어그램은 블록 크기와 데이터 이동 빈도가 성능·메모리 트레이드오프에 어떤 영향을 미치는지, 그리고 온라인 추적을 통해 정규화 상태를 유지하면서 부분 합을 누적하는 방식이 가능함을 근거로 보여준다.

로컬 SRAM에 적재된 활성 서브블록은 Compute Engine으로 전달되어 s_ij = Q_i @ K_j^T 연산이 수행되고, 이후 Softmax 정규화와 V 가중합 누적 단계로 연결된다. 이 과정은 입력 청크를 순차적으로 스트리밍하면서 각 블록마다 내적을 계산하고 소프트맥스 트래커로 정규화 계수를 온라인으로 관리하는 일련의 단계로 동작한다. 다이어그램에 적힌 'online tracking' 표기는 중간 정규화 상수와 부분 합을 유지해 전체 행렬을 한 번에 처리하지 않고도 정규화된 출력을 얻는 방법을 가리키며, 이 방식은 SRAM 용량 한계 내에서 정밀도를 유지하면서 연산을 분할할 수 있다는 실무적 근거를 제공한다.

연산 결과는 로컬에서 누적된 부분 출력으로부터 최종 O_t를 구성한 뒤 HBM으로 쓰기 복귀하는 루프를 통해 합쳐진다. 이 쓰기 단계는 각 청크의 계산이 끝난 뒤 결과를 HBM에 업데이트하는 형태로 동작하며, 다이어그램은 쓰기 복귀가 파이프라인의 마지막 단계로서 데이터 일관성 유지와 전체 출력 조합에 중요한 역할을 한다고 나타낸다. 따라서 블록 단위 스트리밍과 쓰기 누적 전략은 메모리 이동을 최소화하면서도 최종 출력 조합을 보장하므로 대규모 모델 추론에서 효율적인 배치 처리를 가능하게 한다는 의미를 가진다.

다이어그램 전반은 하드웨어 계층 특성에 따른 설계 트레이드오프를 강조하고 있으며, 고속의 로컬 SRAM을 통해 연산 레이턴시를 낮추는 대신 HBM과의 빈번한 데이터 이동을 관리해야 하는 현실적 제약을 드러낸다. 구체적으로 로컬에 너무 큰 블록을 올리면 SRAM 한계를 초과하고 너무 작은 블록은 데이터 이동 오버헤드를 키우는 상충 관계가 존재한다는 점이 암시되어 있다. 이 구조적 이해는 실제 구현에서 블록 크기와 스트리밍 빈도, 온라인 소프트맥스 추적 방식의 조정이 성능과 자원 사용에 직접적인 영향을 미친다는 실무적 결론으로 이어진다.

용어 해설

어텐션 메커니즘(Attention Mechanism): — 입력 토큰 쌍 간 연관성을 가중치로 계산해 응답을 생성하는 연산 방식으로, Query와 Key의 내적을 통해 유사도를 계산하고 Softmax로 정규화한 뒤 Value와 가중합하여 최종 출력을 만든다. 이 과정에서 연산 복잡도와 메모리 사용량이 토큰 수의 제곱에 비례하기 때문에 대규모 문맥 처리를 위해 연산·메모리 최적화가 중요하다.
하이밴드위스 메모리(HBM): — 고대역폭 메모리로서 대용량 데이터를 빠르게 주고받는 데 유리하지만 지연과 에너지 관점에서 로컬 SRAM보다 불리한 특성이 있어 대규모 텐서 저장에 주로 사용된다. FlashAttention 구조에서는 전체 Q/K/V 텐서를 HBM에 두고 작은 청크만 SRAM으로 옮겨 처리하는 방식이 핵심 역할을 한다.
로컬 SRAM(SRAM): — 작고 고속의 온칩 메모리로서 반복 연산에서 높은 대역폭과 낮은 지연을 제공하기 때문에 짧은 길이의 활성 서브블록을 유지하며 연산을 수행하는 임시 저장소로 활용된다. FlashAttention 계열 구현에서는 로컬 SRAM에 청크 단위의 Q/K/V를 올려 온라인 소프트맥스 추적과 부분 합산을 수행한다.
블록 타일링(Block Tiling): — 큰 행렬 연산을 메모리 친화적인 작은 블록으로 분할해 로컬 메모리에 순차적으로 로드하고 계산한 뒤 결과를 누적하여 최종 출력을 얻는 기법으로, 데이터 이동을 줄이고 캐시 활용도를 높여 처리량과 메모리 효율을 개선한다. FlashAttention 다이어그램에서는 HBM에서 SRAM으로 청크를 스트리밍하는 방식이 이 개념을 따른다.

FlashAttention 핵심 해독을 위한 펜앤페이퍼 수학

TL;DR

섹션별 상세

용어 해설

관련 토론

댓글

관련 기사

FlashAttention이 거대한 어텐션 행렬을 메모리에 저장하지 않고 계산하는 원리.

원리 기반 LLM 추론 성능 모델링

Group Query Attention(GQA)의 작동 원리와 설계 의도

Grouped Query Attention의 작동 원리와 KV cache 최적화

하드웨어 친화적 LLM 설계를 위한 모델-하드웨어 공동 설계