分块计算Tiling
让小块数据尽可能留在较快存储层中完成计算。
- 核心约束
- 分块后仍需保持原始运算的数据依赖与结果语义。
- 适用问题
- 矩阵计算、Attention 与存储局部性优化。
- 边界与取舍
- 块大小受硬件容量与并行度约束。
常见操作#
以矩阵乘法 / 注意力为例；块大小受缓存或片上存储容量限制。
| 操作 | 时间复杂度 | 条件与说明 |
|---|---|---|
| 分块矩阵乘法 | 仍有 标量算术主项 | 经典 n×n GEMM 通过重用块减少慢速存储读写，不等于改变三次算术阶数。 |
| 分块注意力 | 精确全注意力仍有 算术量 | N 为序列长度，d 为头维度；在线 softmax 避免保存完整 N×N 中间矩阵。 |
时间与空间#
- 时间复杂度
- 算术工作、跨层数据传输和工作集分别核算
- 辅助空间
- 由块大小、累加器与需要保留的输出决定
好的块大小取决于寄存器、共享内存、向量指令和占用率。FlashAttention 的改进重点是 HBM I/O 与中间存储，不是把全注意力无条件改为线性算法。
工程中的应用#
1 个案例 · 附源码通过分块、融合与重计算，让注意力计算更适应 GPU 存储层次。
附源码片段 · 内存管理
基础原理参考资料
Dao et al. · FlashAttention第 3 节的分块、在线 softmax、重计算与 I/O 分析。
NVIDIA CUDA · Shared memory matrix multiplication矩阵块复用减少全局内存传输的具体实现。