TPU

TPU

SGLang은 Google Cloud TPU를 위해 특별히 최적화된 SGLang-JAX 백엔드를 통해 고성능 TPU 추론을 지원해요. JAX 기반 구현은 TPU 하드웨어의 LLM(Large Language Model) 서빙 작업부하에서 뛰어난 처리량과 낮은 지연을 제공해요.

출처: 문서

본문

TPU 관련 문제나 기능 요청은 sglang-jax GitHub issues 페이지를 방문하세요.

참고: SGLang TPU 지원은 https://github.com/sgl-project/sglang-jax에 별도 저장소로 유지되는 전용 JAX 기반 추론 엔진인 SGLang-JAX 백엔드를 통해 구현돼요.

시스템 요구 사항 (System Requirements)

지원 TPU 하드웨어 (Supported TPU Hardware)

소프트웨어 요구 사항 (Software Requirements)

• Python: 3.12 이상
• JAX: TPU 지원 최신 버전
• 환경: Google Cloud TPU VM 또는 호환 TPU 런타임
• 선택: 간편한 클라우드 배포를 위한 SkyPilot

기능 지원 매트릭스 (Feature Support Matrix)

SGLang-JAX는 프로덕션 LLM 서빙을 위한 포괄적인 TPU 최적화 기능을 제공해요:

어텐션 백엔드 비교 (Attention Backend Comparison)

참고: 우수한 메모리 효율성과 성능 때문에 프로덕션 작업부하에는 FlashAttention 백엔드를 권장해요.

최적화 모델 목록 (Optimized Model List)

SGLang-JAX는 다음 모델 패밀리를 지원해요. 배포 지침은 연결된 TPU cookbook을 사용하세요.

설치 (Installation)

방법 1: PyPI 사용 (권장)

pip install sglang-jax

방법 2: 소스에서 (From Source)

git clone https://github.com/sgl-project/sglang-jax
cd sglang-jax
uv venv --python 3.12 && source .venv/bin/activate
uv pip install -e "python[all]"

방법 3: Docker 사용

참고: TPU용 Docker 지원은 현재 개발 중이에요. PyPI 또는 소스 설치 방법을 사용하세요.

방법 4: SkyPilot으로 Cloud TPU

SkyPilot은 Google Cloud TPU에서 간편한 배포를 제공해요:

1. SkyPilot을 설치하고 GCP 접근을 구성하세요 (SkyPilot documentation 참고).

2. SkyPilot 구성 파일을 만드세요:

# sglang-jax.sky.yaml
resources:
   accelerators: tpu-v6e-4
   accelerator_args:
      tpu_vm: True
      runtime_version: v2-alpha-tpuv6e

run: |
  git clone https://github.com/sgl-project/sglang-jax.git
  cd sglang-jax
  uv venv --python 3.12
  source .venv/bin/activate
  uv pip install -e "python[all]"

3. TPU 클러스터를 실행하세요:

# 표준 배포
sky launch -c sglang-jax sglang-jax.sky.yaml --infra=gcp

# 비용 절감을 위한 spot 인스턴스
sky launch -c sglang-jax sglang-jax.sky.yaml --infra=gcp --use-spot

서빙 엔진 실행 (Launch of the Serving Engine)

기본 예시: Qwen-7B

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server \
    --model-path Qwen/Qwen-7B-Chat \
    --trust-remote-code \
    --dist-init-addr=0.0.0.0:10011 \
    --nnodes=1 \
    --tp-size=4 \
    --device=tpu \
    --random-seed=3 \
    --node-rank=0 \
    --mem-fraction-static=0.8 \
    --max-prefill-tokens=8192 \
    --download-dir=/tmp \
    --dtype=bfloat16 \
    --skip-server-warmup \
    --host 0.0.0.0 \
    --port 30000

주요 파라미터 설명:

1. JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache — 이후 실행에서 서버 시작을 가속화하기 위해 JIT 컴파일 캐싱 활성화
2. --tp-size=4 — 텐서 병렬 처리 크기; TPU 코어 수(보통 1, 4, 8)와 일치시켜야 함
3. --device=tpu — TPU 디바이스 지정 (sglang-jax의 기본값)
4. --dtype=bfloat16 — TPU가 최적화된 bfloat16 정밀도 사용
5. --mem-fraction-static=0.8 — TPU HBM의 80%를 정적 메모리에 할당 (0.2에서 0.9까지 조정 가능)
6. --max-prefill-tokens=8192 — prefill 단계에서 처리하는 최대 토큰 수

고성능 구성: Qwen3-8B

최적 처리량을 위한 프로덕션 작업부하용:

python3 -u -m sgl_jax.launch_server \
    --model-path Qwen/Qwen3-8B \
    --trust-remote-code \
    --tp-size=4 \
    --device=tpu \
    --mem-fraction-static=0.8 \
    --chunked-prefill-size=2048 \
    --dtype=bfloat16 \
    --max-running-requests=256 \
    --page-size=128 \
    --attention-backend=fa

고급: 투기적 디코딩 (EAGLE3)

투기적 디코딩은 호환 모델의 처리량을 20-40% 향상시킬 수 있어요:

python3 -u -m sgl_jax.launch_server \
    --model-path Qwen/Qwen3-32B \
    --trust-remote-code \
    --device=tpu \
    --tp-size=4 \
    --mem-fraction-static=0.8 \
    --max-prefill-tokens=4096 \
    --attention-backend=fa \
    --dtype=bfloat16 \
    --port=30000 \
    --host=0.0.0.0 \
    --disable-overlap-schedule \
    --speculative-algorithm=EAGLE3 \
    --speculative-draft-model-path=AngelSlim/Qwen3-32B_eagle3 \
    --page-size=64 \
    --speculative-eagle-topk=1 \
    --speculative-num-steps=3 \
    --speculative-num-draft-tokens=4

참고: 투기적 디코딩은 현재 Qwen3와 LLaMA 모델 패밀리를 지원해요. 자세한 구성 지침은 Speculative Decoding documentation을 참고하세요.

다중 노드 분산 서빙 (Multi-Node Distributed Serving)

여러 TPU VM이 필요한 큰 모델용:

# Node 0 (코디네이터)
python3 -m sgl_jax.launch_server \
    --model-path MODEL_PATH \
    --dist-init-addr=NODE0_IP:10011 \
    --nnodes=2 \
    --node-rank=0 \
    --tp-size=8 \
    [other parameters...]

# Node 1 (워커)
python3 -m sgl_jax.launch_server \
    --model-path MODEL_PATH \
    --dist-init-addr=NODE0_IP:10011 \
    --nnodes=2 \
    --node-rank=1 \
    --tp-size=8 \
    [other parameters...]

요청으로 벤치마킹 (Benchmarking with Requests)

처리량 테스트 (Throughput Testing)

기본 처리량 벤치마크:

python3 -m sgl_jax.bench_serving \
    --backend sgl-jax \
    --dataset-name random \
    --num-prompts=100 \
    --random-input=512 \
    --random-output=128 \
    --max-concurrency=8 \
    --random-range-ratio=1 \
    --warmup-requests=0

지연 테스트 (Latency Testing)

단일 배치 지연 측정:

python3 -m sgl_jax.bench_one_batch_server \
    --base-url http://127.0.0.1:30000 \
    --model-path Qwen/Qwen-7B-Chat \
    --batch-size=32 \
    --input-len=256 \
    --output-len=32

종합 벤치마크 스크립트 (Comprehensive Benchmark Script)

다양한 구성에서 체계적인 성능 평가용:

#!/bin/bash
set -e

backend=${1:-sgl-jax}
num_prompts_per_concurrency=3
input_seq_lens=(1024 4096 8192)
output_seq_lens=(1 1024)
max_concurrencies=(8 16 32 64 128 256)

for input_seq_len in "${input_seq_lens[@]}"; do
    for output_seq_len in "${output_seq_lens[@]}"; do
        echo "======================================="
        echo "Testing ISL/OSL: $input_seq_len/$output_seq_len"
        echo "======================================="
        for max_concurrency in "${max_concurrencies[@]}"; do
            num_prompts=$((num_prompts_per_concurrency * max_concurrency))
            python3 -m sgl_jax.bench_serving \
                --backend ${backend} \
                --dataset-name random \
                --num-prompts ${num_prompts} \
                --random-input ${input_seq_len} \
                --random-output ${output_seq_len} \
                --max-concurrency ${max_concurrency} \
                --random-range-ratio 1 \
                --disable-ignore-eos \
                --warmup-requests 0
        done
    done
done

모든 벤치마크 파라미터의 자세한 도움말:

python3 -m sgl_jax.bench_serving --help

고급 벤치마킹 기법과 JAX Profiler 프로파일링은 Benchmark and Profiling Guide를 참고하세요.

성능 최적화 (Performance Optimization)

메모리 최적화 (Memory Optimization)

메모리 사용량 줄이기:

• --mem-fraction-static 낮추기 (0.8 → 0.5 → 0.3)
• --max-prefill-tokens 줄이기 (16384 → 8192 → 4096)
• --max-running-requests 줄이기

OOM 오류 처리:

• 보수적인 메모리 설정으로 시작 (--mem-fraction-static=0.5)
• 최적의 균형을 찾을 때까지 점진적으로 증가
• 더 나은 메모리 국소성(locality)을 위해 --page-size 증가 (1 → 16 → 64 → 128)

처리량 최적화 (Throughput Optimization)

초당 토큰을 최대화하려면:

• FlashAttention 백엔드 사용: --attention-backend=fa
• Qwen3 모델에서 투기적 디코딩(EAGLE3) 활성화 (20-40% 향상)
• --max-running-requests를 256+로 증가
• --mem-fraction-static을 0.8+로 설정 (메모리가 허락한다면)
• 더 큰 페이지 크기 사용 (64-128)
• Chunked prefill 활성화: --chunked-prefill-size=2048

지연 최적화 (Latency Optimization)

첫 토큰 시간(TTFT)과 토큰 간 지연을 최소화하려면:

• --page-size를 1-4로 줄이기
• 더 작은 배치를 위해 --max-running-requests 낮추기 (16-32)
• --chunked-prefill-size 줄이기
• GC 일시정지를 피하려면 보수적인 메모리 설정 사용

TPU 특화 최적화 (TPU-Specific Optimizations)

1. JIT 컴파일 캐시:

   export JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache
   

   컴파일된 커널을 캐시하고 서버 시작을 가속화하려면 항상 이 환경 변수를 설정하세요.

2. 데이터 타입 최적화: TPU 네이티브 최적화를 위해 --dtype=bfloat16 사용. TPU는 bfloat16 계산을 위해 특별히 설계됐어요.

3. 텐서 병렬 처리: 최적의 모델 분산을 위해 --tp-size를 TPU 코어 구성(1, 4, 8)과 일치시키세요.

4. 어텐션 백엔드: 프로덕션 작업부하에는 항상 --attention-backend=fa(FlashAttention)를 사용하세요.

문제 해결 (Troubleshooting)

OOM (Out of Memory) 오류

메모리 부족 오류가 발생하면:

1. --mem-fraction-static을 0.8에서 0.5 이하로 낮추기
2. --max-prefill-tokens을 8192에서 4096 또는 2048로 줄이기
3. 동시 배치 크기를 줄이기 위해 --max-running-requests 낮추기
4. 더 나은 메모리 레이아웃 효율을 위해 --page-size 증가

컴파일 시간이 길 때 (Compilation Long-Time)

서버 시작이 너무 오래 걸리면:

1. JAX_COMPILATION_CACHE_DIR이 제대로 설정됐는지 확인
2. 첫 실행은 JIT 컴파일을 필요로 한다는 점 이해 (정상적인 동작)
3. 캐시된 컴파일 덕분에 이후 실행은 훨씬 빠름
4. 첫 요청까지 컴파일을 미루려면 --skip-server-warmup 사용 고려

낮은 처리량 (Low Throughput)

예상 처리량을 얻지 못한다면:

1. --tp-size가 TPU 코어 구성과 일치하는지 확인
2. --attention-backend=fa가 활성화됐는지 확인
3. 더 큰 배치 형성을 위해 --max-running-requests 증가
4. 호환 모델에서 투기적 디코딩 활성화 고려
5. 메모리 설정이 충분한 배치 크기를 허용하는지 확인

연결 문제 (Connection Issues)

클라이언트가 서버에 연결할 수 없으면:

1. 외부 접근(127.0.0.1이 아닌)을 위해 --host=0.0.0.0 확인
2. 방화벽 규칙이 지정된 포트(기본: 30000)의 트래픽을 허용하는지 확인
3. 서버 프로세스가 실행 중인지 확인: curl http://localhost:30000/health

고급 기능 (Advanced Features)

투기적 디코딩 (Speculative Decoding)

SGLang-JAX는 Qwen3와 LLaMA 모델 패밀리를 위한 EAGLE 및 EAGLE3 투기적 디코딩 알고리즘을 지원해요. 투기적 디코딩은 출력 품질에 영향을 주지 않고 처리량을 20-40% 향상시킬 수 있어요.

자세한 구성과 지원 모델 조합은 Speculative Decoding documentation을 참고하세요.

Chunked Prefill

더 나은 TPU 활용을 위해 혼합 prefill-decode 배칭을 활성화하세요:

--chunked-prefill-size=2048 --enable-mixed-chunk

이를 통해 스케줄러가 같은 배치에서 prefill 연산과 decode 연산을 혼합해 전체 처리량을 향상시킵니다.

커스텀 어텐션 백엔드 (Custom Attention Backends)

SGLang-JAX는 플러그인 기반 어텐션 백엔드 시스템을 지원해요. 특정 사용 사례에 최적화된 커스텀 어텐션 커널을 구현할 수 있어요.

구현 세부 사항은 Attention Backend documentation을 참고하세요.

환경 검증 (Environment Verification)

배포 전에 TPU 설정을 검증하세요:

python -c "from sgl_jax import check_env; check_env.check_env()"

이 명령은 다음을 확인해요:

• 설치된 패키지 버전
• TPU 디바이스 가용성과 사양
• 시스템 리소스와 구성
• 설정의 호환성

기여 (Contributing)

SGLang-JAX의 TPU 지원을 개선하는 데 기여를 환영해요!

기여 영역 (Areas for Contribution)

계획된 기능과 새로운 기능 기여 기회를 보려면 Development Roadmap을 확인하세요.

현재 기여 영역:

• 특정 TPU 세대의 성능 최적화
• 추가 모델 아키텍처 지원
• 문서 개선과 예시
• 버그 리포트와 수정
• 벤치마크 결과와 성능 분석

기여 방법 (How to Contribute)

1. sglang-jax 저장소 방문
2. 기여 가이드 읽기
3. 토론을 위해 SGL-JAX Slack 커뮤니티 가입
4. sglang-jax/issues에서 문제 보고

TPU에서 테스트 (Testing on TPU)

테스트를 위해 TPU 접근이 필요한 기여자:

• TPU 하드웨어 접근 정보는 TPU Resources Guide 참고
• 비용 효율적인 테스트를 위해 spot 인스턴스로 SkyPilot 사용
• 성능 검증은 Benchmark and Profiling Guide 따르기

참고 자료 (References)

문서 (Documentation)

• SGLang-JAX 저장소
• SGLang-JAX 설치 가이드
• Qwen 모델 퀵스타트
• 벤치마크와 프로파일링 가이드
• 투기적 디코딩

외부 리소스 (External Resources)

• JAX 문서
• Google Cloud TPU 문서
• SkyPilot 문서

더 알아보기 (Learn more)

• SGLang-JAX 저장소
• 하드웨어 플랫폼 개요