gemm-kernel-optimization
优化面向 NVIDIA 和 AMD GPU 的 Triton 稠密矩阵乘（GEMM）kernel。 覆盖分块 blocking、L2 cache grouping、tensor core 利用，以及双平台 autotune。 适用于编写或优化 matmul、线性层，以及任何基于 GEMM 的 kernel。
来源信息
- 仓库
- sss-hust/awesome-kernel-skills-cn
- 最近来源活动
- 2026年4月14日 13:11
- 检测到的 SKILL.md 语言
- 中文
- 星标
- 0
- 分支
- 0
安装方式
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
检查来源文件
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
文件资源管理器
3 个文件正在显示 SKILL.md
SKILL.md
来源说明 · 只读预览- name
- gemm-kernel-optimization
- description
- 优化面向 NVIDIA 和 AMD GPU 的 Triton 稠密矩阵乘（GEMM）kernel。 覆盖分块 blocking、L2 cache grouping、tensor core 利用，以及双平台 autotune。 适用于编写或优化 matmul、线性层，以及任何基于 GEMM 的 kernel。
# GEMM Kernel 优化
## 概览
通用矩阵乘 `C = A @ B` 是深度学习中最典型、也最吃算力的操作之一。
GEMM 通常是 **compute-bound**。关键指标是 **TFLOPS**，
目标一般是达到硬件峰值的 60% 以上。
| 瓶颈类型 | 指标 | 目标（H100） | 目标（4090） | 目标（MI300X） |
|-------|--------|---------------|---------------|-----------------|
| 计算 | TFLOPS | >500 FP16 | >150 FP16 | >600 FP16 |
## 核心技巧
### 算法：分块 GEMM
1. 把输出 `C` 切成 `[BLOCK_M, BLOCK_N]` 的 tile，每个 program instance 负责一个 tile
2. 沿 K 维按 `BLOCK_K` 分块遍历，并使用 FP32 做累加
3. 每一步加载一个 A tile 和一个 B tile，然后调用 `tl.dot`
4. 最后把结果 tile 写回 C，并转换成输出 dtype
### L2 Cache 分组
朴素的行优先 block 排列会导致过多 HBM 读取。
通过把 `GROUP_SIZE_M` 个行块分成一组、再切换列方向，
可以让 B tile 更久地留在 L2 中。
经验上，`GROUP_SIZE_M=8` 在 A100/H100 上常常能带来 10-15% 的提升。
### Tensor Core 利用
- `tl.dot(a, b, acc)` 会根据架构映射到 WMMA、WGMMA 或 MFMA
- 为了精度，accumulator 应该使用 FP32
- 对 NVIDIA tensor core 路径，`BLOCK_K` 通常至少要 `16`
- 在 MI300X 上，MFMA 通常更偏好 `BLOCK_K=32`，并配合 wavefront-64
## Autotune 配置
### NVIDIA（Ampere/Ada/Hopper）
```python
NVIDIA_CONFIGS = [
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': 128, 'BLOCK_SIZE_N': 128, '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': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
]
```
### AMD MI300X（CDNA3，gfx942）
```python
AMD_CONFIGS = [
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=2),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=2),
triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=2),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'waves_per_eu': 8}, num_warps=4, num_stages=2),
]
```
AMD 侧关键差异：
- wavefront 是 64 线程，不是 32
- 在 HIP 上 `num_stages` 必须 `>=1`
- `waves_per_eu` 有助于调 occupancy：compute-bound 可试 `2-4`，memory-bound 可试 `4-8`
## 验证
```bash
python skills/kernels/gemm/test_gemm.py
```
- 正确性：对 FP16 使用 `torch.allclose(triton_result, torch.matmul(A, B), atol=1e-2, rtol=1e-2)`
- 性能：TFLOPS = `2*M*N*K / latency_s / 1e12`
## 常见陷阱
1. **K 边界缺少 mask**：当 `K` 不能被 `BLOCK_K` 整除时，会读到垃圾数据
2. **累加器 dtype 错误**：用 FP16 做累加会导致严重精度损失
3. **`GROUP_SIZE_M=1`**：等于关闭 L2 grouping，在大矩阵上常常会掉 10-20%
4. **CDNA 上 `BLOCK_K` 太小**：会让 MFMA 效率变差
5. **写回前忘记 `.to(output_dtype)`**：会产生静默精度问题
## Agent 指令
1. 从 `triton_template.py` 开始
2. 先 profile，确认它是不是真的 compute-bound
3. 如果是 compute-bound 且低于峰值的 60%，优先检查 tensor core 利用率，并考虑增大 `BLOCK_K`
4. 如果偏 memory-bound，就检查 L2 命中率，并调整 `GROUP_SIZE_M` 或尝试 Split-K
5. 每次修改后都运行 `test_gemm.py`
## 参考资料
- [Triton Matrix Multiplication Tutorial](https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html)
- [AutoKernel program.md Tier 1](https://github.com/RightNow-AI/autokernel)
- [CUTLASS: Persistent Kernels and Stream-K](https://research.colfax-intl.com/cutlass-tutorial-persistent-kernels-and-stream-k/)
在 GitHub 查看