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

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

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

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