🤖 machine learning

FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval

FlashTrie는 비트 압축 트라이 레이아웃과 협력적 CUDA 커널을 채택하여 CPU 병목 현상을 제거함으로써 생성적 검색을 위한 제약된 빔 서치를 최적화하는 GPU 가속 시스템으로, 대규모 상업적 검색 애플리케이션에서 최대 24배의 속도 향상과 0.71%의 매출 증대를 달성했습니다.

원저자: Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao

게시일 2026-07-14

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 방금 들은 질문을 바탕으로 비밀 코드(예: "DocID: 4592")의 목록을 작성하려고 노력하는 아주 똑똑한 로봇이라고 상상해 보세요. 하지만 함정이 하나 있습니다. 당신은 오직 8억 개의 유효한 항목이 들어 있는 거대하고 사전 승인된 전화부 속에 실제로 존재하는 코드만 작성할 수 있습니다. 만약 전화부에 없는 코드를 추측하면 실패하게 됩니다.

오랫동안 로봇들은 매우 빠르고 조직적인 사서(표준 컴퓨터 칩 또는 CPU에서 실행되는)에게 모든 추측치를 확인해 달라고 요청하며 이 작업을 수행해 왔습니다. 하지만 추측 리스트가 늘어남에 따라 사서는 과부하 상태에 빠졌습니다. 전화부를 확인하는 과정이 교통 체증처럼 변해 속도를 늦추게 된 것입니다. 로봇은 자신의 추측이 허용되는지 확인하기 위해 한 단계씩 줄을 서서 기다려야 했습니다.

여기 FlashTrie가 등장합니다. Microsoft와 Nvidia의 연구진은 사서를 해고하고, 8억 개의 항목이 담긴 전화부 전체를 로봇의 초고속 고성능 메모리(GPU)로 직접 옮기기로 결정했습니다. 하지만 그들은 단순히 책을 옮긴 것이 아니라, 그 구조를 새로 만들었습니다.

"비트 패킹(Bit-Packed)" 전화부의 마법

기존의 전화부를 모든 책이 넓고 빈 공간이 많은 큰 방에 보관되어 있는 거대한 도서관이라고 생각해 보세요. FlashTrie는 이 책들을 압축합니다. 이들은 "비트 압축(bit compression)"이라는 영리한 기술을 사용하여 정보를 꽉 짜내어, 마치 여행 가방을 효율적으로 싸는 것처럼 8억 개의 키워드를 단 3.1 GB의 공간 안에 밀어 넣습니다. 이 용량은 로봇의 고속 메모리 안에 완전히 들어갈 만큼 작기 때문에, 느린 외부 하드 드라이브에서 페이지를 가져오느라 기다릴 필요가 없습니다.

협력적인 댄스

기존 시스템에서는 로봇이 추측을 하고, 사서에게 확인을 요청하고, 답변을 기다린 다음, 다시 다른 추측을 하고 이를 반복했습니다. 이는 외롭고 순차적인 과정이었습니다.

FlashTrie는 게임의 판도를 바꿉니다. 이 시스템은 **512명의 무용수(스레드)**가 완벽한 싱크로 맞춰 함께 춤을 추는 거대한 댄스 플로어와 같은 "협력적 CUDA 커널(cooperative CUDA kernel)"을 사용합니다.

확장(The Expansion): 한 사람이 하나의 추측을 확인하는 대신, 수백 명의 무용수가 동시에 수천 개의 추측을 확인합니다.
검증(The Validation): 이들은 "병렬 이진 탐색(parallel binary search)"(무언가를 찾는 매우 빠른 방법)을 사용하여 추측이 전화부와 일치하는지 확인합니다.
가지치기(The Pruning): 만약 추측이 잘못되었다면 즉시 버립니다. 만약 올바르다면 유지합니다.

모든 일이 (데이터를 가져오기 위해) 메인 컴퓨터(CPU)와 대화하기 위해 멈출 필요 없이 댄스 플로어(GPU) 위에서 일어나기 때문에 과정이 믿을 수 없을 정도로 빨라집니다.

결과: 속도와 지능

연구진은 8억 개의 키워드 라이브러리를 대상으로 테스트를 진행했습니다.

속도: 추측의 수("빔 폭", beam width)를 1,000까지 늘렸을 때, 기존 CPU 시스템은 약 **46 밀리초(ms)**가 걸렸고 리스트가 늘어남에 따라 더 느려졌습니다. 반면 FlashTrie는 시간을 3 밀리초 미만으로 유지했습니다(구체적으로 평균은 1.91 ms였고, 가장 느린 상위 1%도 3.31 ms 미만이었습니다).
성능 향상: 이는 FlashTrie가 고도로 최적화된 CPU 버전보다 최대 24배 더 빠름을 의미합니다.
품질: 결정적으로, 빨라졌다고 해서 정확도가 떨어지는 것은 아니었습니다. FlashTrie는 느린 시스템만큼이나 많은 정확한 코드를 찾아냈습니다. 사실, FlashTrie는 매우 빠르기 때문에 시간 제한을 어기지 않으면서도 기존의 200개 대신 600개의 추측을 확인할 수 있었습니다.

실질적인 영향: 돈 테스트

연구진은 단순히 실험실에서만 테스트하지 않았습니다. 그들은 실제 상용 검색 엔진(당신이 인터넷에서 무언가를 찾을 때 사용하는 것과 같은 종류)에서 FlashTrie를 테스트했습니다. 그들은 여러 국가에서 16일 동안 실험을 수행했습니다.

검색 엔진이 더 많은 추측을 확인하도록 FlashTrie를 사용함으로써, 더 나은 광고를 보여주었습니다.
이는 광고 수익(매출)의 0.71% 증가로 이어졌습니다.
또한 영어 쿼리에 대해서는 클릭률이 0.17%, 비영어 쿼리에 대해서는 0.20% 증가했습니다.
중요한 점은, 광고의 품질이 떨어지지 않았다는 것입니다. "결함률(defect rate, 잘못된 광고가 노출되는 비율)"은 동일하게 유지되었습니다.

FlashTrie가 아닌 것

이 논문이 여기서 작동하지 않거나 필요하지 않다고 명시한 부분들을 주목하는 것이 중요합니다. 연구진은 기존 방식의 "포인터 기반(pointer-based)" 라이브러리를 GPU에서 사용하는 것은 너무 혼란을 주고 무용수들의 속도를 늦히기 때문에 명시적으로 배제했습니다. 또한, 데이터 구조를 재설계(예: "선형 조사(Linear-probe)" 방식)하지 않고 단순히 기존 시스템을 GPU로 옮기기만 하는 것은 그들의 새로운 방식보다 71배에서 209배 더 느릴 것임을 보여주었습니다. 속도 향상은 단순히 더 빠른 하드웨어를 사용했기 때문이 아니라, 전화부와 댄스를 설계하는 특정한 디자인 덕분입니다.

결론

FlashTrie는 속도와 정확성 사이에서 하나를 선택할 필요가 없다는 것을 증명합니다. "전화부"를 저장하는 방식과 "확인"이 일어나는 방식을 재설계함으로써, 그들은 느리고 순차적인 병목 현상을 번개처럼 빠른 병렬 파티로 바꾸었습니다. 이를 통해 로봇은 실시간 인터넷 검색에 필요한 엄격한 시간 제한을 준수하면서도, 더 크게 생각하고(더 많은 옵션 확인) 더 빠르게 생각할 수 있게 되었습니다. 이 시스템의 코드는 검토 과정이 끝난 후 공개될 예정이므로, 다른 이들도 이 새로운 검색 방식을 시도해 볼 수 있습니다.

기술 요약: FlashTrie

문제 정의

생성적 검색(Generative retrieval)은 문서 검색을 쿼리를 문서 식별자(docID)로 직접 매핑하는 시퀀스-투-시퀀스(sequence-to-sequence) 작업으로 재정의합니다. 이러한 패러다임은 엔드-투-엔드 모델링을 압축적으로 제공하지만, 온라인 서빙은 엄격한 지연 시간 예산(latency budget)에 의해 제약을 받습니다. 디코딩 절차, 특히 비자기회적(non-autoregressive, NAR) 디코더의 경우, 유효하지 않은 식별자를 빈번하게 생성하여 이를 필터링하기 위한 제약 기반 탐색이 필요합니다.

현재의 최첨단 방식들은 생성된 식별자가 사전 정의된 라이브러리와 일치하도록 보장하기 위해 **트라이 제약 빔 서치(trie-constrained beam search)**를 활용합니다. 그러나 기존 구현체들은 (표준 MARISA와 같은) 포인터 기반 트라이 표현에 의존하며 CPU에서 실행됩니다. 이러한 설계는 다음과 같은 문제를 야기합니다:

불규칙한 메모리 접근: 하드웨어 활용도를 저하시키고 지연 시간을 높입니다.
제한된 병렬성: 빔 너비(beam width)가 증가함에 따라, CPU 기반의 트래벌(traversal) 및 후보 검증이 서빙의 주요 병목 현상이 됩니다.
GPU에서의 비효율성: 전통적인 억제된(succinct) 트라이(MARISA와 같은)는 문자 단위 키에 최적화되어 있어 GPU 아키텍처와 잘 맞지 않으며, 워프 분기(warp divergence) 및 비결합(uncoalesced) 메모리 접근을 유발합니다.
오케스트레이션 오버헤드: 제약 강제화가 핵심 GPU 디코딩 경로 외부에서 구현되는 경우가 많아, 디코딩 단계마다 반복적인 호스트-디바이스 동기화가 필요합니다.

본 논문은 세 가지 핵심 질문을 다룹니다: 억제된 트라이를 공간 효율성을 희생하지 않으면서 토큰 단위 어휘집에 적응시킬 수 있는가? 제약된 빔 서치를 완전히 GPU 상에서 실행할 수 있는가? 이러한 시스템이 검색 품질을 저하시키지 않으면서 프로덕션 지연 시간 제약을 충족할 수 있는가?

방법론: FlashT리 (FlashTrie)

FlashTrie는 제약 표현과 디코딩 계산을 공동으로 최적화하는 GPU 네이티브 제약 디코딩 프레임워크입니다. 이는 CPU 기반의 오케스트레이션을 완전한 온-디바이스(on-device) 협력형 CUDA 커널로 대체합니다.

1. 데이터 구조: 정수 인식형 억제된 트라이 (Integer-Aware Succinct Trie)

FlashTrie는 GPU 고대역폭 메모리(HBM)에 적합한 작은 메모리 풋프린트를 유지하면서, 대규모 정수 토큰 어휘집(100만 개 이상의 토큰)을 지원하도록 MARISA 트라이 레이아웃을 재설계했습니다.

Narrow LOUDS: 저자들은 LOUDS(Level-Order Unary Degree Sequence) 표현을 32비트 정수 토큰으로 확장했습니다.
비트 압축: 메모리 사용량을 줄이기 위해, 레이아웃은 노드 레이블과 접미사 인덱스를 2바이트 저위 부분(기본 배열에 저장)과 비트 패킹된 고위 부분(필요한 노드를 위한 보조 배열에 저장)으로 분리합니다.
이점: 이 "내로우잉(narrowing)" 기법은 지배적인 배열 크기를 절반으로 줄이고, 제약 인덱스를 축소하며(8k00M 키워드를 약 3GB에 수용), 액세스당 페치(fetch)되는 비트 수를 줄여 메모리 지역성을 개선합니다.

2. 알고리즘: 협력형 온-디바이스 빔 서치 (Cooperative On-Device Beam Search)

FlashTrie는 전체 제약 빔 서치를 단일 지속적(persistent) CUDA 커널 내에서 실행하여, 단계별 호스트 오케스트레이션을 제거합니다.

2단계 병렬성:
- 빔-병렬(Beam-Parallel): 쿼리에 대한 활성 빔들을 서로 다른 스레드 블록(CTA)에 매핑합니다.
- 제안-병렬(Proposal-Parallel): 각 빔 내에서, 스레드들은 정렬된 자식 레이블에 대한 **병렬 이진 탐색(parallel binary search)**을 사용하여 언어 모델의 상위- $K$ 제안을 병렬로 검증합니다.
커널 워크플로우:
1. 확장(Expansion): 활성 빔들이 상위- $K$ 제안을 사용하여 확장됩니다.
2. 검증(Validation): 스레드들이 트라이 자식 레이블에서 이진 탐색을 수행하여 접두사를 검증합니다.
3. 가지치기(Pruning): 병렬 top- $B$ 선택 커널이 전통적인 CPU 힙(heap)을 대체합니다. 이는 추가 전용(append-only) 후보 버퍼와 병렬 정렬을 사용하여 락 경합(lock contention)과 포인터 체이싱 없이 빔을 목표 너비로 가지치기합니다.
동기화: 그리드 단위 배리어(grid-wide barrier)가 단계를 동기화하여, 커널이 단 한 번의 런칭으로 $T$ 디코딩 단계를 수행할 수 있도록 합니다.

주요 기여

GPU 상주 제약 인덱스: 수십억 규모의 제약 라이브러리를 GPU HBM에 완전히 유지하여 메모리 스톨을 줄이는 새로운 비트 압축 트라이 레이아웃을 제시했습니다.
협력형 CUDA 커널: 빔 확장, 검증 및 가지치기를 온-디바이스에서 수행하는 단일 커널 실행 모델을 통해 반복적인 커널 런칭과 호스트 오케스트레이션 오버헤드를 제거했습니다.
GPU 인식 프리미티브: 불규칙한 포인터 기반 트래벌을 병렬 이진 탐색과 힙 없는 top- $K$ 선택으로 교체하여, 워프 활용도를 높이고 분기를 줄였습니다.
확장성: 단일 GPU에서 8억 개의 키워드까지 제약 디코딩을 확장하며, 이는 이전에는 오프라인 설정에서만 가능했던 영역입니다.

실험 결과

시스템은 NVIDIA A100 80GB GPU에서 실행되는 2.2M 토큰 어휘집과 800M 키워드 제약 트라이를 가진 프로덕션 NAR 모델을 대상으로 평가되었습니다.

성능 지표

지연 시간(Latency): FlashTrie는 빔 너비($BW$)가 1000인 경우에도 트라이 탐색 지연 시간을 3 ms 미만으로 줄입니다.
- $BW=1000$에서 FlashTrie는 평균 1.91 ms의 지연 시간을 달성하며, 이는 최적화된 CPU 베이스라인(MARISA-Opt)의 46.30 ms와 비교됩니다.
- 이는 고도로 최적화된 멀티스레드 CPU 베이스라인 대비 24배 빠른 속도를 나타냅니다.
처리량(Throughput): 배치 서빙( $b=8$ ) 환경에서 FlashTrie는 $BW=1000$일 때 피크 1,621 쿼리/초를 기록하며, CPU 베이스라인보다 32배~71배 높은 처리량을 제공합니다.
구축(Construction): FlashTrie는 MARISA-Opt보다 인덱스를 3.4배 빠르게 구축하며, 22% 적은 디스크 공간을 사용합니다.

검색 품질

CPU와의 동일성: FlashTrie는 모든 빔 너비에 대해 모든 정밀도 구간에서 CPU 베이스라인(MARISA-Opt)과 0.001 이내의 정밀도로 일치하는 품질을 보여줍니다.
빔 너비를 통한 품질 향상: FlashTrie는 낮은 지연 시간을 유지하기 때문에, CPU에서는 불가능한 훨씬 더 큰 빔 너비(최대 1000)를 사용할 수 있게 해줍니다. 빔 너비를 100에서 1000으로 늘리면 Precision@100이 0.53에서 0.78로 향상됩니다.

온라인 A/B 테스트

상업용 스폰서 검색 엔진에서의 대규모 온라인 실험 결과:

구성: 처치 그룹(treatment group)은 $BW=600$인 FlashTrie를 사용했고, 대조 그룹(control group)은 지연 시간 제약 내에서 가능한 최대치인 $BW=200$인 MARISA-Opt를 사용했습니다.
결과: FlashTrie는 광고 품질 저하 없이(결함률은 유의미하지 않음) 통계적으로 유의미한 +0.71%의 매출 상승과 0.17–0.20%의 클릭률(CTR) 증가를 달성했습니다.
지연 시간: 처치 그룹은 엄격한 지연 시간 예산을 유지하면서, 대조 그룹의 제한치 대비 p95 지연 시간을 32% 감소시켰습니다.

의의 및 주장

본 논문은 제약 디코딩이 반드시 CPU의 병목 현상으로 남을 필요는 없다고 주장합니다. 병렬 실행을 위해 데이터 구조를 재설계하고 제약 강제화 로직 전체를 GPU로 이동함으로써, FlashTrie는 다음을 가능하게 합니다:

실시간 확장성: 지연 시간이 중요한 애플리케이션(예: 스폰서 검색)에서 생성적 검색 시스템이 지연 시간 예산을 위반하지 않고 빔 크기를 최대 5배까지 확장할 수 있게 합니다.
품질-성능 트레이드오프 해결: 검색 품질과 지연 시간 제약을 분리하여, 시스템이 더 나은 후보 선택을 위해 더 큰 빔 너비를 활용하면서도 낮은 지연 시간을 유지할 수 있도록 합니다.
범용적 적용 가능성: 이 접근 방식은 엔티티 링킹(entity linking)이나 제품 카탈로그 검색과 같이 대규모 정적 어휘집을 사용하는 모든 제약 생성 작업에 적용될 수 있습니다.

저자들은 FlashTrie가, 데이터 구조를 배치 및 병렬 실행에 적합하도록 재설계할 때 전통적으로 CPU에서 구현되던 데이터 구조들이 어떻게 GPU에서 매우 효과적일 수 있는지를 보여주는 광범위한 패턴을 입증한다고 결론짓습니다.