20260625:很多新内容
This commit is contained in:
69
concepts/hardware-aware-algorithm.md
Normal file
69
concepts/hardware-aware-algorithm.md
Normal file
@@ -0,0 +1,69 @@
|
||||
---
|
||||
title: "Hardware-Aware Algorithm (Mamba)"
|
||||
created: 2026-06-18
|
||||
updated: 2026-06-18
|
||||
type: concept
|
||||
tags: ["gpu-optimization", "mamba", "parallel-scan", "kernel-fusion"]
|
||||
sources: ["https://arxiv.org/abs/2312.00752"]
|
||||
---
|
||||
|
||||
# Hardware-Aware Algorithm
|
||||
|
||||
## 定义
|
||||
|
||||
Hardware-Aware Algorithm 是 Mamba 为高效计算 [[selective-state-space|S6(选择性 SSM)]]而设计的 GPU 优化算法。核心思想是**利用 GPU 内存层次结构**(SRAM vs HBM),通过 kernel fusion 和并行扫描(parallel associative scan)避免将扩展状态写入慢速 HBM。
|
||||
|
||||
## 为什么需要
|
||||
|
||||
S6 的选择机制使 SSM 变为时间变化 → **不能用卷积**(卷积要求 LTI)。唯一的计算方式是循环模式:
|
||||
|
||||
```
|
||||
h_t = A_bar_t * h_{t-1} + B_bar_t * x_t (需要 scan)
|
||||
```
|
||||
|
||||
朴素循环是 O(n) 顺序的,对训练不可接受。Mamba 通过以下方案解决。
|
||||
|
||||
## 关键技术
|
||||
|
||||
### 1. 并行关联扫描(Parallel Associative Scan / Blelloch Scan)
|
||||
|
||||
将循环计算的序列依赖展开为前缀和(prefix sum)形式,利用关联性在 GPU 上并行完成:
|
||||
|
||||
```
|
||||
思路: h_t = A_t * h_{t-1} + B_t * x_t
|
||||
= f(f(...f(h_0, x_1), x_2), ..., x_t)
|
||||
|
||||
展开为关联操作,用 Blelloch scan 在 O(log n) 并行步完成
|
||||
```
|
||||
|
||||
### 2. IO 感知 Kernel Fusion
|
||||
|
||||
| 位置 | 操作 |
|
||||
|------|------|
|
||||
| HBM(慢) | 存储输入 x、参数 |
|
||||
| SRAM(快) | 加载 → 离散化 → scan → 写回 |
|
||||
| HBM(慢) | 存储输出 y |
|
||||
|
||||
关键:扩展后的状态 `A_bar, B_bar` 维度为 (B, L, D, N),不在 HBM 中物化——在 SRAM 中计算完毕即写回压缩结果。
|
||||
|
||||
### 3. 重计算(Recomputation)
|
||||
|
||||
反向传播不保留前向的中间状态,在 backward pass 中从输入和参数重新计算 → 减少内存占用。
|
||||
|
||||
## 效果
|
||||
|
||||
- 比所有基于卷积的 SSM(S4、H3、Hyena)快 **3×**(A100 GPU)
|
||||
- 理论复杂度:O(BLDN) vs 卷积 SSM 的伪线性
|
||||
- 使选择性 SSM 的训练在实践中可行
|
||||
|
||||
## 相关概念
|
||||
|
||||
- [[selective-state-space]] — 该算法服务的选择机制
|
||||
- [[parallel-scan]] — 核心计算原语
|
||||
- [[structured-state-space-models]] — 卷积模式的替代路径
|
||||
- [[gu-mamba|Mamba 论文]]
|
||||
|
||||
## 参考
|
||||
|
||||
- [[gu-mamba|Mamba]] (Gu & Dao, 2024) Section 3.3
|
||||
- FlashAttention (Dao et al., 2022) — 同作者的 IO 感知优化风格
|
||||
Reference in New Issue
Block a user