Design of Tensor Core Archiecture for Softmax Acceleration

Large Language Model(LLM)은 점점 더 많은 token을 처리할 수 있는 방향으로 발전하고 있다. LLM은 Multi-Head Attention을 주요 연산으로 활용하는데, 이에 따라 token 개수에 비례하여 행렬곱 및 softmax의 연산량이 증가한다. GPU는 행렬곱 가속을 위해 Tensor Core를 지원하지만, softmax 연산은 범용 연산기인 CUDA Core에서 수행된다. Softmax 연산은 지수 함수 계산과 모든 입력값의 총합 연산을 필요로 하며, CUDA Core에서 이를 수행할 경우 두 가지 주요 한계가 존재한다. CUDA Core는 SIMT 구조를 기반으로 여러 thread를 동시에 실행하지만, softmax 연산의 합산 과정에서는 thread 간 데이터 공유가 필요하기 때문에 warp 단위로 동기화가 요구된다. 둘째, CUDA Core는 Tensor Core 대비 연산 처리량이 낮아 softmax 연산의 성능이 제한될 수 있다. 본 논문은 이러한 한계를 극복하기 위해, 높은 연산 처리량을 제공하고 warp 단위로 연산을 수행하는 Tensor Core에 지수 함수 연산을 위한 Lookup Table(LUT) 및 나눗셈 연산을 위한 reciprocal unit을 추가한 Soft-Tensor Core 구조를 제안한다. 이 기법을 통해 기존 행렬곱 가속뿐만 아니라 softmax 연산도 가속하여, GPU에서의 LLM 연산 성능을 더욱 향상 시킬 수 있을 것으로 기대된다.