myWiki/concepts/parallel-scan.md

title, created, updated, type, tags, sources

title	created	updated	type	tags	sources
Parallel Scan (Parallel Associative Scan)	2026-06-18	2026-06-18	concept	gpu-optimization parallel-algorithm prefix-sum	https://arxiv.org/abs/2312.00752

Parallel Scan

定义

Parallel Scan（并行扫描，也称 Parallel Associative Scan / Blelloch Scan）是一种在 GPU 上并行计算前缀和（prefix sum）的算法。在 Mamba 中，它被用于将 selective-state-space 的顺序循环更新展开为可并行的关联操作，是 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

参考

• Blelloch (1990) — 经典并行 scan 算法
• gu-mamba (Gu & Dao, 2024)