GQLA: ハードウェア適応型大規模言語モデルデコードのためのグループクエリ潜在アテンション

要旨

マルチヘッド潜在注意（MLA）は、DeepSeek-V2/V3で採用されている注意機構であり、キーと値を低ランク潜在表現に統合圧縮し、H100のルーフラインにほぼ完全に適合する。しかし、その学習済み重みが公開する復号パスは、吸収型MQA形式という1系統のみであり、効率的な推論をH100クラスの計算帯域比に依存させ、ヘッド軸に沿ったテンソル並列化を無効化し、輸出規制対象のH20などの汎用推論GPUではマルチトークン予測（MTP）による利得が得られない。本稿では、MLAの最小限の修正であるグループクエリ潜在注意（GQLA）を提案する。GQLAの学習済み重みは、同一パラメータ上で代数的に等価な2つの復号パス、すなわちMLAと同一のMQA吸収パスと、グループごとにキャッシュを拡張したGQAパスを公開する。実行時には、対象ハードウェアに適合するパスを選択するため、再学習もカスタムカーネルも不要であり、単一のGQLA重みセットでH100（MQA吸収、s_q=1）とH20（GQA+MTP、s_q=2）の両方のルーフラインに適合し、かつGQAパス上で最大8方向の冗長性ゼロテンソル並列化をサポートする。ゼロからの事前学習を避けるため、TransMLAをTransGQLAに拡張し、事前学習済みGQAチェックポイントをGQLAモデルに変換する。LLaMA-3-8Bにおいて、この変換によりMQA吸収パス上のトークンあたりKVキャッシュがGQAベースラインの28.125%に圧縮され、かつグループパス上のGQAレベルのトラフィックが構造的に維持される。

English

Multi-head Latent Attention (MLA), the attention used in DeepSeek-V2/V3, jointly compresses keys and values into a low-rank latent and matches the H100 roofline almost perfectly. Its trained weights, however, expose only one decoding path - an absorbed MQA form - which ties efficient inference to H100-class compute-bandwidth ratios, forfeits tensor parallelism along the head axis, and yields no Multi-Token Prediction (MTP) gain on commodity inference GPUs such as the export-restricted H20. We propose Group-Query Latent Attention (GQLA), a minimal modification of MLA whose trained weights expose two algebraically equivalent decoding paths over the same parameters: an MQA-absorb path identical to MLA's, and a GQA path with a per-group expanded cache. The runtime picks the path that matches the target hardware - no retraining, no custom kernels - so a single set of GQLA weights pins the rooflines of both H100 (MQA-absorb, s_q=1) and H20 (GQA + MTP, s_q=2), while supporting up to 8-way zero-redundancy tensor parallelism on the GQA path. To avoid pretraining from scratch we extend TransMLA into TransGQLA, which converts a pretrained GQA checkpoint into a GQLA model; on LLaMA-3-8B it compresses the per-token KV cache to 28.125% of the GQA baseline on the MQA-absorb path while structurally preserving GQA-level traffic on the per-group path.