안녕하세요, 옵트에이아이 박성재입니다.
이 글은 GPU용으로 제안된 FlashDecoding++의 핵심 아이디어를 정리하는 이론편입니다. 디바이스 이식의 구체적인 이야기는 다음 편으로 미루고, 여기서는 "긴 KV 캐시에서 softmax를 빠르고 메모리 효율적으로 계산한다"는 목표가 어떤 원리로 달성되는지에 집중합니다.
1. 왜 온디바이스에서 FlashDecoding++인가
디코딩 단계의 attention은 매 토큰마다 동일한 연산을 반복합니다. Query 길이는 1로 고정된 반면 Key/Value는 KV 캐시를 따라 계속 늘어나므로, 스코어 계산과 softmax는 컨텍스트 길이 K에 비례하는 O(K) 비용을 매 스텝 지불합니다. GPU에서든 NPU에서든 이 구간이 길어질수록 연산보다 데이터를 읽어오는 시간이 지배적으로 변합니다.
*디코딩에서는 길이가 1인 Query가 계속 길어지는 Key/Value 전체를 순차적으로 읽어야 한다*
*Flash-Decoding은 KV 캐시를 여러 split으로 쪼개 병렬 처리한다. 이때 각 split의 부분 결과를 합치려면 블록 간 softmax 통계를 맞추는 과정이 필요해진다 (출처: PyTorch Blog — Flash-Decoding for long-context inference)*
FlashDecoding++는 이 문제를 세 방향에서 공략합니다. 우리의 관심사는 이 중 온디바이스 attention에 가장 직접적으로 기여하는 비동기 softmax입니다.
*FlashDecoding++ 추론 데이터플로우 개요. (a) prefill은 주로 GEMM, decode는 주로 GEMV/Flat GEMM 연산이다. (b) 통합 최댓값 기반 비동기 softmax, (c) 이중 버퍼링 Flat GEMM 최적화, (d) 휴리스틱 데이터플로우 — FlashDecoding++가 재설계하는 세 가지 병목과 각 해법 (출처: Hong et al., arXiv:2311.01282)*
2. 세 가지 병목
(1) 동기화된 부분 Softmax. 긴 시퀀스를 병렬 처리하기 위해 softmax를 여러 블록으로 나누면, 전체 확률을 구하기 위해 각 블록의 최댓값과 합을 서로 공유·갱신하는 동기화가 필요합니다. 이 과정만으로 attention 연산의 약 20% 오버헤드가 발생합니다.
(2) Flat GEMM의 낮은 활용도. 디코딩 단계의 행렬은 M이 매우 작은 납작한(flat) 형태를 띱니다. 기존 라이브러리는 하드웨어 성능을 끌어내기 위해 0으로 패딩해 크기를 키우는데, 이 과정에서 유효하지 않은 연산이 절반 이상을 차지하기도 합니다.
(3) 정적 데이터플로우로 인한 손실. 입력 형태나 하드웨어 구성에 따라 최적의 커널은 달라지지만, 단일 정적 경로로만 연산하면 다양한 GEMM 형태에 유연하게 대응하지 못해 성능이 크게 손실됩니다.
3. 핵심 해법: 통합 최댓값을 이용한 비동기 Softmax
FlashDecoding++는 통계적으로 안전한 통합 최댓값(unified max value) φ를 미리 설정해 각 블록이 독립적으로 계산하도록 만듭니다. 블록 간 동기화가 사라지고 마지막에 단순 합만 남으므로 프리필 1.18배, 디코딩 1.14배의 속도 향상을 얻습니다.
일반적인 online softmax는 실행 중 최댓값을 추적합니다.
softmax(x_i) = exp(x_i - max_j x_j) / Σ_j exp(x_j - max_j x_j)통합 최댓값 방식은 running max를 사전에 고정한 상수 φ로 대체합니다.
softmax(x_i) = exp(x_i - φ) / Σ_j exp(x_j - φ)블록마다 최댓값을 다시 찾을 필요가 없어집니다. 다만 어떤 x_i가 φ 기준을 넘어 오버플로우 위험이 생기면, 해당 블록의 비동기 계산을 중단하고 기존 online softmax로 되돌아가 안전하게 재계산(recomputation)합니다. 즉 φ는 분포를 사전에 관찰해 설정하는 값이며, 분포가 φ를 벗어날 만큼 크면 안전 경로로 fallback하는 구조입니다.
4. 나머지 두 해법
이중 버퍼링을 적용한 Flat GEMM 최적화. 패딩을 최소화하고 이중 버퍼링을 도입해, 데이터를 읽어오는 동안 연산을 수행하는 latency hiding으로 메모리 병목을 우회합니다.
하드웨어 자원 적응형 휴리스틱 데이터플로우. LLM의 선형 레이어 연산은 결국 네 가지(Q/K/V 생성, O projection, 두 FFN)로 귀결되며, 가중치 차원 K·N은 상수이고 오직 M(시퀀스 길이·배치)만 변합니다. 오프라인에서 만들어 둔 룩업 테이블로 M 구간에 따라 커널을 골라, 정적 데이터플로우 대비 손실을 회복합니다.
세 해법 중 온디바이스 attention에 가장 직접 기여하는 것은 비동기 softmax입니다. 다만 이 알고리즘을 실제 하드웨어로 내리면 양자화와 유닛 할당이 결과를 크게 흔듭니다. 다음 편에서 그 이야기를 다룹니다.
참고 문헌
- Hong, K., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
- Dao, T., et al. (2023). Flash-Decoding for long-context inference. PyTorch Blog.
- Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS.