TikTok GQA：Query Head 分组与 KV Cache 规模

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练h_q个query head、h_kv个KV head且h_q可被h_kv整除，head维d；模型转换/训练按结构另定。

解题路线

1. MHA每query head配一组K/V，MQA全部query共享一组，GQA将query heads分组，每组共享一组K/V。共享不是把query也平均成同一个向量。
2. 每个query head仍计算自身scores和softmax，只是取同组K/V。head到group索引映射必须固定，错误broadcast可让所有heads误读同组或形状对不上。
3. decode KV缓存按layers×batch×tokens×2×h_kv×d×每元素字节；与h_q相等的MHA相比缩减因子h_q/h_kv。query输出维度可仍保持h_qd。
4. 节省主要来自K/V投影及缓存读写，不把整个attention FLOPs都按同因子减少，因为query heads仍各有打分和输出。带宽瓶颈decode常受益，prefill行为要分别评估。
5. 共享可能损失head专属key/value表达，质量需验证；从MHA checkpoint转GQA通常需权重转换及继续训练，直接任意删head不保证功能不变。
6. 测head映射、缓存append位置、不同batch长度和full-vs-incremental输出；量化、分页cache与tensor parallel的布局须配套。

手动走一遍

h_q=32、h_kv=8，每4个query heads共享一组K/V，理论KV字节降为MHA的1/4；仍有32份query attention输出，不是只算8份query。

规模与取舍

cache O(LBT h_kv d)，query-score工作仍含h_q；常规全序列attention二次token关系不因GQA变线性。

易错点与追问

容易失分的边界

• 共享K/V也合并query。
• KV省4倍声称所有FLOPs省4倍。
• h_q/h_kv映射错。

继续深挖

• MQA与GQA质量如何评估？
• checkpoint转换怎样初始化KV？
• 分页cache如何组织groups？