softmax-kernel-optimization

# Softmax Kernel 优化 ## 概览 Softmax 定义为： `out[i] = exp(x[i] - max(x)) / sum(exp(x - max(x)))` 在 LLM 常见负载下，Softmax 通常是 **memory-bound**。 关键指标是 **GBps**。 融合 kernel 会在每行的一次 program launch 中完成全部步骤， 避免 max、exp-sum、normalize 三次独立遍历。 ## 核心技巧 ### 安全 Softmax（数值稳定） 1. 把整行加载到寄存器里，并用 FP32 表示 2. 求整行最大值：`m = max(x)` 3. 计算平移后的指数：`e = exp(x - m)` 4. 计算总和：`s = sum(e)` 5. 归一化：`out = e / s` ### 多行处理 对于较短的行宽（例如 `<2048`），每个 program 只处理一行会浪费 GPU 资源。 让一个 program 处理多行通常能提高 occupancy。 经验启发式： `ROWS_PER_PROG = max(1, 4096 / n_cols)` ### `BLOCK_SIZE` 选择 `BLOCK_SIZE = next_power_of_2(n_cols)`， 并且必须保证一个 program 可以覆盖整行。 对于非常大的行（例如 `>32K`），应考虑多次遍历或分块处理。 ## Autotune 说明 - `BLOCK_SIZE` 由 `n_cols` 决定，不是自由参数 - `num_warps`：`n_cols<=2048` 时常用 `4`，`n_cols<=8192` 时常用 `8`，更大时可用 `16` - `NUM_ROWS_PER_PROGRAM`：通常在 `1-8` 之间，取决于行宽 ### AMD 特定说明 - 由于 wavefront-64，`num_warps=4` 就表示 256 线程 - `num_stages=2` 是比较稳妥的默认值；在 HIP 上避免使用 `0` ## 验证 ```bash python skills/kernels/softmax/test_softmax.py ``` - 正确性：与 `torch.softmax(x, dim=-1)` 对比 - 性能：`GBps = (read_bytes + write_bytes) / latency_s / 1e9` ## 常见陷阱 1. **省略减最大值**：会导致 `exp(x)` 溢出 2. **用 FP16 做归约**：指数和很容易溢出 3. **`BLOCK_SIZE < n_cols`**：只能覆盖部分行，结果一定错误 4. **小行宽时每个 program 只处理一行**：很容易把 occupancy 做崩 ## Agent 指令 1. 先确认它是 memory-bound 2. 让 `BLOCK_SIZE` 对齐到 `next_power_of_2(n_cols)` 3. 行宽较小时使用多行处理 4. 目标是超过 HBM 峰值带宽的 60% 5. 遇到超大词表时，考虑分块或 online 变体 ## 参考资料 - [Triton Fused Softmax Tutorial](https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html) - [Online Normalizer Calculation (Milakov & Gimelshein, 2018)](https://arxiv.org/abs/1805.02867)