單線程版本#
逐元素相加：
Triton 實現#
在 Triton 中，向量加法內核透過將向量劃分為多個塊（blocks），並在每個 Grid 中的線程（threads）並行計算，實現高效的向量加法操作。每個線程負責加載兩個向量中對應位置的元素，進行相加並存儲結果。
核心步驟#
- 線程並行計算：每個 Grid 中的線程獨立處理向量中的一部分元素。
- 加載元素：每個線程加載向量 A 和向量 B 中對應位置的元素。
- 元素相加：將加載的元素進行相加。
- 存儲結果：將相加後的結果存儲到輸出向量中。
tl.constexpr 的使用#
tl.constexpr 用於聲明編譯時常量。這意味著使用這個修飾符的變量的值在編譯時就已經確定，而不是在運行時。編譯器可以基於這些常量值進行更 aggressive 的優化來提升內核的執行效率。
@triton.jit
def kernel_vector_addition(a_ptr, b_ptr, out_ptr,
num_elems: tl.constexpr,
block_size: tl.constexpr):
# 內核代碼
上述代碼中，num_elems 和 block_size 被聲明為編譯時常量，使得 Triton 可以在編譯階段優化內核代碼。
確定當前塊與 Program ID#
每個線程塊（block）在 Triton 中都有一個唯一的 Program ID，用於標識當前線程所在的塊。透過 tl.program_id，我們可以確定當前線程所在的塊，從而計算處理的數據偏移量。
pid = tl.program_id(axis=0)
block_start = pid * block_size
處理最後一個 Block#
由於向量長度可能不被塊大小整除，最後一個塊可能只有部分線程需要工作。透過掩碼操作，可以確保只有有效的線程進行計算，避免無效的內存訪問和計算。
掩碼的作用#
Triton 提供了掩碼（mask）操作，用於屏蔽那些不需要工作的線程（最後一個 Grid 中 NA 的線程）。
mask = thread_offsets < num_elems
a_pointers = tl.load(a_ptr + thread_offsets, mask=mask, other=0.0)
b_pointers = tl.load(b_ptr + thread_offsets, mask=mask, other=0.0)
ceil_div 函數的作用#
ceil_div 函數用於計算塊的數量，確保即使向量長度不被塊大小整除，也能覆蓋所有元素。例如 vec_size=10，block_size=3，ceil_div(10, 3)=4，這樣就能確保所有 10 個元素都被處理。
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
說白了，該函數的作用就是高效實現 “向上取整”。
數值精度驗證#
在實現向量加法內核後，驗證數值精度是確保內核正確性的關鍵步驟。透過與 PyTorch 的內置加法操作進行對比，可以確認 Triton 實現的準確性。
def verify_numerics() -> bool:
torch.manual_seed(2020) # seed both cpu and gpu
vec_size = 8192
a = torch.rand(vec_size, device='cuda')
b = torch.rand_like(a)
torch_res = a + b
triton_res = vector_addition(a, b)
fidelity_correct = torch.allclose(torch_res, triton_res)
print(f"{fidelity_correct=}")
return fidelity_correct
驗證了解到我們的 Triton 實現與 PyTorch 原生的數值精度一致，可以進行後面的操作了。
下面是完整的 Kernel 實現：
@triton.jit
def kernel_vector_addition(a_ptr, b_ptr, out_ptr,
num_elems: tl.constexpr,
block_size: tl.constexpr,):
pid = tl.program_id(axis=0)
# tl.device_print("pid", pid)
block_start = pid * block_size # 0 * 2 = 0, 1 * 2 = 2,
thread_offsets = block_start + tl.arange(0, block_size)
mask = thread_offsets < num_elems
a_pointers = tl.load(a_ptr + thread_offsets, mask=mask)
b_pointers = tl.load(b_ptr + thread_offsets, mask=mask)
res = a_pointers + b_pointers
tl.store(out_ptr + thread_offsets, res, mask=mask)
def ceil_div(x: int,y: int) -> int:
return (x + y - 1) // y
def vector_addition(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
output_buffer = torch.empty_like(a)
assert a.is_cuda and b.is_cuda
num_elems = a.numel()
assert num_elems == b.num_elems() # todo - handle mismatched sizes
block_size = 1024
grid_size = ceil_div(num_elems, block_size)
grid = (grid_size,)
num_warps = 8
k2 = kernel_vector_addition[grid](a, b, output_buffer,
num_elems,
block_size,
num_warps=num_warps
)
return output_buffer
基準測試與性能調優#
為了評估 Triton 向量加法內核的性能，下面進行基準測試並探討性能調優的方法。
Benchmark API 介紹#
Triton 提供了豐富的基準測試 API，允許用戶測量內核的執行時間和吞吐量。以下代碼是使用 triton.testing.perf_report 得到性能報告的一個示例：
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=['size'], # 用作圖表x軸的參數名
x_vals=[2**i for i in range(10, 28)], # `x_name`的可能取值
x_log=True, # x軸使用對數刻度
line_arg='provider', # 圖表中不同線條對應的參數名
line_vals=['triton', 'torch'], # `line_arg`的可能取值
line_names=["Triton", "Torch"], # 線條的標籤名
styles=[('blue', '-'), ('green', '-')], # 線條顏色和樣式
ylabel='GB/s', # y軸標籤
plot_name='vector-add-performance', # 圖表名稱，也用作保存文件的文件名
args={}, # 不在`x_names`和`y_name`中的函數參數的值
)
)
def benchmark(size, provider):
x = torch.rand(size, device='cuda', dtype=torch.float32)
y = torch.rand(size, device='cuda', dtype=torch.float32)
quantiles = [0.5, 0.2, 0.8] # 設定分位數
# 根據 provider 選擇不同的計算實現
if provider == 'torch':
ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles)
if provider == 'triton':
ms, min_ms, max_ms = triton.testing.do_bench(lambda: vector_addition(x, y), quantiles=quantiles)
# 計算GB/s
def gbps(ms):
return 12 * size / ms * 1e-06
# 返回中位數、最大值和最小值對應的GB/s
return gbps(ms), gbps(max_ms), gbps(min_ms)
性能報告：
性能對比：
總的來說，vector add 只是個相對簡單的 kernel，相較於複雜的內核更難獲得 Triton 實現帶來的優勢（大部分常用操作 PyTorch 已經透過 CUDA/cuBLAS 等優化到極質了）
調優參數：Num Warps & Block size#
調優內核性能的關鍵在於合理配置 Warp 數量和塊大小。Warp 是 GPU 中的基本執行單元，合理的 Warp 數量和塊大小能夠充分利用 GPU 的並行計算能力，提升內核的執行效率。
block_size = 1024 # 決定每個線程塊處理的元素數量,較大的塊大小可以減少塊的數量，但可能增加每個塊的計算負擔。
grid_size = ceil_div(num_elems, block_size)
grid = (grid_size,)
num_warps = 8 # 每個塊中包含的 Warp 數量,合理配置 Warp 數量可以優化線程的調度和資源利用。
上節 (Softmax in OpenAI Triton)便給出了透過驅動程序來動態調整參數的方式：
# 計算 block_size，為大於等於 cols 的最小 2 的幂
block_size = triton.next_power_of_2(cols)
# 根據 block_size 動態調整 num_warps
num_warps = 4 # 每個 warp 有 32 個線程
if block_size > 2047:
num_warps = 8
if block_size > 4095:
num_warps = 16
參考資料#
感謝：
- 本文內容主要基於這位老師的視頻教程系列 SOTA Deep Learning Tutorials - YouTube
- Triton’s documentation
- o1-preview 語言模型，畫 svg 的能手