Improving Attention Parallelism via Delayed Value Generation and Softmax Offloading

최근 대규모 언어 모델이 다양한 분야에 활용되면서, 연산 성능 최적화에 대한 연구가 주목받고 있다. 이러한 모델은 Attention Layer, Feed-Forward Layer, 그리고 Layer Normalization으로 구성되며, 이 중 Attention Layer가 전체 실행 시간의 약 48%를 차지한다. Attention Layer는 다섯 단계의 순차적인 연산으로 이루어지며, 다음과 같다. 먼저, Query, Key, 그리고 Value를 생성하는 단계가 있으며 생성한 Query와 Key를 이용하여 Score를 계산한다. 그리고 Score 값은 Softmax 연산을 통해 다시 Probablity를 생성하고 이를 다시 첫 단계에서 생성한 Value와 곱하여 Attention Matrix를 만든다. 그 후 최종적으로 가중치를 곱하여 출력값을 생성하는 순차적인 다섯 단계를 가진다. 그러나 초기 단계에서 생성된 Value는 네 번째 단계까지 사용되지 않아 비효율이 발생한다. 본 연구는 이러한 연산 순서의 비효율성을 분석하고, 이를 개선하기 위한 최적화 기법을 제안한다. 첫째, Value 생성을 초기 단계에서 분리하여 Attention Matrix 생성 직전에 수행함으로써 캐시 활용도를 높이고 첫 단계의 연산 시간을 줄인다. 둘째, 연산량이 적은 Softmax는 Data Processing Unit (DPU)에 오프로딩하고, 동시에 GPU에서는 Value 생성을 병렬로 수행하여 처리량을 향상시킨다. 이를 통해 Attention Layer의 연산 순서를 최적화하고 병렬성을 확보하여 전체 성능을 개선한다.