Google, TPU용 PyTorch 네이티브 vLLM·SGLang 추론 엔진 발표

두 엔진은 TorchTPU 위에 구축되어 있다. TorchTPU는 TPU를 PyTorch의 네이티브 디바이스로 다루며, 텐서는 래퍼 타입 없이 표준 PyTorch 텐서로 관리된다. 코드는 eager 실행이나 XLA를 사용한 torch.compile로 동작하며, 집합 통신은 torch.distributed를 거친다. 워커 모델, 스케줄러, KV 캐시 매니저, API 서버는 각 프레임워크의 것을 그대로 유지하고, 그 아래에 TPU 구현을 더하는 설계이다. Pallas로 작성한 커스텀 커널은 컴파일된 그래프 안의 커스텀 오퍼레이터로서 PyTorch에서 호출된다.

지원 기능은 폭넓다. 자동 프리픽스 캐싱, 호스트 메모리나 Managed Lustre로의 KV 캐시 오프로딩, 도구 호출, MTP와 EAGLE-3 등을 이용한 추측적 디코딩, expert parallelism을 동반한 data-parallel attention, 컨텍스트 병렬 처리, prefill과 decode를 서로 다른 호스트로 나누는 분리형 서빙 등이 언급되었다. KV 캐시 이동에는 오픈소스 C++ 라이브러리 TPU Sync를 사용하며, 호스트 측 추가 복사 없이 PyTorch 텐서 뒤의 디바이스 버퍼를 직접 읽는다.

성능 면에서는 SemiAnalysis InferenceX에서 공개 벤치마크를 진행했으며, 에이전트형 장문맥 및 멀티턴을 평가하는 AgentX 트랙에서도 새로운 결과를 냈다고 밝혔다. Ironwood TPU가 Pareto frontier의 대부분에서 앞선다고 한다. 두 엔진은 프리뷰 고객의 프로덕션 환경에서 가동 중이며, SGLang-TorchTPU는 Meta에서 프로덕션으로 운영되고 있다고 Meta 측 담당자가 코멘트했다.

vLLM 측에서는 Inferact와 함께 2.8T 파라미터의 Kimi K3용 Pallas 메가커널을 통합했으며, decode 단계 전체를 하나의 커널로 실행해 Ironwood TPU에서 단일 사용자에게 700토큰/초 이상을 제공했다고 밝혔다. vLLM은 또한 두 엔진에 적극적으로 기여하는 개발자에게 제공할 TPU 컴퓨팅 리소스의 무상 풀을 마련했다. Qwen, DeepSeek, GLM, Kimi, GPT-OSS, Gemma4 등 주요 오픈 모델을 지원하며, 신규 모델의 Day 0 지원에도 힘쓴다. 기존의 vLLM-Jax와 SGLang-JAX도 계속 유지보수된다.