vLLM TPU 임베딩 운영법: 긴 컨텍스트 검색 파이프라인 설계

긴 문서를 검색하는 RAG 시스템에서 병목은 종종 LLM 답변 생성이 아니라 임베딩 생성입니다. 문서가 많고, 한 문서가 길고, 이미지나 표까지 섞이면 임베딩 파이프라인의 처리량과 비용이 전체 제품 경험을 결정합니다. Google Cloud가 vLLM serving engine에 TPU 지원을 통합하고 긴 컨텍스트 멀티모달 임베딩 추론을 다룬 글은 이 지점을 겨냥합니다.

이 글은 “vLLM TPU embedding”, “long context embedding inference”, “RAG embedding pipeline”을 찾는 개발자를 위한 운영 가이드입니다. 특정 클라우드 기능 홍보보다, 긴 컨텍스트 검색 시스템을 만들 때 어떤 병목과 검증 기준을 봐야 하는지에 초점을 둡니다.

긴 컨텍스트 임베딩이 어려운 이유

짧은 텍스트 임베딩은 비교적 단순합니다. 문단을 적당히 자르고, batch로 모델에 넣고, 벡터 DB에 저장하면 됩니다. 하지만 15K token 이상의 긴 컨텍스트를 다루면 문제가 달라집니다.

첫째, 메모리 사용량이 커집니다. 입력 길이가 늘면 attention과 prefill 비용이 늘고, batch 크기를 키우기 어렵습니다. 처리량을 높이려다가 OOM이 나면 작업 큐가 밀립니다.

둘째, chunking 전략이 검색 품질을 좌우합니다. 문서를 너무 잘게 자르면 맥락이 사라지고, 너무 크게 넣으면 비용과 지연시간이 증가합니다. 긴 컨텍스트 모델을 쓴다고 chunking이 필요 없어지는 것은 아닙니다.

셋째, numerical parity가 중요합니다. GPU 기준 결과와 TPU 결과가 미세하게 달라질 수 있습니다. 검색 시스템에서는 작은 벡터 차이가 top-k 순서를 바꿀 수 있습니다. 운영 전에는 기준 모델과의 차이를 측정해야 합니다.

Google 글은 TPU-specific tensor alignment, JAX/XLA compilation pre-warming, hybrid StepPool architecture for chunked prefill 같은 최적화를 언급합니다. 개발자가 가져가야 할 핵심은 “하드웨어를 바꾸면 serving 레이어와 검증 방식도 같이 바뀐다”입니다.

RAG 파이프라인에서 임베딩 병목을 찾는 법

먼저 지표를 나눠야 합니다. 문서 수집 시간, 파싱 시간, chunk 생성 시간, 임베딩 추론 시간, 벡터 DB upsert 시간, 검색 latency를 따로 봅니다. 전체 처리 시간이 느리다는 말만으로는 개선할 수 없습니다.

임베딩 추론에서는 tokens/sec, requests/sec, batch latency, p95 latency, queue wait, 실패율, 재시도율을 봐야 합니다. 특히 batch latency 평균만 보면 위험합니다. 긴 문서 몇 개가 tail latency를 만들고 큐를 막을 수 있습니다.

운영에서는 online embedding과 offline embedding을 분리하는 편이 좋습니다. 사용자가 업로드한 문서를 즉시 검색해야 하는 경우는 online path입니다. 대량 문서 재색인, nightly sync, 권한 변경 반영은 offline path입니다. 두 경로를 같은 큐에 넣으면 대량 배치가 사용자 요청을 막을 수 있습니다.

TPU나 고성능 serving을 쓰기 전 정해야 할 것

고성능 인프라를 붙이기 전에 먼저 입력 정책을 정해야 합니다. 지원할 최대 토큰 수, 파일당 chunk 수, 사용자당 동시 작업 수, 재색인 우선순위, 실패 재시도 횟수를 정하지 않으면 비용이 예측되지 않습니다.

긴 컨텍스트 모델을 쓸 때도 모든 문서를 그대로 넣지 마세요. 문서 구조를 먼저 분석하고 제목, 섹션, 표, 코드블록, 이미지 캡션을 나눠야 합니다. 검색 품질은 모델 크기보다 문서 전처리에서 더 많이 갈릴 때가 많습니다.

멀티모달 문서라면 텍스트 추출 결과와 이미지 설명을 별도 필드로 저장하는 것이 좋습니다. 하나의 벡터만 저장하면 “표 안의 숫자”, “이미지 캡션”, “본문 설명”이 섞여 검색 이유를 설명하기 어렵습니다.

정확도 검증은 top-k 비교로 해야 합니다

임베딩 serving을 바꿀 때는 단순히 API 응답이 200인지 보면 안 됩니다. 기존 GPU serving 결과와 새 TPU serving 결과를 같은 질의 세트로 비교해야 합니다. 봐야 할 것은 cosine similarity 평균이 아니라 검색 결과의 안정성입니다.

예를 들어 200개 질의와 정답 문서 세트를 준비합니다. 기존 파이프라인과 새 파이프라인에서 top-5, top-10 검색 결과를 비교합니다. recall@5, recall@10, nDCG, MRR을 계산합니다. 운영 로그에서 자주 나오는 질의도 포함해야 합니다.

벡터 값이 조금 달라도 검색 결과가 같으면 문제는 작습니다. 반대로 평균 cosine 차이는 작아도 중요한 질의에서 top result가 바뀌면 배포를 보류해야 합니다. 특히 고객지원, 법률, 의료, 금융 문서 검색은 top-k 안정성이 중요합니다.

비용을 줄이는 chunking과 cache 전략

긴 컨텍스트 임베딩 비용을 줄이는 가장 현실적인 방법은 재계산을 줄이는 것입니다. 문서 전체를 매번 다시 임베딩하지 말고, content hash를 저장하세요. 섹션 단위 hash가 바뀐 chunk만 다시 계산하면 됩니다.

또한 metadata 변경과 content 변경을 분리해야 합니다. 권한, 태그, 소유자만 바뀌었다면 벡터를 다시 만들 필요가 없습니다. 벡터 DB metadata만 갱신하면 됩니다.

chunk 크기는 모델과 문서 종류에 따라 다르지만, 실험 없이 정하면 안 됩니다. 512, 1,024, 2,048 token chunk를 비교하고, overlap도 0%, 10%, 20%로 테스트하세요. 긴 컨텍스트 모델은 큰 chunk에서 강하지만, 검색 질의가 짧으면 작은 chunk가 더 잘 맞을 수 있습니다.

실행 체크리스트

• 문서 수집, 파싱, chunking, 임베딩, upsert, 검색 latency를 따로 측정한다.
• online embedding과 offline reindex 큐를 분리한다.
• 파일당 최대 토큰 수, chunk 수, 동시 작업 수, 재시도 횟수를 제한한다.
• 문서 구조를 보존해 제목, 섹션, 표, 이미지 캡션을 metadata로 남긴다.
• serving 변경 전후로 recall@5, recall@10, nDCG, MRR을 비교한다.
• content hash로 변경된 chunk만 다시 임베딩한다.
• metadata 변경은 벡터 재생성 없이 처리한다.
• chunk 크기와 overlap은 실제 질의 세트로 A/B 테스트한다.

출처: Google Developers Blog, “Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU”