Metadata-Version: 2.5
Name: flash-swiglu-mlp
Version: 0.4.2
Summary: A fused Triton SwiGLU MLP kernel (gate_proj + up_proj + SiLU fused, TMA-based) for training on Hopper/Blackwell GPUs.
Project-URL: Homepage, https://github.com/Pearblossom-M/flash-swiglu-mlp
Project-URL: Repository, https://github.com/Pearblossom-M/flash-swiglu-mlp
Project-URL: Issues, https://github.com/Pearblossom-M/flash-swiglu-mlp/issues
Author: Pearblossom-M
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: cuda,fused-kernel,gpu-kernel,llm-training,swiglu,tma,triton
Classifier: Development Status :: 3 - Alpha
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Requires-Dist: torch>=2.4
Requires-Dist: triton>=3.4
Provides-Extra: bench
Requires-Dist: liger-kernel; extra == 'bench'
Provides-Extra: dev
Requires-Dist: build>=1.0; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: ruff>=0.5; extra == 'dev'
Description-Content-Type: text/markdown

**简体中文** | [English](./README.en.md)

# Flash SwiGLU MLP

一个针对 SwiGLU MLP（LLaMA / Mistral / Qwen 类模型的前馈模块）做**全流程协同优化**的 Triton 融合算子：不止融合 SwiGLU 激活本身，而是围绕前向与反向的完整数据流，对 kernel 融合、张量生命周期、显存布局和 kernel launch 开销做端到端优化。

计算对象：`down_proj(silu(gate_proj(x)) * up_proj(x))`

完整的设计思路、性能与正确性测试见技术报告：[flash-swiglu-mlp-technical-report.pdf](./flash-swiglu-mlp-technical-report.pdf)

## 适用设备

| 架构                                                         | 支持情况                                                     |
| ------------------------------------------------------------ | ------------------------------------------------------------ |
| **Hopper（SM90）及以上：H100 / H200 / Blackwell（SM100、SM120，如 RTX 50 系）** | ✅ **推荐**。本实现基于 Triton 的 host-side TMA（`TensorDescriptor`）API，Hopper+ 有硬件 TMA 支持，可发挥全部性能 |
| **Ampere（SM80）：A100 / A800 / RTX 30 系**                  | ✅ 可以运行。Triton 会在无 TMA 硬件的架构上对 descriptor 做兼容处理，但没有硬件 TMA 加速，性能可能有折扣 |
| 更老的架构 / CPU                                             | ❌ 不支持                                                     |

> 本项目目前的性能验证基于单卡 RTX 5060 Ti（Blackwell，SM120）。

## 环境要求

- Python >= 3.9
- `torch >= 2.4`、`triton >= 3.4`（TMA host API 在 triton 3.3.1 尚不可用）
- 实测环境：RTX 5060 Ti（SM120）+ Driver 580.159.03 / CUDA 13.1 / torch 2.13.0 / triton 3.7.1

## 安装

```bash
pip install flash-swiglu-mlp
```

从 GitHub 直接安装：

```bash
pip install git+https://github.com/Pearblossom-M/flash-swiglu-mlp.git
```

从本地源码安装（开发模式）：

```bash
git clone https://github.com/Pearblossom-M/flash-swiglu-mlp.git
cd flash-swiglu-mlp
pip install -e ".[dev]"
```


## 快速开始

提供两种用法：`FlashSwiGLUMLP` 模块（自带权重，可直接作为模型一层）和 `flash_swiglu` 函数式 API（权重自备，嵌入已有模型）。

### 模块用法：FlashSwiGLUMLP

```python
import torch
from flash_swiglu_mlp import FlashSwiGLUMLP

# hidden_size: 模型隐藏层大小；intermediate_size: 前馈中间维度
mlp = FlashSwiGLUMLP(hidden_size=4096, intermediate_size=11008, device="cuda", dtype=torch.bfloat16)

x = torch.randn(2, 1024, 4096, device="cuda", dtype=torch.bfloat16, requires_grad=True)

out = mlp(x)            # [2, 1024, 4096]
out.sum().backward()    # 梯度正常回传到 x 与三个 proj 的权重
```

- 内部由三个 `nn.Linear(bias=False)` 组成：`gate_proj`、`up_proj`（权重 `[intermediate_size, hidden_size]`）和 `down_proj`（权重 `[hidden_size, intermediate_size]`），初始化方式即 `nn.Linear` 默认的 Kaiming uniform；
- 子模块命名与 LLaMA / Mistral / Qwen 的 HuggingFace 实现一致，state_dict 可零改名互相迁移；
- 与 autograd 完全兼容，可直接放入任意 `nn.Module`、被优化器正常更新；
- 在 `torch.no_grad()` 下前向时自动切换到推理专用 kernel（不保存 `G/U`，显存占用更低）。

### 函数式 API：flash_swiglu

已有权重张量时（例如替换现有模型的 MLP），可直接调用函数式接口：

```python
from flash_swiglu_mlp import flash_swiglu

# input: [B, S, hidden_size]
# gate_weight / up_weight: [intermediate_size, hidden_size]
# down_weight: [hidden_size, intermediate_size]
out = flash_swiglu(input, gate_weight, up_weight, down_weight)
```

输入是否需要梯度决定执行路径：任一输入 `requires_grad` 时走训练路径（保存反向所需中间量），否则自动走推理路径。

### 替换 HuggingFace 模型的 MLP（以 LLaMA 为例）

`FlashSwiGLUMLP` 的子模块命名（`gate_proj` / `up_proj` / `down_proj`）与 transformers 中 LLaMA / Mistral / Qwen 等模型的 MLP 完全一致，因此可以零改名加载原有权重，直接原地替换：

```python
import torch
from transformers import AutoModelForCausalLM
from flash_swiglu_mlp import FlashSwiGLUMLP

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B", torch_dtype=torch.bfloat16, device_map="cuda"
)

for layer in model.model.layers:
    new_mlp = FlashSwiGLUMLP(
        hidden_size=model.config.hidden_size,
        intermediate_size=model.config.intermediate_size,
        dtype=torch.bfloat16,
        device="cuda",
    )
    # 键名一一对应，无需任何重命名
    new_mlp.load_state_dict(layer.mlp.state_dict())
    layer.mlp = new_mlp
```

替换后训练与推理均正常：前向走融合 kernel，梯度照常回传；`torch.no_grad()` 下自动切换到推理 kernel。其他 SwiGLU 结构相同的模型（Mistral、Qwen 等）同理。

### Falcon H1 风格模型

本实现同样支持带缩放系数的 Falcon H1 风格 MLP：`down_multiplier * down_proj(gate_multiplier * silu(gate_proj(x)) * up_proj(x))`。只需在构造模块或调用函数时给 `gate_multiplier` 和 `down_multiplier` 赋值即可：

```python
# 模块用法
mlp = FlashSwiGLUMLP(
    hidden_size=2048,
    intermediate_size=4864,
    device="cuda",
    dtype=torch.bfloat16,
    gate_multiplier=1.125,  # 按目标模型的配置填写
    down_multiplier=2.0,
)

# 函数式用法
out = flash_swiglu(
    input, gate_weight, up_weight, down_weight,
    gate_multiplier=1.125, down_multiplier=2.0,
)
```

两个参数默认值均为 `1.0`：其他模型（LLaMA / Mistral / Qwen 等）无需任何改动，用法与之前完全一致。

## 支持的配置

|          | 说明                                                         |
| -------- | ------------------------------------------------------------ |
| dtype    | `float16`、`bfloat16`（性能与正确性均已验证）；`float32` 代码层面支持，未做性能验证 |
| Bias     | 不支持（与 LLaMA / Mistral / Qwen 的 SwiGLU 一致，本就无 bias） |
| 输入布局 | 连续的 `[B, S, dim]`                                         |
| 维度对齐 | 每行数据需 16 字节对齐（TMA 要求），即 fp16/bf16 下 `dim` 与 `hidden_dim` 需为 8 的倍数；建议取 64 的倍数以获得峰值性能 |

## 设计要点

- **前向融合**：单个 Triton kernel 内完成 `gate_proj + up_proj + SiLU + 门控`，输入 tile 在两个投影间复用，`G/U` 留在寄存器中直接参与激活计算；`down_proj` 保留为 cuBLAS GEMM。
- **AGU 连续布局**：`A、G、U` 存于单个 `[B, S, 3H]` 连续缓冲区，减少分配开销、缓解显存碎片，并为反向的大 GEMM 铺路。
- **反向四阶段**：`dA/dG/dU` 融合为一个 kernel（`dA` 不物化）；`dG/dU` 原地覆写 `G/U`，降低峰值显存；`dW_g/dW_u` 合并为单次 GEMM；`dI` 用单累加器融合计算，降低寄存器压力。
- **训练/推理路径分离**：推理不保存 `G/U`，只分配 `A`，自动切换。
- **分桶 Autotune**：按 `(dim, hidden_dim, M 分桶)` 做 key，避免每个 shape 重复调优。
- **L2 Swizzle**：matmul 类 kernel 使用 group-major tile 调度提升 L2 命中。

## 性能与正确性

在 RTX 5060 Ti 上，对比 `torch.compile`（max-autotune-no-cudagraphs）与 Liger-Kernel（`LigerSwiGLUMLP`）：

- **训练**：前向、反向、端到端全面领先，提升约 **5%~15%**，大维度 + 长序列下 TFLOPS 逼近设备算力上限；
- **峰值中间显存**：推理峰值中间显存在全部测试规模下全面优于 torch.compile 与 Liger Kernel。训练时，本实现在长序列场景下，峰值中间显存占用更低，但短序列场景占用反而更高，这主要是因为对比基线在不同数据规模下的行为差异（编译/内存规划策略不同），短序列场景使用了更激进的内存复用优化，显存占用更低；
- **推理**：prefill（大 M）场景领先；decoding（小 M）场景由于 host-side TMA descriptor 的 CPU 开销占比高，目前与 `torch.compile` 相当甚至略低，配合 CUDA Graph 才能体现优势；
- **正确性**：FP16/BF16、训练/推理、小/大形状下，前向输出与全部四个梯度均通过 `torch.allclose` 校验。

详细数据与图表见[技术报告](./flash-swiglu-mlp-technical-report.pdf)。

## 已知限制

1. 不支持 `retain_graph=True` / `create_graph=True` 的多次反向与高阶求导（原地覆写 `G/U` 的代价），常规训练不受影响；
2. 分布式训练（DDP/TP）在结构上兼容，但尚未在多卡环境实测；
3. 首次遇到新的 `(dim, hidden_dim, M 分桶)` 组合会触发一次 autotune，有一定预热开销。

## 测试与基准

```bash
# 单元测试（无可用 GPU 时自动跳过）
pip install -e ".[dev]"
pytest tests/ -v

# 性能 / 显存基准（需要 liger-kernel）
pip install -e ".[bench]"
python benchmarks/bench_flash_swiglu_mlp.py --dtype bfloat16 --only all
```

## 引用

```bibtex
@misc{flashswiglumlp,
  title  = {Flash SwiGLU MLP: a fully-pipeline optimized fused Triton SwiGLU MLP kernel},
  author = {Pearblossom-M},
  year   = {2026},
  url    = {https://github.com/Pearblossom-M/flash-swiglu-mlp}
}
```

## License

Apache License 2.0，见 [LICENSE](./LICENSE)。