52 lines
1.6 KiB
Markdown
52 lines
1.6 KiB
Markdown
---
|
||
title: "Parallel Scan (Parallel Associative Scan)"
|
||
created: 2026-06-18
|
||
updated: 2026-06-18
|
||
type: concept
|
||
tags: ["gpu-optimization", "parallel-algorithm", "prefix-sum"]
|
||
sources: ["https://arxiv.org/abs/2312.00752"]
|
||
---
|
||
|
||
# Parallel Scan
|
||
|
||
## 定义
|
||
|
||
Parallel Scan（并行扫描，也称 Parallel Associative Scan / Blelloch Scan）是一种在 GPU 上并行计算前缀和（prefix sum）的算法。在 Mamba 中，它被用于将 [[selective-state-space|S6]] 的顺序循环更新展开为可并行的关联操作，是 [[hardware-aware-algorithm]] 的核心计算原语。
|
||
|
||
## 为什么需要
|
||
|
||
S6（选择性 SSM）的序列依赖：
|
||
|
||
```
|
||
h_t = A_bar_t * h_{t-1} + B_bar_t * x_t
|
||
```
|
||
|
||
朴素实现是 O(n) 顺序的，无法并行训练。但该操作满足结合律（associativity），可以重新排列为 parallel scan。
|
||
|
||
## 算法直觉
|
||
|
||
将序列分为两半，递归计算：
|
||
- 第一半的完整结果
|
||
- 第二半需要第一半的"最终状态"作为初始条件
|
||
|
||
Blelloch 的两阶段方法：
|
||
1. **Up-sweep（reduce）**：构建局部聚合
|
||
2. **Down-sweep**：将聚合结果分发到各元素
|
||
|
||
复杂度：O(n) work, O(log n) depth → GPU 友好。
|
||
|
||
## 在 Mamba 中的实现
|
||
|
||
Mamba 的硬件感知算法将 scan 与离散化融合在单个 GPU kernel 中，直接在 SRAM 上操作，避免中间结果写入 HBM。
|
||
|
||
## 相关概念
|
||
|
||
- [[hardware-aware-algorithm]] — Mamba 中 scan 的完整实现
|
||
- [[selective-state-space]] — 需要 scan 的原因（时间变化 SSM 不能卷积）
|
||
- [[gu-mamba|Mamba 论文]]
|
||
|
||
## 参考
|
||
|
||
- Blelloch (1990) — 经典并行 scan 算法
|
||
- [[gu-mamba|Mamba]] (Gu & Dao, 2024)
|