原理精讲

Tokenization 基础 BPE 算法详解 GPT 系列 Tokenizer BPE 训练工程化

Attention 机制详解

GPU 编程基础

GPU 架构基础张量布局 Triton 入门：向量加法

Flash Attention 原理详解从朴素实现到 Auto-Tuning Block Pointer 与多维支持 Causal Masking 优化 Grouped Query Attention 反向传播实现

动手训练

系统工程FlashAttention

反向传播实现

会员专享

实现 Flash Attention 的梯度计算，通过 Recomputation 实现内存高效的训练。

登录以继续阅读

这是一篇付费内容，请登录您的账户以访问完整内容。

Grouped Query Attention

实现 GQA/MQA 支持,让多个 Query Head 共享 KV,优化 KV Cache 内存占用。

目录

为什么需要自定义反向传播？

PyTorch 自动微分的局限性

Recomputation 策略

Attention 反向传播的数学原理

前向传播回顾

\frac{\partial \mathcal{L}}{\partial \mathbf{V}}

\frac{\partial \mathcal{L}}{\partial \mathbf{P}}

\frac{\partial \mathcal{L}}{\partial \mathbf{S}}

(Softmax 反向传播)

\frac{\partial \mathcal{L}}{\partial \mathbf{Q}}

\frac{\partial \mathcal{L}}{\partial \mathbf{K}}

完整的梯度计算流程

代码实现解析

Forward Kernel 的修改

Backward Kernel 实现

关键实现细节

1. 循环顺序的反转

2. Atomic Add 处理 dQ

3. 重新计算 P 而非保存

torch.autograd.Function 封装

数值正确性测试

内存占用对比

设计权衡与优化方向

Recomputation 的开销

进一步优化方向