1. TPU는 행렬곱 기계다

TPU(Tensor Processing Unit)는 행렬곱에 특화된 연산 코어인 TensorCore에 고대역폭 메모리(HBM, High Bandwidth Memory — 칩 옆에 쌓아 붙인 큰 저장소)를 붙인 칩이다. 저자들은 TensorCore 안에서 셋을 짚는다. ① MXU(Matrix Multiply Unit)는 시스톨릭 배열(systolic array — 데이터를 격자 모양 회로 사이로 흘려보내며 계산하는 방식)을 써서 8사이클마다 `bf16[8,128]` 행렬과 `bf16[128,128]` 행렬을 곱해 `f32[8,128]` 결과를 낸다. TPU v5e는 1.5GHz에서 MXU 하나가 초당 `5e13`번의 bf16 연산을 하고, 칩 한 장에 MXU가 2~4개 있어 칩 전체로는 초당 `2e14`번(int8이면 `4e14`번)이다. 트릴리엄(TPU v6e)의 MXU는 256×256으로, v5e까지 써 온 128×128보다 커서 사이클당 4배 많은 FLOPs를 낸다. ② VPU(Vector Processing Unit)는 ReLU 같은 일반 연산과 벡터 사이의 덧셈·곱셈, 합산을 맡는다. ③ VMEM(Vector Memory)은 TensorCore 바로 옆에 둔 온칩 저장소로 v5e 기준 128MiB에 불과하지만 MXU로 가는 대역폭은 HBM보다 훨씬 높다. CPU의 L1·L2 캐시와 동작이 비슷하지만 훨씬 크고, 무엇을 올려 둘지는 프로그래머가 직접 정한다. HBM에 있는 데이터를 TensorCore가 쓰려면 먼저 VMEM으로 옮겨야 한다.

TPU v5p는 저자들이 "지금까지 나온 것 중 가장 강력한 축에 든다"고 부르는 세대다. 코어 하나가 초당 `2.5e14`번, 칩 하나가 초당 `5e14`번의 bf16 연산을 하고, 칩 8,960개를 묶은 포드 하나는 초당 4엑사플롭스(bf16)를 낸다. 이만한 컴퓨터를 구글은 여러 대 갖고 있다.

2. 파이프라인 — MXU를 놀리지 않는다

TPU의 모든 연산은 파이프라인으로 겹쳐 돈다. 행렬 X와 A를 곱해 Y를 낼 때 TPU는 먼저 A와 X의 조각을 HBM에서 VMEM으로 복사하고, 이를 MXU에 실어 X는 8×128 조각으로 A는 128×128 조각으로 곱한 뒤 결과를 다시 조각씩 HBM으로 돌려보낸다. 이 복사를 MXU 연산과 겹쳐 두면 MXU는 메모리 전송을 기다리지 않고 계속 돌아가고, 계산은 메모리가 아니라 연산량에 묶인다(compute-bound).

연산 강도(arithmetic intensity — 옮긴 1바이트당 몇 번 연산하는지)가 낮은 계산일수록 이 겹침의 값어치가 커진다. VMEM 대역폭은 HBM보다 약 22배 높아서, VMEM에서 읽고 쓰는 MXU 연산은 연산 강도가 10~20만 되어도 최대 성능에 닿는다. 반대로 가중치를 HBM에서 곧장 읽으면 같은 성능을 내는 데 훨씬 큰 배치 크기가 있어야 한다. 가중치를 VMEM에 미리 올려 두면(프리페치) 다음 계산의 적재 비용을 앞 계산 시간 뒤에 숨길 수 있지만, VMEM이 128MiB로 워낙 작아서 층 하나가 통째로 들어갈 만큼 나눠 두어야만 가능하다고 저자들은 짚는다.

3. 칩 한 장 안 — 메가코어와 트레이

TPU 칩 한 장은 보통(전부는 아니지만) 코어 두 개가 메모리를 나눠 쓰며 사실상 연산량이 두 배인 가속기 하나처럼 움직인다 — 저자들은 이를 메가코어(megacore) 구성이라 부른다. v4·v5·v6 세대가 이 구조를 쓰고, v7은 메가코어 대신 코어 두 개 사이에 고대역폭 링크를 따로 둔다. 더 오래된 v3 이전 칩은 코어마다 메모리가 따로 있어 별개의 가속기 두 개로 취급된다. 추론용으로 만든 v5e는 칩당 코어가 하나뿐이다.

칩은 넷씩 묶여 트레이(tray) 하나를 이루고, 트레이는 PCIe로 CPU 호스트에 붙는다. 흔히 보는 구성은 트레이 하나(칩 4장, 코어 8개, 보통 메가코어 4개로 다룸)를 코랩이나 TPU-VM 하나로 노출하는 형태다. v5e 같은 추론용 칩은 호스트 하나에 트레이가 둘이라 칩 8장이 곧 코어 8개다. PCIe 대역폭은 넉넉하지 않다 — v4 기준 방향마다 초당 16GB로 HBM보다 100배쯤 느리다. 호스트(CPU) 램으로 데이터를 옮기거나 되불러올 수는 있지만 빠르지는 않다.

4. 칩과 칩 사이 — ICI가 만드는 토러스

칩들은 ICI(Inter-Chip Interconnect) 네트워크로 포드(pod) 안에서 서로 이어진다. 오래된 세대(v2·v3)와 추론용 칩(v5e), 트릴리엄(v6e)은 ICI가 가장 가까운 4개 칩과 이어져 2차원 토러스(고리 모양으로 감아 끝과 끝을 이은 격자)를 이루고, v4와 v5p는 가장 가까운 6개와 이어져 3차원 토러스를 이룬다. 이 연결은 호스트를 거치지 않고 칩과 칩이 직접 잇는다. 고리 구조는 두 지점 사이 최대 거리를 N에서 N/2로 줄여 통신을 빠르게 하고, TPU는 여기에 토러스를 뫼비우스 띠처럼 한 번 더 비트는 "뒤틀린 토러스" 배선을 더해 평균 거리를 한층 줄인다.

ICI로 묶은 포드는 아주 커질 수 있다. 최대 포드 크기(슈퍼포드)는 v4가 16×16×16, v5p가 16×20×28이다. 이 큰 포드는 광학 배선으로 다시 연결할 수 있는 4×4×4 칩짜리 정육면체를 이어 붙여 만든다. v5e와 트릴리엄 포드는 16×16 2차원 토러스 하나로 끝이고 그 이상은 못 늘어나지만, 포드끼리는 데이터센터 네트워크(DCN)로 통신할 수 있다.

TPU와 GPU가 갈리는 지점이 이 연결 방식이다. GPU는 스위치 계층을 여러 겹 쌓아 사실상 모든 GPU 사이를 점대점으로 잇는다 — 노드 안(H100은 8개, B200 NVL72는 72개까지)은 직접 연결되고 그보다 큰 토폴로지는 GPU마다 O(log N)번의 홉을 거친다. GPU는 그래서 적은 홉으로 임의의 데이터를 보낼 수 있지만, TPU는 NVLink 스위치 같은 비싼 장비 없이 배선이 단순하고 장치당 링크 수·대역폭이 일정해 훨씬 큰 토폴로지로 늘어난다.

5. 대역폭 넷 — HBM·ICI·PCIe·DCN

TPU v5p의 대역폭 넷 (칩 하나 기준, 공표치)

PCIe와 DCN은 세대마다도 갈린다 — PCIe는 v6e가 초당 `3.2e10`바이트로 v4의 두 배이고, DCN은 v6e·TPU7x가 초당 `1.25e10`바이트, v5e가 초당 `3.125e9`바이트로 더 낮다. ICI는 DCN보다는 훨씬 빠르지만 HBM 대역폭에는 못 미친다. 모델을 여러 칩에 나눠 쪼갤 때는 이 순서(HBM보다 느린 ICI, ICI보다 느린 DCN)를 넘어서는 통신이 MXU를 굶기지 않도록 조심해야 한다. ICI로 이은 TPU 묶음을 저자들은 슬라이스(slice)라고 부르는데, 서로 다른 슬라이스는 DCN으로 잇는다. DCN은 호스트와 호스트를 잇는 연결이라, 슬라이스 사이에서 버퍼를 옮기려면 PCIe로 호스트까지 나갔다가 네트워크를 타고 상대 호스트로 들어간 뒤 다시 PCIe로 HBM까지 들어가야 한다.

저자들이 꼽는 핵심 요지 둘. ① 가중치 행렬은 두 축 모두 최소 128(v6e는 256)까지 채워야 MXU를 온전히 쓴다 — 더 작은 축은 128까지 패딩된다. ② 저정밀도 행렬곱이 대체로 더 빠르다. int8이나 int4는 세대에 따라 bf16보다 2배·4배 빠르지만 VPU 연산은 여전히 fp32로 돈다.

6. 스펙표 — 세대마다 얼마나 다른가

TPU 세대별 스펙 (칩 하나 기준, 공표치)

호스트 크기는 CPU 호스트 하나에 붙은 TPU의 토폴로지다. v5e는 호스트 하나에 4×2로 TPU 8개가 붙는다. 표에서 눈에 띄는 것은 세대가 오를수록 HBM 대역폭과 FLOPs/s가 나란히 오르지만은 않는다는 점이다. v5e는 추론에 맞춰 HBM 용량을 줄인 대신 포드를 16×16 2차원 토러스로 넓게 늘렸고, v5p는 HBM 용량을 키우고 3차원 토러스로 학습에 맞췄다. TPU7x는 이 표 안에서 칩당 HBM 용량과 bf16 FLOPs/s 둘 다 가장 크다.

7. 숫자로 풀어보기 — 두 문제

저자들이 낸 연습문제 중 둘이 이 장의 숫자를 실제로 어떻게 쓰는지 보여준다.

문제 1은 200B 파라미터 모델을 bf16으로 TPU v4p 32개에 나눠 놓고 샘플링할 때 HBM에서 시스톨릭 배열로 파라미터를 전부 읽어 오는 데 얼마나 걸리는지 묻는다. 옮길 바이트는 `sizeof(bf16) * 200e9 = 400e9`, 칩 32개로 나누면 칩당 `12.5e9`바이트이고 칩당 HBM 대역폭이 `1.23e12`바이트/초이니 적재에 약 10밀리초가 걸린다. 저자들은 이 10ms를 그 모델을 샘플링하는 지연시간의 현실적인 하한으로 본다 — 샘플링 한 걸음마다 파라미터를 HBM에서 전부 읽어야 하니 배치 크기가 작을 때는 실제 속도가 이 하한에 가깝게 붙는다.

문제 6은 여러 대역폭을 한 번에 엮는다. int8 행렬 A(128×1,024 제곱, 약 16GB)를 TPU v5e 4×4 슬라이스에 고르게 나눠 각 칩의 호스트 DRAM에 얹어 두고, 이를 전부 TPU{0,0}로 모아 bf16 벡터와 곱하는 문제다. 저자들이 짚은 경로는 넷이다. ① PCIe로 각 호스트가 가진 조각(8GB씩, 링크 16개)을 자기 TPU로 올리는 데 약 63ms가 걸린다. ② ICI로 TPU{0,0}가 나머지 15GB를 두 축·링크당 `4.5e10`바이트/초로 받아오는 데 약 167ms가 걸린다. ③ HBM에서 MXU로 `16e9`바이트를 옮기는 데 약 20ms가 걸린다. ④ 실제 곱셈(`2.7e11` FLOPs, `1.97e14` FLOPs/s)에는 약 1.4ms가 걸린다. TPU가 이 넷을 최대한 겹쳐 돌린다고 보면 전체 시간은 가장 느린 조각인 ICI 복사, 167ms에 가깝게 잡힌다 — 겹침이 완벽하지 않으면 200ms에 가까워진다.