Google TPU (Tensor Processing Unit)
Google TPU 아키텍처: MXU/HBM/ICI 토폴로지, Gasket/ACCEL 드라이버, PCIe/커스텀 인터커넥트, Edge TPU/Coral, XLA/JAX 컴파일 파이프라인, TPU Pod 확장
TPU (Tensor Processing Unit)는 Google이 설계한 ASIC 기반 도메인 특화 가속기로, 행렬 연산(텐서 연산)에 최적화되어 머신러닝 학습(Training)과 추론(Inference) 워크로드를 처리합니다. Cloud TPU는 Google 데이터센터에서 PCIe 또는 자체 호스트 인터페이스로 연결되며, Edge TPU는 PCIe M.2 또는 USB 폼 팩터로 임베디드 환경에서 추론 가속을 제공합니다.
drivers/staging/gasket/ (Gasket 프레임워크)와 drivers/accel/ (ACCEL 서브시스템, 6.2+)을 통해 가속기 드라이버를 지원합니다.TPU 세대별 하드웨어 아키텍처
| 세대 | 연도 | 용도 | 연산 성능 | HBM | 인터커넥트 | 공정 |
|---|---|---|---|---|---|---|
| TPU v1 | 2015 | 추론 전용 | 92 TOPS (INT8) | — (DRAM 8 GB) | PCIe 3.0 x16 | 28nm |
| TPU v2 | 2017 | 학습 + 추론 | 45 TFLOPS (bf16) | HBM2 16 GB | ICI (4링크) | 16nm |
| TPU v3 | 2018 | 학습 + 추론 | 123 TFLOPS (bf16) | HBM2 32 GB | ICI (6링크) | 16nm (수냉) |
| TPU v4 | 2021 | 학습 + 추론 | 275 TFLOPS (bf16) | HBM2e 32 GB | ICI (6링크, 3D Torus) | 7nm |
| TPU v5e | 2023 | 효율 최적화 | 197 TFLOPS (bf16) | HBM2e 16 GB | ICI | N/A |
| TPU v5p | 2023 | 대규모 학습 | 459 TFLOPS (bf16) | HBM2e 95 GB | ICI | N/A |
| TPU v6e (Trillium) | 2024 | 차세대 효율 | 918 TFLOPS (bf16) | HBM 32 GB | ICI | N/A |
| Edge TPU | 2019 | 추론 전용 | 4 TOPS (INT8) | — | PCIe / USB | N/A |
MXU (Matrix Multiply Unit) — 시스톨릭 어레이
TPU의 핵심 연산 유닛은 MXU (Matrix Multiply Unit)으로, 시스톨릭 어레이(systolic array) 구조를 사용합니다. TPU v2/v3는 칩당 2개의 TensorCore를 가지며, 각 TensorCore에 128×128 시스톨릭 어레이 MXU가 탑재됩니다.
| 항목 | 값 | 설명 |
|---|---|---|
| PE 격자 | 128 × 128 | 16,384개 Processing Element |
| MAC/사이클 | 16,384 | 각 PE가 1 MAC 연산 수행 |
| FLOPS/사이클 (bf16) | 32,768 | MAC = multiply + add → ×2 |
| TPU v4 피크 | ~275 TFLOPS (bf16) | ~1.05 GHz × 32,768 × 4 MXU × 2 (bf16 packing) |
TPU 메모리 아키텍처
ICI (Inter-Chip Interconnect)
TPU v2 이후 칩 간 통신을 위한 ICI (Inter-Chip Interconnect)를 제공합니다. ICI는 PCIe를 사용하지 않는 Google 자체 설계 고속 직렬 링크로, 다수의 TPU 칩을 직접 연결하여 Pod 단위의 대규모 분산 학습을 가능하게 합니다.
| 세대 | ICI 링크 수 | 토폴로지 | 링크당 대역폭 | Pod 규모 |
|---|---|---|---|---|
| TPU v2 | 4 | 2D Torus (16×16) | ~496 Gbps | 256 칩 (64 TF/pod) |
| TPU v3 | 6 | 2D Torus (32×16) | ~656 Gbps | 1,024 칩 |
| TPU v4 | 6 | 3D Torus (4×4×4 ~ 큐브) | ~2.4 Tbps (총) | 4,096 칩 |
| TPU v5p | 다수 | 3D Torus | 향상 | 8,960 칩 |
TPU의 PCIe 인터페이스
TPU는 호스트 시스템과 PCIe 버스를 통해 연결됩니다. Edge TPU (Coral)는 표준 PCIe 엔드포인트로, Cloud TPU는 커스텀 호스트 인터페이스 칩을 통해 PCIe로 연결됩니다.
| 디바이스 | PCIe 규격 | Vendor ID | Device ID | BAR 구성 |
|---|---|---|---|---|
| Edge TPU (Apex) | PCIe Gen2 x1 | 1ac1 (Google) | 089a | BAR0: 레지스터, BAR2: 펌웨어/데이터 |
| Cloud TPU v4 | PCIe Gen3 x16 | 1ae0 (Google) | (비공개) | 다수 BAR: CSR, DMA, HBM 윈도우 |
# Edge TPU (Coral PCIe Accelerator) 확인
lspci -nn | grep -i google
# 03:00.0 System peripheral [0880]: Global Unichip Corp. [1ac1:089a]
# 상세 정보
lspci -vvv -s 03:00.0
# Region 0: Memory at f7000000 (64-bit, prefetchable) [size=16M]
# Region 2: Memory at f8000000 (64-bit, prefetchable) [size=2M]
# Capabilities: [80] MSI-X: Enable+ Count=128 Masked-
# Edge TPU BAR 매핑 구조
# BAR0 (16 MB): Apex CSR (Control/Status Registers)
# - 인터럽트 상태/마스크 레지스터
# - DMA 디스크립터 링 베이스/크기
# - 펌웨어 상태 레지스터
# BAR2 (2 MB): 명령어 버퍼 / 데이터 교환 영역
Gasket 프레임워크 (커널 TPU 드라이버)
Google은 TPU 계열 가속기를 위한 커널 드라이버 프레임워크로 Gasket (Google ASIC Software, Kernel Extensions, and Tools)을 개발했습니다. 이 프레임워크는 drivers/staging/gasket/에 위치하며, PCIe 가속기 디바이스의 공통 기능을 추상화합니다.
| 소스 파일 | 역할 |
|---|---|
gasket_core.c | PCI probe/remove, char device 등록, BAR 매핑, ioctl 디스패치(Dispatch) |
gasket_ioctl.c | 표준 ioctl 처리 (페이지 테이블 매핑, 이벤트 fd, 디바이스 리셋) |
gasket_page_table.c | 2단계 디바이스 페이지 테이블 관리 (가상주소 → DMA 주소) |
gasket_interrupt.c | MSI-X 인터럽트 설정 및 eventfd 연동 |
gasket_sysfs.c | sysfs 속성 노출 (펌웨어 버전, 디바이스 상태 등) |
apex_driver.c | Edge TPU (Apex) 전용 드라이버 — Gasket 위에 구현 |
/* Gasket 프레임워크 핵심 구조체 (drivers/staging/gasket/) */
/* gasket_driver_desc — 디바이스별 드라이버 기술자 */
struct gasket_driver_desc {
const char *name; /* 드라이버 이름 */
const char *driver_version; /* 드라이버 버전 문자열 */
/* PCI BAR 기술 */
struct gasket_bar_desc bar_descriptions[GASKET_NUM_BARS];
int num_page_tables; /* 디바이스 페이지 테이블 수 */
struct gasket_page_table_config page_table_configs[...];
/* 인터럽트 기술 */
int num_interrupts;
struct gasket_interrupt_desc interrupts[...];
/* 콜백 */
int (*device_open_cb)(struct gasket_dev *dev);
int (*device_close_cb)(struct gasket_dev *dev);
int (*device_reset_cb)(struct gasket_dev *dev);
enum gasket_status (*device_status_cb)(struct gasket_dev *dev);
long (*ioctl_handler_cb)(...); /* 디바이스별 커스텀 ioctl */
};
/* gasket_dev — 디바이스 인스턴스 (probe 시 생성) */
struct gasket_dev {
struct pci_dev *pci_dev;
struct device *dev;
struct cdev cdev; /* char device */
int dev_idx; /* /dev/apex_N 인덱스 */
/* BAR 매핑 */
struct gasket_bar_data bar_data[GASKET_NUM_BARS];
/* 페이지 테이블 (디바이스 IOVA → 호스트 물리 주소) */
struct gasket_page_table *page_table[...];
/* 인터럽트 */
struct gasket_interrupt_data *interrupt_data;
const struct gasket_driver_desc *driver_desc;
struct mutex mutex;
int ownership_owned; /* 독점 소유권 */
};
Edge TPU (Apex) 드라이버 상세
Edge TPU 드라이버(apex_driver.c)는 Gasket 프레임워크 위에 구현된 구체적인 디바이스 드라이버로, Google Coral PCIe/M.2 가속기를 지원합니다.
/* Edge TPU (Apex) 드라이버 — Gasket 기반 구현 */
/* drivers/staging/gasket/apex_driver.c */
#define APEX_DRIVER_NAME "apex"
#define APEX_DRIVER_VERSION "1.0"
/* PCI ID 테이블 */
static const struct pci_device_id apex_pci_ids[] = {
{ PCI_DEVICE(0x1ac1, 0x089a) }, /* Edge TPU */
{ 0, }
};
/* Apex BAR 기술 */
enum {
APEX_BAR_INDEX = 0, /* BAR0: CSR 레지스터 (16 MB) */
APEX_BAR2_INDEX = 2, /* BAR2: 펌웨어/데이터 (2 MB) */
};
/* 주요 CSR 오프셋 */
#define APEX_BAR2_REG_SCU_BASE 0x1A300 /* SCU (System Control) */
#define APEX_BAR2_REG_KERNEL_HIB 0x1A400 /* Host Interface Block */
/* Apex 드라이버 기술자 — Gasket 프레임워크에 등록 */
static const struct gasket_driver_desc apex_desc = {
.name = APEX_DRIVER_NAME,
.driver_version = APEX_DRIVER_VERSION,
.major = 120,
.minor = 0,
.module = THIS_MODULE,
.bar_descriptions = {
/* BAR0: 레지스터 공간 */
{ .size = 0x1000000, .permissions = 0660 },
/* BAR2: 데이터 공간 */
{ .size = 0x200000, .permissions = 0660 },
},
.num_page_tables = 1,
.page_table_configs = { {
.id = 0,
.mode = GASKET_PAGE_TABLE_MODE_NORMAL,
.total_entries = 8192, /* 4K × 8192 = 32 MB 매핑 가능 */
.base_reg = 0x1A488, /* 페이지 테이블 베이스 레지스터 */
} },
.num_interrupts = 1,
.interrupts = { {
.index = 0,
.reg = 0x1A318, /* 인터럽트 상태 CSR */
.packing = GASKET_IRQ_UNPACKED,
} },
.device_reset_cb = apex_reset,
.device_status_cb = apex_get_status,
};
/* Apex probe — Gasket 프레임워크가 호출 */
static int apex_pci_probe(struct pci_dev *pci_dev,
const struct pci_device_id *id)
{
int ret;
/* Gasket 프레임워크에 디바이스 등록 */
ret = gasket_pci_add_device(pci_dev, &apex_desc);
if (ret)
return ret;
/* /dev/apex_N 디바이스 노드 생성 */
/* BAR 매핑, MSI-X 설정, 펌웨어 로드 등 Gasket이 처리 */
return 0;
}
Gasket 디바이스 페이지 테이블
TPU는 자체 디바이스 페이지 테이블을 가지며, 호스트 사용자 공간(User Space) 가상 주소(Virtual Address)를 디바이스가 접근 가능한 DMA 주소로 변환합니다. Gasket 프레임워크가 이 페이지 테이블의 생명주기를 관리합니다.
/* 페이지 테이블 매핑 ioctl */
struct gasket_page_table_ioctl {
u64 page_table_index; /* 페이지 테이블 ID */
u64 size; /* 매핑 크기 (바이트) */
u64 host_address; /* 사용자 공간 가상 주소 */
u64 device_address; /* 디바이스 측 주소 (IOVA) */
};
/* Gasket 페이지 테이블 매핑 흐름:
*
* 1. 사용자 → ioctl(GASKET_IOCTL_MAP_BUFFER, &pt_ioctl)
* 2. gasket_page_table_map():
* a. get_user_pages_fast() — 사용자 페이지 핀
* b. dma_map_page() — 각 페이지의 DMA 주소 획득
* c. 디바이스 페이지 테이블 엔트리에 DMA 주소 기록
* d. BAR의 페이지 테이블 베이스 레지스터 갱신
* 3. TPU DMA 엔진이 디바이스 페이지 테이블을 참조하여
* 호스트 메모리에 scatter-gather DMA 수행
*/
TPU DMA 데이터 전송 흐름
TPU 인터럽트 처리 (MSI-X / eventfd)
/* Gasket 인터럽트 처리 구조 */
/* 1. MSI-X 벡터 할당 (probe 시) */
static int gasket_interrupt_init(struct gasket_dev *dev)
{
int ret;
/* MSI-X 벡터 할당 */
ret = pci_alloc_irq_vectors(dev->pci_dev,
dev->num_interrupts,
dev->num_interrupts,
PCI_IRQ_MSIX);
if (ret < 0)
return ret;
/* 각 벡터에 ISR 등록 */
for (int i = 0; i < dev->num_interrupts; i++) {
int irq = pci_irq_vector(dev->pci_dev, i);
request_irq(irq, gasket_interrupt_handler,
0, dev->driver_desc->name, dev);
}
return 0;
}
/* 2. ISR — eventfd를 통해 사용자 공간에 통지 */
static irqreturn_t gasket_interrupt_handler(int irq, void *ctx)
{
struct gasket_dev *dev = ctx;
struct eventfd_ctx *eventfd;
/* 인터럽트 상태 레지스터 읽기 및 클리어 */
u64 status = readq(dev->bar_data[0].virt_base + IRQ_STATUS_REG);
writeq(status, dev->bar_data[0].virt_base + IRQ_CLEAR_REG);
/* 등록된 eventfd에 시그널 전달 */
eventfd = dev->interrupt_data->eventfds[0];
if (eventfd)
eventfd_signal(eventfd, 1);
return IRQ_HANDLED;
}
/* 3. 사용자 공간에서 eventfd 등록 (ioctl) */
/*
* 사용자 프로그램:
* int efd = eventfd(0, EFD_NONBLOCK);
* struct gasket_interrupt_eventfd ie = {
* .interrupt = 0, // 인터럽트 인덱스
* .event_fd = efd, // eventfd 파일 디스크립터
* };
* ioctl(tpu_fd, GASKET_IOCTL_SET_EVENTFD, &ie);
*
* // 비동기 완료 대기
* struct pollfd pfd = { .fd = efd, .events = POLLIN };
* poll(&pfd, 1, timeout_ms);
*/
TPU 펌웨어 로딩
/*
* Edge TPU 펌웨어 로딩 과정
*
* 1. 드라이버 probe 시 request_firmware() 호출
* - 펌웨어 파일: /lib/firmware/google/apex_firmware.bin
* - 또는 유저가 sysfs를 통해 커스텀 펌웨어 로드
*
* 2. 펌웨어를 BAR2 영역에 DMA 전송
* - BAR2에 매핑된 디바이스 메모리에 직접 기록
* - 또는 DMA coherent 버퍼를 통해 전송
*
* 3. CSR에 펌웨어 실행 명령 기록
* - SCU 레지스터에 리셋 해제 + 실행 비트 설정
*
* 4. 펌웨어 준비 상태 폴링
* - HIB (Host Interface Block) 상태 레지스터 확인
* - 타임아웃 내에 READY 상태 전환 확인
*/
static int apex_reset(struct gasket_dev *gasket_dev)
{
struct apex_dev *apex = gasket_dev_get_drvdata(gasket_dev);
/* 1. 칩 소프트 리셋 */
writeq(SCU_RESET_VALUE,
gasket_dev->bar_data[2].virt_base + APEX_BAR2_REG_SCU_BASE);
/* 2. 리셋 완료 대기 */
usleep_range(100, 200);
/* 3. 펌웨어 로드 (이미 호스트 메모리에 로드된 경우) */
memcpy_toio(gasket_dev->bar_data[2].virt_base + FW_LOAD_OFFSET,
apex->fw_data, apex->fw_size);
/* 4. 펌웨어 실행 시작 */
writeq(HIB_START_VALUE,
gasket_dev->bar_data[2].virt_base + APEX_BAR2_REG_KERNEL_HIB);
/* 5. READY 상태 폴링 */
return readq_poll_timeout(
gasket_dev->bar_data[2].virt_base + HIB_STATUS_REG,
status, status == FW_STATUS_READY,
100, /* 100 µs 간격 */
5000000); /* 5초 타임아웃 */
}
ACCEL 서브시스템 (Linux 6.2+)
Linux 6.2부터 도입된 ACCEL (Accelerator) 서브시스템은 ML/AI 가속기를 위한 표준 커널 프레임워크입니다. DRM 코어 인프라를 재활용(Recycling)하면서 /dev/accel/accelN 별도 디바이스 노드를 제공하여 GPU와 가속기를 명확히 분리합니다.
| 구분 | DRM (/dev/dri/) | ACCEL (/dev/accel/) |
|---|---|---|
| 용도 | GPU 렌더링/디스플레이 | ML/AI 연산 가속 |
| KMS | 필요 (디스플레이 출력) | 불필요 (연산 전용) |
| GEM | GPU 버퍼 관리 | 가속기 메모리 관리(Memory Management) |
| 스케줄러(Scheduler) | drm_sched (GPU 작업) | 디바이스별 자체 스케줄러 |
| 주요 드라이버 | i915, amdgpu, nouveau | habanalabs (Gaudi), qaic, ivpu |
| 디바이스 노드 | /dev/dri/cardN, /dev/dri/renderDN | /dev/accel/accelN |
/* ACCEL 드라이버 등록 예시 (DRM 기반) */
#include <drm/drm_accel.h>
#include <drm/drm_drv.h>
static const struct drm_driver my_accel_driver = {
.driver_features = DRIVER_COMPUTE_ACCEL, /* 핵심: ACCEL 플래그 */
.name = "my_tpu",
.desc = "My TPU Accelerator",
.date = "20240101",
.major = 1,
.minor = 0,
/* GEM 오브젝트 콜백 */
.gem_create_object = my_gem_create_object,
/* ioctl 테이블 (디바이스별 커스텀 명령) */
.ioctls = my_accel_ioctls,
.num_ioctls = ARRAY_SIZE(my_accel_ioctls),
/* 파일 오퍼레이션 */
.fops = &my_accel_fops,
};
/*
* DRIVER_COMPUTE_ACCEL 플래그의 효과:
* - /dev/accel/accelN 디바이스 노드 자동 생성 (/dev/dri/ 대신)
* - KMS 관련 ioctl 비활성화 (모드 설정 불필요)
* - DRM 코어의 메모리 관리/ioctl 인프라는 그대로 활용
* - sysfs: /sys/class/accel/accel0/
*/
/* PCI probe에서 ACCEL 디바이스 초기화 */
static int my_tpu_probe(struct pci_dev *pdev,
const struct pci_device_id *id)
{
struct drm_device *drm;
int ret;
/* DRM/ACCEL 디바이스 할당 */
drm = drm_dev_alloc(&my_accel_driver, &pdev->dev);
if (IS_ERR(drm))
return PTR_ERR(drm);
/* PCI 리소스 설정 */
ret = pcim_enable_device(pdev);
pci_set_master(pdev);
dma_set_mask_and_coherent(&pdev->dev, DMA_BIT_MASK(64));
/* 디바이스 등록 → /dev/accel/accel0 생성 */
ret = drm_dev_register(drm, 0);
return ret;
}
TPU 소프트웨어 스택 (커널-사용자 공간 인터페이스)
TPU 주요 ioctl 인터페이스
| ioctl | 기능 | 설명 |
|---|---|---|
GASKET_IOCTL_RESET | 디바이스 리셋 | TPU 칩 소프트 리셋 + 펌웨어 재로드 |
GASKET_IOCTL_MAP_BUFFER | 버퍼 매핑 | 사용자 버퍼를 디바이스 페이지 테이블에 등록 |
GASKET_IOCTL_UNMAP_BUFFER | 버퍼 해제 | 디바이스 페이지 테이블 엔트리 제거 + 페이지 언핀 |
GASKET_IOCTL_SET_EVENTFD | 이벤트 등록 | MSI-X 인터럽트에 eventfd 연결 |
GASKET_IOCTL_CLEAR_EVENTFD | 이벤트 해제 | eventfd 연결 해제 |
GASKET_IOCTL_NUMBER_PAGE_TABLES | 페이지 테이블 수 | 지원하는 디바이스 페이지 테이블 수 조회 |
GASKET_IOCTL_PAGE_TABLE_SIZE | 테이블 크기 | 특정 페이지 테이블의 최대 엔트리 수 조회 |
GASKET_IOCTL_SIMPLE_PAGE_TABLE_SIZE | 단순 테이블 크기 | 단순(1단계) 페이지 테이블 크기 조회 |
TPU sysfs 인터페이스
# Edge TPU sysfs 경로
ls /sys/class/apex/apex_0/
# 디바이스 속성
cat /sys/class/apex/apex_0/device/vendor # 0x1ac1 (Google)
cat /sys/class/apex/apex_0/device/device # 0x089a (Edge TPU)
cat /sys/class/apex/apex_0/temp # 칩 온도 (밀리도)
cat /sys/class/apex/apex_0/status # 디바이스 상태
# ACCEL 서브시스템 (Linux 6.2+)
ls /sys/class/accel/accel0/
cat /sys/class/accel/accel0/device/vendor
cat /sys/class/accel/accel0/device/device
# PCIe 링크 상태 확인
lspci -vvv -s 03:00.0 | grep -E "LnkSta|LnkCap"
# LnkCap: Port #0, Speed 5GT/s, Width x1
# LnkSta: Speed 5GT/s, Width x1
# DMA 상태 확인
cat /sys/kernel/debug/dma-buf/bufinfo # DMA-BUF 사용 현황
cat /proc/iomem | grep -i apex # BAR 매핑 주소
Edge TPU 드라이버 설치 및 확인
# Edge TPU (Coral) PCIe 드라이버 — Gasket 모듈 로드
# 1. 커널 모듈 확인/로드
modprobe gasket # Gasket 코어 프레임워크
modprobe apex # Edge TPU (Apex) 드라이버
lsmod | grep apex
# apex 28672 0
# gasket 77824 1 apex
# 2. 디바이스 노드 확인
ls -la /dev/apex_*
# crw-rw---- 1 root apex 120, 0 ... /dev/apex_0
# 3. udev 규칙 (권한 설정)
# /etc/udev/rules.d/65-apex.rules
# SUBSYSTEM=="apex", MODE="0660", GROUP="apex"
# 4. dmesg 로그 확인
dmesg | grep -i apex
# [ 2.345] apex 0000:03:00.0: Apex device found
# [ 2.346] apex 0000:03:00.0: enabling device (0000 -> 0002)
# [ 2.347] apex 0000:03:00.0: successfully loaded firmware
# 5. 추론 테스트 (TensorFlow Lite + Edge TPU Runtime)
# Python에서:
# import tflite_runtime.interpreter as tflite
# interpreter = tflite.Interpreter(
# model_path='model_edgetpu.tflite',
# experimental_delegates=[
# tflite.load_delegate('libedgetpu.so.1')])
TPU 관련 Kconfig 옵션
| 옵션 | 설명 |
|---|---|
CONFIG_STAGING | Staging 드라이버 활성화 (Gasket 포함) |
CONFIG_STAGING_GASKET_FRAMEWORK | Gasket 코어 프레임워크 |
CONFIG_STAGING_APEX_DRIVER | Edge TPU (Apex) 드라이버 |
CONFIG_DRM_ACCEL | ACCEL 서브시스템 (Linux 6.2+, DRM 기반) |
CONFIG_DRM | DRM 코어 (ACCEL 의존성) |
CONFIG_PCI_MSI | MSI/MSI-X (TPU 인터럽트) |
CONFIG_IOMMU_SUPPORT | IOMMU (DMA 보안 격리) |
CONFIG_DMA_SHARED_BUFFER | DMA-BUF (버퍼 공유) |
CONFIG_FW_LOADER | 펌웨어 로딩 인프라 |
TPU vs GPU: 커널 드라이버 관점 비교
| 특성 | GPU (DRM/KMS) | TPU (Gasket/ACCEL) |
|---|---|---|
| 디스플레이 출력 | KMS로 제어 (CRTC, Encoder, Connector) | 없음 (연산 전용) |
| 메모리 관리 | GEM/TTM (GPU VRAM) | 디바이스 페이지 테이블 (HBM) |
| 작업 제출 | 커맨드 링 버퍼 (GPU 명령어) | DMA 디스크립터 + Doorbell CSR |
| 스케줄링 | drm_sched (GPU 컨텍스트 전환) | 하드웨어 자체 스케줄링 |
| 사용자 API | OpenGL/Vulkan → DRM ioctl | XLA/PJRT → Gasket ioctl |
| 칩 간 통신 | NVLink/xGMI (제한적) | ICI (대규모 3D Torus) |
| 데이터 타입 | fp32/fp16/int (다목적) | bf16/int8 (ML 특화) |
| 프로그래밍 | 범용 셰이더 | 시스톨릭 어레이 고정 파이프라인(Pipeline) |
TPU 관련 커널 소스 구조
| 경로 | 설명 |
|---|---|
drivers/staging/gasket/ | Gasket 프레임워크 (Edge TPU 지원) |
drivers/staging/gasket/gasket_core.c | PCI probe, BAR 매핑, char device 관리 |
drivers/staging/gasket/gasket_page_table.c | 디바이스 페이지 테이블 (DMA 주소 변환) |
drivers/staging/gasket/gasket_interrupt.c | MSI-X 설정, eventfd 통지 |
drivers/staging/gasket/apex_driver.c | Edge TPU (Apex) 드라이버 |
drivers/accel/ | ACCEL 서브시스템 (Linux 6.2+) |
drivers/accel/habanalabs/ | Intel Gaudi 가속기 드라이버 |
drivers/accel/qaic/ | Qualcomm Cloud AI 100 드라이버 |
drivers/accel/ivpu/ | Intel VPU (NPU) 드라이버 |
include/drm/drm_accel.h | ACCEL 서브시스템 API 헤더 |
TPU Pod 토폴로지 — 3D Torus
TPU v4부터 칩 간 연결에 3D Torus ICI 토폴로지를 채용합니다. TPU v4 Pod는 4×4×4 = 64개 칩을 하나의 큐브(Cube)로 구성하며, 각 축의 양 끝이 연결되어(wrap-around) 토러스 구조를 형성합니다. 이 구조 덕분에 임의의 두 칩 사이 최대 홉(hop) 수가 절반으로 줄어들어, 대규모 분산 학습에서 집합 통신(Collective Communication) 지연이 크게 감소합니다.
TPU v5p는 이 토러스를 더 큰 규모로 확장하여 최대 8,960개 칩을 하나의 Pod로 연결할 수 있습니다. Pod 슬라이스(Pod Slice)는 전체 Pod 중 사용자에게 할당되는 연속적인 칩 부분 집합을 의미하며, Google Cloud에서 TPU를 프로비저닝할 때 슬라이스 단위로 요청합니다. 전체 Pod(Full Pod)는 하나의 물리적 토러스 네트워크를 공유하는 모든 칩의 집합입니다.
Allreduce 최적화: Torus 토폴로지에서는 각 축 방향으로 독립적인 링(Ring)을 구성할 수 있습니다. 3D Torus의 경우 X, Y, Z 세 축에 대해 병렬로 reduce-scatter와 all-gather를 수행하여, 단순 Ring Allreduce 대비 통신량을 ~3배 줄일 수 있습니다. XLA 런타임은 토폴로지 정보를 기반으로 최적의 집합 통신 패턴을 자동 선택합니다.
| 개념 | 설명 |
|---|---|
| Pod Slice | 전체 Pod 중 사용자에게 할당되는 연속 칩 부분집합 (예: v4-32 = 32칩 슬라이스) |
| Full Pod | 하나의 물리적 3D Torus를 공유하는 모든 칩 (TPU v4 = 최대 4,096칩) |
| ICI 대역폭 | TPU v4 기준 링크당 ~24 GB/s, 칩당 6링크 = 144 GB/s 양방향 |
| Torus 최대 홉 | N×N×N Torus에서 최대 홉 = 3×⌊N/2⌋ (Mesh 대비 절반) |
XLA/JAX 컴파일 파이프라인
TPU에서 실행되는 모든 연산은 XLA (Accelerated Linear Algebra) 컴파일러를 통해 TPU 네이티브 명령어로 변환됩니다. JAX, TensorFlow, PyTorch/XLA 등 프론트엔드 프레임워크가 생성한 연산 그래프는 HLO(High Level Operations) IR로 표현되며, XLA가 이를 최적화하고 TPU 하드웨어에 맞게 낮춥니다(lowering).
PJRT (Pretty-much Just a Runtime)는 프레임워크(JAX, TensorFlow, PyTorch/XLA)와 하드웨어 백엔드(TPU, GPU, CPU) 사이의 표준화된 런타임 인터페이스입니다. PJRT를 통해 프레임워크는 디바이스 종류에 관계없이 동일한 API로 버퍼 할당, 컴파일, 실행을 요청할 수 있으며, 커널 드라이버와의 상호작용(DMA 전송, 인터럽트 대기 등)은 PJRT 백엔드가 처리합니다.
SPMD 파티셔닝: XLA의 GSPMD(Generalized SPMD) 파티셔너는 사용자가 지정한 샤딩 어노테이션(jax.sharding)을 기반으로 텐서를 여러 칩에 분할합니다. 분할된 각 칩은 동일한 프로그램을 실행하되 서로 다른 데이터 파티션을 처리하며, 필요 시 ICI를 통해 집합 연산(all-reduce, all-gather)을 수행합니다.
BF16/FP8 혼합 정밀도: XLA는 행렬곱 연산에 BF16을 사용하고 누적(accumulation) 및 가중치 갱신에는 FP32를 사용하는 혼합 정밀도 전략을 자동 적용합니다. TPU v5p 이후에서는 FP8(E4M3/E5M2) 데이터 타입도 지원하여 MXU 처리량을 추가로 향상시킵니다.
TPU VM 아키텍처
Google Cloud의 TPU 접근 방식은 두 가지입니다: 레거시 TPU Node와 TPU VM. TPU VM은 TPU v2-alpha 이후 기본 모드로, 사용자가 TPU 칩에 직접 연결된 호스트 머신에 SSH로 접근할 수 있습니다.
TPU Node (레거시) 모드에서는 사용자의 GCE VM과 TPU 호스트가 분리되어 있어, gRPC를 통해 원격으로 TPU에 접근했습니다. 이 방식은 네트워크 지연이 추가되고, TPU 호스트의 파일 시스템이나 디바이스 노드에 직접 접근할 수 없는 제약이 있었습니다.
TPU VM 모드에서는 사용자가 TPU 칩에 물리적으로 연결된 호스트 머신에 직접 SSH 접속합니다. 이를 통해 /dev/accel* 디바이스 노드에 직접 접근하고, sysfs를 통해 TPU 상태를 모니터링하며, 필요 시 커스텀 커널 모듈을 로드할 수 있습니다. JAX/TensorFlow 라이브러리를 호스트에서 직접 실행하므로 gRPC 오버헤드가 제거됩니다.
Multislice 학습: 단일 Pod 슬라이스의 칩 수를 초과하는 대규모 모델 학습 시, 여러 Pod 슬라이스를 Google의 DCN(Data Center Network)으로 연결하는 Multislice 구성을 사용합니다. 각 슬라이스 내부는 ICI로 고속 통신하고, 슬라이스 간에는 DCN을 통해 집합 연산을 수행합니다. XLA의 GSPMD 파티셔너가 ICI/DCN 경계를 인식하여 통신 패턴을 최적화합니다.
# TPU VM 접속 및 상태 확인
$ gcloud compute tpus tpu-vm ssh my-tpu-node --zone=us-central1-a
# TPU 디바이스 노드 확인
$ ls /dev/accel*
/dev/accel0 /dev/accel1 /dev/accel2 /dev/accel3
# TPU 칩 상태 확인 (sysfs)
$ cat /sys/class/accel/accel0/device/status
alive
# Multislice 환경 변수 예시
$ export TPU_WORKER_ID=0
$ export TPU_WORKER_HOSTNAMES="10.0.0.1,10.0.0.2,10.0.0.3,10.0.0.4"
전력 및 냉각 관리
TPU의 전력 소비와 냉각 방식은 세대별로 크게 달라집니다. 특히 TPU v3부터 수냉(Liquid Cooling)이 도입되어 고밀도 배치에서도 안정적인 열 관리를 보장합니다.
| 세대 | TDP (칩당) | 냉각 방식 | 비고 |
|---|---|---|---|
| TPU v1 | ~40 W | 공냉 (Air) | 추론 전용, PCIe 카드 폼 팩터 |
| TPU v2 | ~250 W | 공냉 (Air) | 학습 + 추론, 커스텀 보드 |
| TPU v3 | ~450 W | 수냉 (Liquid) | 최초 수냉 도입 — 전용 냉각수 순환 시스템 |
| TPU v4 | ~275 W | 수냉 (Liquid) | 7nm 공정으로 전력 효율 개선 |
| TPU v5p | ~400 W (추정) | 수냉 (Liquid) | 성능 대비 전력 효율 최적화 |
| Edge TPU | ~2 W | 공냉/수동 | 임베디드, USB/PCIe M.2 폼 팩터 |
데이터센터 수냉 시스템: TPU v3 이후 칩은 직접 수냉 방식을 사용합니다. 각 TPU 보드에 냉각수 유로(cold plate)가 부착되어 칩에서 발생하는 열을 냉각수로 직접 전달합니다. 데이터센터 수준에서는 냉각 타워(cooling tower)와 열교환기를 통해 냉각수 온도를 관리하며, 이를 통해 공냉 대비 PUE(Power Usage Effectiveness)를 크게 개선합니다.
Edge TPU 열 관리 — 커널 연동: Edge TPU(Coral)는 임베디드 환경에서 동작하므로 리눅스 커널의 thermal 프레임워크와 연동됩니다. Apex 드라이버는 sysfs를 통해 칩 온도와 쓰로틀링(throttling) 상태를 노출하며, 설정 온도 초과 시 자동으로 클록 주파수를 낮추거나 추론 요청을 거부합니다.
# Edge TPU 온도 확인 (Coral USB Accelerator)
$ cat /sys/class/apex/apex_0/temp
62
# Edge TPU 쓰로틀링 상태 확인
$ cat /sys/class/apex/apex_0/trip_point0_temp
85000
$ cat /sys/class/apex/apex_0/trip_point1_temp
100000
# 커널 thermal zone 연동 확인
$ ls /sys/class/thermal/thermal_zone*/type
# apex_thermal 등의 이름으로 등록됨
# Edge TPU 전력 상태 확인 (PCIe M.2 모듈)
$ lspci -vv -s 03:00.0 | grep -i power
Capabilities: [80] Power Management version 3
Status: D0 PME-Enable+ DSel=0 DScale=0 PME-