Kiến trúc chip AI: Sự trỗi dậy của các dòng chip chuyên dụng từ TPU, GPU đến LPU

Tại Hội nghị chuyên đề quốc tế về Kiến trúc máy tính năm 2018, John Hennessy và David Patterson đã có bài thuyết trình Turing với tiêu đề: "Kỷ nguyên vàng mới cho kiến trúc máy tính".

Vào những năm 1980, khi Hennessy và Patterson thực hiện nghiên cứu đạt giải thưởng Turing của họ, hiệu năng CPU đơn luồng tăng 52% mỗi năm. Đến năm 2018, với sự kết thúc của Định luật Moore và Dennard Scaling, tốc độ này chỉ còn 3%.

Nhu cầu về các kiến trúc chuyên biệt theo miền (DSAs) đã xuất hiện. Ví dụ thực tế mà họ đưa ra là TPU v1 của Google, vốn đã được đưa vào sản xuất: đạt thông lượng gấp 29 lần CPU trong suy luận mạng thần kinh, với hiệu quả năng lượng tốt hơn gấp 80 lần. Dự đoán kết luận của họ là: "thập kỷ tới sẽ chứng kiến một sự bùng nổ kỷ Cambri của các kiến trúc máy tính mới lạ."

Dự đoán này đã trở thành hiện thực. Ngày nay, chúng ta có hàng chục kiến trúc đang được phát triển nghiêm túc. GPUs, TPUs, LPUs, NPUs, DPUs, ASICs, wafer-scale engines, reconfigurable dataflow, neuromorphic, photonic, analog. Đặc biệt, các kiến trúc này tập trung vào khả năng tính toán cho AI.

Các kiến trúc đã giành được sự triển khai thực tế cho đến nay bao gồm: GPUs (NVIDIA, AMD), bộ tăng tốc mảng tâm thu (TPU, Trainium), Cerebras Wafer-Scale Engine và Groq LPU.

NVIDIA là đơn vị dẫn đầu rõ ràng; AMD theo sau với cam kết 6 GW từ cả OpenAI và Meta. Các TPU huấn luyện Gemini và sẽ phục vụ Anthropic với tối đa một triệu chip; Anthropic cũng chạy Claude trên hơn một triệu chip Trainium. Cerebras hiện đang phục vụ suy luận cho OpenAI; Groq LPU đã được sáp nhập vào NVIDIA thông qua một thương vụ mua lại nhân sự trị giá 20 tỷ USD.

Bài viết này nhằm mục đích khảo sát các phương pháp tiếp cận đa dạng này - triết lý, kiến trúc, phương pháp mở rộng (scale-up và scale-out) và ngăn xếp phần mềm (cách lập trình chip) của chúng.

Vấn đề

Tính toán AI bị thống trị bởi phép nhân ma trận. Một transformer là một chuỗi các phép nhân ma trận (matmuls): phép chiếu Q/K/V, attention, phép chiếu đầu ra, FFN - xen kẽ với các phép toán theo phần tử: chuẩn hóa, kích hoạt, cộng phần dư. Việc huấn luyện một mô hình tiên phong thực hiện 10^25 phép tính nhân-tích lũy (các phép nhân ma trận là một chuỗi các phép nhân-tích lũy).

Hình dạng của các phép nhân ma trận đó phụ thuộc vào khối lượng công việc. Quá trình huấn luyện đẩy một loạt các chuỗi về phía trước qua từng lớp, lan truyền ngược lỗi và cập nhật trọng số, với hàng ngàn token chảy qua cùng một ma trận trọng số cùng một lúc. Prefill là giai đoạn nạp prompt của quá trình suy luận: toàn bộ chuỗi đầu vào được chiếu qua mô hình trong một lần truyền, trước khi token đầu ra đầu tiên được tạo ra. Cả huấn luyện và prefill đều xếp chồng nhiều token lên cùng một ma trận trọng số, vì vậy toán học của mỗi lớp là một phép nhân ma trận-ma trận lớn (GEMM), với cường độ tính toán cao (compute-bound). Decode là quá trình tự hồi quy: mô hình phát ra từng token một, mỗi token phụ thuộc vào mọi token trước đó, và token N+1 không thể bắt đầu cho đến khi token N được tạo ra. Chỉ một token được chiếu mỗi bước, vì vậy mọi phép nhân ma trận trở thành phép nhân ma trận-vectơ (GEMV). Việc tạo ra một token đòi hỏi một lần truyền đầy đủ qua mọi trọng số trong mô hình, cộng với việc đọc đầy đủ KV Cache cho attention. Cường độ tính toán giảm theo nhiều bậc so với prefill.

Các hệ thống suy luận khôi phục một phần cường độ đó bằng cách nhóm các token để nâng cấp các GEMV đó trở lại thành GEMM: continuous batching xếp chồng các bước giải mã của nhiều người dùng, speculative decoding xếp chồng K token dự thảo cho mỗi yêu cầu và xác minh chúng trong một lần truyền, và multi-token prediction áp dụng thủ thuật tương tự bên trong chính mô hình. Điều này đạt được mức sử dụng cao hơn của các đơn vị nhân ma trận và đẩy chỉ số Ops/B lên cao. Đối với continuous batching, yêu cầu của mỗi người dùng vẫn đọc KV Cache riêng của nó, vì vậy quá trình giải mã với ngữ cảnh dài chuyển từ giới hạn băng thông trọng số sang giới hạn băng thông KV.

Vấn đề kiến trúc ở đây là di chuyển các con số đến nơi thực hiện phép nhân ma trận đủ nhanh. Điều này được gọi là bức tường bộ nhớ (memory wall): khả năng tính toán đã tăng theo cấp số nhân, nhưng băng thông bộ nhớ thì không.

Mỗi kiến trúc đề xuất một chiến lược khác nhau để chiến thắng trong trò chơi di chuyển dữ liệu. Hiểu về một con chip quy về bốn câu hỏi: dữ liệu nằm ở đâu, nó di chuyển đến các đơn vị tính toán như thế nào, các đơn vị tính toán trông như thế nào và các con chip nói chuyện với nhau như thế nào ở quy mô lớn.

NVIDIA GPU

NVIDIA GPU là một bộ xử lý song song quy mô lớn. Triết lý ở đây là một con chip có thể lập trình với hàng ngàn luồng, được điều phối bởi một CPU chủ và được hiển thị thông qua CUDA, là cỗ máy phù hợp để chạy các khối lượng công việc có thể song song hóa. Mỗi thế hệ đều bổ sung các nguyên hàm tăng tốc vào các Streaming Multiprocessors có thể lập trình mà không làm thay đổi mô hình lập trình. Cùng một con chip đó huấn luyện các transformer, phục vụ suy luận, kết xuất đồ họa và chạy mô phỏng khoa học (tính toán tăng tốc).

Phả hệ

2006

2010

Kiến trúc tính toán thực sự đầu tiên: bộ nhớ đệm L1/L2 hợp nhất, bộ lập lịch warp kép, IEEE-754 FP64.

2012

2014

SM được thiết kế lại với hiệu năng trên mỗi watt cao hơn khoảng 2 lần so với Kepler.

2016

NVLink 1.0, HBM2, thông lượng FP16 gốc; GPU đầu tiên được thiết kế rõ ràng cho học sâu.

2017

2018

2020

2022

2024

2025

Làm mới giữa chu kỳ: thông lượng FP4 tăng khoảng 1,5 lần, 288 GB HBM3e. Được tinh chỉnh cho suy luận ngữ cảnh dài.

2026

2027

Gói GPU 4-die, 1 TB HBM4e mỗi gói. Được triển khai trong các tủ rack NVL576 Kyber 600 kW với 100 PetaFLOPS FP4 mỗi GPU.

Kiến trúc

NVIDIA GPU là một nhóm các lõi hướng thông lượng, một hệ thống phân cấp bộ nhớ sâu để cung cấp dữ liệu cho chúng, cộng với đủ silicon lập lịch để duy trì hàng ngàn luồng đang hoạt động. Các lõi là Streaming Multiprocessors, được sao chép hơn 100 lần mỗi gói: 80 trên V100, 108 trên A100, 132 trên H100, 148 trên B200, 160 trên B300, 224 trên Rubin. Bên trong mỗi SM là cùng một công thức: bốn SM Sub-Partitions, mỗi phần có bộ lập lịch warp riêng, đơn vị điều phối, tệp thanh ghi 16k×32-bit, các làn CUDA Core vô hướng, một Special Function Unit cho các hàm siêu việt và một cổng riêng vào Tensor Cores của SM. Bốn phân vùng chia sẻ một khối L1/bộ nhớ chia sẻ và TMA. Các luồng được nhóm thành các warp gồm 32 luồng thực thi theo kiểu SIMT lock-step; hàng chục warp thường trú mỗi phân vùng cho phép bộ lập lịch ẩn các lần dừng bộ nhớ/số học bằng cách chuyển đổi giữa chúng.

CUDA Cores là thông lượng tính toán ban đầu, và đối với AI, chúng vẫn đảm nhận mọi thứ không phải là phép nhân ma trận: kích hoạt, cộng phần dư, chuẩn hóa, số học địa chỉ. Tuy nhiên, một khối transformer chiếm khoảng 99% FLOPs của phép nhân ma trận, vì vậy thông lượng tính toán áp đảo đến từ Tensor Cores.

Các lõi này thực hiện phép nhân-tích lũy ma trận hợp nhất trên các ô ma trận nhỏ, D=A⋅B+C. Phép nhân ma trận đầy đủ được chia thành các ô đầu ra: để tạo ra một ô đầu ra, một kernel đi qua chiều trong K, lấy A từ một dải hàng của ma trận đầu vào bên trái và B từ một dải cột của ma trận bên phải, và gộp từng tích phần vào một bộ tích lũy đang chạy. C giữ tổng phần cho đến nay, D là giá trị cập nhật được chuyển sang bước tiếp theo. Sau khi vòng lặp trong hoàn tất, D là một ô đã hoàn thành của ma trận đầu ra đầy đủ; toàn bộ phép nhân ma trận được xây dựng từ nhiều ô MMA như vậy.

Hình dạng ô được viết là M × N × K, M×N là kích thước ô đầu ra, và K là lượng chiều trong mà lệnh thực hiện trong một lần chạy; phần còn lại của trục K của phép nhân ma trận được đi qua bởi vòng lặp trong của kernel. Bộ tích lũy được giữ nguyên trong suốt vòng lặp đó: đầu ra D của mỗi MMA trở thành đầu vào C của MMA tiếp theo, vì vậy phương trình thực sự là C←A⋅B+C tại chỗ: các lệnh liên tiếp gộp các tích phần của chúng vào cùng một bộ lưu trữ cho đến khi trục K được đi qua hoàn toàn.