Blog on Lorn

浅谈低精度浮点数

Wed, 10 Jun 2026 20:25:01 +0800

Motivation

最近在搭建 VLM real quantization baseline。最初的实验环境选在 RTX PRO 6000 Blackwell 上，但在实际部署过程中发现，当前主流推理框架对 Blackwell 设备上传统整数低比特量化的支持还不够成熟，尤其是常见的 W4A8、W4A16 等 INT 量化配置，在 vLLM、SGLang、TensorRT-LLM 等推理栈中的可用性、kernel 覆盖和性能表现，都与 Hopper/Ada 等架构存在明显差异。

进一步调研后发现，Blackwell 架构的低精度推理路线正在发生变化：相比传统 INT4/INT8，硬件与软件生态开始更多地围绕低精度浮点格式展开，例如 MXFP4、NVFP4 等 micro-scaling FP4 格式。这类格式在保持 4-bit 存储与计算效率的同时，通过局部 scale、浮点指数以及更细粒度的动态范围适配，试图缓解 INT 量化中常见的离群值敏感和动态范围不足问题。

与此同时，近两年的低精度量化研究也逐渐从“如何把 INT4 做稳”转向“如何利用硬件原生支持的 FP4 格式”。例如，围绕 MXFP4/NVFP4 的推理量化、FP4 全量训练、FP4 误差分析以及格式感知的旋转/缩放方法，已经在 ICLR、ICML、NeurIPS 以及 arXiv 上陆续出现。因此，我准备写下这篇 blog，作为对低精度浮点量化格式的学习笔记，也作为后续在 Blackwell 设备上构建 VLM real quant baseline 的技术背景梳理。

Note

IEEE 754 标准浮点数

我们先来看看在内存中，浮点数是如何按照 IEEE 754 标准存储的。

本质上，浮点数可以看作是用科学计数法来表示一个数：

$$ V = (-1)^S \times M \times R^E $$

其中 S 表示符号位，M 表示尾数，R 表示基数，E 表示指数。对于计算机中的二进制浮点数而言，基数 R = 2。

以 IEEE 754 单精度浮点数，也就是 float32 为例，它一共占 32 bit，具体划分如下：

符号位 S：1 bit
指数字段 Exponent：8 bit
尾数字段 Fraction：23 bit

为了在有限的位宽内尽可能扩大表示范围并提高精度，IEEE 754 对指数和尾数做了特殊规定：

对于规格化浮点数，尾数的整数部分总是 1，因此这个最高位的 1 不需要显式存储，称为隐藏位。于是，尾数字段虽然只有 23 bit，但实际可以表示 24 bit 的有效数字。
指数字段采用偏置（bias）表示。对于 float32，指数字段占 8 bit，存储值范围为 0 ~ 255，偏置值为 127。对于规格化数，真实指数为：

$$ E = \text{Exponent} - 127 $$

需要注意的是，Exponent = 0 和 Exponent = 255 有特殊含义，分别用于表示 0、非规格化数、无穷大和 NaN。因此，规格化单精度浮点数的真实指数范围是 $-126 \sim 127$。

单精度浮点数在内存中的布局如下：

 bit index
 31 30 23 22 0
 │ │ │ │ │
 ▼ ▼ ▼ ▼ ▼
┌───────┬────────────────────────────┬──────────────────────────────────┐
│ S │ Exponent │ Fraction │
│ 1 bit │ 8 bits │ 23 bits │
└───────┴────────────────────────────┴──────────────────────────────────┘
 │ │ │
 │ │ └─ fraction bits of mantissa
 │ └─ biased exponent, bias = 127
 └─ sign bit: 0 = positive, 1 = negative

例如，我们想用单精度浮点数表示 25.125。

首先将它转换为二进制：

整数部分：25(D) = 11001(B)
小数部分：0.125(D) = 0.001(B)

因此：

25.125(D) = 11001.001(B)

将其写成二进制科学计数法：

11001.001(B) = 1.1001001(B) × 2^4

所以：

S = 0
M = 1.1001001
Fraction = 10010010000000000000000
Exponent = 4 + 127 = 131(D) = 10000011(B)

因此，25.125 的 IEEE 754 单精度表示为：

S | Exponent | Fraction
0 | 10000011 | 10010010000000000000000

完整的 32 bit 表示为：

01000001110010010000000000000000

对应的十六进制表示为：

0x41C90000

浮点数表示会带来一定的精度损失，主要原因是并非所有十进制小数都能被有限位宽的二进制小数精确表示。例如，0.2 转换为二进制后是一个无限循环小数：

0.2(D) = 0.0011001100110011...(B)

由于计算机的位宽有限，存储时只能保留有限位，并按照 IEEE 754 的规则进行舍入，因此会产生精度误差。

在 IEEE 754 单精度浮点数下，常用的有限数表示范围大约为：

-3.4 × 10^38 ~ 3.4 × 10^38

更准确地说，单精度浮点数能够表示的最大有限正数约为：

3.4028235 × 10^38

最小正规格化数约为：

1.17549435 × 10^-38

如果考虑非规格化数，还可以表示更接近 0 的数，但精度会进一步降低。

在如今的 LLM 训练推理过程中，更通用的浮点数格式是 FP16 与 BF16。它们与 FP32 的区别在于：

FP16 的指数位为 5，尾数位为 10
BF16 的指数位为 8，尾数位为 7

我们知道，在总位宽固定的情况下，尾数越多则指数越少，能够表示的范围越小，但相应精度越高。因此 FP16 比 BF16 精度更高，但表示范围更小。

正因如此，BF16 通常用在训练场景（更大的动态范围更利于反向传播中的梯度表示），而 FP16 通常用在推理场景（主要是因为目前很多主流推理框架对 FP16 的 kernel 优化做得更好）。

FP8

FP8 是 Hopper 架构设备开始支持的低精度浮点数，通常有两种表示形式：E4M3 和 E5M2（其中 E 表示指数位数，M 表示尾数位数）。

同样地，E4M3 精度更高、范围更小，而 E5M2 范围更大、精度更低。

在支持 FP8 的硬件（Hopper 架构显卡）上，已经支持 FP8 格式的矩阵乘法（E4M3×E4M3、E5M2×E5M2、E4M3×E5M2、E5M2×E4M3）。

我们接下来从精度和性能两个方面，对比分析 FP8 与 INT8 在量化上的表现。

首先分析二者在精度上的表现。

由浮点数的定义可知，浮点数可以看作是在相邻 2 的幂（$2^E$）之间均匀采样的 $2^M$ 个样本。例如在 E5M2 中，2（一次幂）和 4（二次幂）之间有 4 个样本，4 和 8（三次幂）之间也有 4 个样本；而在 E4M3 中，2 和 4 之间有 8 个样本。由此可知：

E5M2 的量化误差比 E4M3 更大；
FP8 表示的数值是非均匀的，浮点量化的误差会随着数值的增大而增大。

这正是 FP8 与 INT8 量化的最大差别：对于 FP8 而言，越靠近 0 分布越稠密，越远离 0 分布越稀疏。

接下来从数学角度分析 FP8 与 INT8 在量化上的表现。

对于 INT8，量化可以写作：

$$ \hat x = s \cdot \text{round}\left(\frac{x}{s}\right) $$

其中 $s$ 是 scale。如果使用对称 INT8，那么通常有：

$$ q \in [-127, 127] $$

假设一组待量化值的最大绝对值是 $\alpha$，那么 scale 大约是：

$$ s = \frac{\alpha}{127} $$

量化绝对误差满足：

$$ |x - \hat x| \leq \frac{s}{2} = \frac{\alpha}{254} $$

而相对误差为：

$$ \frac{|x - \hat x|}{|x|} \leq \frac{\alpha}{254\,|x|} $$

可以看到，无论从绝对误差还是相对误差来看，INT8 量化都更容易受到 outlier 的影响。

而对于 FP8，假设尾数有 $M$ bit，那么在区间 $[2^e, 2^{e+1})$ 内约有 $2^M$ 个采样点，间隔为：

$$ \Delta_e = \frac{2^e}{2^M} = 2^{e-M} $$

因此量化绝对误差大约满足：

$$ |x - \hat x| \leq \frac{1}{2}\,2^{e-M} $$

而相对误差近似满足：

$$ \frac{|x - \hat x|}{|x|} \leq 2^{-(M+1)} $$

也就是说，虽然 FP8 的绝对误差也会随着数值变大而变大，但相对误差大致稳定。

如果是 E4M3，相对误差不会超过 $6.25\%$；而 E5M2 则不会超过 $12.5\%$。

若令 INT8 的相对量化误差等于 E4M3：

$$ \frac{\alpha}{254\,|x|} = \frac{1}{16} $$

可以得到：

$$ |x| \approx 0.063\,\alpha $$

也就是说，当待量化值与最大值的差距不超过约 16 倍时，INT8 的量化效果往往更好；一旦超出这个范围（即数值相对 $\alpha$ 过小），FP8 则是更优的选择。

而我们知道，权重和激活值虽然都存在 outlier，但二者的数值范围并不相同：权重的 outlier 一般是普通 token 的数十倍，而激活值的 outlier 则不一定，甚至可以达到数百倍。

因此在实际部署时，一个理想的选择是权重用 INT 量化、激活值用 FP8 量化，也就是常说的 W4AFP8（在 TensorRT-LLM 中记作 W4A8）。目前的推理框架普遍推荐这种量化方式，例如 TensorRT-LLM 中：

QUANTIZATION METHODS	PERFORMANCE IMPROVEMENT (BATCH SIZE <= 4)	PERFORMANCE IMPROVEMENT (BATCH SIZE >= 16)	ACCURACY IMPACT	CALIBRATION TIME**
FP8 (W8A8)	Medium	Medium	Very Low	Minutes
Int8 SQ (W8A8)	Medium	Medium	Medium	Minutes
Int8 weight-only (W8A16)	Medium	Low	Low	Not Required
Int4 weight-only (W4A16)	High	Low	High	Not Required
Int4 AWQ (W4A16)	High	Low	Low	Tens of Minutes
Int4 GPTQ	High	Low	Low	Tens of Minutes
Int4-FP8 AWQ (W4A8)	High	Medium	Low	Tens of Minutes

可以看到，W4AFP8 在性能提升最大的同时，还能做到较低的精度影响。

接下来从性能角度比较 FP8 与 INT8。

在正式比较之前，需要先明确一点：低精度 Tensor Core 的核心，并不是把整个矩阵乘法都维持在 8 bit 精度下完成，而是采用 mixed-precision MMA 的方式。也就是说，输入矩阵以 FP8 或 INT8 这样的低精度格式参与乘法，但乘积通常会在更高精度的累加器中累加，最后再根据需要转换为目标输出精度。

这样设计的原因在于，GEMM 中每个输出元素本质上都是一个长度为 K 的点积，需要经历大量乘加操作。如果累加阶段仍然使用过低的精度，那么许多较小的增量会因为舍入而被直接吞掉，累加误差也会随着 K 的增大迅速放大。因此，低精度矩阵乘法的“低精度”主要体现在输入和乘法阶段，而不是简单地意味着整个计算链路都只有 8 bit。

从公开硬件规格来看，在支持 FP8 的 NVIDIA GPU 上，FP8 Tensor Core 和 INT8 Tensor Core 往往具有相同的标称峰值吞吐。这说明在产品级实现中，厂商可以通过专门的数据通路设计、流水线化、tile 级并行以及 mixed-precision 累加机制，使 FP8 和 INT8 在峰值计算能力上达到相同水平。

但这并不意味着 FP8 和 INT8 的底层硬件代价完全相同。INT8 的计算本质上是整数/定点乘加，数据格式规则，乘法和累加逻辑相对简单；而 FP8 虽然同样只有 8 bit，但它包含符号位、指数位和尾数位，计算时需要处理指数对齐、尾数运算、规格化、舍入以及高精度浮点累加等逻辑。因此，从单个 MAC 单元或数据通路的实现复杂度来看，FP8 通常比 INT8 更复杂。

所以更准确的理解是：理论上，FP8 运算单元的硬件实现成本高于 INT8；但在 Hopper、Blackwell 这类专门面向 AI 负载设计的 Tensor Core 中，NVIDIA 通过专门的硬件设计把 FP8 和 INT8 的峰值吞吐做到了相同水平。也就是说，FP8 的优势并不是“硬件更简单”，而是在相近吞吐下提供了更大的动态范围，从而更适合大模型中激活值波动大、outlier 明显的场景。（简单来说：放心大胆地用 FP8 吧，NVIDIA 帮你兜底了。）

FP4 / MXFP4 / NVFP4

在 4-bit 下，想要像 FP8 那样直接使用 FP4 是不现实的，因为 FP4 能表示的数值范围实在太小了。在 OCP 标准中，FP4 一般取 E2M1，它能表示的值大概是：

$$ 0,\ \pm 0.5,\ \pm 1,\ \pm 1.5,\ \pm 2,\ \pm 3,\ \pm 4,\ \pm 6 $$

也就是说，单单使用 FP4，能够表示的数值范围相当有限。因此实际应用时，通常不会只存储一个 FP4 值，还会额外存储一个 scale：

$$ \text{real num} = \text{Encoding of FP4} \times \text{scale} $$

scale 的不同对应着目前两种主流的 FP4 格式（Blackwell 确实支持 FP4 这种数值格式，但没有针对 scale 的专用硬件加速）。

我们先来看 MXFP4，它的 scale 格式是每 32 个值共享一个 E8M0 scale，也就是 2 的幂形式的 scale。

接下来按照与 FP8 相同的思路分析它的量化误差（取 block 内最大绝对值为 $A$）：

$$ \hat x_i = s_{mx} \cdot Q_{\text{E2M1}}\left(\frac{x_i}{s_{mx}}\right) $$

一般而言理想的 scale 是 $\frac{A}{6}$；然而在 MXFP4 中 scale 必须是 E8M0 格式，因此需要将其量化为最近的 2 的幂：

$$ s_{mx} = 2^{\text{round}\left(\log_2 \frac{A}{6}\right)} $$

在 TensorRT 中一般会向上取整。为了计算方便，我们定义：

$$ \rho_{mx} = \frac{s_{mx}}{A/6} $$

由于 E8M0 是 2 的幂，向上取整时有：

$$ 1 \leq \rho_{mx} < 2 $$

由于任意量化器的量化误差不会超过一个量化间隔，因此我们有：

$$ |e_i| \leq s_{mx} = \rho_{mx}\,\frac{A}{6} $$

而 $\rho_{mx}$ 与 block 内最大值和 2 的幂的对齐程度相关：当 $A$ 恰好为 2 的幂时 $\rho_{mx}$ 接近 1，误差较小；当 $A$ 略大于某个 2 的幂时 $\rho_{mx}$ 接近 2，误差较大。

例如 $A/6 = 1.99$ 与 $A/6 = 2.01$：前者会被向上取整为 2，此时 $\rho_{mx}$ 接近 1；后者会被向上取整为 4，此时 $\rho_{mx}$ 约为 2。

若对齐得好，则有：

$$ |e_i| \leq \frac{A}{6} $$

否则为：

$$ |e_i| < \frac{A}{3} $$

除了绝对误差外，我们再看看它的归零阈值（即低于该值会被量化为 0）：

$$ |x| < 0.25\,s_{mx} = \rho_{mx}\,\frac{A}{24} $$

于是有：

$$ \frac{A}{24} \leq \text{zero threshold} \leq \frac{A}{12} $$

大约是 $4.17\% A$ 到 $8.33\% A$。

接下来看看 NVFP4，它与 MXFP4 的区别在于：

block 更小：每 16 个元素共享一个 scale；
scale 更细：NVFP4 的 block scale 采用 FP8 E4M3，而不是 E8M0。

此外，NVFP4 还会额外使用一个全局的（per-tensor）FP32 scale 来避免 overflow。一般而言：

$$ s_{\text{global}} = \frac{\text{global}_{amax}}{6 \times 448}\\[4pt] s_{\text{block}} = \frac{\text{block}_{amax}/6}{s_{\text{global}}} $$

其中 6 是 E2M1 的最大可表示值，448 是 E4M3 的最大可表示值，二者相乘保证 block scale 落在 E4M3 的表示范围内。

同样地，我们可以把 NVFP4 写作：

$$ \hat x_i = s_{\text{global}} \cdot s_{\text{block}} \cdot Q_{\text{E2M1}}\left(\frac{x_i}{s_{\text{global}} \cdot s_{\text{block}}}\right) $$

令

$$ s_{nv} = s_{\text{global}} \cdot s_{\text{block}} $$

那么同样地有：

$$ \rho_{nv} = \frac{s_{nv}}{A/6} $$

又因为此时 scale 为 E4M3，其相邻正规格化数的比例最多约为：

$$ 1 + \frac{1}{8} = \frac{9}{8} $$

因此取向上取整时：

$$ 1 \leq \rho_{nv} \leq \frac{9}{8} $$

于是 NVFP4 的绝对误差上界约为：

$$ |e_i| \leq \rho_{nv}\,\frac{A}{6} \leq \frac{3}{16}A $$

约 $18.75\% A$，这明显好于 MXFP4 在未对齐时的最坏情况。

而 NVFP4 的归零阈值为：

$$ \text{zero threshold} \leq 4.69\% A $$

至于 INT4，可以通过与 INT8 类似的方式分析：对称量化下 $q \in [-7, 7]$，scale 为 $s = A/7$，因此绝对误差上界为 $|e_i| \leq s/2 = A/14$，归零阈值同样为 $A/14 \approx 7.14\% A$。

需要注意的是，INT4 的绝对误差上界（约 $7.14\% A$）看似比 NVFP4（约 $18.75\% A$）更小，但这是均匀分布在整个量程上的误差；而对于神经网络中占绝大多数的小数值，FP4 的误差与数值大小成比例，这些小值上的实际误差远小于 INT4 的均匀误差。再结合 FP4 对长尾异常值更友好的动态范围，可以得出结论：在 4-bit 量化中，面对神经网络这种“大量小值 + 长尾异常值”的分布，NVFP4 的综合表现最好。

而从性能上分析，虽然 MXFP4 由于采用 32 元素 block 和 E8M0 power-of-two scale，理论上具有更低的 scale 存储与缩放复杂度，但在真实的 LLM/VLM 推理中，性能并不只由单个 GEMM 的格式开销决定。NVFP4 通过 16 元素 micro-block 和 E4M3 scale 显著降低了量化误差，使更多层能够稳定进入 FP4 路径，减少混合精度 fallback 和额外补偿。因此，在当前 Blackwell + TensorRT-LLM 生态下，NVFP4 往往是更实际的高性能 FP4 选择；而 MXFP4 更适合在对精度不敏感、或有专门误差补偿方法的场景中追求更高的硬件效率。

此外，目前主流推理框架都对 NVFP4 做了专门优化，是性能—精度权衡下的最佳选择。

Reference

Triton

Tue, 14 Apr 2026 21:00:00 +0800

Motivation

CS336课程中在介绍实现FlashAttention2的时候介绍了Triton，我发现这是一个比CUDA更加方便的GPU编程语言。

另一方面许多大厂高性能计算岗位包括但不限于算子开发都逐渐趋向于使用Triton，因此目前看来学习Triton是一件从投资未来的角度来看性价比很高的事情

Introduction

Triton 是 OpenAI 开发的一种面向深度学习与高性能计算的编程语言及编译器，旨在简化 GPU 高性能算子的开发过程，并帮助开发者更容易地写出高效代码。

Triton的官方文档指出,其项目动机来自于基于分块算法的编程范式可以促进构建用于神经网络的高性能计算内核。与传统 CUDA 编程更强调线程级的执行方式不同，Triton 提供了更高层次的块级抽象，开发者可以直接围绕数据块来组织计算，而不必过多关注底层线程调度细节。

我们可以通过矩阵乘法的例子来更加直观的理解。

CUDA编程模型:

#pragma parallel
for(int m = 0; m < M; m++)
#pragma parallel
for(int n = 0; n < N; n++){
 float acc = 0;
 for(int k = 0; k < K; k++)
 acc += A[m, k] * B[k, n];

 C[m, n] = acc;
}

Triton编程模型:

#pragma parallel
for(int m = 0; m < M; m += MB)
#pragma parallel
for(int n = 0; n < N; n += NB){
 float acc[MB, NB] = 0;
 for(int k = 0; k < K; k += KB)
 acc += A[m:m+MB, k:k+KB]
 @ B[k:k+KB, n:n+NB];
 C[m:m+MB, n:n+NB] = acc;
}

GPU

在具体介绍Triton之前，我们先简单了解一下GPU相关的知识，以便减少后续专业术语带来的知识bias。

以下内容来自于本人于南开大学计算机体系结构期末突击的笔记。

相较于CPU而言，GPU的设计思路可以概括为三点:

简化流水线，增加核数：具体而言，对流水线进行瘦身，去掉乱序执行，分支预测等复杂逻辑，节省的空间用来增加大量的计算核心，从而能够同时处理大量数据
多个计算单元共用一条指令：GPU 擅长处理“很多数据做同一种计算”的场景，例如对向量中的每个元素做相同变换。为此，它通常让一组线程(这样的一组线程称为Warp)按照相同的程序逻辑并行执行，只是每个线程处理的数据不同。这种执行方式通常称为 SIMT。
驻留大量线程：GPU 会同时准备很多线程。当一部分线程因为访存而需要等待时，硬件可以立即切换去执行另一部分线程，从而减少计算单元空闲的时间。

我们接下来以数组加法来介绍GPU编程模型中的几个基本概念，这些概念有助于理解后续的并行执行，分块计算和访存优化。

假设我们要计算数组加法:

for i in range(1024):
	C[i] = A[i] + B[i]

那么一个GPU Kernel会以大量线程的形式启动。所有线程执行相同的Kernel代码，但每个线程会根据自己的编号处理不同位置的数据。在本例中，每个线程可以分别负责不同的i

Thread:即线程，这是程序员视角下最基本的并行单位。每个线程执行同一份kernel代码，但处理的数据不同
Block:多个Thread会组成一个Block。其意义在于协作，即同一个block内的线程可以通过共享内存交换数据，也可以使用同步原语协调执行不同 block 之间通常不能直接同步或共享局部数据。
Grid:一次Kernel启动产生的所有线程构成一个grid。grid只是一次kernel启动时的逻辑组织方式，本身并不对应某个具体的硬件结构
Warp:在硬件执行层面，线程通常不会以单个 thread 为单位独立调度，而是会进一步组成 warp。以 NVIDIA GPU 为例，一个 warp 通常包含 32 个线程。warp 内的线程会执行相同的指令流程，但处理各自的数据。这意味着，虽然我们在编程时使用的是 thread 和 block 这些抽象概念，但在硬件上，block 内的线程最终仍会被拆分为多个 warp 来执行。
SM:流处理器，GPU上的block会被分配到SM上执行。一个SM可以同时驻留多个warp，甚至是多个block。当某个warp因访存等高延迟操作暂时停顿时，硬件调度器可以快速切换到其他就绪的 warp 继续执行，从而隐藏延迟、提升整体吞吐。

Weighted Sum

我们接下来通过“加权求和”的例子来具体了解Triton的知识以及它是如何与Pytorch进行协同工作。

Forward Pass

“加权求和”的前向传播过程可以通过如下形式进行描述:

给定一个输入矩阵x,我们将它的每个元素乘以一个按列加权的向量w,然后对每一行求和，最终得到矩阵x与向量w的加权和。

用python代码描述则是:

def weighted_sum(x,weight):
	return (weight*x).sum(axis=-1)

而我们想把这个过程用Triton进行并行化，那么首先需要分析可并行性，我们可以发现矩阵与向量的乘积在最后求和得到结果之前各个元素的运算都是相互独立的，因此每个元素的计算都是可并行的。

因此在编写Triton Kernel时，我们会让每个program instance(Triton中一组执行同一段程序的线程块)负责x中某一块行的加权和，并把对应的标量输出写入输出张量。

例如，假设x的形状是[ROWS,D],weight的形状是[D]，我们会让每个program instance负责x中若干行的加权和的计算，得到每个块行的加权和结果，最后将结果均写入输出张量。

与直接把tensor作为参数不同，我们会在kernel中传入:

指向张量首元素的指针
每个tensor的stride，其含义是告诉我们如何沿着各个维度移动

我们可以利用这些 stride，结合 program ID，将工作分配给不同实例（例如，第 i 个实例处理 x 的第 i 个行块），从而加载当前实例所负责的 x 的一块数据。

在本例中，Triton的Forward Pass和PyTorch的主要区别在于:Triton需要显式地进行指针运算以及load/store操作。

我们来看具体的代码:

import triton
import triton.language as tl
@triton.jit 
def weighted_sum_fwd(
	x_ptr,weight_ptr,			#输入指针
	output_ptr,					#输出指针
	x_stride_row,x_stride_dim,	#stride告诉我们在tensor的每个轴上移动一个元素需要跳多远
	weight_stride_dim,			#通常是1
	output_stride_dim,			#通常是1
	ROWS,D,
	ROWS_TILE_SIZE,D_TILE_SIZE,	#tile形状必须在编译期已知
):
	#每个instance负责计算x中一个行块的加权和
	#`tl.program_id`用于查看当前运行的是哪个线程块
	row_tile_idx = tl.program_id(0)
	
	#block pointer 允许我们从一个N维内存区域中选取数据块
	#并在这个区域中移动所选的数据块
	#在使用block pointer时必须知道:
	#- tensor 第一个元素的指针
	#- tensor 的整体形状
	#- 每个维度的stride
	#- 起始block的N维坐标
	#- 每次load/store的block形状
	#- 内存中维度从主到次的顺序
	x_block_ptr = tl.make_block_ptr(
		x_ptr,
		shape=(ROWS,D),
		strides=(x_stride_row,x_stride_dim),
		offsets=(row_tile_idx * ROW_TILE_SIZE,0),
		block_shape=(ROW_TILE_SIZE,D_TILE_SIZE),
		order(1,0),
	)
	weight_block_ptr=tl.make_block_ptr(
		weight_ptr,
		shape=(D,),
		strides=(weight_stride_dim,),
		offsets=(0,),
		block_shape=(D_TILE_SIZE,),
		order(0,),
	)
	output_block_ptr=tl.make_block_ptr(
		output_ptr,
		shape=(ROWS,),
		strides=(output_stride_row,),
		offsets=(row_tile_idx*ROWS_TILE_SIZE,),
		block_shape=(ROWS_TILE_SIZE,),
		order(0,),
	)
	#初始化一个用于写入的buffer
	output = tl.zeros((ROWS_TILE_SIZE,),dtype=tl.float32)
	for i in range(tl.cdiv(D,D_TILE,SIZE)):
		#加载当前block pointer指向的数据块
		#因为ROW_TILE_SIZE可能不整除ROWS，
		#D_TILE_SIZE 也可能不能整除 D，
		#所以两个维度都需要进行边界检查
		row = tl.load(x_block_ptr,boundary_check=(0,1),padding_option="zero")
		weight = tl.load(weight_block_ptr,boundary_check=(0,),padding_option="zero")
		#计算当前行块的加权和
		output += tl.sum(row * weight[None,:],axis=1)
		#将指针移动到下一个tile,这里参数时(行偏移，列偏移)
		x_block_ptr = x_block_ptr.advance(0,D_TILE_SIZE)	#在最后一个维度上前进D_TILE_SIZE
		weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))
	# 将输出写入 output block pointer（每行一个标量）
 # 因为 ROWS_TILE_SIZE 可能不能整除 ROWS，所以需要边界检查
	tl.store(output_block_ptr,output,boundary_check(0,))

接下来我们将这个实现的Kernel封装到一个Pytorch AutoGrad函数中，使其能够与PyTorch生态协同工作(即接收 Tensor 输入、输出 Tensor，并在 backward pass 中和 autograd 引擎协作):

class WeightedSumFunc(torch.autograd.Function):
	@staticmethod
	def forward(ctx,x,weight):#这里ctx指上下文
		#将x和weight缓存起来，以便在backward中使用
		#backward时我们只会接收到输出张量的梯度
		# 需要利用这些缓存来计算 x 和 weight 的梯度
		D,output_dims = x.shape[-1],x.shape[:-1]
		input_shape = x.shape
		x = rearrange(x,"... d -> (...) d")
		ctx.save_for_backward(x,weight)
		
		assert len(weight.shape) == 1 and weight.shape[0] == D, "维度不匹配"
 assert x.is_cuda and weight.is_cuda, "期望输入是 CUDA tensor"
 assert x.is_contiguous(), "我们的指针运算默认 x 是连续的"
 
 ctx.D_TILE_SIZE = triton.next_power_of_2(D) # 16 #大致让dim维度循环16次
 ctx.ROWS_TILE_SIZE = 16						# 每个线程一次处理16个batch元素
 ctx.input_shape = input_shape
 # 需要初始化一个空的结果 tensor
 # 注意：这里的元素未必初始化为 0
 y = torch.empty(output_dims,device = x.device)
 
 # 以 1D launch grid 启动 kernel
 n_rows = y.numel()
 weighted_sum_fwd[(cdiv(n_rows,ctx.ROWS_TILE_SIZE),)](	#和CUDA的<<<>>>有点像，作用是定义grid大小
 	x,weight,
 	y,
 	x.stride(0),x.stride(1),
 	weight.stride(0),
 	y.stride(0),
 	ROWS = n_rows,D=D,
 	ROWS_TILE_SIZE=ctx.ROWS_TILE_SIZE, D_TILE_SIZE=ctx.D_TILE_SIZE,
 )
 return y.view(input_shape[:-1])

Backward Pass

由于我们定义了自己的kernel，因此也需要自己编写backward函数。

在forward pass中，我们拿到了层的输入，并计算其输出。而在backward pass中我们会拿到目标函数对该层输出的梯度，然后需要计算目标函数对每个输入的梯度。

在这个例子中，我们的操作有两个输入:

矩阵$x \in \mathbb{R}^{n\times h}$
权重向量$w \in \mathbb{R}^h$

记我们的操作为$f(x,w)$，其输出属于$\mathbb{R}^n$

若给定损失函数L对该层输出的梯度$\nabla_{f(x,w)}L$,则根据多元链式法则，可以得到关于x和w的梯度:

$$ (\nabla_xL)_{ij}=\sum_{k=1}^n\frac{\partial f(x,w)_k}{\partial x_{ij}}(\nabla_{f(x,w)}L)_k = w_j \cdot (\nabla_{f(x,w)}L)_i\\ (\nabla_w L)_{j} = \sum_{i=1}^n\frac{\partial f(x,w)_i}{\partial w_j}(\nabla_{f(x,w)}L)_i=\sum_{i=1}^nx_{ij}\cdot(\nabla_{f(x,w)}L)_i $$

这里给出一个很简单的backward计算公式。

为了得到关于x的backward结果，我们根据上式计算w和$\nabla f(x,w)$的外积。而为了得到关于w的backward结果(即$(\nabla_w L)_j$),我们必须把输入与对应输出的梯度逐行相乘并求和。

按照同样的思路，我们来实现这个kernel:

@triton.jit
def weighted_sum_backward(
	x_ptr,weight_ptr,
	grad_output_ptr,
	grad_x_ptr,partial_grad_weight_ptr,
	stride_xr,stride_xd,
	stride_wd,
	stride_gr,
	stride_gxr,stride_gxd,
	stride_gwb, stride_gwd,
 NUM_ROWS, D,
 ROWS_TILE_SIZE, D_TILE_SIZE,
):
	row_tile_idx = tl.program_id(0)
	n_row_tiles = tl.num_programs(0)
	#输入
	grad_output_block_ptr = tl.make_block_ptr(
		grad_output_ptr,
		shape=(NUM_ROWS,),strides=(stride_gr,),
		offsets=(row_tile_idx*ROW_TILE_SIZE,),
		block_shape=(ROWS_TILE_SIZE,),
		order=(0,),
	)
	x_block_ptr = tl.make_block_ptr(
		x_ptr,
		shape=(NUM_ROWS,D,),strides=(stride_xr,stride_xd),
		offsets=(row_tile_idx * ROW_TILE_SIZE,0),
		block_shape=(ROWS_TILE_SIZE,D_TILE_SIZE),
		order=(1,0),
	)
	weight_block_ptr = tl.make_block_ptr(
		weight_ptr,
		shape=(D,),strides=(stride_wd,),
		offsets=(0,),block_shape=(D_TILE_SIZE,),
		order=(0,),
	)
	grad_x_block_ptr = tl.make_block_ptr(
		grad_x_ptr,
		shape=(NUM_ROWS,D,),strides=(stride_gxr,stride_gxd),
		offsets=(row_tile_idx*ROW_TILE_SIZE,D_TILE_SIZE),
		block_shape=(ROW_TILE_SIZE,D_TILE_SIZE),
		order=(1,0),
	)
	partial_grad_weight_block_ptr = tl.make_block_ptr(
		partial_grad_weight_ptr,
		shape=(n_row_tiles,D,),strides=(stride_gwb,stride_gwd),
		offsets=(row_tile_idx,0),
		block_shape=(1,D_TILE_SIZE),
		order=(1,0),
	)
	for i in range(tl.cdiv(D,D_TILE_SIZE)):
		grad_output = tl.load(
			grad_output_block_ptr,
			boundary_check=(0,),
			padding_option="zero",
		)
		weight = tl.load(
			weight_block_ptr,
			boundary_check=(0,),
			padding_option="zero",
		)
		grad_x_row = grad_output[:,None] * weight[None,:]
		tl.store(grad_x_block_ptr,grad_x_row,boundary_chekc=(0,1))
		
		row = tl.load(
			x_block_ptr,
			boundary_check=(0,1),
			padding_option="zero",
		)
		grad_weight_row = tl.sum(row * grad_output[:,None],axis=0,keep_dims=True)
		tl.store(partial_grad_weight_block_ptr,grad_weight_row,boundary_check=(1,))
		
		x_block_ptr = x_block_ptr.advance((0,D_TILE_SIZE))
		weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))
		partial_grad_weight_block_ptr = parital_grad_weight_block_ptr.advance((0,D_TILE_SIZE))
		grad_x_block_ptr = grad_x_block_ptr.advance((0,D_TILE_SIZE))

每个 kernel instance 只负责处理 x 的一个行块，但我们现在需要沿着 x 的行方向求和。因此，在 backward kernel 内，我们不直接完成这个总和，而是假设 partial_grad_weight_ptr 是一个形状为 n_row_tiles × H 的矩阵，其中第一维表示来自每个行块的局部归约结果。

也就是说：

kernel 内部：只在当前行块内完成局部归约
kernel 外部：再用 torch.sum 把所有行块的结果加起来，得到最终的 $\nabla_w$

这样一来，autograd.Function 的最后一部分就相对简单了：

class WeightSumFunc(torch.autograd.Function):
	@staticmethod
	def forward(ctx,x,weight):
		#这里省略，前面定义过了
	@staticmethod
	def backward(ctx,grad_out):
		x,weight = ctx.saved_tensors
		ROWS_TILE_SIZE,D_TILE_SIZE = ctx.ROWS_TILE_SIZE,ctx.D_TILE_SIZE
		n_rows,D = x.shape()
		
		#我们的策略是：每个线程块先写入一个partial buffer，然后对这个buffer做归约得到最终梯度
		parital_grad_weight = torch.empty(
			(cdiv(n_rows,ROWS_TILE_SIZE),D),
			device=x.device,
			dtype=x.dtype
		)
		grad_x = torch.empty_like(x)
		weighted_sum_backward[(cdiv(n_rows,ROWS_TILE_SIZE),)](
			x,weight,
			grad_out,
			grad_x,partial_grad_weight,
			x.stride(0),x.stride(1),
			weight.stride(0),
			grad_out.stride(0),
			grad_x.stride(0), grad_x.stride(1),
 partial_grad_weight.stride(0), partial_grad_weight.stride(1),
 NUM_ROWS=n_rows, D=D,
 ROWS_TILE_SIZE=ROWS_TILE_SIZE, D_TILE_SIZE=D_TILE_SIZE,
		)
		grad_weight = partial_grad_weight.sum(axis=0)
		return grad_x,grad_weight

最后我们就可以得到一个函数，它的使用方法和torch.nn.functional里的函数很相似:

f_weightedsum = WeightedSumFunc.apply

我们可以对两个PyTorch tensor x和w调用这个函数，就会得到类似如下的结果:

tensor([ 90.8563, -93.6815, -80.8884, ..., 103.4840, -21.4634, -24.0192],
 device='cuda:0', grad_fn=<WeightedSumFuncBackward>)

请注意输出 tensor 上附带的 grad_fn —— 这表明 PyTorch 已经知道：当这个 tensor 出现在计算图中并需要做 backward pass 时，应该调用什么函数。

至此，我们就完成了这个 weighted sum 操作的 Triton 实现。

对于其它算子，我们可以按照同样的思路进行。实际上重点还是在于我们应该分析出哪些是可以并行的，program instance该如何划分。

简单来说，在设计Triton算子的时候我们需要考虑三个问题:

输出的自然分块是什么？例如在Weighted Sum中Y[0],Y[1]相互独立按行产生，那么Program Instance可以负责生成一个ROW_TILE_SIZE大小的行块的结果。
一个块的输出需要哪些输入？
这些输入能不能以连续，重复利用的方式被加载？

这三个问题虽然简单，但是在Triton编程中很重要，在实际问题中也很难立即给出标准答案，需要我们不断试错积累经验。

建议参考仓库(Assignment2-systems/Triton/WeightedSum.py)中的代码，不要直接复制粘贴，不然会出现一些奇奇怪怪的错误，这里的代码是参考了CS336 Assignment 2中的代码.

Triton Puzzles (Lite)

为了更深一步的学习Triton,我在网上查找到了一份通过做题来学习Triton的repo。

原始的repo是srush/Triton-Puzzles: Puzzles for learning Triton

但是有一个更好的实现版本是:SiriusNEO/Triton-Puzzles-Lite: Puzzles for learning Triton, play it with minimal environment configuration!

因此我们还是以这个仓库为题面来进行解答。

另外需要强调的一点是，正如Assignment2实验手册中提到的那样，目前大多教程没有使用更新的，更方便的block pointer抽象。因此在此基础上，在本blog中，我将以block pointer抽象来完成这些任务。

之所以说block pointer更加方便，是因为在Triton中，我们的Program Instance的操作对象是划分后的Tile，直观来看就是一个矩形区域。在使用block pointer之前，我们需要得到这个处理对象需要通过Program ID以及SIZE用公式去计算。例如我们想知道处理一个x[ROWS,D]的tile，我们需要写:

row_offsets = row_tile_idx * ROWS_TILE_SIZE + tl.arrange(0,ROWS_TILE_SIZE)
col_offsets = tl.arrange(0,D_TILE_SIZE)
#计算出当前块负责哪些行和哪些列
x_ptrs = x_ptr + row_offsets[:,None] * x_stride_row + col_offsets[None,:] * x_stride_dim
#然后把他们拼接为二维地址
mask = (row_offsets[:, None] < ROWS) & (col_offsets[None, :] < D)
#然后再load
x = tl.load(x_ptrs, mask=mask, other=0.0)

而使用block pointer的话，我们就可以写:

x_block_ptr = tl.make_block_ptr(...)
x = tl.load(x_block,boundary_check=(0,1),padding="zero")

这样做就省去了大量的地址计算，一方面更加便捷，另一方面代码的可读性也大大提升了。接下来我们看看具体的问题以及如何解决吧？

Puzzles 1/2 Constant Add

问题描述:

对向量中的每个元素都加上常数10.使用一个程序ID轴。

问题解答：

我们按照上面Weighted Sum中make_block_ptr的方式定义好x_block_ptr以及z_block_ptr，然后tl.load出x，进行运算得到z最后进行tl.store即可。

由于N0不一定能被B0整除，因此需要进行boundary check

@triton.jit
def add_kernel(x_ptr, z_ptr, N0, B0: tl.constexpr):
 row_tile_idx = tl.program_id(0)
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(row_tile_idx * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(row_tile_idx * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 x = tl.load(x_block_ptr,boundary_check=(0,),padding_option='zero')
 z = x + 10.0
 tl.store(z_block_ptr, z, boundary_check=(0,))
 # Finish me!
 return

Puzzles 3/4 Outer Vector Add

问题描述:

把两个向量按广播方式相加，得到一个二维结果矩阵z。

$$ z_{j,i} = x_i+y_j \quad i=1,\dots,B_0,j=1,\dots,B_1 $$

问题解答:

这个可以利用Python中的broadcast机制:x[None,:]+y[:,None]。在计算时会先将x的形状变为(1,N0),y变为(N1,1)然后相加的时候触发广播机制，两者的 shape 均广播到 (N1, N0)。

在Triton中也有同样的机制，我们可以利用这个得到最终的结果。

考虑到N0,N1,B0,B1的差异，我们的grid相应地也要变成二维的，第一个维度是走x的tile，第二个维度走y的tile

@triton.jit
def add_vec_block_kernel(
 x_ptr, y_ptr, z_ptr, N0, N1, B0: tl.constexpr, B1: tl.constexpr
):
 x_tile_idx = tl.program_id(0)
 y_tile_idx = tl.program_id(1)

 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(x_tile_idx * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 y_block_ptr = tl.make_block_ptr(
 base=y_ptr,
 shape=(N1,),
 strides=(1,),
 offsets=(y_tile_idx * B1,),
 block_shape=(B1,),
 order=(0,)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N1, N0),
 strides=(N0, 1),
 offsets=(y_tile_idx * B1, x_tile_idx * B0),
 block_shape=(B1, B0),
 order=(0, 1)
 )
 x = tl.load(x_block_ptr,boundary_check=(0,),padding_option='zero')
 y = tl.load(y_block_ptr,boundary_check=(0,),padding_option='zero')
 z = x[None, :] + y[:, None]
 tl.store(z_block_ptr, z, boundary_check=(0,1))
 # Finish me!
 return

Puzzles 5 Fused Outer Multiplication

问题描述:

将行向量x与列向量y做外积，然后对结果矩阵的每个元素应用ReLU.

$$ z_{j,i} = \max(0,x_i\cdot y_j),\quad i = 1,\dots,N_0,j=1,\dots,N_1 $$

问题解答：

这一问实际上是与Puzzle 3/4类似的，只有符号从加法变成了乘法。对于ReLU操作直接将z与0取max就行:z = tl.maximum(0,z)

Puzzles 6 Fused Outer Multiplication - Backwards

问题描述：

对如下函数操作进行反向传播:矩阵x与向量y按行相乘，再经过ReLU。

$$ f(x,y) = \text{relu}(x_{j,i}\times y_j),\quad i = 1 \dots N_0,j = 1 \dots N_1\\ dx_{j,i} = f'_x(x,y)_{j,i}\times dz_{j,i} = 1(x_{j,i}\cdot y_j >0)\cdot y_j \cdot dz_{j,i} $$

其中$1(\cdots)$表示指示函数，括号内条件满足时为1，否则为0

问题解答：

按照公式实现即可，没有较为复杂的逻辑，因此不贴代码了。

Puzzles 7 Long Sum

问题描述:

给定一个二维张量x，对其每一行进行求和。

问题解答:

这是一个二维的张量，我们在分块计算的时候虽然直观的想法是每行进行并行计算，但是考虑到每行有T个元素，而我们设计的块的大小是$B_0\times B_1$，而B1远小于B0,因此我们需要沿着第二维分块累加，这也是为什么题目提醒我们需要使用for循环。

当我们使用block pointer时，我们可以直接用tl.advance(pointer,(stride_dim1,...))的方式移动指针。

@triton.jit
def sum_kernel(x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr):
 # Finish me!
 tile_id = tl.program_id(0)
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0, T),
 strides=(T, 1),
 offsets=(tile_id * B0, 0),
 block_shape=(B0, B1),
 order=(1, 0)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(tile_id * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 z = tl.zeros((B0,), dtype=tl.float32)
 for i in range(0, T, B1):
 x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option='zero')
 z += tl.sum(x,1)
 x_block_ptr = tl.advance(x_block_ptr,(0,B1))
 tl.store(z_block_ptr, z, boundary_check=(0,))
 return

Puzzles 8 Long Softmax

问题描述:

对一批logits做Softmax。需要保证数值稳定，即减去最大值后求Softmax

$$ z_{i,j} = \text{softmax}(x_{i,1},\dots,x_{i,T}) \quad i = 1,\dots,N_0 $$

另外需要注意在Triton中建议不要直接用exp而是用exp2.$\text{exp}(x) = 2^{\log_2(e)x}$

问题解答:

我们先来做没有优化的三个for的版本。实际上我们可以将任务做一个拆分。因为我们是分行块处理的，但是tile的大小[B0,B1]其中$B1 < T$，因此我们需要for循环来得到:

每行的max
每行的sum
每行的结果

因此需要3个for循环。因为后续我们将通过Online Softmax算法来实现更加高效实用的版本，因此这里就不贴出具体的代码了，思路和Puzzle 7类似。

Online Softmax的思路其实比较简单，就是第一个for和第二个for可以合并!

考虑两个tile，第一个tile的max我们记为$m_1$,那么此时局部求和的结果为:

$$ sum = \sum_{i = 0}^{B_0-1} \exp(x_i - m_1) $$

第二个tile的max记为$m_2(m_2 > m_1)$,我们记$a = m_1 - m_2$,那么此时我们可以更新局部求和结果为:

$$ sum' = sum \times \exp(a) + \sum_{B_0}^{2B_0-1}\exp(x_i - m_2) $$

由此我们可以通过一个for循环，便得到了求最终结果需要的max以及sum。

@triton.jit
def softmax_kernel(
 x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr
):
 """3 loops ver."""
 block_id_i = tl.program_id(0)
 log2_e = 1.44269504
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0, T),
 strides=(T, 1),
 offsets=(block_id_i * B0, 0),
 block_shape=(B0, B1),
 order=(1, 0)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0, T),
 strides=(T, 1),
 offsets=(block_id_i * B0, 0),
 block_shape=(B0, B1),
 order=(1, 0)
 )
 z = tl.zeros((B0, B1), dtype=tl.float32)
 x_max = tl.full((B0,), float("-inf"), dtype=tl.float32)
 sum = tl.zeros((B0,), dtype=tl.float32)
 x_block_ptr1 = x_block_ptr
 for i in range(0,T,B1):
 x = tl.load(x_block_ptr1, boundary_check=(0, 1), padding_option='zero')
 new_x_max = tl.maximum(x_max, tl.max(x,1))
 if i == 0:
 x_max = new_x_max
 sum = tl.sum(tl.exp2(log2_e * (x - x_max[:, None])),1)
 else:
 scale = tl.exp2(log2_e * (x_max - new_x_max))
 sum = sum * scale + tl.sum(tl.exp2(log2_e * (x - new_x_max[:, None])),1)
 x_max = new_x_max
 x_block_ptr1 = tl.advance(x_block_ptr1,(0,B1))
 for i in range(0,T,B1):
 x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option='zero')
 x = x - x_max[:, None]
 x_exp = tl.exp2(log2_e * x)
 z = x_exp / sum[:, None]
 tl.store(z_block_ptr, z, boundary_check=(0, 1))
 x_block_ptr = tl.advance(x_block_ptr,(0,B1))
 z_block_ptr = tl.advance(z_block_ptr,(0,B1))
 # Finish me!
 return

Puzzles 9 Simple FlashAttention

问题描述:

不使用 program。大小 B0 表示在总共 N0 个 q 中，本次要处理的一批。序列长度是 T。每次处理 B1 < T 个元素（k、v），其中 B1 是某个块大小。

$$ z_i = \sum_{j=1}^T \text{softmax}(q_ik_1,\dots,q_ik_T)_j v_j $$

这个问题可以通过类似于Online Softmax的思路来解决。

问题解答:(为便于理解，推荐读者用纸笔画一下)

考虑q的一个大小为[B0=2]的tile中元素为[q_0,q_1],k和v的一个大小为[B1=2]的tile中的元素为[k_1,k_2],[v_1,v_2]，那么对应的注意力矩阵元素为一个大小为[B0,B1]的二维矩阵，其中元素为[s_11,s_12,s_21,s_22]。

仅考虑结果中的第一个元素$z_1$,此时其累加和为:$z_1 = \text{softmax}(s_{11})v_1+\text{softmax}(s_{12})v_2 = \frac{\exp(s_{11}-m_1)}{l_1}v_1+\frac{\exp(s_12-m_1)}{l_1}v_2$,其中$m_1$为此时的局部最大值，$l_1$为此时的累计指数和.

计算完毕后，指针沿着axis=1的维度进行移动(步长为B1)，此时k和v的一个大小为[B1=2]的tile中的元素为[k_3,k_4],[v_3,v_4]，那么对应的注意力矩阵元素为一个大小为[B0,B1]的二维矩阵，其中元素为[s_13,s_14,s_23,s_24]。

同样考虑结果中的第一个元素$z_1$,此时最大值更新为$m_2$,指数累计和应更新为$l_2 = l_1\times(\text{scale}=\exp(m_1-m_2))+\sum_{j=3}^4 \exp(s_{1j}-m_2)$.

那么我们分两部分考虑此时的结果$z_1$,一部分是新的注意力权重$s_{13},s_{14}$与$v_3,v_4$的加权和，这部分直接和之前一样相加就行，而另一部分则为之前计算出来的贡献，这部分由于m和l的更新，也要对应更新。

更新可以从分子和分母考虑。分子上和Online Softmax类似，需要乘以$\exp(m_1-m_2)$,而分母则需要进行替换为$l_2$,因此最终我们可以更新局部加权和$z_1$为:

$$ z_1' = z_1\times scale\times \frac{l_1}{l_2} + \sum_{j=3}^4\frac{\exp(s_{1j}-m_2)}{l_2}v_j $$

由此不断迭代，便能得到最终结果。

@triton.jit
def flashatt_kernel(
 q_ptr, k_ptr, v_ptr, z_ptr, N0, T, B0: tl.constexpr, B1: tl.constexpr
):
 block_id_i = tl.program_id(0)
 log2_e = 1.44269504
 myexp = lambda x: tl.exp2(log2_e * x)
 q_block_ptr = tl.make_block_ptr(
 base=q_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(block_id_i * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 k_block_ptr = tl.make_block_ptr(
 base=k_ptr,
 shape=(T,),
 strides=(1,),
 offsets=(0,),
 block_shape=(B1,),
 order=(0,)
 )
 v_block_ptr = tl.make_block_ptr(
 base=v_ptr,
 shape=(T,),
 strides=(1,),
 offsets=(0,),
 block_shape=(B1,),
 order=(0,)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(block_id_i * B0,),
 block_shape=(B0,),
 order=(0,)
 )

 qk_max = tl.full((B0,), float("-inf"), dtype=tl.float32)
 qk_sum = tl.zeros((B0,), dtype=tl.float32)
 o = tl.zeros((B0,), dtype=tl.float32)
 for i in range(0, T, B1):
 q = tl.load(q_block_ptr, boundary_check=(0,), padding_option='zero')
 k = tl.load(k_block_ptr, boundary_check=(0,), padding_option='zero')
 v = tl.load(v_block_ptr, boundary_check=(0,), padding_option='zero')
 qk = q[:, None] * k[None, :]
 new_qk_max = tl.maximum(qk_max,tl.max(qk, 1))
 if i == 0:
 qk_max = new_qk_max
 qk_exp = myexp(qk - qk_max[:,None])
 qk_sum = tl.sum(qk_exp, 1)
 o = tl.sum(qk_exp * v[None, :] / qk_sum[:,None], 1)
 else:
 qk_exp = myexp(qk - new_qk_max[:,None])
 scale = myexp(qk_max - new_qk_max)
 scale1 = scale * qk_sum
 qk_sum = qk_sum * scale + tl.sum(qk_exp, 1)
 o = o * scale1 / qk_sum + tl.sum(qk_exp * v[None, :] / qk_sum[:,None], 1)
 qk_max = new_qk_max
 k_block_ptr = tl.advance(k_block_ptr,(B1,))
 v_block_ptr = tl.advance(v_block_ptr,(B1,))
 tl.store(z_block_ptr, o, boundary_check=(0,))
 # Finish me!
 return

Puzzles 10 Two Dimensional Convolution

问题描述:

实现一个带批处理的二维卷积。使用一个program id 轴。块大小为B0表示在N0中一次处理多少个batch。图像x的大小为H x W，并且只有一个通道；卷积核k的大小为kH x kW

$$ z_{i, j, l} = \sum_{oj, ol}^{j+oj\le H, l+ol\le W} k_{oj,ol} \times x_{i,j + oj, l + ol} \text{ for } i = 1\ldots N_0 \text{ for } j = 1\ldots H \text{ for } l = 1\ldots W $$

问题解决：(为了方便理解，最好对照题面的图来理解)

我们考虑如何分块，题目中要求的是按照批维度进行分块拆分并行。那么块的第一维就是B0,注意到卷积运算每次只会用到图像矩阵中KW*KH大小的子矩阵，因此块最合适的大小就是[B0,KH,KW]。这个块与卷积核进行卷积运算，得到一个标量，因此结果的块的大小就是[B0,1,1]，那么这样就能写出block pointer的定义。

接下来模拟二维卷积运算就行了，根据相关知识，我们知道卷积运算就是卷积核在图像矩阵上进行平移(这里默认步长为1).那么我们用两个for循环来模拟这个平移计算的流程就行，然后在内层for进行一次后沿着dim=1平移一格，在外层for进行一次后沿dim=0平移一格。

@triton.jit
def conv2d_kernel(
 x_ptr, k_ptr, z_ptr, N0, H, W, KH: tl.constexpr, KW: tl.constexpr, B0: tl.constexpr
):
 block_id_i = tl.program_id(0)
 k_block_ptr = tl.make_block_ptr(
 base=k_ptr,
 shape=(KH, KW),
 strides=(KW, 1),
 offsets=(0, 0),
 block_shape=(KH, KW),
 order=(1, 0),
 )
 k = tl.load(k_block_ptr, boundary_check=(0, 1), padding_option="zero")

 x_row_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0, H, W),
 strides=(H * W, W, 1),
 offsets=(block_id_i * B0, 0, 0),
 block_shape=(B0, KH, KW),
 order=(2, 1, 0),
 )
 z_row_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0, H, W),
 strides=(H * W, W, 1),
 offsets=(block_id_i * B0, 0, 0),
 block_shape=(B0, 1, 1),
 order=(2, 1, 0),
 )

 for _ in range(H):
 x_block_ptr = x_row_block_ptr
 z_block_ptr = z_row_block_ptr
 for _ in range(W):
 x = tl.load(x_block_ptr, boundary_check=(0, 1, 2), padding_option="zero")
 z = (x * k[None, :, :]).sum(1).sum(1)
 tl.store(z_block_ptr, z[:, None, None], boundary_check=(0, 1, 2))
 x_block_ptr = tl.advance(x_block_ptr, (0, 0, 1))
 z_block_ptr = tl.advance(z_block_ptr, (0, 0, 1))
 x_row_block_ptr = tl.advance(x_row_block_ptr, (0, 1, 0))
 z_row_block_ptr = tl.advance(z_row_block_ptr, (0, 1, 0))

 return

Puzzles 11 Matrix Multiplication

问题描述：

使用三条program id轴。块大小B2表示在N2中要处理的Batch数量。块大小B0表示在N0中要处理的行数，块大小B1表示在N1中要处理的列数。

中间维度的大小为MID

$$ z_{i, j, k} = \sum_{l} x_{i,j, l} \times y_{i, l, k}\quad \text{for } i = 1\ldots N_2,\ j = 1\ldots N_0,\ k = 1\ldots N_1 $$

Hint:可以使用tl.dot，它可以计算一个更小规模的矩阵乘法

问题解答:

我们知道两个矩阵的形状分别为shape(x) = [N2,N0,MID],shape(y) = [N2,MID,N1].我们分块的话，显然需要在MID维度上进行拆分，那么两个矩阵对应的块大小为:[B2,B0,B_MID],[B2,B_MID,B1],然后将这些小矩阵的乘积进行累加，便能得到最终的结果。

@triton.jit
def dot_kernel(
 x_ptr,
 y_ptr,
 z_ptr,
 N0,
 N1,
 N2,
 MID,
 B0: tl.constexpr,
 B1: tl.constexpr,
 B2: tl.constexpr,
 B_MID: tl.constexpr,
):
 block_id_j = tl.program_id(0)
 block_id_k = tl.program_id(1)
 block_id_i = tl.program_id(2)
 # Finish me!
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N2,N0,MID),
 strides=(N0 * MID,MID,1),
 offsets=(block_id_i * B2, block_id_j * B0, 0),
 block_shape=(B2, B0, B_MID),
 order=(0, 1, 2),
 )
 y_block_ptr = tl.make_block_ptr(
 base=y_ptr,
 shape=(N2,MID,N1),
 strides=(MID * N1,N1,1),
 offsets=(block_id_i * B2, 0, block_id_k * B1),
 block_shape=(B2, B_MID, B1),
 order=(0, 1, 2),
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N2,N0,N1),
 strides=(N0 * N1,N1,1),
 offsets=(block_id_i * B2, block_id_j * B0, block_id_k * B1),
 block_shape=(B2, B0, B1),
 order=(0, 1, 2),
 )
 z = tl.zeros((B2, B0, B1), dtype=tl.float32)
 for i in range(0, MID, B_MID):
 x = tl.load(x_block_ptr, boundary_check=(0,1,2),padding_option='zero')
 y = tl.load(y_block_ptr, boundary_check=(0,1,2),padding_option='zero')
 z += tl.dot(x,y)
 x_block_ptr = tl.advance(x_block_ptr,(0,0,B_MID))
 y_block_ptr = tl.advance(y_block_ptr,(0,B_MID,0))
 tl.store(z_block_ptr, z, boundary_check=(0,1,2))
 return

Puzzles 12 Quantized Matrix Mult

问题描述:

我们将对一个矩阵乘法进行量化压缩，具体而言是将权重矩阵以更低精度存储，并额外配合一个偏移项和缩放项。

在本问题中，我们的weight将以4bit存储。一个32位整数中可以存放FPINT个这样的值。此外，对每个连续group个权重，我们还会额外存储:

1个scale浮点值
1个shift的4bit值

这些scale和shift是按weight的列来存储的。而activation则单独以普通浮点数格式存储。

$$ z_{j, k} = \sum_{l} sc_{j, \frac{l}{g}} (w_{j, l} - sh_{j, \frac{l}{g}}) \times y_{l, k} \quad \text{for } j = 1\ldots N_0,\ k = 1\ldots N_1 $$

其中g表示分组大小。

问题解答:

本题的主要难点，不在于计算，而是张量的形状处理。我们先来梳理以下完整的操作流程。

首先我们还是按照正常的矩阵乘法在MID维度上进行拆分，然后由于权重是以INT8格式存储的，我们还需要将其提取出来(利用位运算即可)。然后每个GROUP我们都会存储一个浮点型的scale和INT4类型的shift.

为了方便后续理解，我们看看各个分块的大小(仅看输入):

激活值,这个最好理解，它是以浮点型存储的，然后因为是在MID维度做了拆分，因此块大小为[B_MID,B_1],需要注意访问的顺序是列优先的，因此order = (0,1)
权重矩阵，这个需要考虑到INT32存储INT8带来的形状差异，我们记一个INT32可以存储FPINT个INT4,那么权重矩阵实际的大小为[N0,MID//FPINT],那么对应的块大小就是[B0,B_MID//FPINT]
scale,它是浮点型存储的，并且是一个GROUP(在MID维度分组)一个,那么其形状就是[N0,GROUP],对应块大小为[B0,GROUP]
offset,它与Scale类似，是一个GROUP一个，但是它是以INT4存储的，因此也需要考虑FPINT的事情，它的实际大小就是:[B0,GROUP//FPINT]

分完块之后就需要把用int4的块提取为int32格式的，方便后续计算。

然后需要考虑形状的问题了，我们提取出来的INT32格式的Weight的张量形状为[B0,B_MID // FPINT,FPINT],但是我们的scale的形状为[B0,GROUP]，我们无法将其直接与Weight进行运算得到dequant的权重，因此我们需要将其reshape为[B0,GROUP,B_MID//GROUP].offsets也同理，需要从[B0, GROUP // FPINT, FPINT]处理为[B0,GROUP]。之后借助广播进行dequant后，再进行矩阵乘法。

@triton.jit
def quant_dot_kernel(
 scale_ptr,
 offset_ptr,
 weight_ptr,
 activation_ptr,
 z_ptr,
 N0,
 N1,
 MID,
 B0: tl.constexpr,
 B1: tl.constexpr,
 B_MID: tl.constexpr,
):
 block_id_j = tl.program_id(0)
 block_id_k = tl.program_id(1)

 def extract(x):
 over = tl.arange(0,8) * 4
 mask = 2**4 - 1
 return (x[:,:, None] >> over) & mask

 activation_block_ptr = tl.make_block_ptr(
 base = activation_ptr,
 shape = (MID,N1),
 strides = (N1,1),
 offsets = (0, block_id_k * B1),
 block_shape = (B_MID, B1),
 order = (0,1)
 )
 weight_block_ptr = tl.make_block_ptr(
 base = weight_ptr,
 shape = (N0,MID // FPINT),
 strides = (MID // FPINT,1),
 offsets = (block_id_j * B0, 0),
 block_shape = (B0, B_MID // FPINT),
 order = (1,0)
 )
 scale_block_ptr = tl.make_block_ptr(
 base = scale_ptr,
 shape = (N0, GROUP),
 strides = (GROUP,1),
 offsets = (block_id_j * B0, 0),
 block_shape = (B0, GROUP),
 order = (1,0)
 )
 offset_block_ptr = tl.make_block_ptr(
 base = offset_ptr,
 shape = (N0,GROUP // FPINT),
 strides = (GROUP // FPINT,1),
 offsets = (block_id_j * B0, 0),
 block_shape = (B0, GROUP // FPINT),
 order = (1,0)
 )
 z_block_ptr = tl.make_block_ptr(
 base = z_ptr,
 shape = (N0,N1),
 strides = (N1,1),
 offsets = (block_id_j * B0, block_id_k * B1),
 block_shape = (B0, B1),
 order = (1,0)
 )
 z = tl.zeros((B0, B1), dtype=tl.float32)
 for i in range(0, MID, B_MID):
 scale_fp32 = tl.load(scale_block_ptr, boundary_check=(0,1), padding_option='zero')#[B0,GROUP]
 
 offset_int32 = extract(tl.load(offset_block_ptr, boundary_check=(0,1), padding_option='zero'))#[B0, GROUP // FPINT, FPINT]
 offset_int32 = offset_int32.reshape(B0, GROUP)#[B0, GROUP]
 
 weight_int32 = extract(tl.load(weight_block_ptr, boundary_check=(0,1), padding_option='zero'))#[B0, B_MID // FPINT, FPINT]
 weight_int32 = weight_int32.reshape(B0, B_MID)#[B0, B_MID]
 weight_int32 = weight_int32.reshape(B0,GROUP,B_MID // GROUP)#[B0, GROUP, B_MID // GROUP]

 activation_fp32 = tl.load(activation_block_ptr, boundary_check=(0,1), padding_option='zero')
 
 weight_fp32 = scale_fp32[:,:,None] * (weight_int32 - offset_int32[:,:,None])
 weight_fp32 = weight_fp32.reshape(B0,B_MID)

 z += tl.dot(weight_fp32, activation_fp32)
 scale_block_ptr = tl.advance(scale_block_ptr,(0, GROUP))
 offset_block_ptr = tl.advance(offset_block_ptr,(0, GROUP // FPINT))
 weight_block_ptr = tl.advance(weight_block_ptr,(0, B_MID // FPINT))
 activation_block_ptr = tl.advance(activation_block_ptr,(B_MID, 0))
 tl.store(z_block_ptr, z, boundary_check=(0,1))

 # Finish me!
 return

Reference

namoe的解答
Lite作者的解答
先进编译实验室的讲解
CS336 Assignment2-systems的实验指导手册
CSE559M Note

Transformer中的位置编码的探索与思考

Tue, 14 Apr 2026 20:48:13 +0800

Motivation

最近在完成CS336的Assignment1，从头搭建一个大模型，其中实现使用的是主流的旋转位置编码(RoPE)。虽然对照指导书完成了相关实验，但是对位置编码的选择以及原理存在一定的知识缺漏，因此写本博客以加强理解。

Note

我们在下面的数学推到中使用列向量推导

为什么Transformer需要位置编码?

设输入序列长度为L，模型的维度为d，那么我们输入的序列第i个token的表示为$x_i \in \mathbb{R}^{d}$,将整个序列拼接起来我们有:

$$ X = [x_1,x_2,\dots,x_L]\in \mathbb{R}^{d\times L} $$

那么Transformer Block中的MHA(由于最终结果是每一个Head通过相同的过程得到的不同结果进行concat，因此我们这里只考虑一个Head)在计算Attention Score时首先通过线性映射得到:

$$ Q = W_Q X \quad K = W_KX\quad V = W_V X $$

其中:

$$ W_Q,W_K,W_V \in \mathbb{R}^{d\times d} $$

因此:

$$ Q,K,V \in \mathbb{R}^{d\times L} $$

接下来会计算第i个query和第j个key的相似度:

$$ s_{ij} = \frac{q_i^\top k_j}{\sqrt{d}} $$

写成矩阵形式就是:

$$ S = \frac{Q^\top K}{\sqrt{d}}\in \mathbb{R}^{L\times L} $$

然后进行Softmax操作(沿key维度):

$$ A = \text{softmax}(S)\in \mathbb{R}^{L\times L} $$

最后进行加权求和:

$$ \text{Attention(Q,K,V)} = VA^\top \in \mathbb{R}^{d\times L} $$

令$P \in \mathbb{R}^{L\times L}$为置换矩阵，表示对序列顺序的重排。我们有:

$$ X' = XP $$

那么我们重新计算Attention:

$$ Q' = W_Q X' = W_QXP = QP \quad K' = W_KX'=KP \quad V' = VP $$$$ S' = \frac{Q'^\top K'}{\sqrt{d}} = \frac{P^\top Q^\top K P}{\sqrt{d}} $$$$ A' = \text{softmax}(S') = P^\top AP $$$$ \text{Attention}(Q',K',V') = VPA'^\top = VPP^\top A^\top P = VA^\top P = \text{Attention}(Q,K,V)P $$

上式表明，在不引入任何位置信息的情况下，Attention计算具有置换同变性。但是我们需要注意到，P是一个置换矩阵，它只是索引的重定向，我们计算出来的Attention的数值没有发生任何变化。如果我们把Attention输出结果看作带有语义信息的Embedding向量，那么它的语义将不会携带位置语义的信息。

举一个具体的例子:

我们输入的文本序列是["猫","吃","鱼"]，对应的输入序列为[x_1,x_2,x_3],而对应的Attention输出为[y1,y2,y3],而我们进行交换得到:["鱼","吃","猫"]，对应的输入序列为[x_3,x_2,x_1],根据我们上面的证明,Attention的输出将是[y3,y2,y1]，在数值上没有任何变换，只是简单的位置交换，也就是说Attention的输出没有提取出在语义结构(如主谓宾等)的信息。

位置编码为什么可以解决这个问题？

我们现在需要解决的问题是打破这种数值不变性，也就是交换了输入位置后，我们希望最终的输出为[y3',y2',y1'].我们将Attention视为一个输出为集合的函数f,那么上述不变性用数学语言表述就是:

$$ f(\dots,x_m,\dots,x_n,\dots) = f(\dots,x_n,\dots,x_m,\dots) $$

因此，我们要做的事情就是打破这种不变性，比如在每个位置都加上一个不同的编码向量:

$$ \hat f(\dots,x_m,\dots,x_n,\dots) = f(\dots,x_m+p_m,\dots,x_n+p_n,\dots) $$

一般来说，只要每个位置的编码向量不同，那么这种全对称性就被打破了，即可以用f~代替f来处理有序的输入。

我们写作矩阵的形式:

我们将位置编码矩阵定义为$E = [e_1,e_2,\dots,e_L]\in \mathbb{R}^{d\times L}$。每一个列向量$e_i$仅与位置索引i有关。

融入后的输入矩阵为:

$$ X_{pos} = X + E $$

此时我们重新推导Q,K,V的生成过程(以Q为例):

$$ Q = W_Q(X + E) = W_Q X + W_Q E = Q_X + Q_E $$

那么同理:

$$ K = K_X + K_E,V = V_X + V_E $$

我们可以将$\{K,Q,V\}_{E}$视为位置信息的“特征表达”

接下来我们计算$S = \frac{Q^\top K}{\sqrt{d}}$,我们具体来看其中一个点积项$s_{ij}$:

$$ s_{ij} = \frac{1}{\sqrt{d}}(q_{x,i}+q_{e,i})^\top(k_{x,j}+k_{e,j})\\ =\frac{1}{\sqrt{d}}(q^\top_{x,i}k_{x,j}+q^\top_{x,i}k_{e,j}+q^\top_{e,i}k_{x,j}+q^\top_{e,i}k_{e,j}) $$

那么此时就多出了位置-内容，位置-位置的信息。

我们再按照同样的推理,令$X' = XP$,那么此时:

$$ Q' = W_Q(X'+E) = W_Q(XP+E) = Q_XP+Q_E $$

同理:

$$ K' = K_XP+K_E \quad V' = V_XP+V_E $$

那么:

$$ S' = \frac{Q'^\top K'}{\sqrt{d}} = \frac{(Q_XP+Q_E)^\top(K_XP+K_E)}{\sqrt{d}}\\ =\frac{1}{\sqrt{d}}(P^\top Q_X^\top K_XP + P^\top Q_X^\top K_E + Q_E^\top K_XP+Q_E^\top K_E) $$

显而易见,$S' \neq P^\top SP$,且其中的$Q_E^\top K_E$项没有被P作用，也就意味着原本的置换对称性也就被打破了，因此我们可以用$\hat f $来代替f来处理有序的输入。

也就是说，位置编码的引入可以解决我们说的Attention置换数值不变性的问题。

怎样的位置编码是好的？

我们现在想要进一步分析位置编码的性质，从而设计更好的位置编码。我们将$\hat f$展开至二阶项(为了简化考虑，写作矩阵形式):

$$ \hat f(X) =f(X+E) \approx f(X) + \nabla f(X)\cdot E+\frac{1}{2}E^\top H_f(X)E $$

那么我们来看与位置编码有关的项:

一阶项$\nabla f(X)\cdot E = \sum_{i=1}^L e^\top_i \frac{\partial f}{\partial x_i}$:这项依赖于单一位置，所以是绝对位置信息
二阶项$\frac{1}{2}E^\top H_{f(X)}E = \sum_{i=1}^L\sum_{j=1}^L e_i^\top \frac{\partial^2 f}{\partial{x_i}\partial x_j}e_j$:这项与包含了任意两个位置的交互，所以是相对位置信息

对于一阶项，只要位置编码在每个位置都是独特的就能表述绝对位置的信息。而对于二阶信息，对于一个理想的位置编码而言，应该让这个二阶项满足某种平移不变性，即对于任意位移k，位置i与位置i+k的交互模式应该是稳定的。

我们先从最简单的情况入手，假设$\mathbf{H} = I$为单位矩阵，那么此时$E^\top E$是两个位置编码的内积，我们希望在这个简单的例子中该项表达的是相对位置信息，即存在某个函数g使得:

$$ \langle p_m,p_n\rangle = g(m-n) $$

这里的$p_m,p_n$为d维向量，这里我们从最简单的$d = 2$入手。我们称上式是一个位置编码为一个合理位置编码的条件式。

对于2维向量，我们借助复数来推导，视向量[x,y]为复数$x+yi$,那么我们有:

$$ \langle p_m,p_n\rangle=a_xb_x+a_yb_y = \text{Re}(p_m\hat p_n) $$

其中$\hat w$为$w$的共轭复数。

为了满足上式，我们可以假设存在复数$q_{m-n}$,使得:

$$ p_m\hat p_n = q_{m-n} $$

这样两边取实部就得到条件式。为了解这个方程，我们可以使用复数的指数形式，假设$p_m = r_me^{i\phi_m},\hat p_n = r_ne^{-i\phi_n},q_{m-n} = R_{m-n}e^{i\Phi_{m-n}}$,那么有:

$$ r_mr_ne^{i(\phi_m-\phi_n)}=R_{m-n}e^{i\Phi_{m-n}} $$

于是我们得到等式:

$$ \left\{ \begin{array}{l} r_m r_n = R_{m-n} \\ \phi_m - \phi_n = \Phi_{m-n} \end{array} \right. $$

对于第一个方程，带入m = n，可以得到$r_m^2 = R_0$,即$r_m$为一个常数，为了简单，我们令其为1；

对于第二个方程，显然等差数列满足上述性质(令m = 0有$\Phi_m = \phi_m$)，设公差为$\theta$，则通项为$\phi_m = \Phi_m = m\theta$,由此我们得到二维情况下的位置编码的解:

$$ p_m = e^{im\theta} \to p_m = \left(\begin{align}\cos m\theta\newline \sin m\theta\end{align}\right) $$

由于内积满足线性叠加性，我们可以由二维的情况直接扩展到更高偶数维的情况:

$$ p_m = \left(\begin{matrix}e^{im\theta_0}\\e^{im\theta_1}\\\vdots\\e^{im\theta_{d/2-1}}\end{matrix}\right) \to p_m = \left(\begin{matrix} \cos m\theta_0 \\ \sin m\theta_0 \\ \cos m\theta_1 \\ \sin m\theta_1 \\ \vdots\\ \cos m\theta_{d/2-1} \\ \sin m\theta_{d/2-1} \\ \end{matrix}\right) $$

这样我们就求出了满足条件式的一组解，显然解不唯一。

此外，一个好的位置编码应该满足远程衰减的性质，即随着$|m-n|$的增大,$\langle p_m,p_n\rangle$有趋于0的趋势。

那么有:

$$ \langle p_m,p_n\rangle = \text{Re}[e^{i(m-n)\theta_0} + e^{i(m-n)\theta_1}+\dots+e^{i(m-n)\theta_{d/2-1}}]\\ =\sum_{j=0}^{d/2-1}\cos (k\theta_j) \quad k = m-n $$

由于在LM中d通常为768，是一个较大值，因此我们可以将离散的索引j映射到连续变量$t\in[0,1]$上。设$\theta_j$是某个光滑单调函数$f(t)$生成的，即$\theta_j = f(2j/d)$。利用Euler-Maclaurin的一阶近似，我们可以将求和转化为积分:

$$ \langle p_m,p_n\rangle \approx \frac{d}{2}\int_0^1 \cos(k\cdot f(t))dt $$

那么现在的问题就转化为了，寻找一个函数$f(t)$,使得上述震荡积分在k很大时具有较好的衰减性质(足够快)。

根据黎曼-勒贝格引理，只要频率分布函数f(t)满足光滑且严格单调的条件，当$k\to \inf$时，被积函数的高频震荡将导致正负面积互相抵消，使得积分值必然趋近于0.

在Transformer2017的论文中的Sinusoidal位置编码选择的是$\theta_t = 10000^{-t}$.

由此，我们便推导出了Sinusoidal位置编码的形式:

$$ \left\{ \begin{array}{l} p_{k,2i} = \sin(k/10000^{2i/d})\\ p_{k,2i+1} = \cos(k / 10000^{2i/d}) \end{array} \right. $$

我们只能说明它的合理性，但是无法说明它的最优性，因为它并不一定最优[Lol]

事实上，一个可行的方案是将位置编码中的$\theta_i$设为可学习的参数，其初始值为$\theta_i = 10000^{-2i/d}$.

在上述推导中，都是基于H = I这个简单情况，对于一般的H，使用上述Sinusoidal位置编码，还能具备我们理想的性质吗？

事实上，有研究表明$^{[3]}$ 在网络规模足够大的情况下，其Hessian矩阵将呈现出块对角的形式，因此我们考虑H是一个对角阵的情况，此时:

$$ p_m^\top H p_n = \sum_{i=1}^{d/2}H_{2i,2i}\cos m\theta_i \cos n\theta_i+H_{2i+1,2i+1}\sin m\theta_i\sin n\theta_i $$

由和差化积有:

$$ \sum_{i=1}^{d/2}\frac{1}{2}(H_{2i,2i}+H_{2i+1,2i+1})\cos(m-n)\theta_i + \frac{1}{2}(H_{2i,2i}-H_{2i+1,2i+1})\cos(m+n)\theta_i $$

可以看到其中是包含了相对位置项(m-n)的，只是会出现m+n项。

因此我们可以认为Sinusoidal位置编码是一个有效的位置编码。

为什么RoPE$^{[4]}$比Sinusoidal位置编码更好？

在先前分析为什么需要位置编码的过程中，我们知道位置编码的作用是在计算Attention Score时，让QK点积中包含文本的结构语义信息，具体而言是下式中的后三项包含位置信息:

$$ S' = \frac{Q'^\top K'}{\sqrt{d}} = \frac{(Q_XP+Q_E)^\top(K_XP+K_E)}{\sqrt{d}}\\=\frac{1}{\sqrt{d}}(P^\top Q_X^\top K_XP + P^\top Q_X^\top K_E + Q_E^\top K_XP+Q_E^\top K_E) $$

我们可以发现，除了最后一项是相对位置信息外，另外两项则是绝对位置信息与文本语义信息的耦合，这实际上是一种冗余，因为绝对位置信息的作用相较于相对位置信息而言一方面具有误导性，例如在短文本中结尾出现在绝对位置10中，而10在长文本中可能仅仅是文本开头；另一方面模型还会额外去学习已经得到的相对位置信息造成资源的浪费.

为了解决这个问题，工业界进行了许多尝试$^{[5]}$,但是这些尝试普遍会带来大量的额外计算存储开销。由于绝对位置编码具有实现简单，计算速度快的特点，并且在Sinusoidal位置编码中我们也能看到通过绝对位置编码在一定程度上是可以得到相对位置信息的,如果可以通过绝对位置编码的方式实现相对位置编码，那么就是“集各家之所长”，“鱼和熊掌兼得”了。

为了实现这个目标，我们假设通过下述运算来给q,k添加绝对位置信息:

$$ \hat q_m = f(q,m)\quad \hat k_n = f(k,n) $$

而我们希望得到如下恒等关系:

$$ \langle f(q,m),f(k,n)\rangle = g(q,k,m-n) $$

为了求解方便，我们令:

$$ f(q,0) = q \quad f(k,0) = k $$

求解思路与之前推导Sinusoidal类似，我们先考虑二维的情况，然后借助复数来求解。

$$ \text{Re}[f(q,m)\hat f(k,n)] = g(q,k,m-n) $$

设:

$$ \begin{array}{rcl} f(q,m) &=& R_f(q,m)e^{i\theta_f(q,m)}\\ \hat f(k,n) &=& R_f(k,n)e^{-i\theta_f(k,n)}\\ g(q,k,m-n) &=& R_g(q,k,m-n)e^{i\theta_g(q,k,m-n)} \end{array} $$

那么带入方程求解得到:

$$ \left\{ \begin{array}{l} R_f(q,m)R_f(k,n) = R_g(q,k,m-n)\\ \theta_f(q,m) - \theta_f(k,n) = \theta_g(q,k,m-n) \end{array} \right . $$

对于第一个方程，令m = n有:

$$ R_f(q,m)R_f(k,n) = R_g(q,k,0) = R_f(q,0)R_f(k,0)=||q||||k|| $$

那么我们可以直接令$R_f(q,m) = ||q||$,即它不依赖于m。

对于第二个方程，同样地令m=n有:

$$ \theta_f(q,m)-\theta_f(k,n)=\theta_g(q,k,0)=\theta_f(q,0)-\theta_f(k,0) = \theta_q-\theta_k $$

这里的$\theta_q,\theta_k$是q,k本身的辐角。

由此我们得到:

$$ \theta_f(q,m) - \theta_q = \theta_f(k,n) - \theta_k $$

因此$\theta_f(q,m)-\theta_q$应该是一个只与m有关的而与q无关的函数，因为我们希望恒等式对任意的q,k都成立，所以右侧必须与q,k无关，因此记为常数，记为$\theta$，该式记为$\phi(m)$,即

$$ \theta_f(q,m) = \theta_q + \phi(m) $$

令n = m - 1 有:

$$ \phi(m) - \phi(m-1) = \theta_g(q,k,1)+\theta_k-\theta_q $$

故$\{\phi(m)\}$为等差数列，公差为$\theta$,得到$\phi(m) = m\theta$

由此，我们得到了二维情况下用复数表示的RoPE:

$$ f(q,m) = R_f(q,m)e^{i\theta_f(q,m)} = ||q||e^{i(\theta_q+m\theta)} = \vec{q}e^{im\theta} $$

根据复数乘法的集合意义，该变换实际上对应着向量的旋转，我们可以将其写为矩阵形式:

$$ f(q,m) = \left( \begin{matrix} \cos m\theta \quad -\sin m\theta\\ \sin m \theta \quad \cos m\theta \end{matrix} \right) \left( \begin{matrix} q_0\\ q_1 \end{matrix} \right) $$

由内积的线性叠加性，我们可以得到任意偶数维度的RoPE:

$$ \left( \begin{matrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \dots & 0 & 0 & \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \dots & 0 & 0 & \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \dots & 0& 0\\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \dots &\cos m\theta_{d/2-1} & -\sin m \theta_{d/2-1}\\ 0 & 0 & 0 & 0 & \dots & \sin m\theta_{d/2-1} & \cos m \theta_{d/2-1} \end{matrix} \right) \left( \begin{matrix} q_0\\q_1\\q_2\\q_3\\ \vdots\\q_{d-2}\\q_{d-1} \end{matrix} \right) $$

我们称左侧的矩阵为旋转矩阵记为$R_m$.也就是说给位置为m的向量q乘上矩阵$R_m$,位置为n的向量k乘上矩阵$R_n$,用变换后的Q，K序列做Attention，那么Attention就自动包含相对位置信息了，因为恒等式成立:

$$ (R_mq)^\top(R_nk) = q^\top R_m^\top R_nk = q^\top R_{n-m}k $$

值得指出的是，Rm是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性。

此外，在具体实现时，我们并不会拿这个大矩阵去乘以向量，而是采用逐位相乘的方式:

$$ \left( \begin{matrix} q_0\\q_1\\q_2\\q_3\\ \vdots\\q_{d-2}\\q_{d-1} \end{matrix} \right)\otimes \left( \begin{matrix} \cos m\theta_0\\\cos m\theta_0 \\\cos m\theta_1 \\ \cos m\theta_1\\ \vdots \\ \cos m\theta_{d/2-1} \\ \cos m\theta_{d/2-1} \end{matrix} \right)+ \left( \begin{matrix} -q_1\\q_0\\-q_3\\q_2\\ \vdots\\-q_{d-1}\\q_{d-2} \end{matrix} \right)\otimes \left( \begin{matrix} \sin m\theta_0\\\sin m\theta_0\\\sin m\theta_1\\\sin m\theta_1\\ \vdots\\\sin m\theta_{d/2-1}\\\sin m\theta_{d/2-1} \end{matrix} \right) $$

对于$\theta$的选择，作者选择了与Sinusoidal位置编码一样的$\theta_i = 10000^{-2i/d}$,从而带来远程衰减性。

至此，我们完成了对RoPE的推导，并通过推导成功说明了为何RoPE相较于Sinusoidal位置编码更好，因为它通过绝对位置的注入方式，实现了相对位置的注入，而不带来其它冗余。

最后，下面是我在CS336中实现的一个RoPE:

from einops import rearrange, einsum
import torch
import torch.nn as nn

class RotaryPositionalEmbedding(nn.Module):
 def __init__(self,theta,d_k,max_seq_len,device):
 """
 d_k: int, 维度大小，必须为偶数
 theta: float, RoPE中的\Theta值
 max_seq_len: int, 最大序列长度
 device: torch.device, 设备
 """
 super().__init__()
 assert d_k % 2 == 0, "d_k must be even"
 self.theta = theta
 self.d_k = d_k
 self.max_seq_len = max_seq_len
 self.device = device

 #一共有d / 2个频率
 half_dk = d_k // 2
 k = torch.arange(0,half_dk,device=device).float()
 inv_freq = 1.0 / (self.theta ** (2.0 * k / d_k))

 positions = torch.arange(0,max_seq_len,device = device).float()

 angles = einsum(positions,inv_freq,"max_seq_len,half_dk->max_seq_len half_dk")
 cos = torch.cos(angles)
 sin = torch.sin(angles)

 self.register_buffer("cos",cos,persistent = False)
 self.register_buffer("sin",sin,persistent = False)

 def forward(self,x,token_positions):
 """
 inputs:
 x: ...,seq_len,d_k
 token_positions:...,seq_len
 returns:
 x_rotated: ...,seq_len,d_k
 """
 cos = self.cos[token_positions] # ...,seq_len,half_dk
 sin = self.sin[token_positions] # ...,seq_len,half_dk

 x_even = x[...,0::2]
 x_odd = x[...,1::2]

 x_rot_even = x_even * cos - x_odd * sin
 x_rot_odd = x_even * sin + x_odd * cos

 out = torch.empty_like(x)
 out[...,0::2] = x_rot_even
 out[...,1::2] = x_rot_odd
 return out

参考

Hessian矩阵系列串讲

Mon, 24 Nov 2025 13:28:52 +0800

定义

假设有一实值函数$f(x_1,x_2,\dots,x_n)$，若$f$的所有二阶偏导数都存在且在定义域里连续，那么我们定义函数$f$的Hessian矩阵为如下一个$n\times n$的方阵:

$$ \mathbf{H} = \left[ \begin{matrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1\partial x_2} & \dots&\frac{\partial^2 f}{\partial x_1 \partial x_n}\\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} & \dots& \frac{\partial^2 f}{\partial x_2 \partial x_n}\\ \vdots & \vdots & \ddots & \vdots\\ \frac{\partial^2 f}{\partial x_n \partial x_1} & \frac{\partial^2 f}{\partial x_n \partial x_2} & \dots & \frac{\partial^2 f}{\partial x_n^2} \end{matrix} \right] $$

或使用下标标记表示为:

$$ \mathbf{H}_{ij} = \frac{\partial^2 f}{\partial x_i \partial x_j} $$

一般性质

对称性:

对于机器学习中遇到的大多数函数(特别是那些具有连续二阶偏导数的函数),混合偏导数的求导顺序无关紧要。这被称为克莱罗定理或施瓦茨定理,它说明了混合偏导数的相等性:

$$ \frac{\partial^2 f}{\partial x_i\partial x_j} = \frac{\partial^2 f}{\partial x_j \partial x_i} $$

这意味着Hessian矩阵是对称的，即$\mathbf{H} = \mathbf{H}^\top$

正定性与局部最优值的相关性

通过Hessian矩阵的正定性，我们可以判断该点是局部最小值，局部最大值还是鞍点

正定:若Hessian矩阵$H(x^*)$是正定的(意味着$\forall v \neq 0,v^\top\mathbf{H}v>0$)，则函数在$x^*$处有局部最小值。函数在该点周围向所有方向向上弯曲，类似于碗的底部
负定:若Hessian矩阵$H(x^{*})$是负定的(意味着$\forall v \neq 0,v^\top\mathbf{H}v<0$),则函数在$x^*$处有局部最大值。函数在该点周围向所有方向向下弯曲，类似于圆顶的顶部
不定:若Hessian矩阵$H(x^{*})$是不定的,则函数在$x^*$处有鞍点。函数在某些方向向上弯曲，在另一些方向向下弯曲，就像马鞍

Hessian在神经网络中的计算

Hessian矩阵在神经网络计算的许多方面有着重要作用，包括:

⼀些⽤来训练神经⽹络的⾮线性最优化算法是基于误差曲⾯的⼆阶性质的，这些性质

由Hessian矩阵控制(比如牛顿法和拟牛顿法)
对于训练数据的微⼩改变，Hessian矩阵构成了快速重新训练前馈⽹络的算法的基础
Hessian矩阵的逆矩阵⽤来鉴别神经⽹络中最不重要的权值，这是⽹络“剪枝”算法的⼀部分(LeCun的OBD)

对于Hessian矩阵的众多应用而言，一个重要的需要考虑的问题是计算效率。在神经网络中有$W$个参数(包括权值和偏置),那么Hessian矩阵的大小就是$W\times W$,那么计算Hessian矩阵的计算量为$O(W^2)$。这在具有大量参数的神经网络中是难以接受的，因此我们需要进行一些高效的近似。

对角近似

这个方法最早由Yan LeCun在OBD剪枝方法中提出

我们只保留Hessian矩阵的对角线元素$H_{i,i}$，把非对角线元素置为0.(这个方法在很大程度可以可以方便求逆)

对于神经网络中的第j个神经元的输入$a_j$对应权重为$w_{ji}$:

$$ \frac{\partial^2 E_n}{\partial w_{ji}^2} = \frac{\partial^2 E_n}{\partial a_j^2}z_i^2 $$

其中$z_i$是上一层神经元的输出。

而$\frac{\partial^2 E_n}{\partial a_j^2}$可以通过链式法则递归计算(类似于反向传播):

$$ \frac{\partial^2 E_n}{\partial a_j^2} = \underbrace{\frac{\partial}{\partial a_j}[h'(a_j)\sum_{k}w_{kj}\frac{\partial E_n}{\partial a_k}]}_{链式法则}=h'(a_j)^2 \sum_{k,k'}w_{kj}w_{k'j}\frac{\partial^2 E_n}{\partial w_k\partial w_{k'}} + h''(a_j)\sum_k w_{kj}\frac{\partial E_n}{\partial a_n} $$

忽略二阶导中的非对角线项$k\neq k'$:

$$ \frac{\partial^2 E_n}{\partial a_j^2}\approx \underbrace{h'(a_j)^2\sum_k w_{kj}^2 \frac{\partial^2 E_n}{\partial a_k^2}}_{链式法则} + \underbrace{h''(a_j)\sum_{k}w_{kj}\frac{\partial E_n}{\partial a_k}}_{链式法则} $$

从而一次反向传播便可以计算出来，时间复杂度为$O(W)$

但是如你所见，在这个方法下，Hessian矩阵完全退化为了一个diag，这在神经网络中是不合理的，因为非线性层的引入会让Hessian矩阵的交叉项不为0，更重要的是这种误差会随着网络层数的堆砌不断放大。

因此针对这个问题在对角近似上近年来不断有相关文章发表，下面对我阅读过的一些进行简单介绍:

HesScale:在LeCun的方法的骨架上，针对常见网络的最后一层的softmax+CE结构可以通过一个简单的公式求出diag(H)的精确值:
$$ \frac{\partial^2 E_n}{\partial^2 a_j^2} = p - p\circ p $$
其中p为$\text{softmax}(a_j)$,$\circ$表示向量的逐元素乘积,为方便书写我们将$\{a_1,\dots,a_n\}$写作z

证明如下:

预测$\mathbf{p} = \text{softmax}(z)$,即$p_i = \frac{\exp(z_i)}{\sum_{j = 1}^k \exp(z_j)}$,设目标分布为$t\in \Delta^{K-1},\sum_{i}t_i = 1$，对于单样本而言，其损失可以写作:
$$ \mathcal{L}(z) = -\sum_{j=1}^Kt_i\log p_i = -t^\top z + \log \sum_{j=1}^K \exp(z_j) $$
求一阶导我们有:
$$ \nabla_z \mathcal{L} = -t + \nabla_z \log \sum_{j=1}^K \exp{z_j} = -t + \mathbf{p} $$
求二阶导相当于$\mathbf{p}$对z求导。我们知道$\text{softmax}$的Jocobian为
$$ \frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} - p_j),J_{\text{softmax}}(z) = \text{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top $$
其中$\delta_{i,j}$为Kronecker函数，写成矩阵形式，则只有对角线元素为1，其余为0.

因此我们有:
$$ \nabla^2_z\mathcal{L} = \text{diag}(\mathbf{p}) - \mathbf{pp}^\top $$
于是我们可以得到:
$$ \frac{\partial^2 L}{\partial^2 a_j} = p_j(1-p_j) $$
AdaHessian:

具体而言，AdaHessian在进行Hessian矩阵的对角近似的时候利用了Hutchinson估计，它通常也被用来对矩阵的迹进行估计:

对任意矩阵$\mathbf{A}\in\mathbb{R}^{d\times d}$,若随机向量$z = (z_1,\dots,z_d)^\top$满足:
$$ \mathbb{E}[z_i] = 0\\ \mathbb{E}[z_iz_j] = \delta_{ij} $$
那么有:
$$ \mathbb{E}[z\odot (Az)] = \text{diag}(A) $$
下面进行证明:

记估计量$d = z\odot (Az)$的第i个分量:
$$ d_i = z_i(Az)_i = z_i\sum_{j=1}^d A_{ij}z_j = \sum_{j = 1}^d A_{ij}z_iz_j\\ \mathbb{E}[d_i] = \sum_{j=1}^dA_{ij}\mathbb{E}[z_iz_j] = \sum_{j=1}^d A_{ij}\delta_{ij} = A_{ii} $$
由此我们证明了$\mathbb{E}[z\odot(Az)] = \text{diag}(A)$

在神经网络中，相较于直接计算完整的Hessian矩阵以及其OBD方式的对角近似，计算其矩阵向量积(HVP,Hessian Vector product)是并不困难的，它只需要一次反向传播：
$$ Hz = \frac{\partial (g^\top z)}{\partial \theta} = \frac{\partial g^\top}{\partial \theta}z + g^\top \frac{\partial z}{\partial \theta} = \frac{\partial g^\top}{\partial \theta}z $$
由此，我们可以通过多次在满足Rademacher分布的向量取样计算$z\odot (Az)$的期望就能得到$\text{diag}(A)$的无偏估计，在实际的应用中，只进行一次取样就能得到较为不错的结果。

外积近似

在神经网络应用于回归问题时，通常采用下面形式的平方和误差

$$ E = \frac{1}{2}\sum_{n=1}^N(y_n-t_n)^2 $$

那么Hessian矩阵可以写成如下形式:

$$ H = \nabla\nabla E = \nabla(\sum_{n=1}^N (y_n-t_n)\nabla y_n) = \sum_{n = 1}^N \nabla y_n(\nabla y_n)^\top + \sum_{n=1}^N(y_n-t)\nabla\nabla y_n $$

在网络已经训练好的情况下，输出$y_n$与$t_n$接近，因此第二项可以忽略，由此我们得到了Hessian矩阵的外积近似

$$ H\approx \sum_{n=1}^N b_n b_n^\top $$

其中，$b_n = \nabla y_n = \nabla a_n$(输出单元的激活函数就是恒等函数)。这种方法中的Hessian矩阵可以跟随反向传播算法在$O(W)$个步骤内高效地求出误差函数地一阶导数。再通过简单地乘法就可以在$O(W^2)$步骤内求出矩阵元素。

Hessian矩阵逆的计算

使用外积近似，我们可以提出一个计算Hessian矩阵的逆的高效办法,首先我们有:

$$ \mathbf{H}_N = \sum_{i=1}^n b_n b_n^\top $$

其中，$b_n = \nabla w a_n$时数据点n产生的输出单元激活对梯度的贡献。我们现在推导一个建立Hessian矩阵的顺序步骤，每次处理一个数据点。假设我们已经使用了前L个数据点得到了Hessian矩阵的逆。通过将第L+1个数据点的贡献单独写出来，我们有:

$$ \mathbf{H}_{L+1} = \mathbf{H}_L + b_{L+1}b_{L+1}^\top $$

为了计算Hessian矩阵的逆，我们考虑下面这个矩阵的恒等式:

$$ (M + vv^\top)^{-1} = M^{-1} - \frac{(M^{-1}v)(v^\top M^{-1})}{1+v^\top M^{-1}v} $$

若我们令$H_L = M$,且$b_{L+1} = v$,我们有:

$$ H_{L+1}^{-1} = H_{L}^{-1} - \frac{H_L^{-1}b_{L+1}b_{L+1}^\top H_{L}^{-1}}{1 + b_{L+1}^\top H_L^{-1}b_{L+1}} $$

通过这种方式，数据点可以依次使用，直到L+1=N,整个数据集处理完毕。于是，这个结果表示一个计算Hessian矩阵的逆的算法。这个算法只需对数据集扫描一次。最开始的矩阵$\mathbf{H_0}$被选为$\alpha I$,其中$\alpha$是一个较小的量，从而算法实际找的是$\mathbf{H} + \alpha I$的逆。结果对于$\alpha$的精确值不敏感。

Hessian矩阵的逆计算的K-FAC分解方法：

在特定场景下，Hessian矩阵的逆有一个利用Kronecker分解的近似方法，它利用了在loss为log-like形式下$\mathbb{E}[H] = F$的特点，其中$F$为Fisher信息矩阵:

$$ F = \mathbb{E}[\nabla \log p(x)\nabla \log p(x)^\top] $$

其中上式中做了如下形式的简写:

$$ \nabla \log{p(x)} = \nabla_\theta \log{(p|\theta)}\\ \mathbb{E}[p(x)] = E_{x\sim p(x|\theta)} $$

首先我们有:

$$ \mathbb{E}[\nabla \log p(x)] = \int (\nabla \log p(x))p(x) dx\\ =\int \frac{\nabla \log p(x)}{p(x)}p(x) dx\\ =\int \nabla \log p(x)dx\\ =\nabla 1 = 0 $$

那么当loss为log-like形式下时:

$$ L = -\log p(x) $$

我们求它的Hessian:

$$ \nabla^2 L = -\nabla^2 \log p(x)\\ =-\nabla \frac{\nabla p(x)}{p(x)}\\ =\frac{\nabla p(x)}{p(x)}^\top \frac{\nabla P(x)}{p(x)} - \frac{P(x)^2}{p(x)}\\ =\nabla \log p(x)^\top \nabla \log p(x) - \frac{\nabla^2 P(x)}{p(x)} $$

对该式求期望，我们有:

$$ \mathbb{E}[H]=\mathbb{E}[\nabla \log p(x)^\top \nabla \log p(x)] - \mathbb{E}[\frac{\nabla^2 p(x)}{p(x)}] = F - \int \frac{\nabla^2 p(x)}{p(x)} p(x) dx\\ =F - \int \nabla^2 p(x)dx\\ =F - \nabla^2 \int p(x)\\ =F $$

由于$\theta$本质为所有层$\mathbf{W}$的拼接，我们有:

$$ d\theta = \nabla_{\theta}L(x)\\ F = \mathbb{E}[\nabla L(x)\nabla L(x)^\top] = \mathbb{E}[d\theta d\theta^\top]\\ \theta = [\text{vec}(W_0)^\top,\text{vec}(W_1)^\top,\dots,\text{vec}(W_n)^\top]^\top $$

带入展开得到:

$$ F_{ij} = \mathbb{E}[\text{vec}(dW_i)\text{vec}(dW_j)^\top] $$

令$a_i,g_i$分别为第i层的前向输入和反向传播梯度，由反向传播算法有:

$$ d W_i = g_ia_i^\top $$

带入有:

$$ \text{vec}(dW_i) = \text{vec}(g_ia_i^\top) = a_i\otimes g_i $$

这里可能不太直观，这是因为平常对kronecker积($\otimes$)的接触较少，具体而言，我们这样定义Kronecker积:

$$ A\otimes B = \left[ \begin{matrix} a_{11}B &\dots &a_{1n}B\\ \vdots & \ddots& \vdots\\ a_{m1}B & \dots & a_{mn}B \end{matrix} \right] $$

在这里，由于我们将$dW_i$即$\nabla W_i$第i层参数的梯度进行了向量化，原本的$\nabla W_i$ 是如下形式:

$$ \nabla W_i = g_i a_{i}^\top = \left[\begin{matrix}g_1a_1 & \dots &g_1a_n\\ \vdots & \ddots & \vdots\\ g_ma_1 &\dots&g_ma_n\end{matrix} \right] $$

按列堆叠(vec)就有:

$$ \text{vec}(dW_i) = \left[ \begin{matrix} g_1a_1\\ \vdots\\ g_1a_n\\ \vdots\\ g_ma_n \end{matrix} \right] = a_i\otimes g_i $$

那么我们可以将Fisher矩阵写作如下形式:

$$ F_{ij} = \mathbb{E}[\text{vec}(dW_i)\text{vec}(dW_j)]\\ =\mathbb{E}[(a_i\otimes g_i)(a_j\otimes g_j)^\top]\\ =\mathbb{E}[(a_ia_j^\top)\otimes (g_ig_j^\top)]\\ \approx \mathbb{E}[a_ia_j^\top]\otimes\mathbb{E}[g_ig_j^\top] $$

这个近似相当于我们忽略了$\text{Cov}_\otimes(a_ia_j^\top,g_ig_j^\top)$,这实际上是合理的，尤其是在网络较深的情况下。这里给出一个链接提供一个比较详细的说明。

虽然直接求$F_{ij}$的复杂度仍然不变，但是利用Kronecker积的性质，我们在求逆时可以得到较大的性能提升:

$$ (A\otimes B)^{-1} = A^{-1}\otimes G^{-1} $$

并且在实际运算中，并不是整个网络的Fisher进行计算，而是按层做块对角运算:

$$ \mathbf{F} = \text{blockdiag}(F_1,F_2,\dots),F_l\approx A_l\otimes G_l $$

这样每层求逆的时候只需要求两个小矩阵的逆。

一些统计量的计算：

除了直接对Hessian矩阵的近似计算，我们有时候仅仅需要对Hessian矩阵的统计量进行计算，例如矩阵的迹，最大特征值等。

在计算这些统计量时，我们需要一个重要的算子:Hv,即Hessian矩阵与任意向量的乘积，这个乘积我们在上面已经证明了通过一次简单的反向传播算法可以得到。(PyHessian是一个Python库，它实现了这个算子的高效计算)

顶部若干特征值

利用幂迭代法 + HVP可以高效地求解特征值。

简单来说幂迭代法的流程很简单:

给定矩阵A以及初始非零向量$x_0$,

迭代:
$$ y_{k+1} = Ax_k\\ x_{k+1} = \frac{y_{k+1}}{||y_{k+1}||}\\ \mu_{k+1} = \frac{x_{k+1}Ax_{k+1}^\top}{x_{k+1}x_{k+1}^\top} $$
具体的收敛性证明见教材
矩阵的迹

在上面近似计算中，我们曾用Hutchinson估计对Hessian矩阵进行了对角近似，事实上，我们还可以进行迹的估计。

简单来说，方法一样，从Rademacher分布中取一随机向量$v$,然后有恒等式:
$$ \text{Tr}(H) = \text{Tr}(HI) = \text{Tr}(H\mathbb{E}[vv^\top]) = \mathbb{E}[\text{Tr}(Hvv^\top)] = \mathbb{E}[v^\top H v] $$
Hessian矩阵在神经网络下的特殊结构
本小节内容主要基于Towards Quantifying the Hessian Structure of Neural Networks,B站上有作者的讲解视频[FAI] 港中深张雨舜 | 浅谈神经网络Hessian矩阵的特殊结构

这篇文章主要说明了在神经网络中，Hessian矩阵往往具有近块对角结构，这表明曲率信息(二阶信息)主要在层内耦合，层与层的二阶交互很弱。这间接的说明了Layer-wise的量化/剪枝的合理性。

此外，文章还进一步从理论和实验上阐释了造成这个现象的原因：

静态力量:即使在随机初始化阶段，即在训练开始之前，神经网络的Hessian矩阵就已经呈现出近块对角结构。这种结构的形成与网络的架构设计有关，因此被称为“静态力量”。具体来说，对于线性模型和单隐藏层网络，无论是使用均方误差（MSE）损失还是交叉熵（CE）损失，Hessian矩阵的对角块和非对角块在随机初始化时就已经表现出明显的差异。
动态力量:在训练过程中，Hessian矩阵的结构会进一步发生变化。特别是在使用CE损失时，训练过程会逐渐消除初始时存在于跨层Hessian分量（Hwv）中的“块循环”（block-circulant）模式，而对角块和非对角块的近块对角结构则保持稳定。这种由训练过程引起的结构变化被称为“动态力量”。

另外一点，文章证明了类别数C是影响Hessian矩阵结构的主要因素。(在实验中，隐藏层Hessian的非对角块与对角块的比值以$\frac{1}{\sqrt{C}}$的速度衰减，输出层Hessian的衰减速率为$\frac{1}{C}$).这个结果对于大模型而言是友好的，因为在神经网络中C的大小往往是$1e3\sim 1e4$级别的，这说明大模型的Hessian是具有强对角块的结构的！这一点实际上在传统的方法上人们意识or无意识的用到了(对角近似)，但是更丰富层面以及基于这个发现的在计算的可行性和性能的Trade-off做的工作是比较少的，在优化器那边做的比较多。

此外还有不少文献揭露了Hessian矩阵具有低秩特征谱，即只有少数的特征值显著大，其余大多接近0.

Blog on Lorn

浅谈低精度浮点数

Motivation

Note

IEEE 754 标准浮点数

FP8

FP4 / MXFP4 / NVFP4

Reference

Triton

Motivation

Introduction

GPU

Weighted Sum

Forward Pass

Backward Pass

Triton Puzzles (Lite)

Puzzles 1/2 Constant Add

Puzzles 3/4 Outer Vector Add

Puzzles 5 Fused Outer Multiplication

Puzzles 6 Fused Outer Multiplication - Backwards

Puzzles 7 Long Sum

Puzzles 8 Long Softmax

Puzzles 9 Simple FlashAttention

Puzzles 10 Two Dimensional Convolution

Puzzles 11 Matrix Multiplication

Puzzles 12 Quantized Matrix Mult

Reference

Transformer中的位置编码的探索与思考

Motivation

Note

为什么Transformer需要位置编码?

位置编码为什么可以解决这个问题？

怎样的位置编码是好的？

为什么RoPE$^{[4]}$比Sinusoidal位置编码更好？

参考

Hessian矩阵系列串讲

定义

一般性质

Hessian在神经网络中的计算

Hessian矩阵在神经网络下的特殊结构