Triton 是一種用於並行程式設計的語言和編譯器。它旨在提供一個基於 Python 的程式設計環境，以高效編寫自定義 DNN 計算核心，並能夠在現代 GPU 硬體上以最大吞吐量執行。
更多 Triton 中文文件可訪問 →https://triton.hyper.ai/
在本教程中，您將編寫一個非常簡短的高效能 FP16 矩陣乘法核心，其效能可以與 cuBLAS 或 rocBLAS 相媲美。
您將具體學習以下內容：
- 塊級矩陣乘法。
- 多維指標算術。
- 為提高 L2 快取命中率而進行的程式重排序。
- 自動效能調優。
動機
矩陣乘法是現代大多數高效能運算系統的關鍵構建塊。
矩陣乘法難以最佳化，因此其實現通常由硬體供應商自行完成，作為所謂「核心庫」（例如 cuBLAS ）的一部分。
這些庫通常是專有的，不能輕易定製以滿足現代深度學習工作負載的需求（例如融合啟用函式）。
在本教程中，您將學習如何藉助一種更易於定製和擴充套件的方法，用 Triton 實現高效的矩陣乘法。
整體來說，我們將編寫的核心將實現以下的分塊演算法，用於計算一個 (M, K) 乘以一個 (K, N) 的矩陣：
# Do in parallel
# 並行進行
for m in range(0, M, BLOCK_SIZE_M):
# Do in parallel
# 並行進行
for n in range(0, N, BLOCK_SIZE_N):
acc = zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=float32)
for k in range(0, K, BLOCK_SIZE_K):
a = A[m : m+BLOCK_SIZE_M, k : k+BLOCK_SIZE_K]
b = B[k : k+BLOCK_SIZE_K, n : n+BLOCK_SIZE_N]
acc += dot(a, b)
C[m : m+BLOCK_SIZE_M, n : n+BLOCK_SIZE_N] = acc其中，每次雙重巢狀的迴圈迭代都由專用的 Triton 程式例項執行。
計算核心
實際上，上述演算法在 Triton 中實現起來相當簡單。
主要困難在於計算內迴圈中必須讀取 A 和 B 塊的記憶體位置。為此，我們需要多維指標算術。
指標算術
因此，對於行主序的二維張量 X，X[i, j] 的記憶體位置由 &X[i, j] = X + i*stride_xi + j*stride_xj 給出。
因此，A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K]和B[k : k+BLOCK_SIZE_K, n:n+BLOCK_SIZE_N] 的指標塊可以用虛擬碼定義為：
&A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K] = a_ptr + (m : m+BLOCK_SIZE_M)[:, None]*A.stride(0) + (k : k+BLOCK_SIZE_K)[None, :]*A.stride(1);
&B[k : k+BLOCK_SIZE_K, n:n+BLOCK_SIZE_N] = b_ptr + (k : k+BLOCK_SIZE_K)[:, None]*B.stride(0) + (n : n+BLOCK_SIZE_N)[None, :]*B.stride(1);這意味著在 Triton 中可以將 A 和 B 塊的指標初始化（即 k=0 ）為以下程式碼。還要注意，當 M 不是 BLOCK_SIZE_M 的倍數或 N 不是 BLOCK_SIZE_N 的倍數時，我們需要額外的取模運算來應對，這種情況下我們可以用一些無用的值填充資料，這些值不會對結果有影響。對於 K 維度，我們將在後面使用掩碼載入語義來處理。
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None]*stride_am + offs_k [None, :]*stride_ak)
b_ptrs = b_ptr + (offs_k [:, None]*stride_bk + offs_bn[None, :]*stride_bn)然後在內迴圈中更新如下：
a_ptrs += BLOCK_SIZE_K * stride_ak;
b_ptrs += BLOCK_SIZE_K * stride_bk;L2 快取最佳化
正如上面提到的，每個程式例項計算 C 的一個 [BLOCK_SIZE_M, BLOCK_SIZE_N] 塊。
重點要記住這些塊的計算順序，因為它會影響我們程式的 L2 快取命中率，而且，簡單的行主序排序是行不通的。
pid = tl.program_id(axis=0)
grid_n = tl.cdiv(N, BLOCK_SIZE_N)
pid_m = pid // grid_n
pid_n = pid % grid_n一種可能的解決方案是以促進資料重用的順序啟動塊。
在轉向下一列之前，可以透過將 GROUP_M 行的塊進行「超級分組」來實現此目的：
# Program ID
# 程式 ID
pid = tl.program_id(axis=0)
# Number of program ids along the M axis
# M 軸上程式 id 的數量
num_pid_m = tl
# Number of programs ids along the N axis
# N 軸上程式 id 的數量
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
# Number of programs in group
# 組中程式數量
num_pid_in_group = GROUP_SIZE_M * num_pid_n
# Id of the group this program is in
# 本程式所在的組 id
group_id = pid // num_pid_in_group
# Row-id of the first program in the group
# 組內第一個程式的行 id
first_pid_m = group_id * GROUP_SIZE_M
# If `num_pid_m` isn't divisible by `GROUP_SIZE_M`, the last group is smaller
# 如果 `num_pid_m` 不能被 `GROUP_SIZE_M` 整除，最後一組會比較小
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
# *Within groups*, programs are ordered in a column-major order
# 在組內，程式按列主序排序。
# Row-id of the program in the *launch grid*
# 程式的行 id
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
# Col-id of the program in the *launch grid*
# 啟動網格中的程式的行 id
pid_n = (pid % num_pid_in_group) // group_size_m例如，在以下的矩陣乘法示例中，每個矩陣都是 9*9 個塊。可以看到，如果按行主序計算輸出，我們需要載入 90 個塊到 SRAM 中來計算前 9 個輸出塊，但如果按組順序計算，我們只需要載入 54 個塊。
實際上，這種做法可以在某些硬體架構上顯著提升我們的矩陣乘法核心效能，例如在 A100 上，效能提升可以超過 10%，從 220 到 245 TFLOPS 不等。
最終結果
import torch
import triton
import triton.language as tl
def is_cuda():
return triton.runtime.driver.active.get_current_target().backend == "cuda"
def is_hip_mi200():
target = triton.runtime.driver.active.get_current_target()
return target.backend == 'hip' and target.arch == 'gfx90a'
def get_cuda_autotune_config():
return [
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3,
num_warps=8),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5,
num_warps=2),
triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5,
num_warps=2),
# Good config for fp8 inputs.
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3,
num_warps=8),
triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3,
num_warps=8),
triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4,
num_warps=4)
]
def get_hip_autotune_config():
return [
triton.Config(
{'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 1, 'waves_per_eu': 2},
num_warps=4, num_stages=0),
triton.Config(
{'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 4, 'waves_per_eu': 2},
num_warps=8, num_stages=0),
triton.Config(
{'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 1, 'waves_per_eu': 2},
num_warps=8, num_stages=0),
triton.Config(
{'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'waves_per_eu': 3},
num_warps=4, num_stages=0),
triton.Config(
{'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 1, 'waves_per_eu': 8},
num_warps=4, num_stages=0),
]
def get_autotune_config():
if is_cuda():
return get_cuda_autotune_config()
else:
return get_hip_autotune_config()
# `triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` decorator, which consumes:
# `triton.jit` 函式可以透過使用 `triton.autotune` 裝飾器進行自動調優，該裝飾器接受以下內容：
# - A list of `triton.Config` objects that define different configurations of
# meta-parameters (e.g., `BLOCK_SIZE_M`) and compilation options (e.g., `num_warps`) to try
# - 一組 `triton.Config` 物件的列表，這些物件定義了不同的元引數配置（例如 `BLOCK_SIZE_M`）和編譯選項（例如 `num_warps`）以供嘗試。
# - An auto-tuning *key* whose change in values will trigger evaluation of all the
# provided configs
# - 一個自動調優的 key，其值的變化將觸發對所有提供的配置進行評估。
@triton.autotune(
configs=get_autotune_config(),
key=['M', 'N', 'K'],
)
@triton.jit
def matmul_kernel(
# Pointers to matrices
# 矩陣指標
a_ptr, b_ptr, c_ptr,
# Matrix dimensions
# 矩陣維度
M, N, K,
# The stride variables represent how much to increase the ptr by when moving by 1
# element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
# by to get the element one row down (A has M rows).
# 這些步幅變數表示在特定維度移動 1 個元素時，`ptr` 應該增加多少。例如，`stride_am` 指示了為了訪問下一行的元素（假設 `A` 有 `M` 行），需要增加多少 `a_ptr`。
stride_am, stride_ak, #
stride_bk, stride_bn, #
stride_cm, stride_cn,
# Meta-parameters
# 元引數
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, #
GROUP_SIZE_M: tl.constexpr, #
ACTIVATION: tl.constexpr #
):
"""Kernel for computing the matmul C = A x B.
A has shape (M, K), B has shape (K, N) and C has shape (M, N)
"""
"""計算矩陣乘法 C = A x B 的核心演算法。
其中，A 的形狀為 (M, K)，B 的形狀為 (K, N)，C 的形狀為 (M, N)。
"""
# -----------------------------------------------------------
# Map program ids `pid` to the block of C it should compute.
# 將程式 ID `pid` 對映到它應計算的 C 塊。
# This is done in a grouped ordering to promote L2 data reuse.
# 這是按組順序進行的，以促進 L2 資料重用。
# See above `L2 Cache Optimizations` section for details.
# 詳細資訊請參見上述的 `L2 快取最佳化` 部分。
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
# ----------------------------------------------------------
# Create pointers for the first blocks of A and B.
# 建立 A 和 B 第一個塊的指標
# We will advance this pointer as we move in the K direction
# and accumulate
# 在沿著 K 方向移動時，我們將推進這個指標並累加
# `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
# `a_ptrs` 是一個 [BLOCK_SIZE_M, BLOCK_SIZE_K] 大小的指標塊
# `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
# `b_ptrs` 是一個 [BLOCK_SIZE_K, BLOCK_SIZE_N] 大小的指標塊
# See above `Pointer Arithmetic` section for details
# 詳細資訊請參見上述的 `指標算術` 部分。
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
# -----------------------------------------------------------
# Iterate to compute a block of the C matrix.
# 迭代計算 C 矩陣的一個塊。
# We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
# of fp32 values for higher accuracy.
# 我們累加到一個 `[BLOCK_SIZE_M, BLOCK_SIZE_N]` 大小的 fp32 值塊，以提高精度。
# `accumulator` will be converted back to fp16 after the loop.
# `accumulator` 在迴圈結束後將轉換回 fp16。
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
# Load the next block of A and B, generate a mask by checking the K dimension.
# 載入 A 和 B 的下一個塊，透過檢查 K 維度生成一個掩碼。
# If it is out of bounds, set it to 0.
# 如果超出邊界設為 0
a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
# We accumulate along the K dimension.
# 透過著 K 維度進行累加。
accumulator = tl.dot(a, b, accumulator)
# Advance the ptrs to the next K block.
# 指標前進到下一個 K 塊。
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
# You can fuse arbitrary activation functions here
# while the accumulator is still in FP32!
# 在累加器仍然是 FP32 的情況下，您可以在這裡融合任意啟用函式！
if ACTIVATION == "leaky_relu":
accumulator = leaky_relu(accumulator)
c = accumulator.to(tl.float16)
# -----------------------------------------------------------
# Write back the block of the output matrix C with masks.
# 寫回帶有掩碼的輸出矩陣 C 的塊。
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, c, mask=c_mask)
# We can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`.
# 我們可以透過在 `matmul_kernel` 中將 `leaky_relu` 作為 `ACTIVATION` 元引數來融合 `leaky_relu`。
@triton.jit
def leaky_relu(x):
return tl.where(x >= 0, x, 0.01 * x)
現在我們可以建立一個方便的 wrapper 函式，只接受兩個輸入張量，並且：(1) 檢查任何 shape 約束；(2) 分配輸出；(3) 啟動上述的核心。
def matmul(a, b, activation=""):
# Check constraints.
# 檢查約束
assert a.shape[1] == b.shape[0], "Incompatible dimensions"
assert a.is_contiguous(), "Matrix A must be contiguous"
M, K = a.shape
K, N = b.shape
# Allocates output.
# 分配輸出
c = torch.empty((M, N), device=a.device, dtype=torch.float16)
# 1D launch kernel where each block gets its own program.
# 1 維啟動核心，其中每個塊都有自己的程式。
grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
matmul_kernel[grid](
a, b, c, #
M, N, K, #
a.stride(0), a.stride(1), #
b.stride(0), b.stride(1), #
c.stride(0), c.stride(1), #
ACTIVATION=activation #
)
return c
單元測試
對自定義矩陣乘法操作進行測試，與 原生 torch 實現（例如 cuBLAS）進行對比。
torch.manual_seed(0)
a = torch.randn((512, 512), device='cuda', dtype=torch.float16)
b = torch.randn((512, 512), device='cuda', dtype=torch.float16)
triton_output = matmul(a, b)
torch_output = torch.matmul(a, b)
print(f"triton_output_with_fp16_inputs={triton_output}")
print(f"torch_output_with_fp16_inputs={torch_output}")
# Bigger tolerance for AMD MI200 devices.
# 對於 AMD MI200 裝置，使用更大的容差。
# MI200 devices use reduced precision fp16 and bf16 and flush input and
# output denormal values to zero. Detailed info is at: https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
# MI200 裝置使用降低精度的 FP16 和 BF16，並將輸入和輸出的非規格化值清零。詳細資訊在以下連結：https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
rtol = 1e-2 if is_hip_mi200() else 0
if torch.allclose(triton_output, torch_output, atol=1e-2, rtol=rtol):
print("✅ Triton and Torch match")
else:
print("❌ Triton and Torch differ")
TORCH_HAS_FP8 = hasattr(torch, "float8_e5m2")
if TORCH_HAS_FP8 and is_cuda():
torch.manual_seed(0)
a = torch.randn((512, 512), device="cuda", dtype=torch.float16)
b = torch.randn((512, 512), device="cuda", dtype=torch.float16)
a = a.to(torch.float8_e5m2)
# pre-transpose b for efficiency.
# 提前轉置 b 提高效率
b = b.T
b = b.to(torch.float8_e5m2)
triton_output = matmul(a, b)
torch_output = torch.matmul(a.to(torch.float16), b.to(torch.float16))
print(f"triton_output_with_fp8_inputs={triton_output}")
print(f"torch_output_with_fp8_inputs={torch_output}")
if torch.allclose(triton_output, torch_output, atol=0.125, rtol=0):
print("✅ Triton and Torch match")
else:
print("❌ Triton and Torch differ")Out:
triton_output_with_fp16_inputs=tensor([[-10.9531, -4.7109, 15.6953, ..., -28.4062, 4.3320, -26.4219],
[ 26.8438, 10.0469, -5.4297, ..., -11.2969, -8.5312, 30.7500],
[-13.2578, 15.8516, 18.0781, ..., -21.7656, -8.6406, 10.2031],
...,
[ 40.2812, 18.6094, -25.6094, ..., -2.7598, -3.2441, 41.0000],
[ -6.1211, -16.8281, 4.4844, ..., -21.0312, 24.7031, 15.0234],
[-17.0938, -19.0000, -0.3831, ..., 21.5469, -30.2344, -13.2188]],
device='cuda:0', dtype=torch.float16)torch_output_with_fp16_inputs=tensor([[-10.9531, -4.7109, 15.6953, ..., -28.4062, 4.3320, -26.4219],
[ 26.8438, 10.0469, -5.4297, ..., -11.2969, -8.5312, 30.7500],
[-13.2578, 15.8516, 18.0781, ..., -21.7656, -8.6406, 10.2031],
...,
[ 40.2812, 18.6094, -25.6094, ..., -2.7598, -3.2441, 41.0000],
[ -6.1211, -16.8281, 4.4844, ..., -21.0312, 24.7031, 15.0234],
[-17.0938, -19.0000, -0.3831, ..., 21.5469, -30.2344, -13.2188]], device='cuda:0', dtype=torch.float16)✅ Triton and Torch matchtriton_output_with_fp8_inputs=tensor([[-21.4375, 13.1719, 6.0352, ..., 28.7031, 8.6719, -40.7500],
[ 10.0000, 37.0000, -5.5664, ..., 20.9844, 46.8125, 30.8281],
[ 19.5625, -3.0078, -20.0469, ..., -2.1309, -8.0625, 12.5625],
...,
[-18.1562, -34.1562, -27.4219, ..., -27.3906, -24.0938, -12.3516],
[ -3.3945, -8.6250, -23.6562, ..., -4.1094, -3.5332, -16.0781],
[-23.9688, -3.2637, -33.6875, ..., 17.3125, -36.6250, 25.8594]], device='cuda:0', dtype=torch.float16)torch_output_with_fp8_inputs=tensor([[-21.4375, 13.1719, 6.0352, ..., 28.7031, 8.6719, -40.7500],
[ 10.0000, 37.0000, -5.5664, ..., 20.9844, 46.8125, 30.8281],
[ 19.5625, -3.0078, -20.0469, ..., -2.1309, -8.0625, 12.5625],
...,
[-18.1562, -34.1562, -27.4219, ..., -27.3906, -24.0938, -12.3516],
[ -3.3945, -8.6250, -23.6562, ..., -4.1094, -3.5332, -16.0781],
[-23.9688, -3.2637, -33.6875, ..., 17.3125, -36.6250, 25.8594]], device='cuda:0', dtype=torch.float16)✅ Triton and Torch match
基準測試
比較核心與 cuBLAS 或 rocBLAS 的效能差異。此處以方陣為例進行講解，也可以可以按需調整指令碼，對其他 matrix shape 進行基準測試。
ref_lib = 'cuBLAS' if is_cuda() else 'rocBLAS'
configs = []
for fp8_inputs in [False, True]:
if fp8_inputs and (not TORCH_HAS_FP8 or not is_cuda()):
continue
configs.append(
triton.testing.Benchmark(
x_names=["M", "N", "K"], # Argument names to use as an x-axis for the plot 作為繪圖 x 軸的引數名
x_vals=[128 * i for i in range(2, 33)], # Different possible values for `x_name` `x_names` 引數的不同可能值
line_arg="provider", # Argument name whose value corresponds to a different line in the plot 對應繪圖中不同線的引數名
# Possible values for `line_arg` `line_arg` 的可能值
# Don't compare to cublas for fp8 cases as torch.matmul doesn't support fp8 at the moment. 在 fp8 情況下不與 cuBLAS 比較，因為 torch.matmul 目前不支援 fp8。
line_vals=["triton"] if fp8_inputs else [ref_lib.lower(), "triton"], # Label name for the lines
line_names=["Triton"] if fp8_inputs else [ref_lib, "Triton"], # Line styles
styles=[("green", "-"), ("blue", "-")],
ylabel="TFLOPS", # Label name for the y-axis y 軸的標籤名稱
plot_name="matmul-performance-" +
("fp16" if not fp8_inputs else "fp8"), # Name for the plot, used also as a file 繪圖名稱，也用作儲存繪圖的檔名 name for saving the plot.
args={"fp8_inputs": fp8_inputs},
))
@triton.testing.perf_report(configs)
def benchmark(M, N, K, provider, fp8_inputs):
a = torch.randn((M, K), device='cuda', dtype=torch.float16)
b = torch.randn((K, N), device='cuda', dtype=torch.float16)
if TORCH_HAS_FP8 and fp8_inputs:
a = a.to(torch.float8_e5m2)
b = b.T
b = b.to(torch.float8_e5m2)
quantiles = [0.5, 0.2, 0.8]
if provider == ref_lib.lower():
ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.matmul(a, b), quantiles=quantiles)
if provider == 'triton':
ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul(a, b), quantiles=quantiles)
perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
return perf(ms), perf(max_ms), perf(min_ms)
benchmark.run(show_plots=True, print_data=True)
Out:
matmul-performance-fp16:
matmul-performance-fp8:
Download Jupyter notebook: 03-matrix-multiplication.ipynb
Download Python source code: 03-matrix-multiplication.py
Download zipped: 03-matrix-multiplication.zip