# 融合 Softmax

## 实现目标

`fused_softmax.py` 将一行 Softmax 所需的最大值归约、指数、求和和归一化融合到
单个 Triton 内核中。与会产生多个中间张量的朴素 PyTorch 写法相比，该实现只需
读取一次输入并写入一次输出。

## 核心数据流

1. `naive_softmax()` 展示数值稳定的参考算法：每行先减去最大值，再计算指数和
   归一化结果。
2. `softmax()` 根据列数选择不小于该列数的最小二次幂作为 `block_size`，并结合
   GPU 属性估算 warp 数量、流水线阶段数和可并发程序数。
3. `_softmax_kernel()` 让每个程序处理若干间隔分布的行。每行只加载一次，在
   SRAM 中完成最大值归约、指数、求和与除法，最后写回 DRAM。
4. 超出真实列数的槽位在加载时填充为负无穷，确保它们对最大值和 softmax 分母
   没有贡献。

## 关键细节与限制

- 单行经过二次幂填充后必须能够放入可用 SRAM，否则包装函数会拒绝启动。
- 输入必须是二维、连续且位于当前 CUDA 设备的张量。
- 启动程序数受行数、SM 数量和内核资源占用共同限制，以便程序通过循环覆盖全部行。
- 融合的主要收益来自减少 DRAM 读写，而不是减少 softmax 的数学运算量。

## 测试与基准

```bash
python3 fused_softmax.py
python3 fused_softmax.py --benchmark
```

第一条命令执行数值测试；第二条命令还会比较 Triton 与 PyTorch 的有效内存带宽。

## 配套视频

[bilibili](https://www.bilibili.com/video/BV1fMyWBgERM?spm_id_from=333.788.videopod.episodes&vd_source=e98b669ccbafff4b5aa59dd6303b722f&p=9)
