vLLM TorchTPU: Google Cloud TPU 추론을 위한 새로운 vLLM 백엔드

vLLM TorchTPU는 vLLM이 Google Cloud TPU에서 모델을 서빙할 수 있게 해주는 새로운 out-of-tree vLLM 플랫폼 플러그인으로, vLLM의 기존 모델 코드, CLI, 스케줄러, KV 캐시 관리자, OpenAI 호환 서버를 변경 없이 그대로 재사용합니다. 이는 vLLM에서 이미 작동하는 대부분의 모델이 모델 정의를 수정하지 않고도 TPU에서 실행될 수 있다는 점, 그리고 attention, Mixture-of-Experts, 양자화된 matmul을 위한 수작업으로 작성된 Pallas 커널을 오픈 소스로 공개한다는 점에서 중요합니다.

vLLM TorchTPU란 무엇인가

vLLM TorchTPU는 vllm-torchtpu 패키지로, Google Cloud TPU에서 vLLM을 실행하기 위한 백엔드입니다. Google의 TorchTPU 스택을 기반으로 하며, 이 스택은 PyTorch를 TPU에서 기본적으로 실행하고 vLLM에 이미 있는 모델 정의를 실행합니다. 그 결과, 오늘날 vLLM에서 작동하는 대부분의 모델은 모델 코드를 변경하지 않고도 TPU에서 서빙될 수 있습니다.

이 프로젝트는 릴리스 한 달 전에 발표된 Inferact와 Google 간의 엔지니어링 파트너십에서 나왔습니다. 릴리스 게시물은 프로젝트의 세 가지 목표를 명시합니다:

  1. vLLM의 기존 모델 코드를 재사용하여 새로운 모델을 TPU에 쉽게 도입할 수 있게 합니다.
  2. attention, Mixture-of-Experts, 양자화된 matmul을 위한 Pallas 커널로 오픈 소스에서 TPU 추론 성능을 끌어올립니다.
  3. vLLM 경험을 동일하게 유지합니다: 다른 하드웨어에서 사용되는 것과 동일한 CLI, OpenAI 호환 서버, 스케줄러, KV 캐시 관리자를 사용합니다.

TPU 백엔드가 vLLM에 연결되는 방식

vllm-torchtpu는 out-of-tree vLLM 플랫폼 플러그인입니다. 설치되면 vLLM은 vllm.platform_plugins 엔트리 포인트를 통해 자동으로 이를 인식하므로, vLLM의 포크나 패치된 복사본이 필요하지 않습니다.

모델은 vLLM의 표준 모델 러너를 통해 실행되며 torch.compile(backend="tpu")로 컴파일되어 그래프를 XLA로 낮춥니다. 각 TPU 장치는 vLLM이 GPU에서 실행되는 것과 동일한 방식으로 자체 워커 프로세스를 가지며, 장치 간 통신은 torch.distributed를 통해 이루어집니다.

실질적인 결과: 스케줄러, 워커, KV 캐시 관리자, API 서버는 모두 수정되지 않은 업스트림 vLLM입니다. TPU 관련 작업은 이러한 구성 요소 아래에 있으며, 대부분 컴파일된 그래프 내에서 호출되는 사용자 정의 연산으로 등록된 JAX/Pallas 커널 집합에 있습니다.

시작하기: 설치 및 모델 서빙

Python 3.12를 실행하는 Cloud TPU VM에서 단일 pip 명령으로 설치할 수 있습니다:

pip install vllm-torchtpu

패키지는 플러그인과 함께 일치하는 vLLM, torch, torch_tpu, libtpu를 설치합니다. 설치 가이드는 uv, 소스 빌드, Docker도 다룹니다.

발표에서 가져온 대표적인 서빙 명령으로, 8개의 TPU 칩에서 FP8 Qwen3-Coder MoE 모델을 expert parallelism 및 FP8 KV 캐시와 함께 사용합니다:

vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --kv-cache-dtype fp8 \
  --max-model-len 16384 \
  --async-scheduling

첫 번째 시작은 구성된 형태에 대해 TPU 그래프를 컴파일하며, 몇 분이 걸릴 수 있습니다. VLLM_XLA_CACHE_PATH를 설정하면 재시작 시 컴파일된 아티팩트를 재사용하여 매 부팅 시 재컴파일을 피할 수 있습니다.

추측 디코딩, 프리픽스 캐싱, 멀티-LoRA, 구조화된 출력, 멀티모달 입력은 모두 일반적인 vLLM 플래그로 활성화됩니다. 퀵스타트는 오프라인 배치 추론, 멀티모달 모델, 멀티-LoRA를 포함한 추가 예제를 다룹니다.

Pallas 커널: TPU용으로 수작업 튜닝된 것

vLLM TorchTPU는 컴파일러에 전적으로 의존하지 않고 TPU에서 수작업 튜닝이 유리한 연산에 사용자 정의 Pallas 커널을 사용합니다. 발표된 커널 세트는 다음을 포함합니다:

  • MHA 및 GQA 모델용 Ragged Paged Attention v3.
  • DeepSeek, Kimi, GLM용 Dense 및 Sparse Multi-head Latent Attention (MLA).
  • 하이브리드 선형 attention 모델용 Gated DeltaNet 및 Kimi Delta Attention.
  • Expert parallelism을 갖춘 Fused Mixture-of-Experts.
  • FP8, INT8, MXFP4, NVFP4 양자화 matmul.

확장: 6가지 병렬 처리 차원

vLLM의 6가지 병렬 처리 차원 모두 단일 호스트 및 다중 호스트 TPU 슬라이스에서 지원되며, 하나의 배포에서 결합할 수 있습니다.

차원 플래그
텐서 병렬 처리 (TP) --tensor-parallel-size
데이터 병렬 처리 (DP) --data-parallel-size
파이프라인 병렬 처리 (PP) --pipeline-parallel-size
전문가 병렬 처리 (EP) --enable-expert-parallel
프리필 컨텍스트 병렬 처리 (PCP) --prefill-context-parallel-size
디코드 컨텍스트 병렬 처리 (DCP) --decode-context-parallel-size

대규모 Mixture-of-Experts 모델의 경우, 프로젝트는 일반적으로 데이터 병렬 attention과 전문가 병렬 처리를 결합할 때 최상의 처리량이 나온다고 보고합니다. 예를 들어 단일 v7x 호스트에서 EP를 활성화한 DP=8, TP=1입니다. 이 구성은 전문가를 슬라이스 전체에 분산하면서 attention을 각 장치에 로컬로 유지합니다. 긴 컨텍스트의 경우 PCP와 DCP가 프리필과 KV 캐시를 장치 간에 분할합니다.

명시된 초점은 오늘날 서빙 시스템에 가장 큰 부담을 주는 워크로드, 특히 컨텍스트가 계속 커지고 도구 출력이 길어질 수 있는 다중 턴 에이전트 및 멀티모달 트래픽에 맞춰져 있습니다.

KV 캐시 오프로딩 및 분리형 서빙

vllm-torchtpu는 Python 런타임 외부에서 TPU 장치 버퍼를 직접 읽을 수 있는 오픈 소스 C++ 라이브러리인 TPU Sync를 기반으로 KV 이동 기능을 구축합니다. vLLM의 KV 커넥터 인터페이스를 통해 연결되어 두 가지 기능을 제공합니다:

  • 호스트 DRAM으로의 제로 카피 KV 캐시 오프로딩으로, 긴 다중 턴 세션의 프리픽스 캐시가 HBM 용량에 제한되지 않습니다.
  • 낮은 KV 전송 오버헤드로 TPU 호스트 간 분리형 프리필/디코드 서빙. 프리필과 디코드는 다른 레이아웃을 사용할 수 있으며(예: 프리필은 PCP=4, TP=2, 디코드는 DP=8, TP=1), KV 캐시는 전송 중에 리샤딩됩니다.

추측 디코딩 및 프리픽스 캐싱

추측 디코딩은 MTP, EAGLE-3, DFlash, DSpark, n-gram, suffix 디코딩을 지원합니다. 자동 프리픽스 캐싱은 Qwen3.5 및 Kimi K3와 같은 하이브리드 모델에서도 작동하며, 선형 attention 상태가 KV 캐시와 함께 캐시됩니다.

지원 모델 및 하드웨어

vLLM TorchTPU는 vLLM이 지원하는 대부분의 모델을 서빙할 수 있으며, Qwen, DeepSeek, Kimi, GLM, Llama, GPT-OSS, Gemma 제품군에 대해 TPU 특화 커널, 양자화, 병렬 처리 작업을 제공합니다.

TPU v7x (Ironwood)는 권장 하드웨어이며 모든 풀 리퀘스트에서 CI로 테스트됩니다. TPU v6e, v5e, v5p는 실험적 단계로 설명됩니다. 야간 지원 매트릭스는 CI의 TPU v7x 정확성 및 벤치마크 결과를 추적하며, Supported Models & Features 페이지는 검증된 체크포인트와 각 기능을 활성화하는 방법을 나열합니다.

로드맵: 다음 단계

명시된 목표는 주요 제품군(Qwen, DeepSeek, GLM, Kimi, Gemma, GPT-OSS)의 새로운 오픈 모델이 출시 당일에 TPU에서 실행되어야 한다는 것입니다. 공개 로드맵은 다음을 나열합니다:

  • 릴리스 후 24시간 내에 새로운 오픈 가중치 모델을 온보딩, 벤치마킹, 지원 매트릭스 게시하는 자동화된 CI 파이프라인.
  • 프리픽스, 캐시, 세션 인지 라우팅 및 llm-d TPU 메모리 계층 라우팅.
  • TorchTPU 생태계의 다른 프로젝트와 공유되는 Tokamax 커널용 PyTorch 인터페이스.
  • vllm-omni를 통한 확산 모델 서빙.

프로젝트의 기여 가이드를 통해 기여를 명시적으로 환영합니다.

기여자 TPU 용량 풀

출시와 함께 프로젝트는 커뮤니티 기여 TPU 용량 풀을 엽니다. vllm-torchtpu를 기반으로 개발하는 개발자는 무료 TPU 시간을 신청할 수 있으며, Inferact의 핵심 vLLM 유지관리자가 설계 및 PR 검토를 도와줍니다. 신청은 발표에 링크된 공개 Google Form을 통해 진행됩니다.

코드 및 문서를 찾을 수 있는 곳

  • 퀵스타트, Cloud TPU 설정, 설치(pip, uv, 소스, Docker) 문서는 vllm-torchtpu 저장소에 있습니다.
  • 배포 리소스에는 TPU v7x (Ironwood) 서빙 레시피, tpu-recipes 저장소의 Ironwood에서 vLLM TorchTPU 레시피, Kubernetes 배포 가이드가 포함됩니다.
  • 토론은 vLLM Developer Slack의 #sig-tpu 채널에서 이루어집니다.

SUMMARY: vLLM TorchTPU(vllm-torchtpu)는 Google의 TorchTPU 스택을 기반으로 구축된 새로운 out-of-tree vLLM 플랫폼 플러그인으로, 모델 코드를 변경하지 않고도 대부분의 기존 vLLM 모델을 Google Cloud TPU에서 서빙하며, 수작업으로 튜닝된 Pallas 커널과 vLLM의 6가지 병렬 처리 차원을 모두 지원합니다.

TITLE: vLLM TorchTPU: Google Cloud TPU 추론을 위한 새로운 vLLM 백엔드

Sources