Google TPU (Tensor Processing Unit)

Google TPU 아키텍처: MXU/HBM/ICI 토폴로지(Topology), Gasket/ACCEL 드라이버, PCIe/커스텀 인터커넥트, Edge TPU/Coral, XLA/JAX 컴파일 파이프라인(Pipeline), TPU Pod 확장

TPU (Tensor Processing Unit)는 Google이 설계한 ASIC 기반 도메인 특화 가속기로, 행렬 연산(텐서 연산)에 최적화되어 머신러닝 학습(Training)과 추론(Inference) 워크로드를 처리합니다. Cloud TPU는 Google 데이터센터에서 PCIe 또는 자체 호스트 인터페이스로 연결되며, Edge TPU는 PCIe M.2 또는 USB 폼 팩터로 임베디드 환경에서 추론 가속을 제공합니다.

TPU 세대별 하드웨어 아키텍처

MXU (Matrix Multiply Unit) — 시스톨릭 어레이

TPU의 핵심 연산 유닛은 MXU (Matrix Multiply Unit)으로, 시스톨릭 어레이(systolic array) 구조를 사용합니다. TPU v2/v3는 칩당 2개의 TensorCore를 가지며, 각 TensorCore에 128×128 시스톨릭 어레이 MXU가 탑재됩니다.

TPU 메모리 아키텍처

ICI (Inter-Chip Interconnect)

TPU v2 이후 칩 간 통신을 위한 ICI (Inter-Chip Interconnect)를 제공합니다. ICI는 PCIe를 사용하지 않는 Google 자체 설계 고속 직렬 링크로, 다수의 TPU 칩을 직접 연결하여 Pod 단위의 대규모 분산 학습을 가능하게 합니다.

TPU의 PCIe 인터페이스

TPU는 호스트 시스템과 PCIe 버스(Bus)를 통해 연결됩니다. Edge TPU (Coral)는 표준 PCIe 엔드포인트로, Cloud TPU는 커스텀 호스트 인터페이스 칩을 통해 PCIe로 연결됩니다.

# 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 프레임워크 핵심 구조체 (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와 가속기를 명확히 분리합니다.

/* 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 인터페이스

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 옵션

TPU vs GPU: 커널 드라이버 관점 비교

TPU 관련 커널 소스 구조

TPU v6e Trillium · v7 Ironwood (2024~2025)

Google은 2024년 5월 Google I/O에서 TPU v6e(코드명 Trillium)를 공개하고 같은 해 10월 프리뷰로 배포했습니다. 2025년 4월 Google Cloud Next에서는 7세대 TPU인 TPU v7(Ironwood)를 발표해 대규모 추론(Inference)과 차세대 학습(Training) 양쪽을 모두 겨냥한 단일 아키텍처를 제시했습니다.

TPU v6e (Trillium)

• 성능: 칩당 918 TFLOPS(bf16), TPU v5e 대비 4.7배 성능 향상.
• 메모리: HBM 용량·대역폭 모두 2배로 확장(32 GB, ~3.58 Tb/s).
• ICI: TPU v5e 대비 2배 대역폭. Pod는 256칩/Pod 구성.
• 물리 인터페이스: 호스트와는 PCIe Gen5를 통해 연결하며, 커널 측은 accel 서브시스템의 TPU 드라이버 경로가 유지됩니다.
• 성숙도: Gemini 1.5/2.0 학습의 주력 플랫폼으로 사용되며, Cloud TPU v6e 런타임은 JAX/PyTorch XLA에서 안정 상태입니다.

TPU v7 (Ironwood)

Ironwood는 Google이 공식적으로 "AI 추론 시대를 위해 처음부터 설계한 TPU"라 소개한 7세대 제품입니다. 다음 특성은 공개된 제품 문서와 키노트 자료에 기반합니다:

Ironwood는 OCP(Open Compute Project) 기반의 액체 냉각 랙과 결합되어, 9,216칩 규모의 단일 Pod가 하나의 일관된 공유 메모리 도메인처럼 동작합니다. 이 규모의 ICI 페브릭(Fabric)은 기존 v5p의 3D Torus를 넘어 3D Torus + 고차 다차원 라우팅(Routing)으로 확장되어, 칩 간 집합 통신의 평균 홉 수를 획기적으로 줄였습니다.

커널/드라이버 관점 영향

• accel 서브시스템 정착: Linux v6.2에서 DRM 산하에 추가된 drivers/accel/이 AI 가속기의 표준 자리가 되었고, Habana Gaudi, Apple NPU, Rivos, 그리고 Google TPU 계열 모두 이 프레임워크 모델을 따릅니다. v6.14에는 AMD XDNA NPU(Ryzen AI) 드라이버(drivers/accel/amdxdna/)가 메인라인에 추가되어 x86 노트북·데스크톱 SoC 내장 NPU가 accel 서브시스템을 통해 사용자 공간에 노출되는 선례를 만들었습니다.
• PCIe 진화 반영: v6e/v7 모두 PCIe Gen5를 채택하여 호스트 쪽 드라이버는 MSI-X/PASID 경로를 그대로 쓰지만, BAR 크기와 MMIO 윈도우는 커집니다. 대용량 BAR(64-bit Resizable BAR, PCIe 2.3+) 지원이 전제됩니다.
• SVA/PASID 상호작용: 유저모드에서 TPU에 ENQCMDS 유사 경로로 작업을 제출할 때 IOMMU SVA가 필수입니다. 2025년에는 IOMMUFD viommu 기능과 결합돼 다중 테넌트 가상화가 더 자연스럽게 이루어집니다.
• 펌웨어 측면: Ironwood용 Gasket 계보 드라이버는 고유 RPC 큐를 유지하되, 이벤트 리포팅을 DRM device ioctl 경로로 일원화하는 방향입니다. 이는 GPU/NPU가 공통적으로 따르는 "metadata via DRM ioctl"(v6.18 Tyr 사례) 흐름과 일치합니다.

커널 6.14부터: AMD XDNA NPU 드라이버(drivers/accel/amdxdna)가 메인라인에 합류하여 accel 서브시스템이 단순 클라우드 AI 가속기를 넘어 온디바이스(On-Device) NPU까지 포괄하는 범용 프레임워크로 자리매김했습니다.

TPU Pod 토폴로지 — 3D Torus

TPU v4부터 칩 간 연결에 3D Torus ICI 토폴로지를 채용합니다. TPU v4 Pod는 4×4×4 = 64개 칩을 하나의 큐브(Cube)로 구성하며, 각 축의 양 끝이 연결되어(wrap-around) 토러스 구조를 형성합니다. 이 구조 덕분에 임의의 두 칩 사이 최대 홉(hop) 수가 절반으로 줄어들어, 대규모 분산 학습에서 집합 통신(Collective Communication) 지연(Latency)이 크게 감소합니다.

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 런타임은 토폴로지 정보를 기반으로 최적의 집합 통신 패턴을 자동 선택합니다.

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 처리량(Throughput)을 추가로 향상시킵니다.

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 상태를 모니터링하며, 필요 시 커스텀 커널 모듈(Kernel Module)을 로드할 수 있습니다. JAX/TensorFlow 라이브러리를 호스트에서 직접 실행하므로 gRPC 오버헤드(Overhead)가 제거됩니다.

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)이 도입되어 고밀도 배치에서도 안정적인 열 관리(Thermal Management)를 보장합니다.

데이터센터 수냉 시스템: 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-

참고자료

• Google Cloud TPU Documentation
• Gasket Driver (GitHub)
• Coral (Edge TPU) Documentation
• TPU v6e (Trillium) — Google Cloud Docs
• TPU7x (Ironwood) — Google Cloud Docs
• Introducing Trillium, sixth-generation TPUs — Google Cloud Blog
• 3 things to know about Ironwood, Google's latest TPU
• This is the Google TPU v7 Ironwood Chip — ServeTheHome
• Linux accel Subsystem Introduction

관련 문서

• PCI / PCIe 서브시스템 — PCIe 아키텍처와 드라이버 기초
• DMA — DMA 매핑 API
• IOMMU — DMA 격리와 주소 변환(Address Translation)
• SR-IOV — PCI 하드웨어 가상화(Virtualization)