forked from ccf-ai-infra/Intro-ops
261 lines
7.8 KiB
Plaintext
261 lines
7.8 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"# 04 — Softmax 算子\n",
|
||
"\n",
|
||
"## 学习目标\n",
|
||
"\n",
|
||
"1. 理解数值稳定性：为什么需要减 max\n",
|
||
"2. 理解 online softmax 算法（一遍扫描 vs 三趟扫描）\n",
|
||
"3. 理解 TileLang 中为什么用 `exp2`/`log2` 替代 `exp`/`log`\n",
|
||
"4. 完成 NVIDIA + TileLang 两个 kernel 的 TODO\n",
|
||
"\n",
|
||
"Softmax 是四个算子中最复杂的——既需要线程间通信（同步归约），又有数值稳定性陷阱。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 概念导入：Softmax 的数值稳定性\n",
|
||
"\n",
|
||
"### 朴素公式\n",
|
||
"\n",
|
||
"$$\\text{softmax}(x_i) = \\frac{e^{x_i}}{\\sum_j e^{x_j}}$$\n",
|
||
"\n",
|
||
"### 问题\n",
|
||
"\n",
|
||
"$e^{88.7} \\approx 1.6 \\times 10^{38}$，接近 FP32 上限。如果 $x_i = 100$，则 $e^{100}$ 溢出为 inf，inf/inf = NaN。\n",
|
||
"\n",
|
||
"### 解决：减最大值\n",
|
||
"\n",
|
||
"$$\\text{softmax}(x_i) = \\frac{e^{x_i - \\max(x)}}{\\sum_j e^{x_j - \\max(x)}}$$\n",
|
||
"\n",
|
||
"分子分母同除 $e^{\\max(x)}$，数学结果不变，但 $e^{x_i - \\max(x)}$ 最大为 $e^0 = 1$，永不溢出。\n",
|
||
"\n",
|
||
"详细图示见 [docs/diagrams/softmax-pipeline.md](../docs/diagrams/softmax-pipeline.md)。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## PyTorch 参考实现"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import torch\n",
|
||
"\n",
|
||
"src = torch.randn(32, 128, device=\"cuda\", dtype=torch.float32)\n",
|
||
"expected = torch.softmax(src, dim=1)\n",
|
||
"\n",
|
||
"print(f\"输入 shape: {src.shape}\")\n",
|
||
"print(f\"输出 shape: {expected.shape}\")\n",
|
||
"print(f\"每行和: {expected.sum(dim=1)[:5]}\") # 应该全为 1.0\n",
|
||
"\n",
|
||
"# 演示溢出问题\n",
|
||
"big = torch.tensor([100.0, 200.0, 300.0], device=\"cuda\")\n",
|
||
"naive = torch.exp(big) / torch.exp(big).sum()\n",
|
||
"print(f\"\\n朴素 softmax([100, 200, 300]): {naive}\") # NaN!\n",
|
||
"\n",
|
||
"stable = torch.softmax(big, dim=0)\n",
|
||
"print(f\"稳定 softmax([100, 200, 300]): {stable}\") # [0, 0, 1]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## NVIDIA CUDA Kernel（三趟扫描）\n",
|
||
"\n",
|
||
"打开 `ops/softmax/nvidia/kernel.cuh`。基础版本使用三趟扫描：\n",
|
||
"\n",
|
||
"1. Pass 1: 求行最大值 `max_val`\n",
|
||
"2. Pass 2: 写 `exp(x - max)` 到 out，同时累加 sum\n",
|
||
"3. Pass 3: 除 sum 归一化"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"KERNEL_SKELETON = \"\"\"\n",
|
||
"__global__ void softmax_rowwise_kernel(\n",
|
||
" float *out, const float *in, int64_t rows, int64_t cols) {\n",
|
||
"\n",
|
||
" // TODO: implement a numerically stable row-wise softmax kernel.\n",
|
||
" //\n",
|
||
" // Suggested steps:\n",
|
||
" // 1. One block per row.\n",
|
||
" // 2. Pass 1: find row max via shared memory reduction.\n",
|
||
" // 3. Pass 2: exp(x - max) → out, accumulate sum via reduction.\n",
|
||
" // 4. Pass 3: out[i] /= sum.\n",
|
||
"}\n",
|
||
"\"\"\"\n",
|
||
"print(KERNEL_SKELETON)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 为什么是三趟？\n",
|
||
"\n",
|
||
"因为 softmax 有两个需要全局信息才能计算的步骤：\n",
|
||
"- 求 max 需要全局信息（所有元素的最大值）\n",
|
||
"- 求 sum 需要全局信息（所有 exp(x-max) 的和）\n",
|
||
"\n",
|
||
"这两个全局信息都通过 shared memory 归约获得，每一步都需要 block 内同步。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 验证 NVIDIA 版本\n",
|
||
"import subprocess, sys\n",
|
||
"result = subprocess.run([\n",
|
||
" sys.executable, \"-m\", \"pytest\",\n",
|
||
" \"tests/op_tests/test_softmax.py\", \"-v\",\n",
|
||
" \"--backend\", \"nvidia\"\n",
|
||
"], capture_output=True, text=True)\n",
|
||
"print(result.stdout)\n",
|
||
"if result.returncode == 0:\n",
|
||
" print(\"✓ NVIDIA softmax kernel 测试全部通过！\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## TileLang Kernel（Online Softmax）\n",
|
||
"\n",
|
||
"打开 `ops/softmax/tilelang/kernel.py`。TileLang 版本使用 **online softmax**——只需两趟扫描。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"TILELANG_SKELETON = \"\"\"\n",
|
||
"@tilelang.jit\n",
|
||
"def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):\n",
|
||
" log2_e = 1.44269504 # log2(e) for exp/log2 conversion\n",
|
||
" N, M = T.const(\"N, M\")\n",
|
||
" dtype = T.float32\n",
|
||
" src: T.Tensor((N, M), dtype)\n",
|
||
" out = T.empty((N, M), dtype)\n",
|
||
"\n",
|
||
" # TODO: implement a tiled row-wise softmax with online algorithm.\n",
|
||
" #\n",
|
||
" # Two-pass approach:\n",
|
||
" # Pass 1: scroll through column tiles, update running log-sum-exp.\n",
|
||
" # Pass 2: scroll again, normalize each tile with final lse.\n",
|
||
"\n",
|
||
" return out\n",
|
||
"\"\"\"\n",
|
||
"print(TILELANG_SKELETON)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 为什么用 `exp2`/`log2`？\n",
|
||
"\n",
|
||
"GPU 硬件对 2 的幂运算 (`2^x`) 有专门的快速指令，比自然指数 (`e^x`) 快。\n",
|
||
"\n",
|
||
"转换公式：$\\exp(x) = 2^{x \\cdot \\log_2(e)}$\n",
|
||
"\n",
|
||
"其中 $\\log_2(e) \\approx 1.44269504$（代码中的 `log2_e` 常量）。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 验证 TileLang 版本\n",
|
||
"result = subprocess.run([\n",
|
||
" sys.executable, \"-m\", \"pytest\",\n",
|
||
" \"tests/op_tests/test_softmax.py\", \"-v\",\n",
|
||
" \"--backend\", \"tilelang\"\n",
|
||
"], capture_output=True, text=True)\n",
|
||
"print(result.stdout)\n",
|
||
"if result.returncode == 0:\n",
|
||
" print(\"✓ TileLang softmax kernel 测试全部通过！\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 对比验证"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from operator_runtime import softmax\n",
|
||
"\n",
|
||
"src = torch.randn(32, 128, device=\"cuda\", dtype=torch.float32)\n",
|
||
"my_out = softmax(src, dim=1, backend=\"nvidia\")\n",
|
||
"torch_out = torch.softmax(src, dim=1)\n",
|
||
"\n",
|
||
"print(f\"我的 kernel 行和: {my_out.sum(dim=1)[:5]}\")\n",
|
||
"print(f\"PyTorch 行和: {torch_out.sum(dim=1)[:5]}\")\n",
|
||
"print(f\"最大误差: {(my_out - torch_out).abs().max().item():.2e}\")\n",
|
||
"print(f\"一致: {torch.allclose(my_out, torch_out, atol=1e-5, rtol=1e-5)}\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 检查清单\n",
|
||
"\n",
|
||
"- [ ] 理解数值稳定性问题：为什么需要减 max\n",
|
||
"- [ ] 理解 online softmax 算法（一遍扫描 vs 三趟扫描）\n",
|
||
"- [ ] 理解 log-sum-exp 的滚动更新逻辑\n",
|
||
"- [ ] 理解 TileLang 中为什么用 `exp2`/`log2` 替代 `exp`/`log`\n",
|
||
"- [ ] NVIDIA `kernel.cuh` TODO 完成\n",
|
||
"- [ ] TileLang `kernel.py` TODO 完成\n",
|
||
"- [ ] 两种后端测试全部通过\n",
|
||
"- [ ] 性能与 PyTorch 参考实现对比\n",
|
||
"- [ ] 尝试 warp-level 优化"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"name": "python",
|
||
"version": "3.12.0"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 5
|
||
}
|