生成式模型的注意力计算-矩阵视角笔记

从矩阵运算的视角拆解生成式模型的注意力机制，能彻底跳出“黑盒直觉”，看清每一步计算的硬件执行逻辑，这也是大模型推理优化的核心底层依据。

📐 基础定义：Q/K/V 全矩阵化生成

生成式模型的输入是形状为 [seq_len, d_model] 的嵌入矩阵 

𝑋

X，其中每一行对应一个 token 的嵌入向量。

注意力计算的第一步，是通过三个可训练的权重矩阵 

𝑊

𝑞

,

𝑊

𝑘

,

𝑊

𝑣

W

q

,W

k

,W

v

，一次性生成整个序列的查询、键、值矩阵：

𝑄

=

𝑋

⋅

𝑊

𝑞

,

𝐾

=

𝑋

⋅

𝑊

𝑘

,

𝑉

=

𝑋

⋅

𝑊

𝑣

Q=X⋅W

q

,K=X⋅W

k

,V=X⋅W

v

三个权重矩阵的形状均为 [d_model, d_k]，其中 

𝑑

𝑘

d

k

 是单头注意力的维度

这一步本质是三次‌通用矩阵乘法（GEMM）‌，GPU 上可以直接通过 Tensor Core 硬件加速，是整个注意力计算中最容易优化的部分

🔢 注意力分数矩阵的生成

得到 Q/K/V 矩阵后，通过矩阵乘法一次性计算出所有 token 两两之间的注意力分数矩阵 

𝐴

A：

𝐴

=

𝑄

⋅

𝐾

𝑇

𝑑

𝑘

A=

d

k

Q⋅K

T

分数矩阵 

𝐴

A 的形状为 [seq_len, seq_len]，其中第 

𝑖

i 行第 

𝑗

j 列的元素 

𝐴

𝑖

𝑗

A

ij

，代表第 

𝑖

i 个 token 对第 

𝑗

j 个 token 的关注权重

除以 

𝑑

𝑘

d

k

 是为了抵消点积结果随维度增长的方差膨胀，保证 Softmax 输入的数值稳定性

这一步的矩阵乘法是 O(n²) 复杂度，当序列长度超过 4096 时，计算量和显存占用会快速上升

🔒 生成式模型专属：因果掩码矩阵

生成式模型是自回归的，当前 token 不能看到未来位置的信息，这一步通过‌下三角因果掩码矩阵‌强制实现：

构造一个和分数矩阵 

𝐴

A 同形状的掩码矩阵，上三角区域（

𝑗

>

𝑖

j>i 的位置）全部填充为 -inf

将掩码矩阵和分数矩阵逐元素相加

后续 Softmax 计算时，-inf 位置的权重会直接变成 0，彻底屏蔽未来 token 的信息泄露

这是生成式模型和双向注意力模型（如BERT）最核心的矩阵层面差异，也是保证自回归生成合法性的关键约束。

🧮 Softmax 归一化与最终输出矩阵

对掩码后的分数矩阵逐行执行 Softmax 操作，得到最终的注意力权重矩阵 

𝐴

′

A

′

，再和 V 矩阵相乘得到注意力输出矩阵 

𝑂

O：

𝑂

=

𝐴

′

⋅

𝑉

O=A

′

⋅V

输出矩阵 

𝑂

O 的形状和 V 完全一致，每一行对应一个 token 融合了全序列上下文信息的新向量

整个计算链路从输入嵌入矩阵开始，全程没有零散的向量级循环，所有操作都是批量矩阵运算，完美适配 GPU 的并行计算特性

🚀 多头注意力的矩阵合并优化

生成式模型的多头注意力，在工程上会把多个头的 Q/K/V 矩阵直接合并成大矩阵运算，避免多次小矩阵乘法的开销：

把 

𝑊

𝑞

W

q

 拆分为 

ℎ

h 个形状为 [d_model, d_k] 的子矩阵，对应 

ℎ

h 个注意力头

输入嵌入矩阵直接和合并后的大权重矩阵相乘，一次性得到所有头的 Q/K/V 结果

最后把所有头的输出矩阵在特征维度拼接，再通过一个输出投影矩阵 

𝑊

𝑜

W

o

 还原回 [seq_len, d_model] 维度

这种矩阵合并策略，是现代大模型推理框架（vLLM、TensorRT-LLM）实现高吞吐的核心优化手段之一。

需要我为你补充‌多头注意力的完整矩阵运算Python实现代码‌吗？可以直接运行验证每一步矩阵形状的变化。