gpumode · 강의 아카이브
《GPU Mode》 L088 2025 Mid priority transcript · failed

TinyTPU — TPU 를 처음부터 만들어보며 이해하기

Google 의 TPU (Tensor Processing Unit) 가 어떻게 동작하는지 — 추상이 아니라 한 개 칩의 systolic array, accumulator, MMU 를 직접 구현하면서 이해하는 학습용 emulator/구현 프로젝트. William Zhang 이 깐 학습 자료의 위치, systolic array 의 한 cycle 단위 동작, GEMM 의 매핑, GPU 와의 비교, 그리고 hardware bring-up 의 학습 가치를 — 자막 없이 외부 자료로 — 재구성한 학습 노트.

TPU systolic array emulator MAC unit accumulator SRAM GEMM
W
Speaker
William Zhang
TinyTPU 저자 · 확인 필요
강의 번호
L088
스피커
William Zhang
학습 우선순위
Mid · 호흡 고르기
자막
failed · 외부 자료
§ 01강의가 풀려는 문제· why TinyTPU

“TPU 가 어떻게 도는지” 를 추상이 아니라 코드로 만지는 길

대부분의 ML 엔지니어는 TPU 를 JAX 의 한 backend 로 만난다. jax.devices('tpu'), pmap, shard_map. 그 위에서 도는 hardware 가 무엇이고 왜 GPU 와 다르게 작동하는지는 — 추상 뒤에 가려져 있다. TinyTPU 는 그 추상을 한 단계 걷어낸다.

강의의 입장은 — 외부 자료에서 추정하건대 — 학습용 프로젝트의 일관된 입장. “이해의 가장 빠른 길은 직접 만들어보는 것”. 작은 TPU 의 핵심 (systolic array MXU + accumulator + SRAM) 을 Python / Verilog / 어떤 형태로든 구현하면서 — 다음 사실들이 손에 잡힌다.

  1. systolic array 가 GEMM 에 자연스러운 이유 — data 가 PE 사이를 한 cycle 씩 흘러가면서 곱하고 더한다. control logic 이 매우 단순.
  2. TPU 의 instruction set 이 작은 이유 — GPU 처럼 일반 compute 를 다루지 않고 matmul 한 가지에 최적화. ISA 가 그만큼 작다.
  3. HBM ↔ SRAM ↔ MXU 의 데이터 흐름 — TPU 의 latency / throughput 이 결정되는 자리.
  4. compiler 의 책임 — XLA 가 GEMM 을 systolic array 의 모양에 맞춰야 하는 이유.
자료 한계 · 확인 필요

강의 자막은 실패. 슬라이드 / repo 폴더가 GPU Mode 측에 미공개. TinyTPU 라는 이름의 프로젝트는 외부에 여러 변형이 있어 (GitHub 검색하면 학습용 TPU emulator 들이 등장) — 이 강의의 정확한 구현 형태는 영상 직접 확인 필요. 노트의 입장은 학습용 TPU emulator 의 일반적 구조 + Google 의 TPU 공식 자료 기반.

“추상은 이해를 빠르게 하지만, hardware 자체의 직관은 — 직접 cycle 단위로 만들어봐야 손에 들어온다.”학습 노트 · 재구성
§ 02TPU 아키텍처 한 눈에· MXU · vector unit · HBM

한 칩의 핵심 블록 — 그리고 그 사이의 데이터 경로

TPU 의 외부 인터페이스는 GPU 와 비슷해 보이지만 — 안의 구성이 매우 다르다. 한 그림에 박아둔다.

FIG · TPU 한 칩의 블록 다이어그램대략적 형태 (TPUv2/v3 기준)
┌────────────────────────────────────────────────┐ │ HBM (high-bandwidth) │ │ 파라미터 · activation │ └─────────────────────┬──────────────────────────┘ │ ▼ (~수백 GB/s) ┌────────────────────────────────────────────────┐ │ VMEM / on-chip SRAM (수십 MB) │ │ activation 과 parameter 의 staging │ └──┬──────────────────┬─────────────────────┬────┘ │ │ │ ▼ ▼ ▼ ┌──────┐ ┌──────────────┐ ┌────────┐ │MXU 1│ │ vector unit │ │ scalar │ │128×128│ │ ReLU, exp, │ │ unit │ │systolic│ │ bias 등 │ │ 제어 │ │array │ │ │ │ │ └──┬────┘ └──────┬───────┘ └────┬───┘ │ │ │ └─── accumulator ───┘ │ │ │ └────── ICI / DMA ───────────┘ (chip 사이 연결)
핵심 — MXU (Matrix Multiply Unit) 가 칩의 80% 이상 면적을 차지하는 systolic array. vector unit 이 elementwise op (ReLU, softmax 보조 등). scalar unit 이 control flow. HBM ↔ VMEM 사이의 DMA 가 데이터 흐름의 핵심 bottleneck.
MXU
128×128 systolic array (TPUv2/v3) 또는 256×256 (TPUv4+). 한 cycle 에 한 번의 matmul-accumulate. 전체 칩의 throughput 의 본업.
VMEM (Vector Memory)
on-chip SRAM. activation / parameter 의 staging. 수십 MB. GPU 의 SRAM 보다 큼.
vector unit
elementwise op (ReLU, sigmoid, exp). MXU 의 출력을 받아 후처리. SIMD-style.
scalar unit
control flow, address 계산, instruction issue. CPU 의 fetch/decode 와 비슷한 역할.
HBM
off-chip 메모리. 수십 GB. parameter 와 activation 의 본가. 한 step 의 데이터가 여기에 있다.
ICI (Inter-Chip Interconnect)
TPU pod 안의 칩들을 연결하는 mesh / torus 네트워크. NVLink 와 비슷한 역할.
GPU 와의 큰 차이

GPU 는 여러 SM (수십~수백) 이 SIMT 모델로 일한다. 같은 칩 안에 일반 compute 가 다양. TPU 는 한 거대한 MXU 가 systolic array 모델로 일한다. 일반 compute 는 vector unit 으로 좁게 다룸. GPU 가 일반 compute 의 가속기, TPU 가 GEMM 의 가속기 — 이 입장 차이가 ISA, 컴파일러, 사용자 코드 모두에 영향.

§ 03systolic array 모델· data flow

한 cycle 씩 데이터가 흘러가면서 자동으로 곱셈-덧셈이 일어난다

systolic array 의 직관은 — 1982년 Kung 의 원래 논문부터 — “data 가 array 를 통과하면서 일이 자동으로 일어난다”. 한 시간 frame 씩 따라가면 직관이 잡힌다.

FIG · 4×4 systolic array 의 4-cycle 진행weight stationary · activation flowing
cycle 1
a₀
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
cycle 2
·
a₀
·
·
a₁
·
·
·
·
·
·
·
·
·
·
·
cycle 3
·
·
a₀
·
·
a₁
·
·
a₂
·
·
·
·
·
·
·
cycle 4 (full)
·
·
·
a₀
·
·
a₁
·
·
a₂
·
·
a₃
·
·
·
한 PE (processing element) 안에는 한 weight 가 stationary 하게 박혀 있다. activation 이 한 cycle 마다 한 칸 오른쪽으로 이동. 각 PE 가 “자기 weight × 들어온 activation” 을 곱해 누적. 4×4 array 가 fully fill 되는 데 4 cycle, 그 후 매 cycle 에 4 개의 결과를 토해낸다.

이 모델의 가치는 — control logic 의 단순성. 각 PE 가 똑같은 일을 한다 — 입력 받기, 곱하기, 더하기, 다음 PE 로 전달. global control 도 데이터 distributor 도 필요 없다. 같은 면적에 더 많은 ALU를 박을 수 있는 이유.

weight stationary
weight 가 PE 에 박혀 있고 activation 이 흐름. TPU 의 표준 매핑. weight reuse 가 자동.
output stationary
결과(C 의 한 element) 가 한 PE 에 누적되고 weight 와 activation 이 둘 다 흐름. 다른 매핑.
input stationary
activation 이 PE 에 박혀 있고 weight 가 흐름. 또 다른 매핑.
pipeline fill
array 가 fully active 되기까지 N cycle (N = array dim). 작은 GEMM 에서 이 비용이 크다.
pipeline drain
마지막 데이터가 array 를 빠져나오는 데 또 N cycle. 같이 작은 GEMM 의 효율을 깎는 자리.
utilization
array 가 fully active 되는 비율. 작은 GEMM 에서 50% 미만으로 떨어지기 쉬움. compiler 의 큰 책임.
§ 04학습 목적 emulator· why build one

“만들어 보면 안다” — emulator 가 가르치는 것

실물 TPU 는 만질 수 없지만, emulator 는 직접 만들 수 있다. TinyTPU 같은 학습용 구현이 가르쳐주는 것들.

cycle-accurate timing
한 cycle 마다 PE 가 어떤 일을 하는지 visualize. timing 의 직관.
memory hierarchy
HBM ↔ VMEM ↔ MXU 의 latency / throughput 차이를 직접 모델링. data movement 의 비용 손에.
utilization tradeoff
작은 GEMM 의 utilization 이 왜 떨어지는지 직접 측정. matrix dim 과 systolic dim 의 관계.
compiler 의 일
한 GEMM 을 array 에 어떻게 매핑할지 — “이 일이 컴파일러의 본업” 임을 직접 코드로 만남.
non-matmul 의 비용
softmax, layer norm 같은 비-matmul op 가 vector unit 에서 어떻게 도는지. relative cost.
디버깅 가능
실물 hardware 의 black box 가 사라짐. 매 cycle 의 state 를 print 가능.
구현의 난이도 layer

학습용 구현은 보통 세 layer 로 진행 — L1: Python functional model(numpy 로 systolic array 의 시뮬레이션), L2: cycle-accurate Python emulator(timing 까지), L3: Verilog / Chisel RTL(FPGA 까지 가능). 강의가 어느 layer 에 초점인지는 — 확인 필요.

§ 05예시: GEMM on TPU· A·B → C 의 매핑

한 GEMM 이 systolic array 위에서 어떻게 도는가

M×K · K×N → M×N 의 GEMM 을 K=128, M=N=K 가 array dim 의 배수라고 가정하고, weight stationary 매핑을 따라가본다.

FIG · 128×128 systolic array 위 GEMM한 tile 의 cycle 단위 흐름
1. load B (128×128) from VMEM into MXU PEs as stationary weights 2. stream A rows through MXU — each row enters at one row of PEs 3. for cycle 1..128: - row_i of A enters at PE[0,_] - shifts right one column per cycle - each PE computes: acc += a_in * w_pe - and forwards a_in to PE[r, c+1] 4. for cycle 129..256: - results stream out of MXU bottom row - go to accumulator (vector register) 5. apply vector unit ops (bias add, ReLU) 6. write C back to VMEM, eventually HBM
한 128×128 tile 이 약 256 cycle (fill 128 + steady 128). 큰 GEMM 은 여러 tile 을 sequence 로 흘려보낸다. weight 를 다시 로드하는 비용이 크므로 — A 의 row 를 N 의 모양에 따라 오래 흘려보내야 효율. 이게 compiler 의 tile size 결정의 핵심.

이 그림에서 잡히는 사실 두 가지.

non-square shape 의 비용

실제 LLM 의 GEMM 은 정확히 128 의 배수가 아니다. M, N, K 가 다르고 작은 차원이 있을 수 있다. 이 자리에서 padding 또는 utilization 손실. compiler (XLA) 가 이걸 잘 매핑하는 게 큰 책임. L001 의 NCU 의 “tail effect” 와 같은 본질의 문제.

§ 06GPU 와 비교· SIMT vs systolic

같은 GEMM, 다른 hardware 패러다임

GPU (SIMT)
  • 일반 compute 가속기 — matmul 만이 아니라 다양한 패턴.
  • SM 안의 thread 들이 같은 instruction 을 다른 데이터에 — SIMT.
  • Tensor Core 가 matmul 가속을 받지만, 여전히 SM 위.
  • warp / block / grid 의 multi-level 추상.
  • register / shared memory / global 의 hierarchy 사용자 측 통제.
  • 비-matmul 도 ALU 위에서 자연스러움.
TPU (systolic)
  • matmul 가속기 — MXU 가 80% 면적.
  • data 가 PE 사이를 흘러가면서 일이 일어남 — control logic 단순.
  • 일반 compute 는 vector unit 에서 별도 — 좁은 set.
  • 한 chip 의 작업 단위가 “MXU 한 번 돌리기” 의 비교적 큰 단위.
  • VMEM 이 사용자 표면에 안 노출 — XLA 가 책임.
  • 비-matmul 은 별도 unit 에서 도므로 fusion 의 모양이 다름.

이 차이가 사용자 코드에 미치는 영향 — GPU 의 사용자는 launch shape, BLOCK_SIZE, shared memory 를 직접 통제한다. TPU 의 사용자는 그런 통제를 거의 하지 않는다 — XLA 가 다 한다. 사용자 표면이 더 추상적인 만큼, hardware 의 모양이 코드에 잘 드러나지 않는다.

“GPU 는 ‘일반 compute 의 도서관’, TPU 는 ‘matmul 의 공장’ — 둘 다 잘 도는 자리가 다르다.”학습 노트 · 재구성
utilization 패턴
GPU 는 작은 커널이 많이 도는 자리에서도 여러 SM 으로 분산 가능. TPU 는 한 큰 MXU — 작은 GEMM 에서 utilization 이 더 떨어짐.
memory pattern
GPU 의 user-managed shared memory vs TPU 의 compiler-managed VMEM. 차이가 큰 자리.
programming model
CUDA / Triton 의 imperative thread-level vs JAX / XLA 의 functional / graph-level.
collective
TPU pod 의 ICI 가 mesh/torus topology, GPU 의 NVLink 가 hub-spoke. 알고리즘 디자인 차이.
debug / profile
GPU 는 NCU 의 풍부한 hint, TPU 는 XLA 의 HLO dump 와 timeline. 도구의 두께가 다름.
생태계
GPU 의 ecosystem 이 훨씬 큰. TPU 는 Google 내부 + JAX 사용자. open source contribution 의 모양이 다름.
§ 07hardware bring-up 학습· RTL · simulation · FPGA

emulator 너머 — 실물 hardware 까지 가는 길

학습용 emulator 가 끝이 아니라면 — 그 다음은 RTL (Verilog/Chisel) 과 FPGA. TinyTPU 같은 프로젝트가 가르치는 hardware bring-up 의 layer.

L1 · functional Python
numpy 로 systolic array 의 결과 계산. timing 무시. 가장 빠른 학습.
L2 · cycle-accurate Python
매 cycle 의 state 를 명시적으로 다룸. simpy 같은 simulation framework 가능. timing intuition.
L3 · RTL (Verilog/Chisel)
실제 hardware description. iverilog / verilator 로 simulate. RTL 단계에서 wire / register 의 구체화.
L4 · FPGA synthesis
RTL 을 FPGA 로 합성. 작은 array (4×4, 8×8) 가 실제 hardware 위에서 도는 경험.
L5 · ASIC tape-out
진짜 칩 만들기. 학습 프로젝트의 영역은 아니지만 — tinygrad TinyChip, OpenROAD 같은 프로젝트가 그 다음 라운드.
학습 가치
L1 ~ L3 까지가 학습용 프로젝트의 sweet spot. RTL 까지 한 번 만져보면 — hardware 의 직관이 굳는다.
왜 이게 ML 엔지니어에게 가치인가

새 hardware 가 매년 등장한다 — H100, B200, Cerebras, Groq, Trainium. 각자의 패러다임이 다르고, 그 위에서 잘 돌게 하는 책임이 ML 엔지니어에게 자주 떨어진다. 한 번이라도 hardware 를 직접 만들어본 사람은 새 hardware 의 spec sheet 를 펼치자마자 “여기서 잘 돌고 여기서 안 돌겠다” 의 직관이 다르게 작동한다.

§ 08채택· where it teaches

TinyTPU 같은 학습 프로젝트가 잘 도는 자리

“ML 엔지니어가 직접 hardware 를 만들 일은 없을 것이다. 그러나 hardware 의 직관 없이 production ML 을 잘 짜기는 점점 어려워진다.”학습 노트 · 재구성
§ 09future· extensions

TinyTPU 의 자연스러운 다음 단계

  1. multi-chip ICI — 한 칩에서 끝나지 않고, mesh/torus 로 연결된 여러 칩의 모델링. TPU pod 의 모양.
  2. vector unit 의 더 깊은 모델 — softmax / layernorm 의 cycle 단위 cost 가 GEMM 대비 어떻게 큰지.
  3. compiler 통합 — small XLA-like compiler 의 직접 작성. HLO 그래프 → systolic mapping.
  4. quantization — INT8 / BF16 / FP8 의 PE 에서의 동작. L084 Numerics 와의 교차.
  5. FPGA 배포 — RTL 까지 짜고 FPGA 위에서 작은 GEMM 이 실제로 도는 데까지. 학습의 완결.
  6. 새 architecture 비교 — Cerebras 의 wafer-scale, Groq 의 LPU, Tenstorrent 의 grayskull. systolic 의 변형들.
§ 10기억할 메모· key takeaways

다시 열었을 때 5분 안에 손에 잡혀야 할 것

TinyTPU 의 위치
학습용 TPU emulator/구현 프로젝트. systolic array, accumulator, VMEM 의 핵심을 직접 만지면서 이해.
TPU 의 핵심
MXU (systolic array) + vector unit + scalar unit + VMEM + HBM. GEMM 가속기.
systolic 의 직관
data 가 PE 사이를 한 cycle 씩 흐르면서 자동으로 곱셈-덧셈. control 단순. weight stationary 가 표준.
GPU 와의 차이
GPU = 일반 compute 가속기 / SIMT. TPU = matmul 가속기 / systolic. 같은 GEMM, 다른 패러다임.
utilization tradeoff
array dim (128) 의 배수가 아닌 GEMM 에서 utilization 손실. fill / drain 비용. compiler 의 책임.
학습 layer
L1 functional Python → L2 cycle-accurate → L3 RTL Verilog → L4 FPGA. 한 layer 씩 진도.
hardware 직관
한 번 만들어보면 — 새 hardware 의 spec sheet 가 다르게 읽힌다. ML 엔지니어의 long-term 가치.
non-matmul 의 자리
vector unit 에서 별도 처리. fusion 의 모양이 GPU 와 다름. softmax 의 비용이 상대적으로 큼.
Slides official repo 에 미공개
Related Google TPU paper (Jouppi 외, ISCA 2017) · In-Datacenter Performance Analysis of a TPU · 확인 필요

손에 새기기 — 학습 시퀀스

  1. numpy 로 systolic array — 4×4 array 의 functional 구현. weight stationary 매핑. 4×4 GEMM 결과를 baseline 과 비교.
  2. cycle-accurate 로 확장 — 매 cycle 의 PE state 를 print. fill / drain 의 cycle 수 직접 측정.
  3. utilization 측정 — 다양한 GEMM 모양 (square / tall / wide / odd dim) 에 대해 utilization 곡선 그리기.
  4. vector unit 추가 — bias add, ReLU 같은 op 를 별도 unit 으로 모델링. fusion 의 효과.
  5. RTL 까지 한 번 — Chisel 또는 Verilog 로 한 PE 의 RTL 짜기. iverilog 로 simulate. 작은 array 의 RTL 까지.
  6. JAX TPU 로 같은 op 실행 — 같은 GEMM 을 JAX/TPU 로 돌리고 HLO dump 와 비교. 자기 emulator 가 본 것이 실제 시스템에 어떻게 등장하는지.
§ 11다른 강의로 이어지는 길· connections

TinyTPU 의 자리에 등장하는 다른 강의들

§ 12열린 질문· open questions

다음에 다시 들었을 때 직접 검증해야 할 것들

검증 메모

이 노트는 강의 자막 실패 + slide 부재 상태에서 — TPU 의 공식 architecture 자료(Jouppi 외 ISCA 2017 paper, Google AI 블로그) + systolic array 의 일반 지식 + 학습용 hardware 프로젝트의 일반 형태로 재구성. TinyTPU 의 정확한 구현 디테일, 강의의 정확한 데모 시퀀스, 스피커의 정확한 입장은 영상 직접 확인 필요.

← Lecture 087 NVSHMEM low-latency Lecture 089 → cuTile (NVIDIA)