Triton và TPU

CUDA C rất mạnh nhưng cũng rất dài dòng. Triton cho phép bạn viết kernel GPU bằng Python. TPU mang lại một lựa chọn thay thế cho GPU với những đánh đổi khác biệt. File này đề cập đến lập trình kernel Triton, Flash Attention như một nghiên cứu điển hình, kiến trúc TPU cùng JAX/Pallas, và cách chọn công cụ phù hợp. Về Vulkan và tính toán GPU đa nền tảng, xem file 07.

• File trước đã dạy lập trình GPU bằng CUDA C. File này leo thêm một bậc trên thang trừu tượng: Triton cho bạn 80% hiệu năng của CUDA với 20% công sức, hoàn toàn bằng Python. TPU và Vulkan cung cấp các mục tiêu phần cứng thay thế cho những trường hợp sử dụng cụ thể.

Triton: Kernel GPU bằng Python

• Triton (của OpenAI) là một ngôn ngữ dựa trên Python dùng để viết kernel GPU. Thay vì phải tính toán từng luồng riêng lẻ (CUDA), bạn thao tác với các khối (blocks) dữ liệu. Trình biên dịch của Triton tự động xử lý ánh xạ luồng, gộp bộ nhớ (memory coalescing), quản lý bộ nhớ dùng chung, và nhiều tối ưu hóa khác.

• Tại sao Triton lại quan trọng: CUDA C đòi hỏi kiến thức sâu về lập lịch warp, xung đột bank trong bộ nhớ dùng chung, áp lực thanh ghi (register pressure), và các mẫu gộp truy cập. Triton che giấu hầu hết những thứ này, giúp việc phát triển kernel GPU trở nên dễ tiếp cận với các nhà nghiên cứu ML — những người biết Python nhưng không rành lập trình hệ thống.

Kernel Triton đầu tiên của bạn

import triton
import triton.language as tl
import torch

@triton.jit
def add_kernel(
    x_ptr, y_ptr, output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,  # hằng số thời gian biên dịch
):
    # Mỗi thể hiện chương trình xử lý một khối gồm BLOCK_SIZE phần tử
    pid = tl.program_id(axis=0)  # tôi là khối thứ mấy?
    block_start = pid * BLOCK_SIZE

    # Các offset cho khối này
    offsets = block_start + tl.arange(0, BLOCK_SIZE)

    # Mặt nạ để xử lý trường hợp n_elements không chia hết cho BLOCK_SIZE
    mask = offsets < n_elements

    # Nạp dữ liệu (có mặt nạ: các truy cập ngoài biên trả về 0)
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)

    # Tính toán
    output = x + y

    # Lưu kết quả
    tl.store(output_ptr + offsets, output, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    output = torch.empty_like(x)
    n_elements = output.numel()

    # Khởi chạy: một chương trình cho mỗi khối
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)

    return output


# Cách dùng
x = torch.randn(1000000, device='cuda')
y = torch.randn(1000000, device='cuda')
z = add(x, y)

• Các điểm khác biệt so với CUDA:
  • Không quản lý luồng một cách tường minh. Bạn tư duy theo khối (chương trình), không phải theo luồng.
  • tl.arange(0, BLOCK_SIZE) tạo ra một vector các offset cho toàn bộ khối. Mọi phép toán trên vector này đều được tự động vector hóa.
  • mask xử lý các điều kiện biên (giống như thanh ghi mặt nạ của AVX-512, file 03). Không cần vòng lặp dọn dẹp scalar.
  • tl.load và tl.store tự động xử lý truy cập gộp (coalesced).
  • @triton.jit biên dịch hàm thành PTX (mã lắp ráp GPU) ở lần gọi đầu, sau đó lưu đệm kernel đã biên dịch.

Kernel Softmax Triton

• Softmax là một ví dụ Triton hay vì nó cần nhiều lượt quét qua dữ liệu (max, trừ, exp, tổng, chia) và có lợi khi giữ dữ liệu trong SRAM (bộ nhớ dùng chung) giữa các lượt:

@triton.jit
def softmax_kernel(
    output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # Mỗi chương trình xử lý một hàng
    row_idx = tl.program_id(0)
    row_start = input_ptr + row_idx * input_row_stride

    # Nạp hàng
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols
    row = tl.load(row_start + col_offsets, mask=mask, other=-float('inf'))

    # Softmax: lấy max để ổn định số học, rồi exp, rồi chuẩn hóa
    row_max = tl.max(row, axis=0)
    numerator = tl.exp(row - row_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = numerator / denominator

    # Lưu kết quả
    output_start = output_ptr + row_idx * output_row_stride
    tl.store(output_start + col_offsets, softmax_output, mask=mask)

• Trong PyTorch, F.softmax(x, dim=-1) khởi chạy 3 kernel riêng biệt (max, exp-và-tổng, chia), mỗi kernel đọc từ và ghi ra bộ nhớ toàn cục. Bản Triton làm mọi thứ trong một kernel duy nhất, giữ dữ liệu trong thanh ghi/SRAM. Chính sự hợp nhất kernel (kernel fusion) này là lý do các kernel Triton tự viết có thể nhanh hơn 2-4 lần so với các phép toán có sẵn của PyTorch.

Tự động tinh chỉnh (Auto-Tuning) của Triton

• Triton hỗ trợ tự động tinh chỉnh: thử nhiều cấu hình và chọn ra cái nhanh nhất:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64}),
    ],
    key=['M', 'N', 'K'],  # tinh chỉnh lại khi các giá trị này thay đổi
)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
    ...

• Triton đo hiệu năng của từng cấu hình trên phần cứng thực tế và chọn ra cái nhanh nhất. Kích thước tile tối ưu phụ thuộc vào kiến trúc GPU, kích thước ma trận, và cách bố trí bộ nhớ — tự động tinh chỉnh tìm ra chúng mà không cần thử nghiệm thủ công.

Triton so với CUDA: Khi nào dùng cái nào

	Triton	CUDA C
Ngôn ngữ	Python	C/C++
Mức trừu tượng	Cấp khối	Cấp luồng
Tốc độ phát triển	Nhanh (10-50 dòng mỗi kernel)	Chậm (100-500 dòng)
Trần hiệu năng	~80-95% CUDA tinh chỉnh thủ công	100% (điều khiển phần cứng đầy đủ)
Bộ nhớ dùng chung	Tự động	Thủ công
Gộp truy cập	Tự động	Thủ công
Nguyên hàm cấp warp	Hạn chế	Đầy đủ (shuffle, vote, v.v.)
Hỗ trợ phần cứng	Chỉ NVIDIA (AMD thử nghiệm)	Chỉ NVIDIA

• Dùng Triton cho: các kernel hợp nhất, các mẫu attention tùy chỉnh, hàm kích hoạt, hầu hết nhu cầu kernel cho nghiên cứu ML.
• Dùng CUDA C cho: hiệu năng tối đa (phần còn lại 5-20%), nguyên hàm cấp warp, tính song song phụ thuộc dữ liệu phức tạp, khi Triton không thể biểu đạt được mẫu của bạn.

Nghiên cứu điển hình: Flash Attention

• Flash Attention (Dao et al., 2022) là kernel tự viết có ảnh hưởng lớn nhất trong ML gần đây. Nó tính attention với bộ nhớ \(O(n)\) thay vì \(O(n^2)\), cho phép xử lý chuỗi dài hơn nhiều.

• Vấn đề: attention chuẩn tính \(\text{softmax}(QK^T / \sqrt{d}) \cdot V\). Ma trận \(QK^T\) có kích thước \(n \times n\) với \(n\) là độ dài chuỗi. Với \(n = 128K\), ma trận này là \(128K \times 128K \times 4\) byte = 64 GB. Nó không vừa với bộ nhớ GPU.

• Cái nhìn sâu: bạn không cần vật chất hóa toàn bộ ma trận \(n \times n\). Hãy tính attention trong các tile: nạp một khối \(Q\), một khối \(K\), tính điểm attention cục bộ, cộng dồn, và chuyển sang khối tiếp theo. Ma trận \(n \times n\) không bao giờ được vật chất hóa hoàn toàn — chỉ một tile tồn tại trong SRAM tại một thời điểm.

• Softmax trực tuyến (online softmax): phần khó là softmax, vốn cần biết giá trị lớn nhất trên toàn bộ hàng (để ổn định số học). Flash Attention dùng thủ thuật softmax trực tuyến: duy trì một giá trị lớn nhất đang chạy và điều chỉnh tỷ lệ các giá trị đã tính trước khi tìm thấy một giá trị lớn nhất mới. Điều này cho phép tính softmax tăng dần, từng tile một.

• Thuật toán:

Với mỗi khối các hàng Q:
    Với mỗi khối các cột K:
        1. Nạp Q_block từ HBM vào SRAM
        2. Nạp K_block từ HBM vào SRAM
        3. Tính S_block = Q_block @ K_block.T (trong SRAM)
        4. Cập nhật max đang chạy, điều chỉnh tỷ lệ các kết quả trước
        5. Tính exp(S_block - max_đang_chạy)
        6. Cập nhật tổng đang chạy và bộ cộng dồn đầu ra
    Nạp V_block và tính đầu ra cuối cùng
    Ghi khối đầu ra trở lại HBM

• Tại sao nó nhanh: vòng lặp trong chỉ hoạt động hoàn toàn trong SRAM (bộ nhớ dùng chung). Bộ nhớ toàn cục (HBM) chỉ được truy cập để nạp các khối Q, K, V và ghi đầu ra cuối cùng. Hệ số tái sử dụng dữ liệu tỷ lệ thuận với kích thước SRAM, vốn nhanh hơn HBM khoảng 100 lần khi truy cập.

• Flash Attention được cài đặt bằng cả Triton và CUDA C. Bản CUDA nhanh hơn (~10% hiệu quả hơn), nhưng bản Triton dễ đọc và dễ sửa đổi hơn nhiều, điều này rất quan trọng cho nghiên cứu các biến thể attention mới.

Kiến trúc TPU

• TPU (Tensor Processing Unit) là các bộ tăng tốc ML tùy chỉnh của Google. Chúng tiếp cận theo hướng hoàn toàn khác so với GPU:

• Systolic array (mảng tuần hoàn): đơn vị tính toán lõi của TPU là một Matrix Multiply Unit (MXU), một mảng tuần hoàn 128×128 hoặc 256×256 tính phép nhân ma trận bằng cách cho dữ liệu chảy qua một lưới các đơn vị nhân-tích lũy (multiply-accumulate). Dữ liệu đi vào từ các cạnh và lan truyền qua mảng, mỗi đơn vị thực hiện một phép nhân-cộng và chuyển kết quả cho đơn vị tiếp theo.

• Không giống như GPU (lên lịch hàng nghìn luồng độc lập), mảng tuần hoàn là một luồng dữ liệu duy nhất, xác định. Không có lập lịch luồng, không có phân kỳ warp, không có dự đoán nhánh. Sự đơn giản này làm cho MXU cực kỳ tiết kiệm năng lượng cho phép nhân ma trận.

• HBM: TPU dùng cùng loại Bộ nhớ Băng thông Cao (High Bandwidth Memory) như GPU. TPU v5e có 16 GB HBM2e mỗi chip; TPU v5p có 95 GB HBM2e.

• ICI (Inter-Chip Interconnect - Kết nối liên chip): các pod TPU kết nối hàng trăm TPU bằng một mạng tốc độ cao tùy chỉnh. Tính song song dữ liệu và song song mô hình (chương 6) trên các pod TPU được JAX hỗ trợ sẵn.

• BFloat16: TPU là thiết bị đầu tiên dùng bfloat16 (chương 13, file 02). BF16 có cùng dải số mũ như float32 (ngăn tràn số khi huấn luyện) nhưng với độ chính xác định mức (mantissa) thấp hơn. Sự đánh đổi này lý tưởng cho ML, nơi các giá trị gradient trải rộng một dải lớn nhưng không cần tới 23 bit độ chính xác.

Lập trình TPU: JAX và Pallas

• TPU được lập trình thông qua JAX và XLA. Bạn viết mã Python/JAX, jax.jit biên dịch nó thành XLA HLO, và XLA biên dịch HLO thành các lệnh riêng cho TPU. Không có CUDA, không có C++.

import jax
import jax.numpy as jnp

@jax.jit
def matmul(a, b):
    return jnp.dot(a, b)

# Đoạn này chạy trên CPU, GPU, hoặc TPU tùy thiết bị
a = jnp.ones((1024, 1024))
b = jnp.ones((1024, 1024))
c = matmul(a, b)

• Pallas là API viết kernel của JAX — tương đương với Triton nhưng của JAX. Nó cho phép bạn viết các kernel mức thấp mà XLA biên dịch cho GPU hoặc TPU:

from jax.experimental import pallas as pl
import jax.numpy as jnp

def add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]

def add_pallas(x, y):
    return pl.pallas_call(
        add_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(x.shape[0] // 128,),
        in_specs=[pl.BlockSpec((128,), lambda i: (i,)),
                  pl.BlockSpec((128,), lambda i: (i,))],
        out_specs=pl.BlockSpec((128,), lambda i: (i,)),
    )(x, y)

• Pallas mới hơn và chưa chín muồi bằng Triton, nhưng nó là cách duy nhất để viết kernel tùy chỉnh cho TPU (vì TPU không hỗ trợ CUDA).

GPU so với TPU

	GPU (NVIDIA)	TPU (Google)
Tính sẵn có	Mọi cloud, on-premise	Chỉ Google Cloud
Lập trình	CUDA C, Triton, PyTorch	JAX/XLA, Pallas
Linh hoạt	Tính toán đa dụng	Tối ưu cho ML nặng về ma trận
Đỉnh FLOPS nhân ma trận	Rất cao (Tensor Cores)	Rất cao (MXU)
Phép toán không nhân ma trận	Tốt	Chậm hơn (chuyển qua đơn vị vector, không qua MXU)
Mở rộng đa chip	NVLink (8 GPU), InfiniBand	ICI (hàng nghìn TPU, tích hợp chặt chẽ hơn)
Hiệu quả chi phí	Cạnh tranh	Thường rẻ hơn cho huấn luyện quy mô lớn
Hệ sinh thái	Lớn nhất (PyTorch, TensorFlow, JAX)	Tập trung vào JAX

• Dùng GPU cho: hầu hết khối lượng công việc ML, nghiên cứu dựa trên PyTorch, phục vụ suy luận, các khối lượng có tính toán ngoài nhân ma trận đáng kể.
• Dùng TPU cho: huấn luyện JAX quy mô lớn (hàng nghìn chip), huấn luyện nhạy cảm về chi phí trên Google Cloud, các khối lượng bị chi phối bởi phép nhân ma trận.

Chọn công cụ đúng

Khối lượng công việc	Công cụ tốt nhất	Lý do
Huấn luyện ML (PyTorch)	NVIDIA GPU + CUDA/Triton	Hệ sinh thái lớn nhất, công cụ tốt nhất
Huấn luyện ML (JAX, quy mô lớn)	TPU hoặc NVIDIA GPU	TPU cho chi phí ở quy mô Google, GPU cho sự linh hoạt
Kernel hợp nhất tùy chỉnh	Triton (Python) hoặc CUDA C	Triton cho tốc độ phát triển, CUDA cho hiệu năng đỉnh
Kernel JAX tùy chỉnh	Pallas	Lựa chọn duy nhất cho TPU, cũng chạy được trên GPU
Suy luận đa nền tảng	Vulkan (file 07) hoặc ONNX Runtime	Chạy trên mọi hãng GPU
Suy luận mobile/edge	Metal (Apple), Vulkan (Android), NNAPI	Bộ tăng tốc riêng theo nền tảng
Suy luận trình duyệt	WebGPU (file 07)	Lựa chọn duy nhất trong trình duyệt
Suy luận chỉ CPU	ONNX Runtime + AVX/NEON	Không cần GPU, dùng SIMD (file 02-03)
Phần cứng mới lạ	SDK riêng của hãng	Mỗi bộ tăng tốc có bộ công cụ riêng

Bài tập lập trình (dùng CoLab với runtime GPU)

1. Viết và chạy một kernel Triton cho phép cộng vector. So sánh hiệu năng của nó với phép cộng có sẵn của PyTorch.

   import triton
   import triton.language as tl
   import torch
   import time
   
   @triton.jit
   def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
       pid = tl.program_id(0)
       offs = pid * BLOCK + tl.arange(0, BLOCK)
       mask = offs < n
       x = tl.load(x_ptr + offs, mask=mask)
       y = tl.load(y_ptr + offs, mask=mask)
       tl.store(out_ptr + offs, x + y, mask=mask)
   
   n = 10_000_000
   x = torch.randn(n, device='cuda')
   y = torch.randn(n, device='cuda')
   
   # Triton
   out_triton = torch.empty_like(x)
   grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
   add_kernel[grid](x, y, out_triton, n, BLOCK=1024)
   
   # PyTorch
   out_torch = x + y
   
   # Kiểm tra tính đúng
   assert torch.allclose(out_triton, out_torch, atol=1e-5)
   
   # Đo hiệu năng
   torch.cuda.synchronize()
   start = time.time()
   for _ in range(1000):
       add_kernel[grid](x, y, out_triton, n, BLOCK=1024)
   torch.cuda.synchronize()
   triton_time = (time.time() - start) / 1000
   
   start = time.time()
   for _ in range(1000):
       out_torch = x + y
   torch.cuda.synchronize()
   torch_time = (time.time() - start) / 1000
   
   print(f"Triton:  {triton_time*1000:.3f} ms")
   print(f"PyTorch: {torch_time*1000:.3f} ms")
   print(f"Tỷ lệ:   {torch_time/triton_time:.2f}x")
   

2. Viết một kernel hợp nhất Triton thực hiện nhân + cộng + ReLU trong một lượt duy nhất. So sánh với ba phép toán PyTorch riêng biệt.

   import triton
   import triton.language as tl
   import torch
   import time
   
   @triton.jit
   def fused_mul_add_relu_kernel(x_ptr, w_ptr, b_ptr, out_ptr, n, BLOCK: tl.constexpr):
       pid = tl.program_id(0)
       offs = pid * BLOCK + tl.arange(0, BLOCK)
       mask = offs < n
       x = tl.load(x_ptr + offs, mask=mask)
       w = tl.load(w_ptr + offs, mask=mask)
       b = tl.load(b_ptr + offs, mask=mask)
       result = tl.maximum(x * w + b, 0.0)  # hợp nhất: nhân + cộng + relu
       tl.store(out_ptr + offs, result, mask=mask)
   
   n = 10_000_000
   x = torch.randn(n, device='cuda')
   w = torch.randn(n, device='cuda')
   b = torch.randn(n, device='cuda')
   
   # Hợp nhất (Triton)
   out_fused = torch.empty_like(x)
   grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
   fused_mul_add_relu_kernel[grid](x, w, b, out_fused, n, BLOCK=1024)
   
   # Không hợp nhất (PyTorch)
   out_unfused = torch.relu(x * w + b)
   
   assert torch.allclose(out_fused, out_unfused, atol=1e-5)
   
   # Đo hiệu năng
   torch.cuda.synchronize()
   start = time.time()
   for _ in range(1000):
       fused_mul_add_relu_kernel[grid](x, w, b, out_fused, n, BLOCK=1024)
   torch.cuda.synchronize()
   fused_time = (time.time() - start) / 1000
   
   start = time.time()
   for _ in range(1000):
       out_unfused = torch.relu(x * w + b)
   torch.cuda.synchronize()
   unfused_time = (time.time() - start) / 1000
   
   print(f"Hợp nhất (Triton):    {fused_time*1000:.3f} ms")
   print(f"Không hợp nhất (PyTorch): {unfused_time*1000:.3f} ms")
   print(f"Tăng tốc:           {unfused_time/fused_time:.2f}x")
   

3. Đo xem trình biên dịch XLA của JAX tự động hợp nhất các phép toán như thế nào. So sánh một chuỗi phép toán có và không có jit.

   import jax
   import jax.numpy as jnp
   import time
   
   def chain_ops(x):
       x = x * 2.0
       x = x + 1.0
       x = jnp.maximum(x, 0.0)  # ReLU
       x = x / jnp.sum(x)
       return x
   
   chain_jit = jax.jit(chain_ops)
   x = jax.random.normal(jax.random.PRNGKey(0), (10000, 1000))
   
   # Khởi động
   _ = chain_jit(x)
   jax.block_until_ready(_)
   
   # Eager (mỗi op là một lượt khởi chạy kernel riêng)
   start = time.time()
   for _ in range(100):
       y = chain_ops(x)
   jax.block_until_ready(y)
   eager_time = (time.time() - start) / 100
   
   # JIT (XLA hợp nhất các phép toán)
   start = time.time()
   for _ in range(100):
       y = chain_jit(x)
   jax.block_until_ready(y)
   jit_time = (time.time() - start) / 100
   
   print(f"Eager: {eager_time*1000:.2f} ms")
   print(f"JIT:   {jit_time*1000:.2f} ms")
   print(f"Tăng tốc: {eager_time/jit_time:.1f}x (XLA hợp nhất 4 phép toán thành 1 kernel)")