피드로 돌아가기
Dev.toAI/ML
원문 읽기
TPU의 저비용 고효율을 가로막는 Static Shape 제약과 아키텍처 비용 분석
Why TPUs Aren't Popular (Even Though They're Cheaper Per Token)
AI 요약
Context
TPU/Trainium은 Token당 비용과 전력 효율이 NVIDIA GPU보다 우수함에도 불구하고 실제 도입률이 낮음. 이는 하드웨어의 Systolic Array 구조가 컴파일 타임에 데이터 형상을 고정해야 하는 Static Shape 제약을 가지기 때문임.
Technical Solution
- SIMT 방식의 GPU와 달리 TPU는 AOT(Ahead-of-Time) 컴파일러(XLA)를 통한 고정형 바이너리 생성 구조 채택
- 데이터 형상 변경 시마다 새로운 바이너리를 생성해야 하는 JIT Recompile Stall 문제 발생
- 가변 길이 입력을 처리하기 위해 고정 크기 텐서에 데이터를 맞추는 Padding 및 Masking 기법 강제
- FlashAttention의 cu_seqlens를 통한 동적 할당 대신 StaticShapeAttention 구조로 불필요한 FLOPs 소모
- 입력 데이터의 형상을 미리 정의하고 이에 맞는 미리 컴파일된 바이너리로 라우팅하는 Bucket Routing 전략 필요
- 고정된 비즈니스 워크플로우나 Long-context Prefill 작업을 통해 Padding 낭비를 최소화하는 최적화 수행
실천 포인트
- 입력 데이터의 Shape이 가변적인 챗봇 서비스라면 NVIDIA/AMD의 SIMT 아키텍처 우선 검토 - 입력 데이터 형상을 제어할 수 있는 CLI나 고정 파이프라인 환경에서 TPU의 Cost-per-token 이점 활용 - TPU 도입 시 단순 칩 비용 외에 Padding 처리 및 바이너리 관리 등을 위한 엔지니어링 공수(OpEx) 산정 - Long-context 워크로드 비중을 높여 Static Array의 상대적 낭비를 줄이는 데이터 패킹 전략 검토