Qrita: 피벗 기반 절단 및 선택을 활용한 GPU용 고성능 Top-k 및 Top-p 알고리즘

Qrita: High-performance Top-k and Top-p Algorithm for GPUs using Pivot-based Truncation and Selection

Top-k와 Top-p는 거대 언어 모델(LLM) 샘플링에서 주로 사용되는 절단(truncation) 연산자입니다. 널리 사용되고 있음에도 불구하고, 방대한 어휘 집합에 대해 이를 효율적으로 구현하는 것은 여전히 중요한 과제입니다. 기존의 접근 방식들은 주로 정렬(sorting)에 의존하여 GPU 상에서 상당한 연산 및 메모리 오버헤드를 발생시키거나, 알고리즘의 출력을 변형시키는 확률적 방식에 의존합니다. 본 연구에서는 피벗 기반 선택 전략을 활용한 효율적인 Top-k 및 Top-p 알고리즘인 Qrita를 제안합니다. 그래프 신경망의 노드 선택을 위해 피벗 기반 탐색을 사용하는 RTop-k에 기초하여, Qrita는 다음 두 가지 핵심 기술을 통해 피벗 기반 탐색의 개념을 Top-k와 Top-p 모두로 확장합니다. 1. 대상 요소들의 탐색 공간을 대폭 줄여주는 가우시안 기반 시그마 절단(Gaussian-based sigma-truncation), 2. 피벗 탐색 반복 횟수를 절반으로 줄이고 결정론적 출력을 보장하는 중복 처리가 포함된 4진 피벗 탐색(Quaternary pivot search)입니다. 우리는 널리 사용되는 GPU 프로그래밍 언어인 Triton을 사용하여 Qrita의 전체 구현을 제공합니다. vLLM, SGLang, Flashinfer와 같은 고성능 LLM 실행 엔진의 Top-k 및 Top-p 커널과 비교 평가한 결과, Qrita는 정렬 기반 알고리즘과 동일한 출력을 제공하면서도 최대 2배의 처리량(throughput)을 달성하고 메모리 사용량을 절반으로 줄인 것으로 나타났습니다.

Top-k and Top-p are the dominant truncation operators in the sampling of large language models. Despite their widespread use, implementing them efficiently over large vocabularies remains a significant challenge. Existing approaches often rely on sorting, which incur significant computation and memory overhead on GPUs, or stochastic approaches, which alter the algorithm output. In this work, we propose Qrita, an efficient Top-k and Top-p algorithm based on a pivot-based selection strategy. Based on RTop-k, which uses a pivot-based search for node selection in graph neural networks, Qrita extends the concept of pivot-based search to both Top-k and Top-p with two key techniques: 1. Gaussian-based sigma-truncation, which greatly reduces the search space of the target elements, and 2. Quaternary pivot search with duplication handling, which halves the pivot search iteration and guarantees deterministic output. We provide the full implementation of Qrita using Triton, a popular GPU programming language. Our evaluation of Qrita against the Top-k and Top-p kernels of high performance LLM execution engines such as vLLM, SGLang, and Flashinfer show that Qrita achieves up to 2 times throughput and half memory use while providing the same output to the the sorting-based algorithms.