TL;DR

TPU Raiden은 TPU 상에서 JAX와 PyTorch용 C++ 확장과 Python 래퍼를 제공하는 개발 중인 프로젝트로, Bazel 8.6.0과 Python 3.12 환경을 요구한다. PyTorch 경로는 ABI에 민감해 정확히 매칭되는 torch/torch_tpu가 필요하며 patchelf를 사용해 빌드된 확장에 링크 정보를 주입해야 정상 동작한다. 운영 측면에서는 /dev/shm 기반의 영속적 KV 캐시를 통해 프로세스 재시작 시 cold start를 방지하는 기능과 TPUVM 빌드 시 디스크 공간 설정 같은 실무 팁이 문서에 구체적으로 적혀 있어 TPU에서 직접 서비스하거나 실험하려는 엔지니어에게 필요한 전제와 절차를 한눈에 확인시켜 준다.

핵심 포인트

• 레포는 TPU에서 동작하는 C++ 확장 라이브러리와 그에 연동되는 Python 래퍼를 제공하며 JAX와 PyTorch 양쪽을 모두 지원하는 구조를 갖고 있다. 소스 빌드는 Bazel 8.6.0으로 컴파일된 공유 라이브러리(_tpu_raiden_jax.so, _tpu_raiden_torch.so)를 생성하고 PyPI용 휠 대신 Google Artifact Registry의 사내 휠이나 로컬 빌드를 통해 배포할 수 있다. 안정된 작업 지점은 루트의 lkg.version 파일에 커밋 해시로 기록되어 있어서 git checkout $(cat lkg.version)으로 검증된 리비전을 재현할 수 있다.
• PyTorch 경로는 ABI 민감성이 명시되어 있어 torch_tpu와 정확히 맞는 torch 버전이 필요하며 문서에 예시로 torch==2.11.0을 언급하고 있다. PyTorch 확장 빌드에는 patchelf가 필수이고, 빌드 스크립트가 생성된 .so 파일에 patchelf --add-needed를 실행해 libpywrap_torch_tpu_common.so 의 NEEDED 엔트리를 주입해 중복 XLA allocator 등록 오류를 예방한다. 따라서 PyTorch 환경을 맞추지 않으면 확장 모듈을 정상적으로 로드할 수 없고 문서에 이런 전제 조건을 분명히 제시하고 있다.
• TPUVM 특성을 반영한 운영 주의사항과 영속적 KV 캐시용 POSIX shared memory (/dev/shm) 지원을 포함하고 있어서 서비스 재시작 시 cold start를 피할 수 있다. RAIDEN_SHM_KEY, RAIDEN_SHM_MODEL_UID, RAIDEN_SHM_SERVER_NAME 같은 환경변수를 통해 네임스페이스와 모델 검증 UID를 지정하고, 초기 실행(cold boot)에는 파일을 생성하며 재시작(warm boot) 시 기존 세그먼트에 재부착하고 TPU DMA 엔진에 페이지를 재등록한다. /dev/shm/raiden_cache_* 파일을 수동으로 삭제해 호스트 메모리를 회수할 수 있다는 운영 절차도 README에 구체적으로 적혀 있다.
• 개발·테스트 워크플로우는 제공된 스크립트로 단순화되어 있는데 ./build.sh [jax|torch|both]로 빌드 범위를 지정하고 ./run_tests.sh [jax|torch|both]로 단위 테스트를 실행할 수 있다. 빌드 스크립트는 워크스페이스로 이동해 Bazel 컴파일, (PyTorch의 경우) patchelf 적용, requirements 설치, 그리고 컴파일된 .so를 대응 프레임워크 패키지로 복사하는 단계로 구성되어 있다. TPUVM의 /tmp 디스크 부족 문제를 피하려면 BAZEL_OUTPUT_BASE를 충분한 디스크가 있는 디렉토리로 설정해야 한다는 운영 팁이 포함되어 있어 TPU 환경에서 빌드할 때 실무적으로 유의해야 할 점을 바로 확인할 수 있다.