GPU 기반의 빠른 임베딩
검색 및 컴퓨터부터 API 플랫폼에 이르기까지 Perplexity의 모든 영역에서 빠르고 정확한 검색은 필수적입니다. 이면에서는 임베딩 및 랭킹 모델이 핵심적인 역할을 수행하여, 시스템이 주어진 쿼리에 가장 관련성이 높은 결과를 식별할 수 있도록 돕습니다. 당사는 최첨단
검색 및 컴퓨터부터 API 플랫폼에 이르기까지 Perplexity의 모든 영역에서 빠르고 정확한 검색은 필수적입니다. 이면에서는 임베딩 및 랭킹 모델이 핵심적인 역할을 수행하여, 시스템이 주어진 쿼리에 가장 관련성이 높은 결과를 식별할 수 있도록 돕습니다. 당사는 pplx-embed과 같은 자체 모델을 학습시키고 서빙함으로써 최고 수준의 품질과 지연 시간을 달성합니다.
이 글에서는 이 특수한 종류의 모델을 위한 Perplexity의 서빙 인프라 내부를 살펴봅니다. AI 네이티브 검색의 추론 요구 사항을 효율적으로 해결하기 위한 기술을 논의하며, 엑사바이트 규모의 검색 인덱스를 구동하는 동시에 모델의 빠른 프로토타이핑과 평가를 가능하게 합니다. 이러한 기술들은 집합적으로 검색 품질과 효율성의 파레토 프론티어를 확장하여, 가장 낮은 비용과 지연 시간으로 에이전트와 사용자에게 최상의 결과를 제공할 수 있게 합니다.
검색을 위한 임베딩
전형적인 검색 설정에서는 인덱싱된 문서를 임베딩 모델을 사용하여 고차원 벡터 스페이스에 매핑하고 벡터 데이터베이스에 저장합니다. 동일한 모델을 사용하여 쿼리를 임베딩함으로써, 쿼리의 벡터에 가장 가까운 벡터를 찾아 유사한 문서를 찾을 수 있습니다. 이로 인해 추론 엔진이 서빙해야 하는 두 가지 서로 다른 트래픽 패턴이 발생합니다.
- 배치 임베딩: 데이터베이스를 구축, 확장 또는 재인덱싱할 때, 대량의 문서를 벡터 스페이스에 임베딩하여 비용을 최소화하기 위해 처리량을 극대화합니다.
벡터 검색에 이어, 처리량과 지연 시간 사이의 균형을 유지하면서 대규모 문서 배치의 점수를 산정해야 합니다.
- 온라인 임베딩: 데이터베이스를 쿼리할 때, 조회를 위해 짧은 쿼리를 임베딩해야 하며 지연 시간을 최소화합니다.
당사는 사용 사례 전반에 걸쳐 가능한 한 많은 공통 구성 요소를 활용할 수 있도록 추론 인프라를 구축했습니다. 임베딩을 생성하기 위해 주로 소형 Transformer 모델을 사용하므로, LLM 추론 코드와 구현의 대부분을 공유합니다. 배치 임베딩은 연산 집약적인 프리필(prefill)과 유사한 반면, 주로 소수의 토큰에서 실행되는 온라인 임베딩은 메모리 집약적인 디코드(decode)와 연산 관점에서 유사합니다. 따라서 당사는 최적화된 프리필 및 디코드 커널을 재사용하여 임베딩 모델을 서빙합니다. 결과적으로, 온라인 임베딩 워크로드의 낮은 지연 시간을 유지하면서 최소한의 추가 엔지니어링 작업으로 대규모 배치 추론 처리량을 달성할 수 있습니다.
Tulips, Roses, 그리고 약간의 Ivy
내부적으로 그리고 API 플랫폼을 통한 외부적으로 표준화된 API를 통해 추론을 제공합니다. 내부적으로는 임베딩 요청 처리에 여러 서비스가 관여합니다.
- Ivy는 Perplexity 서비스가 호출하는 Rust HTTP 게이트웨이입니다.
JSON 파싱, 토큰화, 입력 템플릿화, 배치 분할 등의 요청에 대한 CPU 측 작업을 처리하며, 요청을 다운스트림 서버를 위한 커스텀 gRPC 프로토콜로 변환합니다. 이러한 분리를 통해 더 무거운 추론 인스턴스를 건드리지 않고도 토큰화 및 입력 포맷팅과 관련된 특정 파라미터를 구성할 수 있습니다.
- Tulip은 추론 서버 인터페이스입니다.
Rust, tokio, `tonic`으로 구현된 gRPC 서버입니다. Tulip은 gRPC 추론 요청을 받아 스케줄링과 배치를 처리합니다. 그런 다음 배치를 ROSE 엔진으로 보내고, 완료된 응답을 클라이언트에 반환합니다.
- **ROSE**(Runtime-Optimized Serving Engine)는 모델 추론을 구현합니다.
주로 Python으로 정의되며, 다양한 모델을 위한 커널, 레이어 및 정의를 제공합니다. ROSE는 모델을 통한 순방향 패스를 구현하며, 임베딩에 특화된 CUDA 그래프 관리도 제공합니다. 이 엔진은 배치를 받아 가속기에서 수행하는 연산에 대한 참조를 반환하는 step() 함수를 통해 Tulip과 연결됩니다.
커널 그 이상을 주목하며
Transformer 기반 모델과 기본 Hopper/Blackwell 아키텍처는 모두 성숙한 기술이므로, GPU 측에서의 임베딩 추론은 다양한 추론 엔진 전반에 걸쳐 대개 최적의 구현으로 수렴되었습니다. 그럼에도 불구하고, 당사는 모델을 클라이언트에 엔드투엔드로 노출하는 런타임과 하네스에서 추가적인 개선 기회를 발견했습니다. 특히, CUDA 그래프를 세심하게 관리하고 네이티브 Rust 엔진에서 GPU 측 결과를 비동기적으로 추적하기 위한 LazyTensor 추상화를 구축함으로써 지연 시간을 단축할 수 있음을 확인했습니다. 당사는 ROSE의 모델 구현과 효과적으로 연동되도록 Tulip에 이러한 기능을 구현했습니다.
Tulip
당사는 Tulip을 모델 서빙 상에서 가능한 한 경량화된 인터페이스로 설계했습니다. Tulip은 Tokio 비동기 작업에서 들어오는 요청을 처리하며, 추적하는 요청 풀을 유지하고 가속기로 디스패치할 배치를 스케줄링합니다. Tulip의 스케줄링 메커니즘은 매우 간단합니다. Tulip이 작업을 디스패치하거나 결과를 기다리는 동안 요청들이 누적됩니다. 누적된 요청 중에서 선착순(first-come, first-served) 방식으로 시퀀스를 선택하여 모델을 통해 실행합니다.
간단한 스케줄링 메커니즘은 모델 성능에 대한 관찰에서 동기를 얻었습니다. 소형 임베딩 모델의 경우, 당사가 서빙하는 시퀀스 길이에서 밀집 레이어의 선형 비용이 어텐션의 이차 비용보다 지배적임을 확인했습니다. 따라서 지연 시간은 시퀀스 수가 아니라 주로 토큰 수에 비례합니다. 결과적으로 배치가 GPU를 포화시키기에 충분히 커지면(10억 개 미만의 모델에서 약 512 토큰), 그 안에 더 많은 시퀀스를 패킹해도 효율성이 향상되지 않습니다.
모델과 효과적으로 연동하기 위해, Tulip은 CUDA 그래프와 레이지 결과 추적에 의존하여 GPU와 CPU 작업을 오버랩하고 사용 가능한 리소스를 온전히 활용합니다.
CUDA 그래프 관리
모델의 순방향 패스를 실행하는 것은 CPU 측 작업과 GPU 측 작업을 모두 수반합니다. CPU는 배치 스케줄링과 적절한 파라미터로 커널을 런칭하는 책임을 지며, GPU는 관련 행렬 곱셈, 어텐션, 노름 또는 활성화 커널을 실행합니다. 학습 및 재인덱싱과 같은 고처리량 워크로드의 경우, 배치 크기와 GPU 측 지연 시간이 모두 크기 때문에 CPU 측 오버헤드는 무시할 수 있는 수준입니다. 그러나 더 작은 배치 크기에서는 CPU 측 작업이 GPU 측 작업을 능가할 수 있습니다.
오버헤드를 완화하기 위해, 독립적인 커널을 런칭하는 대신 CUDA 단일 호출로 순방향 패스의 모든 커널을 런칭하는 데 필요한 메타데이터를 캡처하도록 CUDA 그래프를 구축할 수 있습니다. 이를 통해 CUDA 그래프를 캡처할 수 있는 구성을 위해 비용이 많이 드는 Python 및 PyTorch 코드를 다시 실행할 필요가 없어집니다.
각 모델별로 GPU 실행 비용이 CPU 측 커널 런칭 비용보다 비싸지는 최소 토큰 수를 결정하는 변곡점을 추적합니다. 임베딩 모델은 크기가 작기 때문에, 이 변곡점은 수천 개의 토큰과 수십 개의 시퀀스로 이루어진 배치에서 발생함을 관찰했습니다. 일부 어텐션 구현은 커널 런칭을 구성하기 위해 동적 호스트 측 입력에 의존하므로, 전체 모델 프리필/밀집 CUDA 그래프를 사용할 수 없습니다. 당사는 추론 엔진에서 이를 활성화하기 위해 관련 커널에 대한 변경 사항을 업스트림에 반영(upstreamed)했습니다.
오버헤드를 해결하기 위해, 모든 임베딩 모델에 대해 전체 모델 CUDA 그래프를 구축하고 CPU 작업과 GPU 작업을 오버랩합니다. CUDA 그래프는 CPU 측 오버헤드를 최소화하므로, 그래프가 런칭되면 다음 배치가 준비될 때마다 해당 배치의 실행을 시작하고 큐에 넣을 수 있는 여유 시간이 생깁니다. 대기 중인 배치의 결과는 LazyTensor로 추적되며, 이를 통해 Rust의 비동기 작업이 이전 배치의 실행이 완료될 때까지 블록될 수 있습니다. CUDA 그래프는 커널 런칭 비용에 발목 잡히지 않도록 보장하여 저지연 서빙을 돕고, CPU가 다음 배치 작업을 더 빨리 수행할 수 있도록 자유롭게 하므로 고처리량 경우의 스케줄링을 용이하게 합니다.
각각의 서로 다른 구성에 대해 CUDA 그래프를 캡처해야 하며, 임베딩의 경우 이는 시퀀스 개수와 토큰 개수 조합당 하나의 그래프를 의미합니다. 이 그리드가 광범위하므로 토큰 개수를 64 또는 256의 배수인 버킷으로 패딩합니다. 이로 인해 여전히 수천 개의 그래프가 생성되며, 일반적인 모델의 경우 캡처하는 데 수 분이 소요될 수 있습니다. 캡처 비용은 두 가지 원인에서 발생합니다. 커널을 컴파일하고 필요한 다양한 커널의 버퍼를 설정하기 위해 실행되어야 하는 즉시 순방향 패스와, 그 후 Python 코드를 다시 실행하는 캡처 실행입니다.
엔진이 서빙을 수행함에 따라 CUDA 그래프를 지연(lazily) 캡처하여 시작 비용을 완화합니다. 각 구성을 추적하고, 두 번째 히트에서 그래프 캡처 및 재생을 트리거하기 전에 즉시(eager) 웜업 실행을 거치도록 보장합니다. 동일한 그래프 구성의 모든 후속 실행은 CUDA 그래프 재생을 거칩니다. 지연 그래프 캡처는 시작 시 p99 지연 시간에 영향을 미치지만, 수 분간의 즉시 작업을 여러 시간에 걸쳐 분산시키는 데 유용합니다. 더 빠른 시작 시간은 임베딩 배포를 더 잘 확장하고 관리할 수 있도록 해줍니다.
레이지 텐서(Lazy Tensors)
CUDA를 통해 GPU 작업은 비동기식으로 진행됩니다. 커널을 비동기적으로 런칭하면 스트림에 큐에 추가되므로, 호스트 코드는 결과 벡터를 읽기 위해 명시적으로 동기화해야 합니다. 더 높은 수준의 병렬 처리를 용이하게 하고 디바이스에서 이전 작업이 완료되기를 기다리는 동안 향후 배치를 시작할 수 있도록, 값을 추적하기 위해 LazyTensor 추상화에 의존합니다.
LazyTensor는 페이지 잠금 메모리의 호스트 버퍼와 디바이스에서 데이터를 복사하는 이벤트를 통한 cudaMemcpyAsync 연산을 추적합니다. 이는 동일한 스트림에서 순방향 패스가 런칭된 직후에 시작됩니다. 복사 연산은 스트림상의 이전 모든 커널이 실행될 때까지 기다려야 하므로, 연관된 이벤트는 순방향 패스의 완료와 CPU에서의 결과 사용 가능 여부를 모두 추적합니다.
ROSE 인코더 엔진에서 LazyTensor를 활용하여 GPU와 CPU 작업을 오버랩합니다. 각 step() 호출이 CUDA 그래프를 실행하고 완료될 때까지 기다리는 대신, step()은 결과를 비동기적으로 추적하기 위해 LazyTensor를 반환합니다. CUDA 그래프와 결합된 이 방식은 낮은 지연 시간과 더 나은 처리량을 달성하는 데 도움을 줍니다.
ROSE
당사는 원래 LLM 서빙을 위해 구축했던 ROSE 엔진을 수정하여 임베딩 모델의 실행도 처리할 수 있도록 했습니다. 임베딩 모델 지원에 필요한 노력을 최소화하기 위해, ROSE는 LLM과 임베딩 간의 코드를 적극적으로 재사용합니다. 예를 들어, pplx-embed 서빙과 Qwen3.5 LLM 디코딩은 모두 동일한 커널을 거칩니다. 이러한 공유를 통해 프로토타이핑, 평가, 프로덕션 추론을 위해 LLM으로부터 원래 미세 조정된 임베딩 모델을 쉽게 서빙할 수 있습니다.
밀집 레이어의 경우, 토큰 벡터가 독립적으로 처리되므로 임베딩과 LLM 추론이 동일합니다. 어텐션 레이어에서는 LLM에 필요한 페이징된 프리필 및 디코드 설정과 함께, 래기드(ragged) 입력에 대한 지원을 추가하여 차이점을 처리합니다. 임베딩 모델을 서빙할 때, KV 캐시를 인스턴스화하지 않고 패딩을 피하기 위해 래기드 포맷을 지원하는 다양한 어텐션 커널로 디스패치합니다. 지원하는 변환 및 캘리브레이션 루틴도 LLM과 공유됩니다.
Ivy
추론 HTTP 프록시 레이어인 Ivy 또한 성능에서 중요한 역할을 합니다. 프로덕션에서는 요청 페이로드가 다양하므로, 개별 요청을 개별 레플리카에 라우팅하면 부하 불균형이 발생할 수 있습니다. Ivy는 대규모 배치 요청을 청크로 분할하고 레플리카 간에 부하를 분산하여, 사용률을 개선하고 지연 시간을 평탄화합니다. Ivy에 완전히 배포된 사내 유니그램 토크나이저(in-house unigram tokenization)에 대한 당사의 최근 작업은 기성 토크나이저에 비해 지연 시간을 대폭 개선합니다.
...그렇지만 여전히 중요한 커널
ROSE는 다양한 어텐션 백엔드를 지원합니다. 서로 다른 커널이 특정 문제 크기에 적합할 수 있습니다. 시간이 지남에 따라, 당사는 래기드 어텐션을 구현하기 위해 FlashInfer 2, FlashInfer 3, FlashAttention 4 커널을 통합했습니다.
일반적으로 FlashAttention 4가 더 빠릅니다. 그러나 매우 긴 시퀀스 길이에서는 FlashInfer 3이 Qwen 기반 모델에서 더 우수한 성능을 보입니다. 성능과 튜닝은 어텐션 헤드의 수와 차원에 따라 달라질 수 있으므로, 당사는 여러 구성을 지원하며 서빙 시 개별 사례별로 결정을 내립니다.
벤치마크
평가 데이터셋에서 파생된 실제 모델 가중치와 입력에 대해 BF16 정밀도로 추론을 실행하여 vLLM v0.22.0과 벤치마크를 비교했습니다. 모든 타이밍 실행 전에는 코사인 유사도의 편차가 0.1% 이내인지 확인하는 웜업(warmup) 실행이 선행되었습니다.
저지연 임베딩 (p50 / p90 / p99 / 최대 ms)
사전 토큰화된 요청 배치 크기 1, 완전 순차적 요청, 시퀀스 길이 128, 512, 4096 토큰에 대한 런타임을 보고합니다.
저지연 점수 산정 (p50 / p90 / p99 / 최대 ms)
사전 토큰화된 요청 배치 크기 5, 25, 50, 토큰 시퀀스 길이 512.
고성능 임베딩 (emb/s)
요청 배치 크기 100, 요청을 제출하는 동시 프로세스 4개, 시퀀스 길이 512, 1024, 4096 토큰.
고동시성 임베딩 (p50 / p90 / p99 / 최대 ms)
시퀀스 길이 512, 배치 크기 1, 하지만 1, 2, 4, 8, 16개의 동시 요청을 전송합니다. 이 벤치마크에는 Ivy를 통한 토큰화 비용과 Ivy 및 Tulip 간의 네트워킹 오버헤드도 포함됩니다.
결론 및 향후 연구
Ivy, Tulip, ROSE로 구성된 서빙 인프라를 통해 Perplexity는 더 낮은 지연 시간과 더 나은 처리량으로 임베딩을 서빙할 수 있으며, 기성 솔루션에 비해 줄어든 비용으로 더 정확한 검색 결과를 얻을 수 있습니다.
특정 모델에 집중하고 전체 스택에 대한 소유권을 가짐으로써, 성능과 유연성 사이의 효과적인 균형을 잡는 데 필요한 자유를 얻을 수 있으며, 재사용성과 성능이 뛰어난 Rust 프리미티브를 보다 일반적인 Python 모델링 코드와 혼용할 수 있습니다. vLLM, SGLang, TokenSpeed와 같은 많은 오픈소스 추론 엔진들은 Rust와 C++ 같은 언어를 스택에 통합하고 있습니다. 당사는 지난 2년간 Rust에 투자하여 성능과 유지보수성 모두에서 큰 보상을 얻었습니다. 임베딩 구현의 대부분을 LLM 서빙 스택과 공유함으로써, 임베딩 모델 유지보수에 상당한 엔지니어링 노력을 들이지 않고도 처리량 향상 효과를 얻을 수 있습니다.
모델이 진화함에 따라, CPU 바운드 및 GPU 바운드 지연 시간을 모두 줄이기 위해 스택의 각 레이어를 지속적으로 개선할 것입니다. Ivy 및 Tulip 내의 커스텀 gRPC 기반 프로토콜을 통해 통신을 조정하여 네트워크 지연 시간을 줄일 수 있으며, ROSE는 연산 처리량을 개선할 기반을 제공합니다. 또한 생태계 전반에서 자유 스레드(free-threaded) Python에 대한 지원이 성장함에 따라, Python과 Rust 간의 상호 운용성을 더욱 개선하여 오버헤드를 줄일 수 있게 될 것입니다.