# %%
"""
用 Triton 实现分块矩阵乘法 C = A @ B。
建议先记住三个层次：
1. 整体矩阵
A: (M, K) B: (K, N) C: (M, N)
2. 一个 Triton program（程序实例）负责 C 的一个二维小块
C_block: (BLOCK_SIZE_M, BLOCK_SIZE_N)
3. 计算这个 C 小块时，沿 K 方向分批读取
A_block: (BLOCK_SIZE_M, BLOCK_SIZE_K)
B_block: (BLOCK_SIZE_K, BLOCK_SIZE_N)
C_block += A_block @ B_block
例如 BM=2、BN=2、BK=2、K=4 时，一个 C 小块要累加两轮：
C_block = A[:, 0:2] @ B[0:2, :]
+ A[:, 2:4] @ B[2:4, :]
所以分块矩阵乘法没有改变数学结果，只是改变了 GPU 读取和计算数据的方式。
"""
import os
import torch
import triton
import triton.language as tl
# 当前 Triton 后端对应的 PyTorch 设备。
# CUDA 环境通常类似 cuda:0；HIP/ROCm 环境在 PyTorch 中也常通过 torch.cuda 接口暴露。
DEVICE = triton.runtime.driver.active.get_active_torch_device()
def is_cuda():
"""当前 Triton 后端是否为 NVIDIA CUDA。"""
return triton.runtime.driver.active.get_current_target().backend == "cuda"
def get_cuda_autotune_config():
"""返回 CUDA 后端的候选调优配置，由 Triton 实测后选择最快的一套。"""
# BLOCK_SIZE_M / N：一个 program 计算的 C 小块尺寸。
# BLOCK_SIZE_K：每轮点积处理的 K 方向长度。
# GROUP_SIZE_M：只改变 program 的排列顺序，以提高缓存复用，不改变分块尺寸和结果。
# num_warps：一个 program 由多少个线程束协作执行。
# num_stages：软件流水线深度；较深时可以更充分地重叠加载与计算，但会占用更多资源。
#
# 大块可以提高数据复用率，却会增加寄存器占用；小块资源压力低，却可能增加 program 数量。
# 因此没有一套参数适合所有矩阵尺寸和 GPU，需要自动调优。
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),
# K 分块较大的候选项。FP8 元素占用空间较小，部分硬件上可受益于更大的 K 分块。
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():
"""返回 AMD HIP/HCU 后端的候选调优配置。"""
# AMD 与 NVIDIA 的矩阵指令、线程束组织和资源限制不同，因此单独准备候选配置。
sizes = [
{'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 6},
{'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4},
{'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 6},
{'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 6},
{'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4},
{'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4},
{'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4},
{'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 6},
]
# 字典合并运算 size | {...} 不会修改原 size，而是为每套配置统一添加 AMD 参数。
# matrix_instr_nonkdim=16 是 AMD 矩阵指令的非 K 维尺寸提示；它不是 C 的分块尺寸。
return [
triton.Config(
size | {'matrix_instr_nonkdim': 16},
num_warps=8,
num_stages=2,
)
for size in sizes
]
def get_autotune_config():
"""根据当前后端选择 CUDA 或 HIP/HCU 调优配置。"""
return get_cuda_autotune_config() if is_cuda() else get_hip_autotune_config()
# 自动调优过程：
# 1. 第一次遇到某组 (M, N, K) 时编译并测试候选配置；
# 2. 选择最快配置；
# 3. 缓存结果，后续相同 key 通常直接复用。
#
# key 只包含 M/N/K，意味着输入布局或 dtype 改变时未必重新调优。如果实际项目会混用多种
# dtype/stride，应该评估是否需要让它们也进入调优 key，或拆成不同内核。
@triton.autotune(configs=get_autotune_config(), key=['M', 'N', 'K'])
@triton.jit
def matmul_kernel(
# A、B、C 的首地址。
a_ptr, b_ptr, c_ptr,
# 逻辑形状：A(M, K) @ B(K, N) -> C(M, N)。
M, N, K,
# stride 的单位是"元素个数"，而不是字节数。
# A[i, j] 的元素偏移 = i * stride_am + j * stride_ak。
stride_am, stride_ak,
stride_bk, stride_bn,
stride_cm, stride_cn,
# constexpr 表示编译期常量。不同分块配置会产生不同版本的 GPU 内核。
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
ACTIVATION: tl.constexpr,
):
"""计算 C = A @ B；每个 Triton program 负责 C 的一个二维分块。"""
# -------------------------------------------------------------------------
# 第 1 步：把一维 program 编号 pid 映射为二维输出块坐标 (pid_m, pid_n)
# -------------------------------------------------------------------------
#
# "program"不是 Python 进程，也不是单个 GPU 线程；它是 Triton 的一个并行工作单元，
# 内部由 num_warps 指定的多个线程束共同计算一个 C 小块。
pid = tl.program_id(axis=0)
# C 在 M、N 方向各有多少个块。cdiv 是向上取整除法：cdiv(130, 64) == 3。
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
# GROUP_SIZE_M 仅控制 program 顺序，目的是让相邻 program 更可能复用缓存中的 B 数据。
#
# 具体例子：num_pid_m=4、num_pid_n=3、GROUP_SIZE_M=2。
# C 的块坐标为：
#
# N块0 N块1 N块2
# M块0 C00 C01 C02
# M块1 C10 C11 C12
# M块2 C20 C21 C22
# M块3 C30 C31 C32
#
# 下方映射让 pid 顺序变成：
#
# pid 0,1 -> C00, C10 （共同使用 B 的第0列块）
# pid 2,3 -> C01, C11 （共同使用 B 的第1列块）
# pid 4,5 -> C02, C12
# pid 6,7 -> C20, C30
# pid 8,9 -> C21, C31
# pid 10,11 -> C22, C32
#
# 因此 GROUP_SIZE_M=2 不是让一个 program 同时计算两个块，而是让相邻的两个 program
# 优先处理同一个 N 块下的两个 M 块。GROUP_SIZE_M=8 同理，只是每组最多有 8 个 M 块。
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。例如 num_pid_m=10、GROUP_SIZE_M=8 时，
# 第一组有 M块0~7，最后一组只有 M块8~9，所以最后一组的 group_size_m 是 2。
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_in_group = pid % num_pid_in_group
pid_m = first_pid_m + pid_in_group % group_size_m
pid_n = pid_in_group // group_size_m
# 这些 assume 不执行运行时检查，而是把已知范围告诉编译器，帮助优化整数和地址运算。
tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(stride_am > 0)
tl.assume(stride_ak > 0)
tl.assume(stride_bn > 0)
tl.assume(stride_bk > 0)
tl.assume(stride_cm > 0)
tl.assume(stride_cn > 0)
# -------------------------------------------------------------------------
# 第 2 步：生成当前 program 所负责数据的逻辑下标
# -------------------------------------------------------------------------
#
# 若 BM=4 且 pid_m=2，则原始 M 行号为 [8, 9, 10, 11]。
# 若 BN=4 且 pid_n=1，则原始 N 列号为 [4, 5, 6, 7]。
# offs_k 是当前一轮 K 小块内的相对下标，例如 BK=4 时为 [0, 1, 2, 3]。
#
# M/N 方向使用 % 回绕，保证读取 A/B 时地址有效。最后写 C 时使用真实下标和 mask，
# 所以回绕读取产生的"临时结果"不会被写到越界位置。
#
# 例：M=10、BM=4、pid_m=2，原始行是 [8,9,10,11]，取模后为 [8,9,0,1]。
# 前两行是真正结果；后两行只是安全占位，最终 c_mask 会禁止写回它们。
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)
# -------------------------------------------------------------------------
# 第 3 步：由下标和 stride 生成 A、B 小块的二维"指针张量"
# -------------------------------------------------------------------------
#
# offs_am[:, None] 的形状是 (BM, 1)，offs_k[None, :] 的形状是 (1, BK)，
# 广播相加后得到 (BM, BK) 个地址。因此 a_ptrs 的形状为 (BM, BK)。
# 同理，b_ptrs 的形状为 (BK, BN)。
#
# 小例子：offs_am=[2,3]、offs_k=[0,1]，A 连续且每行4个元素，stride_am=4、stride_ak=1：
#
# offs_am[:,None] * 4 = [[ 8], offs_k[None,:] = [[0,1]]
# [12]]
#
# 广播相加得到元素偏移 [[8,9], [12,13]]，即 A[2:4, 0:2] 的四个地址。
#
# 注意：a_ptrs/b_ptrs 此时只是地址集合，真正读数据发生在下面的 tl.load。
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
# 当前 program 的输出小块，使用 FP32 累加以降低长点积的舍入误差。
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
# -------------------------------------------------------------------------
# 第 4 步：沿 K 方向分块加载、矩阵乘加
# -------------------------------------------------------------------------
#
# 数学上，一个输出元素是完整行列点积：
# C[i,j] = sum(A[i,k] * B[k,j], k=0...K-1)
#
# Triton 把这条长点积切成多段再相加。例如 K=6、BK=2：
# 第0轮处理 k=0,1
# 第1轮处理 k=2,3
# 第2轮处理 k=4,5
# 最终结果与一次完成完整点积相同（浮点累加顺序不同可能造成末位微小差异）。
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
# 最后一轮可能不足 BK 个元素。例如 K=5、BK=2 时需要3轮，最后一轮只有 k=4 有效。
# 无效位置用 0 填充，因为补 0 不改变点积结果。
k_remaining = K - k * BLOCK_SIZE_K
a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
# 形状规则：(BM, BK) @ (BK, BN) -> (BM, BN)。
# 这里不是只得到一个数，而是同时完成 BM*BN 个"分段行列点积"。
accumulator = tl.dot(a, b, accumulator)
# 地址移动到下一段 K：A 沿列方向移动 BK，B 沿行方向移动 BK。
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
# -------------------------------------------------------------------------
# 第 5 步：可选融合激活，然后把有效输出写回 C
# -------------------------------------------------------------------------
if ACTIVATION == "leaky_relu":
accumulator = leaky_relu(accumulator)
# 本例固定输出 FP16。转换只影响存储精度，前面的累加仍是 FP32。
c = accumulator.to(tl.float16)
# 写回必须使用未取模的真实 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 + offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
# 二维 mask 的形状是 (BM, BN)：只有真实位于矩阵内部的位置才写入显存。
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, c, mask=c_mask)
@triton.jit
def leaky_relu(x):
"""逐元素 Leaky ReLU：正数不变，负数乘 0.01。"""
return tl.where(x >= 0, x, 0.01 * x)
# %%
# Python 调用接口
def matmul(a, b, activation=""):
"""检查输入、创建输出张量，并启动 Triton 矩阵乘法内核。"""
# A(M,K) @ B(K,N) 的公共维 K 必须相等。
assert a.ndim == 2 and b.ndim == 2, "A 和 B 必须都是二维矩阵"
assert a.shape[1] == b.shape[0], "A 的列数必须等于 B 的行数"
assert a.is_contiguous(), "本示例要求矩阵 A 连续存储"
assert a.device == b.device, "A 和 B 必须位于同一设备"
M, K = a.shape
_, N = b.shape
c = torch.empty((M, N), device=a.device, dtype=torch.float16)
# 每个 program 负责一个 C 小块，所以 program 总数是：
# ceil(M/BM) * ceil(N/BN)
#
# 例：M=130、N=70、BM=64、BN=32，需要 3*3=9 个 program。
# grid 写成 lambda，是因为自动调优会改变 BM/BN，program 数也必须同步重算。
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.manual_seed(0)
a = torch.rand((512, 512), device=DEVICE, dtype=torch.float16) - 0.5
b = torch.rand((512, 512), device=DEVICE, dtype=torch.float16) - 0.5
triton_output = matmul(a, b)
torch_output = torch.matmul(a, b)
print(f"FP16 输入下的 Triton 输出：{triton_output}")
print(f"FP16 输入下的 PyTorch 输出：{torch_output}")
if torch.allclose(triton_output, torch_output, atol=1e-2, rtol=0):
print("✓ Triton 与 PyTorch 的结果一致")
else:
print("✗ Triton 与 PyTorch 的结果不一致")
# 仅当 PyTorch 与 CUDA GPU 都支持 float8_e5m2 时测试 FP8。
# FP8 有效位较少，因此容差比 FP16 更宽松。
TORCH_HAS_FP8 = hasattr(torch, "float8_e5m2")
if TORCH_HAS_FP8 and is_cuda():
torch.manual_seed(0)
a = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
b = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
a = a.to(torch.float8_e5m2)
# 这里是方阵，所以转置后形状仍为 (512,512)。若改成非方阵，必须重新检查逻辑形状。
b = b.T.to(torch.float8_e5m2)
triton_output = matmul(a, b)
torch_output = torch.matmul(a.to(torch.float16), b.to(torch.float16))
if torch.allclose(triton_output, torch_output, atol=0.125, rtol=0):
print("✓ FP8：Triton 与 PyTorch 参考结果一致")
else:
print("✗ FP8：Triton 与 PyTorch 参考结果不一致")
# %%
# 性能基准测试
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"],
x_vals=[128 * i for i in range(2, 33)], # 方阵边长从256增加到4096。
line_arg="provider",
line_vals=["triton"] if fp8_inputs else [ref_lib.lower(), "triton"],
line_names=["Triton"] if fp8_inputs else [ref_lib, "Triton"],
styles=[("green", "-"), ("blue", "-")],
ylabel="TFLOPS",
plot_name="matmul-performance-" + ("fp16" if not fp8_inputs else "fp8"),
args={"fp8_inputs": fp8_inputs},
)
)
@triton.testing.perf_report(configs)
def benchmark(M, N, K, provider, fp8_inputs):
"""测量指定矩阵尺寸下 Triton 与参考库的吞吐量。"""
a = torch.randn((M, K), device=DEVICE, dtype=torch.float16)
b = torch.randn((K, N), device=DEVICE, dtype=torch.float16)
if TORCH_HAS_FP8 and fp8_inputs:
a = a.to(torch.float8_e5m2)
b = b.T.to(torch.float8_e5m2)
# quantiles依次返回中位、较快和较慢耗时。
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
)
elif provider == 'triton':
ms, min_ms, max_ms = triton.testing.do_bench(
lambda: matmul(a, b), quantiles=quantiles
)
else:
raise ValueError(f"未知 provider: {provider}")
# GEMM 约有 2*M*N*K 次浮点运算：每个输出元素含 K 次乘法和约 K 次加法。
# elapsed_ms * 1e-3 把毫秒转成秒，1e-12 把 FLOPS 转成 TFLOPS。
perf = lambda elapsed_ms: 2 * M * N * K * 1e-12 / (elapsed_ms * 1e-3)
return perf(ms), perf(max_ms), perf(min_ms)
os.makedirs("../result", exist_ok=True)
benchmark.run(show_plots=True, print_data=True, save_path="../result")
