AI 推理术语表

模型架构

注意力 head dimension

也称为 attention head dimension、head dimension、head_dim、head dim、注意力头维度

先用大白话

注意力 head dimension 是每个注意力头使用的向量宽度,会影响计算量、缓存占用,以及它与加速器矩阵 tile 的匹配程度。

技术定义

注意力 head dimension 是每个头的 query、key 或 value 表示的特征宽度;并非所有架构中这几个宽度都相等。

工程细节

矩阵引擎按固定 tile 形状计算,kernel 可以通过 padding、打包或改变布局来适配注意力维度。因此,同一个小于矩阵阵列宽度的 head,在一种布局中可能浪费计算,在另一种映射中却能更充分地利用阵列。Head 数、序列长度、精度、切分方式,以及 query-key 与 value 各自的维度都会影响实际利用率。

为什么重要

有利于模型质量或 GPU kernel 的维度,在 TPU 上可能需要不同布局。用 head dimension 除以阵列宽度,不能通用地估计整模型吞吐量,甚至不能概括所有注意力矩阵乘法的利用率。

如何在 InferenceX 中解读

TPU AgentX 文章讨论了 Muse Spark Open 的 head dimension 64,以及模型形状与 TPU 矩阵硬件的匹配。应结合具体 kernel 和布局理解这一例子,不应把维度比直接当作整模型的性能预测。

在真实基准中理解这一概念

全部文章