기술 요약: Qrita - 피벗 기반 절단 및 선택을 활용한 고성능 Top-k 및 Top-p

문제 제기

대규모 어휘 (종종 10 만 개 이상의 토큰을 초과) 를 가진 대규모 언어 모델 (LLM) 에 대한 Top-k 및 Top-p 샘플링 알고리즘의 효율적인 구현은 GPU 에서 여전히 중요한 과제로 남아 있습니다. 기존 접근 방식은 계산 효율성과 출력 정확성 사이의 트레이드오프에 직면해 있습니다:

• 정렬 기반 방법: vLLM 과 같은 엔진에서 사용되며, 정확하고 결정론적인 결과를 제공하지만 전체 어휘 정렬과 스트라이드 메모리 접근으로 인해 높은 계산 및 메모리 오버헤드를 초래합니다.
• 확률적 방법: FlashInfer 및 SGLang 과 같은 라이브러리는 정렬을 피하기 위해 거절 샘플링을 사용하여 높은 처리량을 제공합니다. 그러나 이들은 전체 후보 집합이 아닌 단일 샘플링된 토큰만 반환하므로, 사적 디코딩 검증 및 강화 학습과 같이 결정론적 후보 집합이 필요한 작업과는 호환되지 않습니다.
• 통계적 컷오프 방법: Top-nσ 및 Min-p 와 같은 접근 방식은 평균 및 표준 편차와 같은 통계적 임계값을 사용하여 정렬을 제거합니다. 효율적이지만 샘플링 분포와 출력을 변경하여 생성 품질 및 일관성에 관한 논쟁을 초래합니다.

정확한 Top-k 및 Top-p 의 의미를 훼손하지 않으면서 메모리 효율성과 엄격한 결정론적 행동 (전체 후보 집합 반환) 을 제공하는 알고리즘에 대한 중요한 격차가 존재합니다.

방법론

저자들은 피벗 기반 절단 및 선택을 기반으로 한 GPU 용 고성능 Top-k 및 Top-p 알고리즘인 Qrita를 제안합니다. Qrita 는 그래프 신경망 (GNN) 을 위해 원래 설계된 RTop-k 의 피벗 기반 전략을 확장하여 대규모 LLM 어휘를 처리하면서도 결정론을 보장합니다. 이 방법은 두 가지 핵심 기술에 의존합니다:

1. 가우스 기반 $\sigma$-절단

LLM 로지트가 일반적으로 "노이즈" 배경에 비해 고확률 토큰이 이상치를 형성하는 준 가우스 분포를 따른다는 점을 인식하여, Qrita 는 사전 필터링 단계를 도입합니다.

• 메커니즘: 알고리즘은 로지트 분포를 $Z \sim \mathcal{N}(\mu, \sigma^2)$로 근사합니다. 샘플 블록에서 평균 ($\mu$) 과 표준 편차 ($\sigma$) 를 계산하고 절단 임계값 $t = \mu + \delta \cdot \sigma$를 계산합니다.
• 룩업 테이블: 계수 $\delta$는 정규 분위수 함수에서 유도된 사전 계산된 룩업 테이블을 통해 결정되며, 목표 $k$ 또는 $p$ 값에 맞게 조정됩니다.
• 결과: 이 단계는 어휘의 거대한 "노이즈" 영역을 단일 패스로 절단합니다. 절단된 집합 크기 $|Z_t|$가 목표 $k$를 초과하는 경우 ("절단 히트"), 검색은 이 더 작은 하위 집합에서만 진행됩니다. 그렇지 않으면 전체 로지트 집합으로 되돌아갑니다. 이는 후속 피벗 검색을 위한 검색 공간을 크게 줄입니다.

2. 중복 처리를 위한 4 분 피벗 검색

정확한 피벗 값 $\tau$(정확히 $k$개의 요소가 $\ge \tau$인 값) 를 효율적이고 결정론적으로 찾기 위해:

• 4 분 검색: 이진 검색 대신 Qrita 는 각 반복에서 현재 검색 범위의 25%, 50%, 75% 백분위수에 해당하는 세 개의 피벗을 평가합니다. 이는 이진 검색에 비해 검색 반복 횟수를 절반으로 줄입니다.
• 중복 처리: 표준 피벗 검색은 중복 로지트 값이 존재할 때 (저정밀도 또는 긴 꼬리 분포에서 일반적) 실패하거나 비결정론적일 수 있습니다. Qrita 는 피벗보다 큰 최소 로지트 값 ($z_{mn}$) 과 그 발생 횟수 ($N_{zmn}$) 를 추적합니다.
• 종료 조건: 피벗 $\tau_n$이 발견되어 $\tau_n$보다 엄격하게 큰 요소의 수가 $k$보다 작지만, 최소 중복 값 $z_{mn}$ 이상인 요소의 수가 $k$ 이상일 때 검색이 종료됩니다. 이를 통해 알고리즘은 Top-k 제약을 충족시키기 위해 경계 값의 몇 개의 중복 인스턴스를 유지해야 하는지 정확하게 결정할 수 있으며, 재현 가능하고 결정론적인 출력을 보장합니다.

구현 세부 사항

• 프레임워크: GPU 프로그래밍 언어인 Triton을 사용하여 구현되었습니다.
• 자동 튜닝: "Tiered Triton Block Size Autotune" 전략을 사용합니다. 전체 로지트 집합과 절단된 집합 사이의 벡터 크기가 극적으로 변하기 때문에, 컴파일러가 최적의 블록 크기를 선택할 수 있도록 허용합니다 (예: 전체 집합에는 더 큰 블록, 절단된 집합에는 더 작은 블록) 하드웨어 활용도를 극대화합니다.
• 파이프라인: 알고리즘은 배치를 병렬로 처리하며, 가우스 절단, 4 분 검색, 그리고 최종 출력을 생성하기 위한 중복 처리 마스크 생성을 순차적으로 적용합니다.

주요 기여

1. 결정론적 고성능 샘플링: Qrita 는 FlashInfer/SGLang 과 달리 전체 Top-k/Top-p 후보 집합을 제공하면서도 정렬 기반 알고리즘의 정확한 의미를 유지하여 사적 디코딩 및 RL 에 적합합니다.
2. 검색 공간 축소: 로지트의 통계적 특성을 활용하여 가우스 $\sigma$-절단은 검색 공간을 크게 줄여 전체 어휘 정렬을 방지합니다.
3. 강건한 피벗 검색: 명시적인 중복 처리와 결합된 4 분 검색의 도입은 결정론적 행동을 보장하고 반복 횟수를 줄여 대규모 어휘에서 표준 피벗 기반 방법의 비효율성을 해결합니다.
4. 통합: Qrita 는 vLLM의 GPU 실행 경로에 대한 기본 Top-k 및 Top-p 샘플러로 구현되었습니다.

결과

저자들은 다양한 모델 (Llama, Qwen, GPT-OSS, Gemma) 을 사용하여 NVIDIA H100, RTX4090, AMD MI355X GPU 에서 Qrita 를 단순 PyTorch 정렬, vLLM 기본 샘플러, FlashInfer, SGLang 과 비교 평가했습니다.

• 처리량: Qrita 는 H100 에서 vLLM 기준 대비 최대 1.4$\times$의 엔드투엔드 서비스 처리량 개선을 달성했으며, 일부 구성에서는 RTX4090 에서 최대 2$\times$의 개선을 보였습니다. 정렬 기반 기준보다 훨씬 뛰어난 성능을 발휘하면서도 정렬 기반 알고리즘의 출력과 일치합니다.
• 메모리 사용량: Qrita 는 정렬에 필요한 큰 중간 버퍼 할당을 피함으로써 정렬 기반 방법 대비 메모리 사용량을 약 50% 줄입니다.
• 절단 연구:
  • 가우스 절단: 전체 로지트 집합을 검색하는 것보다 약 74% 의 지연 시간 감소를 제공하여 가장 큰 속도 향상을 제공합니다.
  • 4 분 검색: "절단 히트" 경우 검색 공간이 축소되어 이진 검색보다 약간 느리지만, 전체 어휘를 검색해야 하는 폴백 경우에서 약 2.3ms 의 상당한 지연 시간 감소를 제공합니다.
  • 중복 처리: 최소 오버헤드 (지연 시간 약 3.5% 증가) 를 추가하지만 정확성과 결정론에 필수적입니다.
  • 자동 튜닝: 단일 고정 블록 크기를 사용하는 것보다 실행 지연 시간을 거의 30% 줄입니다.

중요성 및 주장

이 논문은 Qrita 가 대규모 어휘 디코딩을 위한 메모리 효율성과 엄격한 결정론적 행동 사이의 격차를 해소한다고 주장합니다. 전체 어휘 정렬을 피하고 통계적 사전 필터링을 통한 피벗 기반 선택을 활용함으로써 Qrita 는 전체 후보 집합이 필요한 고급 작업 (예: 사적 디코딩) 과 호환되는 LLM 디코딩을 위한 고성능 솔루션을 제공한다고 명시합니다. 저자들은 Qrita 를 Top-k/Top-p 의 통계적 의미를 대체하는 것이 아니라, 휴리스틱 통계적 컷오프 방법이나 확률적 단일 토큰 샘플러와 달리 이를 정확히 보존하는 효율적인 구현으로 위치시킵니다. vLLM 에 대한 기본 샘플러로서의 통합은 생산 환경에서의 실용적 유용성을 강조합니다.