模型架构
注意力 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 和布局理解这一例子,不应把维度比直接当作整模型的性能预测。