왜 임베딩 서빙이 다시 어려워졌을까?
임베딩 모델은 텍스트·이미지·오디오 같은 비정형 데이터를 dense vector로 변환해서 시맨틱 검색, 추천, 의도 분류 같은 다운스트림 태스크를 가능하게 하는 '번역기'예요. 프로토타입 단계에서는 소형 텍스트 임베딩 모델을 그냥 띄우면 끝이지만, 수백만 쿼리 규모로 넘어가는 순간 병목이 완전히 달라집니다.
실무에서 가장 자주 부딪히는 세 가지는 이거예요.
- 가속기 용량의 탄력성: 트래픽이 출렁일 때 컴퓨트를 즉시 늘리고 줄일 수 있어야 함
- 비용/성능 효율: 같은 처리량을 더 싸게 뽑아내야 함
- 수치적 패리티: 이기종 하드웨어 백엔드 간에 임베딩 결과가 수학적으로 동일해야 함 (특히 엔터프라이즈)
특히 차세대 임베딩 모델은 **4K+ 토큰(텍스트)에서 15K+ 토큰(멀티모달)**까지 초장문 컨텍스트를 처리해야 해서, 기존 서빙 스택으로는 HBM이 버티질 않아요. 이 글에서는 Google Cloud가 vLLM에 네이티브 TPU 지원을 넣으면서 해결한 4가지 핵심 최적화를 실무 관점에서 뜯어봅니다. (근거자료: Google Developers Blog 원문)
최소 예제: Qwen3-Embedding-8B를 TPU에서 띄우기
vLLM의 네이티브 pooling runner를 쓰면 TPU에서도 dense pooling 출력을 바로 뽑을 수 있어요. 아래는 공식 레시피에서 발췌한 최소 실행 예제입니다.
from vllm import LLM
# Cloud TPU에서 vLLM의 네이티브 pooling runner로 Qwen3-Embedding-8B 초기화
llm = LLM(
model="Qwen/Qwen3-Embedding-8B",
runner="pooling", # dense pooling 출력 활성화
tensor_parallel_size=2, # TPU 토폴로지 메시에 샤딩
max_model_len=16384,
max_num_batched_tokens=512,
dtype="bfloat16",
trust_remote_code=True
)
# 입력에 대한 dense 벡터 임베딩 추출
prompts = ["Enterprise-grade semantic retrieval on TPUs with vLLM."]
results = llm.embed(prompts)
embedding_vector = results[0].outputs.embedding
핵심 포인트는 runner="pooling"이에요. 이걸 켜야 LLM의 생성 헤드가 아니라 pooling 레이어에서 벡터가 나옵니다. tensor_parallel_size는 TPU 토폴로지 메시 크기와 반드시 맞춰야 하고요.
내부에서 뭘 고쳤길래 돌아갈까?
TPU에서 vLLM을 돌리려면 크게 네 가지 지점을 손봐야 했어요.
- 어휘 행렬 패딩 (Vocabulary Padding) — TPU의 MXU(Matrix Execution Unit)는 텐서 병렬화(TP)로 어휘 행렬을 메시에 샤딩할 때 엄격한 나눗셈 제약을 걸어요. 그래서 All-Gather 실행 중 텐서 정렬이 정확히 맞도록 하드웨어 안전 패딩 전략을 넣었습니다.
- 레이지 로더 호환 (Attribute Promotion) — vLLM은 TPU에서 콜드 스타트 지연과 호스트 메모리 피크를 줄이려고 lazy-loading을 씁니다. 여기서 unquantization 파이프라인에 attribute promotion을 넣어야 초기화 실패가 사라져요.
- JIT 컴파일 캐시 사전 워밍 (Sharding-aware Pre-warming) — 멀티프로세스 배포에서 JAX/XLA 컴파일 트랩을 피하려면 추론 전에 컴파일 캐시를 미리 lock해둬야 합니다.
- Chunked Prefill 상태 유지 (StepPool + CachedRequestState) — 초장문 컨텍스트는 pooling 레이어에서 Chunked Prefill이 필수인데, 스텝 경계를 넘어가면서 pooling 상태가 날아가는 문제가 있어요. 하이브리드 StepPool과 메타데이터 마이그레이션으로 preemption 이후에도 상태가 누적되도록 만들었습니다.
주의사항: 패리티 검증 없이 프로덕션에 올리지 마세요
TPU로 옮기면 성능은 좋아지지만, 수치적 패리티를 반드시 검증해야 해요. 임베딩은 검색 품질과 직결되니까 미세한 오차도 랭킹을 흔듭니다.
검증 방법은 단순합니다. TPU에서 나온 벡터(vTpu)와 XPU 골든 레퍼런스(vRef)의 코사인 유사도를 재면 돼요.
| 입력 유형 | 통과 임계값 |
|---|---|
| 텍스트 | ≥ 0.999 |
| 멀티모달 (텍스트+이미지) | ≥ 0.995 |
1.0에 가까울수록 하드웨어 백엔드 간 수치 패리티가 확보된 겁니다. 멀티모달이 임계값이 살짝 낮은 이유는 이미지 인코더 쪽 연산 경로가 더 길어서 누적 오차 여지가 크기 때문이에요.
이 기술의 한계와 주의점
- TPU 종속성: vLLM-TPU 스택에 최적화가 들어가 있어서, 다른 가속기로 옮길 때는 별도 검증이 필요해요.
- JIT 워밍 비용: pre-warming을 안 하면 첫 요청 레이턴시가 튀어요. 오토스케일링 환경에서는 새 노드가 뜰 때마다 이 비용을 치릅니다.
- 패딩 오버헤드: 어휘 크기가 TPU 메시 크기의 배수가 아니면 패딩으로 약간의 메모리 낭비가 생겨요.
- Chunked Prefill 튜닝:
max_num_batched_tokens를 너무 작게 잡으면 상태 전환 오버헤드가 커집니다.
함께 보면 좋은 글
- CSS 상대 색상(Relative Color)으로 애니메이션 테마링 끝내기 — 프론트엔드 쪽 최적화 감각을 같이 키워보고 싶다면.
- 넷플릭스의 Cassandra 넓은 파티션 전략 — 대용량 시계열 데이터를 다루는 시스템 설계 관점에서 읽어볼 만해요.
한국 개발 생태계에서의 적용 맥락
국내에서는 Cloud TPU를 직접 프로비저닝해서 쓰는 팀이 아직 많지 않아요. 대부분 AWS/GCP의 GPU 인스턴스로 임베딩 서빙을 돌리고 있죠. 그래서 이 글의 내용이 당장 '내일 적용'하기는 어려울 수 있습니다. 다만 다음 시나리오에서는 검토 가치가 있어요.
- RAG 파이프라인에서 임베딩 비용이 GPU 청구서의 큰 비중을 차지하는 경우 — TPU의 가성비가 유효할 수 있어요.
- 국내 SI 프로젝트에서 온프레미스 임베딩 서버를 요구하는 경우 — vLLM이 이미 표준이라 마이그레이션 경로가 열려 있습니다.
- 멀티모달 검색을 자체 구축하는 커머스/미디어 기업 — 15K 토큰 컨텍스트가 필요한 순간이 옵니다.
다음 단계 학습 방향
- vLLM 아키텍처 이해: pooling runner, scheduler, chunked prefill이 어떻게 맞물리는지 먼저 파악하세요.
- TPU 프로그래밍 모델: MXU, HBM, 토폴로지 메시 개념을 익히면 위 최적화가 왜 필요한지 이해됩니다.
- 패리티 테스트 자동화: CI에 코사인 유사도 회귀 테스트를 넣어두면 하드웨어 교체 시 안전해요.
- GKE Custom Compute Classes: 오토스케일링 우선순위 규칙으로 여러 가속기 타입을 넘나드는 전략을 설계해보세요.
결론적으로, 초장문 멀티모달 임베딩을 프로덕션에서 돌리려면 '그냥 모델 띄우기'가 아니라 하드웨어 제약 → 프레임워크 패치 → 수치 검증의 3단계를 반드시 거쳐야 합니다. 이 흐름을 이해해두면 TPU든 GPU든 어떤 백엔드로 옮겨가도 응용할 수 있어요.