왜 지금 TPU에서 PyTorch 네이티브 추론인가

TPU는 수년간 대규모 AI 워크로드를 고성능·고효율로 처리해 왔습니다. 하지만 개발자들은 늘 같은 불만을 토로했습니다. "JAX 백엔드는 좋은데, 우리 서빙 스택을 PyTorch로 커스터마이징하고 디버깅하고 싶다." 에이전틱(agentic) 워크로드와 멀티모달 모델이 본격화되면서, 단순히 '돌아가는' 수준이 아니라 '효율적으로 서빙'하는 게 핵심 경쟁력이 됐죠.

그래서 구글은 TorchTPU를 만들었습니다. PyTorch를 TPU에서 네이티브로 실행하는 레이어입니다. 그리고 그 위에 vLLM(Inferact 협업)과 SGLang(RadixArk 협업) 엔진을 새로 올렸습니다. 목표는 명확합니다. PyTorch 커뮤니티가 익숙한 도구 그대로 모델을 TPU에 올리고 빠르게 서빙하도록 만드는 것.

이 글은 구글 개발자 블로그의 근거자료를 바탕으로 실무 관점에서 재구성했습니다.

TorchTPU의 동작 방식: 래퍼 없는 네이티브 디바이스

TorchTPU는 TPU를 'PyTorch 네이티브 디바이스'로 만듭니다. 핵심은 다음과 같아요.

  • 텐서는 tpu 디바이스 위의 표준 PyTorch 텐서로 관리됩니다. 래퍼 타입이 없습니다.
  • 코드는 즉시 실행(eager)되거나 torch.compile로 XLA 컴파일됩니다.
  • 집합 통신(collective)은 torch.distributed를 그대로 통과합니다.
  • 각 디바이스가 자체 PyTorch 프로세스를 돌리기 때문에, 두 엔진 모두 각 프레임워크의 워커 모델·스케줄러·KV 캐시 매니저·API 서버를 그대로 유지하고 하위만 TPU 구현으로 갈아끼웁니다.

대부분의 모델은 두 엔진에서 표준 업스트림 PyTorch 모델 정의를 그대로 사용합니다. TPU 특화 최적화는 커널과 레이어 구현에만 들어갑니다. Pallas로 작성한 커스텀 커널은 PyTorch의 커스텀 오퍼레이터로 컴파일 그래프 안에서 직접 호출됩니다. 이 '플랫폼-플러그인' 설계 덕분에 프레임워크 재사용은 극대화하고, TPU 최적화는 독립적으로 빠르게 발전시킬 수 있습니다.

프로덕션 기능과 에이전틱 워크로드 대응

두 엔진 모두 프로덕션 전 범위를 겨냥합니다. 특히 에이전틱 워크로드는 여러 기능을 동시에 요구합니다. 에이전트는 확장되는 멀티턴 대화를 반복 재전송하기 때문에, 재계산 대신 캐시된 프리픽스를 찾아야 효율적입니다. 두 엔진은 재사용 가능한 프리픽스를 HBM에 보관하고, 호스트 오프로드를 켜면 오래된 프리픽스를 2차 계층인 호스트 메모리로 자동 이동시킵니다.

여기서 두 기능이 KV 캐시를 '계산한 디바이스 밖'으로 옮깁니다.

  • 디스어그리게이티드 서빙(disaggregated serving): 프리필과 디코드를 별도 TPU 호스트에서 실행. 요청의 KV 캐시가 프리필 호스트에서 디코드 호스트로 이동해야 합니다.
  • 티어드 KV 오프로드(tiered KV offload): HBM에 더 이상 안 들어가는 KV 캐시를 호스트 DRAM과 Google Cloud Managed Lustre 보조 스토리지로 이동하고, 나중에 같은 프리픽스를 재사용할 때 다시 가져옵니다.

어느 쪽이든 전송이 연산과 겹치지 않으면 지연에 그대로 더해집니다. 두 엔진은 이를 TPU Sync로 해결합니다. TPU에서 KV 캐시 전송과 관리를 위한 오픈소스 C++ 라이브러리입니다. PyTorch 텐서 뒤의 디바이스 버퍼를 직접 읽기 때문에 호스트에 별도 스테이징 복사가 필요 없고, Python 밖의 자체 스레드에서 전송을 실행해 프리필·디코드가 계속 진행됩니다. vLLM-TorchTPU는 vLLM의 KV 커넥터 인터페이스로 노출하고, SGLang-TorchTPU는 디스어그리게이션 전송 백엔드 겸 HiCache I/O 백엔드로 사용합니다.

두 엔진은 이미 프리뷰 고객과 함께 프로덕션 규모로 돌고 있습니다. 성능 투명성을 위해 커뮤니티 표준인 SemiAnalysis InferenceX 프로젝트로 공개 벤치마크를 진행합니다.

TPU server rack powering large-scale PyTorch inference engines vLLM and SGLang Dev Environment Setup

벤치마크와 파트너십: 숫자로 보는 현주소

Ironwood TPU(TPU 7x)에 Qwen 3.5를 올린 초기 제출을 바탕으로, 엔지니어링 팀이 스택을 추가 최적화해 백만 토큰당 최저 비용 2.01센트라는 업계 선도 수치를 기록했습니다. 에이전틱 롱컨텍스트·멀티턴 워크로드를 평가하는 AgentX 트랙에도 새로운 SOTA 결과를 공개했고, Ironwood TPU가 파레토 프론티어 대부분에서 앞섭니다.

이 백엔드는 진공에서 나온 게 아닙니다. PyTorch를 받아들이고 오픈소스 메인테이너와 직접 일하면서 업스트림 정렬과 커뮤니티 우선을 유지합니다.

  • SGLang 팀과의 협업: TPU용 엔진을 만들어 메타 프로덕션에서 이미 돌리고 있습니다. 메타의 PyTorch-TPU 리드 Zongwei Zhou는 "PyTorch 네이티브 서빙 스택이 기존 AI 인프라에 자연스럽게 통합된다"고 밝혔습니다.
  • RadixArk: TorchTPU가 다양한 실리콘에 SGLang을 롤아웃하는 개발 속도를 크게 높였다고 평가합니다. Banghua Zhu CTO는 "고객에게 효율 이득과 유연한 컴퓨트 선택을 그대로 전달한다"고 했습니다.
  • Kimi K3 사례: 2.8조 파라미터 Kimi K3용 고성능 Pallas 메가커널을 통합해, 디코드 스텝 전체를 하나의 커널로 실행. Ironwood TPU에서 단일 사용자에게 초당 700토큰 이상을 서빙합니다.
  • Inferact: vLLM-TorchTPU 스택으로 Kimi K3에 대한 첫 AgentX 예비 제출을 주도했습니다.
# TorchTPU 위에서 vLLM 엔진을 사용하는 개념 예시 (의사 코드)
from vllm import LLM, SamplingParams

# TPU 디바이스에서 네이티브 PyTorch 텐서로 동작
llm = LLM(
    model="Qwen/Qwen3.5",
    device="tpu",              # TPU 네이티브 디바이스
    enable_prefix_caching=True, # 에이전틱 멀티턴 대응
    kv_cache_dtype="auto",
)

sampling = SamplingParams(temperature=0.7, max_tokens=512)
out = llm.generate(["TPU에서 PyTorch 네이티브 추론을 설명해줘"], sampling)
print(out[0].outputs[0].text)

위 코드는 개념 이해용 예시입니다. 실제 API는 각 엔진의 공식 문서를 확인하세요.

주의사항과 한계: 도입 전에 짚어야 할 것들

  1. JAX 백엔드와의 공존: vLLM-JAX와 SGLang-JAX는 계속 유지·지원됩니다. 기존 JAX 기반 파이프라인을 무리하게 갈아엎을 필요는 없습니다.
  2. 모델 커버리지: Qwen, DeepSeek, GLM, Kimi, GPT-OSS, Gemma 등 인기 오픈 모델군은 지원되지만, 신규 모델의 Day 0 지원은 계속 확장 중입니다. 사내 파인튜닝 모델은 커널 호환성 검증이 필요합니다.
  3. 디스어그리게이션·오프로드의 복잡도: TPU Sync가 전송을 숨겨주지만, 프리필/디코드 호스트 분리는 네트워크 토폴로지와 스케줄링 설계를 요구합니다. 소규모 워크로드에서는 오히려 오버헤드가 될 수 있습니다.
  4. 벤치마크 해석: 2.01센트/백만 토큰은 특정 모델·TPU 구성·배치 조건에서의 수치입니다. 자사 워크로드에 그대로 대입하지 말고 InferenceX 리더보드의 조건을 확인하세요.

다음 단계 학습 방향

  • TorchTPU 아키텍처: torch.compile + XLA 백엔드가 TPU 그래프를 어떻게 컴파일하는지 공부하면 커스텀 커널 최적화 감이 잡힙니다.
  • Pallas 커널: TPU용 커스텀 커널 작성법을 익히면 메가커널 사례를 이해할 수 있습니다.
  • KV 캐시 관리: 프리픽스 캐싱, 티어드 오프로드, 디스어그리게이션 서빙은 LLM 서빙 엔지니어링의 핵심 주제입니다.
  • vLLM/SGLang 내부 구조: 스케줄러와 KV 커넥터 인터페이스를 뜯어보면 TPU 백엔드가 어디에 끼어드는지 명확해집니다.

국내 대규모 추론 인프라를 운영하는 팀이라면, TPU 도입 시 비용 모델(예약 vs 온디맨드)과 GKE 연동, 사내 PyTorch 파이프라인과의 정합성을 먼저 검토하는 걸 권합니다. 특히 국내 SI 환경에서는 기존 JAX 기반 자산과의 마이그레이션 비용을 반드시 산정하세요.

Cloud infrastructure diagram showing disaggregated prefill and decode on TPU hosts IT Technology Image

vLLM-TorchTPU vs SGLang-TorchTPU 비교

항목vLLM-TorchTPUSGLang-TorchTPU
개발 파트너InferactRadixArk
TPU 네이티브 기반TorchTPUTorchTPU
KV 캐시 전송vLLM KV 커넥터 인터페이스로 TPU Sync 노출디스어그리게이션 전송 백엔드 + HiCache I/O 백엔드
프리픽스 캐싱HBM + 호스트 오프로드 2계층HBM + 호스트 오프로드 2계층
디스어그리게이티드 서빙지원지원
대표 사례Kimi K3 AgentX 예비 제출메타 프로덕션, Kimi K3 메가커널 700+ tok/s
커스텀 커널Pallas 커널을 PyTorch 커스텀 오퍼레이터로 호출동일
생태계 정렬업스트림 vLLM 정렬업스트림 SGLang 정렬
적합한 팀vLLM 기반 서빙 스택 보유SGLang 기반 에이전틱 서빙 중심

두 엔진 모두 JAX 기반 백엔드(vLLM-JAX, SGLang-JAX)와 병행 운영됩니다. 신규 도입이라면 워크로드 특성(에이전틱 멀티턴 vs 단발 대량 배치)과 팀의 기존 프레임워크 숙련도를 기준으로 선택하세요.

AI agent workload running on TPU with PyTorch native backend for high throughput Development Concept Image

결론: TPU 생태계의 'PyTorch 네이티브' 전환점

이번 발표의 본질은 새로운 엔진 두 개가 아니라, TPU가 PyTorch 네이티브 디바이스가 됐다는 점입니다. 래퍼 없는 텐서, torch.compile/XLA, torch.distributed 그대로. 여기에 Pallas 커스텀 커널과 TPU Sync가 더해져 에이전틱 워크로드의 KV 캐시 병목을 정면으로 다룹니다.

실무 관점 조언은 세 가지입니다.

  1. 작게 시작하세요. 기존 PyTorch 모델 하나를 TorchTPU 위 vLLM으로 올려 InferenceX 조건과 자사 지표를 비교하세요.
  2. KV 캐시 전략을 먼저 설계하세요. 프리픽스 재사용 패턴이 명확하면 티어드 오프로드와 디스어그리게이션이 큰 이득이 됩니다.
  3. JAX 자산은 유지하세요. 두 진영 모두 계속 지원되므로, 검증된 워크로드는 그대로 두고 신규 에이전틱 서빙부터 PyTorch 네이티브로 옮기는 하이브리드가 현실적입니다.

함께 보면 좋은 글

본 콘텐츠는 신뢰할 수 있는 출처를 바탕으로 AI 도구를 활용하여 초안이 작성되었으며, 편집자의 검토를 거쳐 발행되었습니다. 전문가의 조언을 대체하지 않습니다.