欢迎光临
我们一直在努力

datawhale--llm-algo-leetcode4️⃣

参考链接:笔记

04. Attention MHA GQA

模型架构总览

本节讲解的 Attention 在完整 Transformer Block 中的位置和内部结构:

#mermaid-svg-R5SmpaAwHULlR2n6{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-R5SmpaAwHULlR2n6 .error-icon{fill:#552222;}#mermaid-svg-R5SmpaAwHULlR2n6 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-R5SmpaAwHULlR2n6 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-R5SmpaAwHULlR2n6 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-R5SmpaAwHULlR2n6 .marker.cross{stroke:#333333;}#mermaid-svg-R5SmpaAwHULlR2n6 svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-R5SmpaAwHULlR2n6 p{margin:0;}#mermaid-svg-R5SmpaAwHULlR2n6 .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 .cluster-label text{fill:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 .cluster-label span{color:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 .cluster-label span p{background-color:transparent;}#mermaid-svg-R5SmpaAwHULlR2n6 .label text,#mermaid-svg-R5SmpaAwHULlR2n6 span{fill:#333;color:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 .node rect,#mermaid-svg-R5SmpaAwHULlR2n6 .node circle,#mermaid-svg-R5SmpaAwHULlR2n6 .node ellipse,#mermaid-svg-R5SmpaAwHULlR2n6 .node polygon,#mermaid-svg-R5SmpaAwHULlR2n6 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-R5SmpaAwHULlR2n6 .rough-node .label text,#mermaid-svg-R5SmpaAwHULlR2n6 .node .label text,#mermaid-svg-R5SmpaAwHULlR2n6 .image-shape .label,#mermaid-svg-R5SmpaAwHULlR2n6 .icon-shape .label{text-anchor:middle;}#mermaid-svg-R5SmpaAwHULlR2n6 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-R5SmpaAwHULlR2n6 .rough-node .label,#mermaid-svg-R5SmpaAwHULlR2n6 .node .label,#mermaid-svg-R5SmpaAwHULlR2n6 .image-shape .label,#mermaid-svg-R5SmpaAwHULlR2n6 .icon-shape .label{text-align:center;}#mermaid-svg-R5SmpaAwHULlR2n6 .node.clickable{cursor:pointer;}#mermaid-svg-R5SmpaAwHULlR2n6 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-R5SmpaAwHULlR2n6 .arrowheadPath{fill:#333333;}#mermaid-svg-R5SmpaAwHULlR2n6 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-R5SmpaAwHULlR2n6 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-R5SmpaAwHULlR2n6 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-R5SmpaAwHULlR2n6 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-R5SmpaAwHULlR2n6 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-R5SmpaAwHULlR2n6 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-R5SmpaAwHULlR2n6 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-R5SmpaAwHULlR2n6 .cluster text{fill:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 .cluster span{color:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-R5SmpaAwHULlR2n6 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-R5SmpaAwHULlR2n6 rect.text{fill:none;stroke-width:0;}#mermaid-svg-R5SmpaAwHULlR2n6 .icon-shape,#mermaid-svg-R5SmpaAwHULlR2n6 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-R5SmpaAwHULlR2n6 .icon-shape p,#mermaid-svg-R5SmpaAwHULlR2n6 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-R5SmpaAwHULlR2n6 .icon-shape .label rect,#mermaid-svg-R5SmpaAwHULlR2n6 .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-R5SmpaAwHULlR2n6 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-R5SmpaAwHULlR2n6 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-R5SmpaAwHULlR2n6 :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

输入 x(batch, seq_len, hidden_size)

RMSNormAttention 前

Multi-Head Attention本节重点

残差连接x + Attention(…)

RMSNormMLP 前

SwiGLU MLP

残差连接h + MLP(…)

输出(batch, seq_len, hidden_size)

结构说明: 本节聚焦 A2(Multi-Head Attention)的内部实现,包括 MHA、GQA 和 KV Cache 机制。

Attention 内部架构(同时覆盖 MHA / GQA)

#mermaid-svg-Ij4rbLMaXM44ib6i{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-Ij4rbLMaXM44ib6i .error-icon{fill:#552222;}#mermaid-svg-Ij4rbLMaXM44ib6i .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-Ij4rbLMaXM44ib6i .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-Ij4rbLMaXM44ib6i .marker{fill:#333333;stroke:#333333;}#mermaid-svg-Ij4rbLMaXM44ib6i .marker.cross{stroke:#333333;}#mermaid-svg-Ij4rbLMaXM44ib6i svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-Ij4rbLMaXM44ib6i p{margin:0;}#mermaid-svg-Ij4rbLMaXM44ib6i .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i .cluster-label text{fill:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i .cluster-label span{color:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i .cluster-label span p{background-color:transparent;}#mermaid-svg-Ij4rbLMaXM44ib6i .label text,#mermaid-svg-Ij4rbLMaXM44ib6i span{fill:#333;color:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i .node rect,#mermaid-svg-Ij4rbLMaXM44ib6i .node circle,#mermaid-svg-Ij4rbLMaXM44ib6i .node ellipse,#mermaid-svg-Ij4rbLMaXM44ib6i .node polygon,#mermaid-svg-Ij4rbLMaXM44ib6i .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-Ij4rbLMaXM44ib6i .rough-node .label text,#mermaid-svg-Ij4rbLMaXM44ib6i .node .label text,#mermaid-svg-Ij4rbLMaXM44ib6i .image-shape .label,#mermaid-svg-Ij4rbLMaXM44ib6i .icon-shape .label{text-anchor:middle;}#mermaid-svg-Ij4rbLMaXM44ib6i .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-Ij4rbLMaXM44ib6i .rough-node .label,#mermaid-svg-Ij4rbLMaXM44ib6i .node .label,#mermaid-svg-Ij4rbLMaXM44ib6i .image-shape .label,#mermaid-svg-Ij4rbLMaXM44ib6i .icon-shape .label{text-align:center;}#mermaid-svg-Ij4rbLMaXM44ib6i .node.clickable{cursor:pointer;}#mermaid-svg-Ij4rbLMaXM44ib6i .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-Ij4rbLMaXM44ib6i .arrowheadPath{fill:#333333;}#mermaid-svg-Ij4rbLMaXM44ib6i .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-Ij4rbLMaXM44ib6i .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-Ij4rbLMaXM44ib6i .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Ij4rbLMaXM44ib6i .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-Ij4rbLMaXM44ib6i .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Ij4rbLMaXM44ib6i .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-Ij4rbLMaXM44ib6i .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-Ij4rbLMaXM44ib6i .cluster text{fill:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i .cluster span{color:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-Ij4rbLMaXM44ib6i .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-Ij4rbLMaXM44ib6i rect.text{fill:none;stroke-width:0;}#mermaid-svg-Ij4rbLMaXM44ib6i .icon-shape,#mermaid-svg-Ij4rbLMaXM44ib6i .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Ij4rbLMaXM44ib6i .icon-shape p,#mermaid-svg-Ij4rbLMaXM44ib6i .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-Ij4rbLMaXM44ib6i .icon-shape .label rect,#mermaid-svg-Ij4rbLMaXM44ib6i .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Ij4rbLMaXM44ib6i .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-Ij4rbLMaXM44ib6i .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-Ij4rbLMaXM44ib6i :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

输出

Attention 计算

KV Cache

RoPE 旋转

多头切分

Q/K/V 投影

输入 x(batch, seq_len, d)

W_qLinear(d, h*D)

W_kLinear(d, kv_h*D)

W_vLinear(d, kv_h*D)

Q heads(B, H, S, D)

K heads(B, KV_H, S, D)

V heads(B, KV_H, S, D)

Q_rotated

K_rotated

历史 K

历史 V

Q·K^T / sqrt(D)

Causal Mask

Softmax

加权 V

合并多头(B, S, H*D)

W_oLinear(h*D, d)

GQA 架构对比

#mermaid-svg-12EUGRYgPBtaobcB{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-12EUGRYgPBtaobcB .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-12EUGRYgPBtaobcB .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-12EUGRYgPBtaobcB .error-icon{fill:#552222;}#mermaid-svg-12EUGRYgPBtaobcB .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-12EUGRYgPBtaobcB .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-12EUGRYgPBtaobcB .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-12EUGRYgPBtaobcB .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-12EUGRYgPBtaobcB .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-12EUGRYgPBtaobcB .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-12EUGRYgPBtaobcB .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-12EUGRYgPBtaobcB .marker{fill:#333333;stroke:#333333;}#mermaid-svg-12EUGRYgPBtaobcB .marker.cross{stroke:#333333;}#mermaid-svg-12EUGRYgPBtaobcB svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-12EUGRYgPBtaobcB p{margin:0;}#mermaid-svg-12EUGRYgPBtaobcB .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-12EUGRYgPBtaobcB .cluster-label text{fill:#333;}#mermaid-svg-12EUGRYgPBtaobcB .cluster-label span{color:#333;}#mermaid-svg-12EUGRYgPBtaobcB .cluster-label span p{background-color:transparent;}#mermaid-svg-12EUGRYgPBtaobcB .label text,#mermaid-svg-12EUGRYgPBtaobcB span{fill:#333;color:#333;}#mermaid-svg-12EUGRYgPBtaobcB .node rect,#mermaid-svg-12EUGRYgPBtaobcB .node circle,#mermaid-svg-12EUGRYgPBtaobcB .node ellipse,#mermaid-svg-12EUGRYgPBtaobcB .node polygon,#mermaid-svg-12EUGRYgPBtaobcB .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-12EUGRYgPBtaobcB .rough-node .label text,#mermaid-svg-12EUGRYgPBtaobcB .node .label text,#mermaid-svg-12EUGRYgPBtaobcB .image-shape .label,#mermaid-svg-12EUGRYgPBtaobcB .icon-shape .label{text-anchor:middle;}#mermaid-svg-12EUGRYgPBtaobcB .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-12EUGRYgPBtaobcB .rough-node .label,#mermaid-svg-12EUGRYgPBtaobcB .node .label,#mermaid-svg-12EUGRYgPBtaobcB .image-shape .label,#mermaid-svg-12EUGRYgPBtaobcB .icon-shape .label{text-align:center;}#mermaid-svg-12EUGRYgPBtaobcB .node.clickable{cursor:pointer;}#mermaid-svg-12EUGRYgPBtaobcB .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-12EUGRYgPBtaobcB .arrowheadPath{fill:#333333;}#mermaid-svg-12EUGRYgPBtaobcB .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-12EUGRYgPBtaobcB .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-12EUGRYgPBtaobcB .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-12EUGRYgPBtaobcB .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-12EUGRYgPBtaobcB .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-12EUGRYgPBtaobcB .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-12EUGRYgPBtaobcB .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-12EUGRYgPBtaobcB .cluster text{fill:#333;}#mermaid-svg-12EUGRYgPBtaobcB .cluster span{color:#333;}#mermaid-svg-12EUGRYgPBtaobcB div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-12EUGRYgPBtaobcB .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-12EUGRYgPBtaobcB rect.text{fill:none;stroke-width:0;}#mermaid-svg-12EUGRYgPBtaobcB .icon-shape,#mermaid-svg-12EUGRYgPBtaobcB .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-12EUGRYgPBtaobcB .icon-shape p,#mermaid-svg-12EUGRYgPBtaobcB .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-12EUGRYgPBtaobcB .icon-shape .label rect,#mermaid-svg-12EUGRYgPBtaobcB .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-12EUGRYgPBtaobcB .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-12EUGRYgPBtaobcB .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-12EUGRYgPBtaobcB :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

KV 头数 H → G

KV 头数 G → 1

MQA

Q 头 1

K 头(全部共享)

V 头(全部共享)

Q 头 2

Q 头 …

Q 头 H

GQA

Q 头 1

K 头 1(共享)

V 头 1(共享)

Q 头 2

Q 头 …

K 头 2(共享)

V 头 2(共享)

Q 头 H

MHA

Q 头 1

K 头 1

V 头 1

Q 头 2

K 头 2

V 头 2

Q 头 …

K 头 …

V 头 …

Q 头 H

K 头 H

V 头 H

本节讲解的模型架构模块

模块章节形状变化核心作用
Q 投影 #2.1, #3.1 [B, S, d] → [B, S, H*D] 生成 Query 向量
K/V 投影 #2.1, #3.1 [B, S, d] → [B, S, G*D] GQA 中只生成 G=num_kv_heads 组 K/V
多头切分 #3.1 [B, S, H*D] → [B, H, S, D] 分离注意力头
KV Cache 拼接 #3.2 [B, G, past_len, D] + [B, G, S_q, D] → [B, G, S_kv, D] 复用历史 K/V,避免为历史 token 重复生成 K/V
repeat_kv #3.2 [B, KV_H, S, D] → [B, H, S, D] GQA 扩充 KV 头
注意力分数 #3.3 Q·K^T / sqrt(D) → [B, H, S_q, S_kv] 每个 Query 位置与可见 Key 位置计算相关性
Softmax #3.3 [B, H, S_q, S_kv] 沿最后一个 S_kv 维归一化
加权 V #3.3 probs·V → [B, H, S_q, D] 按注意力权重聚合 Value
输出投影 #3.4 [B, S, H*D] → [B, S, d] 合并多头

数据流详解(含具体数值)

#mermaid-svg-xrArwJiu7QIREQKm{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-xrArwJiu7QIREQKm .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-xrArwJiu7QIREQKm .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-xrArwJiu7QIREQKm .error-icon{fill:#552222;}#mermaid-svg-xrArwJiu7QIREQKm .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-xrArwJiu7QIREQKm .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-xrArwJiu7QIREQKm .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-xrArwJiu7QIREQKm .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-xrArwJiu7QIREQKm .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-xrArwJiu7QIREQKm .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-xrArwJiu7QIREQKm .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-xrArwJiu7QIREQKm .marker{fill:#333333;stroke:#333333;}#mermaid-svg-xrArwJiu7QIREQKm .marker.cross{stroke:#333333;}#mermaid-svg-xrArwJiu7QIREQKm svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-xrArwJiu7QIREQKm p{margin:0;}#mermaid-svg-xrArwJiu7QIREQKm .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-xrArwJiu7QIREQKm .cluster-label text{fill:#333;}#mermaid-svg-xrArwJiu7QIREQKm .cluster-label span{color:#333;}#mermaid-svg-xrArwJiu7QIREQKm .cluster-label span p{background-color:transparent;}#mermaid-svg-xrArwJiu7QIREQKm .label text,#mermaid-svg-xrArwJiu7QIREQKm span{fill:#333;color:#333;}#mermaid-svg-xrArwJiu7QIREQKm .node rect,#mermaid-svg-xrArwJiu7QIREQKm .node circle,#mermaid-svg-xrArwJiu7QIREQKm .node ellipse,#mermaid-svg-xrArwJiu7QIREQKm .node polygon,#mermaid-svg-xrArwJiu7QIREQKm .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-xrArwJiu7QIREQKm .rough-node .label text,#mermaid-svg-xrArwJiu7QIREQKm .node .label text,#mermaid-svg-xrArwJiu7QIREQKm .image-shape .label,#mermaid-svg-xrArwJiu7QIREQKm .icon-shape .label{text-anchor:middle;}#mermaid-svg-xrArwJiu7QIREQKm .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-xrArwJiu7QIREQKm .rough-node .label,#mermaid-svg-xrArwJiu7QIREQKm .node .label,#mermaid-svg-xrArwJiu7QIREQKm .image-shape .label,#mermaid-svg-xrArwJiu7QIREQKm .icon-shape .label{text-align:center;}#mermaid-svg-xrArwJiu7QIREQKm .node.clickable{cursor:pointer;}#mermaid-svg-xrArwJiu7QIREQKm .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-xrArwJiu7QIREQKm .arrowheadPath{fill:#333333;}#mermaid-svg-xrArwJiu7QIREQKm .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-xrArwJiu7QIREQKm .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-xrArwJiu7QIREQKm .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xrArwJiu7QIREQKm .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-xrArwJiu7QIREQKm .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xrArwJiu7QIREQKm .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-xrArwJiu7QIREQKm .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-xrArwJiu7QIREQKm .cluster text{fill:#333;}#mermaid-svg-xrArwJiu7QIREQKm .cluster span{color:#333;}#mermaid-svg-xrArwJiu7QIREQKm div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-xrArwJiu7QIREQKm .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-xrArwJiu7QIREQKm rect.text{fill:none;stroke-width:0;}#mermaid-svg-xrArwJiu7QIREQKm .icon-shape,#mermaid-svg-xrArwJiu7QIREQKm .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xrArwJiu7QIREQKm .icon-shape p,#mermaid-svg-xrArwJiu7QIREQKm .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-xrArwJiu7QIREQKm .icon-shape .label rect,#mermaid-svg-xrArwJiu7QIREQKm .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xrArwJiu7QIREQKm .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-xrArwJiu7QIREQKm .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-xrArwJiu7QIREQKm :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

Step 5: 输出投影#3.4

W_o: (4096, 4096)

Step 4: Attention 计算#3.3

scores: (2, 32, 16, 21)

softmax

output: (2, 32, 16, 128)

Step 3: repeat_kv#3.2

K: (2, 32, 21, 128)

V: (2, 32, 21, 128)

Step 2: 多头切分 + KV Cache示例 past_len=5

Q: (2, 32, 16, 128)

K: (2, 8, 21, 128)

V: (2, 8, 21, 128)

Step 1: Q/K/V 投影#3.1

q_proj.weight: (4096, 4096)PyTorch [out,in]

k_proj.weight: (1024, 4096)PyTorch [out,in]

v_proj.weight: (1024, 4096)PyTorch [out,in]

输入 x(2, 16, 4096)

输出(2, 16, 4096)


0. 学习本章前先认识这些名词

Attention 的实现涉及多头机制、KV Cache 和分组策略,需要先建立对这三个概念的基本理解。

0.1 注意力基础

名词英文简单理解在本节中的作用
注意力 Attention 让模型聚焦相关部分 核心机制
Query Q 当前查询位置用于“匹配”的向量 寻找相关 Key
Key K 每个可见位置用于“被匹配”的向量 与 Query 计算相关性
Value V 每个可见位置提供给聚合的内容向量 根据注意力权重被加权求和
缩放点积 Scaled Dot-Product Q·K^T / sqrt(d) 计算相关性分数

必须先懂: 注意力机制的基本直觉(Q 提问、K 匹配、V 回答)。

0.2 多头机制

名词英文简单理解在本节中的作用
注意力头 Attention Head 独立学习的注意力子空间 并行计算多组 Q/K/V
多头注意力 Multi-Head Attention 多个注意力头并行 MHA 标准实现
分组查询 Grouped Query 多个 Q 头共享一个 K/V 头 GQA 的核心思想
多查询 Multi-Query 所有 Q 头共享一个 K/V 头 MQA 的极端情况

必须先懂: 为什么需要多个注意力头(不同头学习不同类型的依赖关系)。

0.3 推理优化

名词英文简单理解在本节中的作用
KV Cache KV Cache 缓存历史 K/V 避免重复计算 自回归生成加速
自回归 Autoregressive 逐个生成 token KV Cache 的使用场景
因果掩码 Causal Mask 当前 token 只能看到历史 防止信息泄露

必须先懂: 自回归生成的推理流程(生成第

N

N

N 个 token 需要前

N

−

1

N-1

N−1 个 token 的信息)。

0.4 数学符号

符号含义在本节中的作用

H

H

H

Query 头数 num_heads

G

G

G

KV 头数(分组数) num_kv_heads

D

D

D

每个头的维度 head_dim = hidden_dim / H

S

S

S

序列长度 seq_len

g

=

H

/

G

g = H/G

g=H/G

每个 KV 头对应的 Q 头数 GQA 的核心参数

0.5 维度关系

架构KV 头数KV Cache 大小(单层单 token)代表模型
MHA

H

H

H

2

×

H

×

D

×

b

2 \\times H \\times D \\times b

2×H×D×b bytes

原始 Transformer
GQA

G

G

G(

1

<

G

<

H

1 < G < H

1<G<H)

2

×

G

×

D

×

b

2 \\times G \\times D \\times b

2×G×D×b bytes

LLaMA-2 70B、LLaMA 3 等
MQA

1

1

1

2

×

1

×

D

×

b

2 \\times 1 \\times D \\times b

2×1×D×b bytes

部分 MQA 模型

其中 b 是每个缓存元素的字节数,例如 FP16/BF16 通常为 2 bytes,FP32 为 4 bytes。KV Cache 大小必须和实际 dtype 一起计算。


1. 从 MHA 到 GQA:架构演进

1.1 MHA:标准多头注意力

结构:

Q: [batch, seq, H, D]
K: [batch, seq, H, D] ← 每个 Q 头有独立的 K/V
V: [batch, seq, H, D]

Attention = softmax(Q·K^T / sqrt(D)) · V

优点:

  • 每个 Q 头拥有独立 K/V 头,参数与 KV 表示更充分;但不能简单断言所有任务上“一定最强”
  • 标准实现,兼容性好

缺点:

  • KV Cache 显存占用大(与 Q 头数成正比)
  • 解码阶段常受到 KV Cache 读取带宽和容量的显著影响,尤其在长上下文或较大 batch 下

1.2 MQA:极致压缩

结构:

Q: [batch, seq, H, D]
K: [batch, seq, 1, D] ← 所有 Q 头共享 1 个 K/V
V: [batch, seq, 1, D]

K = repeat(K, H) → [batch, seq, H, D]
V = repeat(V, H) → [batch, seq, H, D]

优点:

  • KV Cache 仅为 MHA 的 1/H
  • 推理速度大幅提升

缺点:

  • K/V 头共享程度最高,表示容量更受约束
  • 相比 MHA/GQA 可能带来质量损失,但损失大小依模型和训练方式而定

1.3 GQA:折中方案

结构:

Q: [batch, seq, H, D]
K: [batch, seq, G, D] ← 共 G 个 KV 头;每个 KV 头对应 H/G 个 Q 头
V: [batch, seq, G, D]

K = repeat(K, H/G) → [batch, seq, H, D]
V = repeat(V, H/G) → [batch, seq, H, D]

优点:

  • KV Cache 为 MHA 的 G/H(如 8/64 = 1/8)
  • 保留更多 KV 表达能力,效果损失小

LLaMA-2/3 配置示例:

模型hidden_dimnum_headsnum_kv_heads压缩比
LLaMA-2 7B 4096 32 32 1x(MHA)
LLaMA-2 70B 8192 64 8 8x
LLaMA-3 8B 4096 32 8 4x
LLaMA-3 70B 8192 64 8 8x

2. 核心公式与张量维度

2.1 Scaled Dot-Product Attention

Attention

(

Q

,

K

,

V

)

=

Softmax

(

Q

K

T

d

k

)

V

\\text{Attention}(Q, K, V) = \\text{Softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V

Attention(Q,K,V)=Softmax(dk​

​QKT​)V

为什么除以

d

k

\\sqrt{d_k}

dk​

​?

防止点积结果过大,导致 Softmax 梯度消失:

  • 假设 Q, K 的元素均值为 0,方差为 1
  • Q

    ⋅

    K

    Q \\cdot K

    Q⋅K 的方差 =

    d

    k

    d_k

    dk​

  • 除以

    d

    k

    \\sqrt{d_k}

    dk​

    ​ 使方差回到 1

2.2 张量维度追踪

输入: x: [batch, seq_len, hidden_dim]

步骤 1:线性投影

数学上把每个 token 表示看成行向量:

Q = x @ W_q → [batch, seq_len, num_heads * head_dim]
K = x @ W_k → [batch, seq_len, num_kv_heads * head_dim]
V = x @ W_v → [batch, seq_len, num_kv_heads * head_dim]

其中数学权重可写为 W_q:[d, H*D]、W_k/W_v:[d, G*D]。

PyTorch nn.Linear(in_features, out_features) 内部的 weight 按 [out_features, in_features] 存储,所以代码等价于:

Q = x @ q_proj.weight.T
K = x @ k_proj.weight.T
V = x @ v_proj.weight.T

数学结构没有改变,只是 PyTorch 的权重存储约定不同。

步骤 2:多头切分

xq = xq.reshape(batch, seq_len, num_heads, head_dim).transpose(1, 2)
→ [batch, num_heads, seq_len, head_dim]

xk = xk.reshape(batch, seq_len, num_kv_heads, head_dim).transpose(1, 2)
→ [batch, num_kv_heads, seq_len, head_dim]

步骤 3:KV Cache 拼接(推理时)

if kv_cache:
xk = torch.cat([k_cache, xk], dim=2)
xv = torch.cat([v_cache, xv], dim=2)

步骤 4:repeat_kv 扩充

xk = repeat_kv(xk, num_heads // num_kv_heads)
→ [batch, num_heads, seq_len, head_dim]

步骤 5:注意力计算

scores = xq @ xk.transpose(-2, -1) / sqrt(head_dim)
→ [batch, num_heads, query_len, kv_len]

probs = softmax(scores, dim=-1)
output = probs @ xv
→ [batch, num_heads, seq_len, head_dim]

步骤 6:输出投影

output = output.transpose(1, 2).reshape(batch, seq_len, -1)
→ [batch, seq_len, num_heads * head_dim]

output = output @ W_o
→ [batch, seq_len, hidden_dim]

数学上 W_o:[H*D, d];PyTorch o_proj.weight:[d, H*D],实际仍是 output @ o_proj.weight.T。


理论结构 vs 本章教学代码: 真实 LLaMA/Qwen 类模型通常会在 Q/K 切头之后对 Q/K 应用 RoPE,并把旋转后的 K写入 KV Cache。本章下面的最小代码为了单独练习 MHA/GQA/KV Cache,没有实现 RoPE;不要把这个省略理解为真实 LLM 没有位置编码。

3. 动手实战

3.1 任务要求

请补全 GroupedQueryAttention.forward 函数中的 TODO 部分。

import torch
import torch.nn as nn
import math
import torch.nn.functional as F

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) –> torch.Tensor:
"""将 KV 头复制 n_rep 次以匹配 Query 头数"""
batch, num_kv_heads, slen, head_dim = hidden_states.shape
if n_rep == 1:
return hidden_states
hidden_states = hidden_states.unsqueeze(2).expand(–1, –1, n_rep, –1, –1)
return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)

class GroupedQueryAttention(nn.Module):
def __init__(self, hidden_dim: int, num_heads: int, num_kv_heads: int = None):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads

assert num_heads % self.num_kv_heads == 0
assert hidden_dim % num_heads == 0, "hidden_dim 必须能被 num_heads 整除"

self.num_queries_per_kv = self.num_heads // self.num_kv_heads
self.head_dim = hidden_dim // num_heads

self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(num_heads * self.head_dim, hidden_dim, bias=False)

def forward(
self,
x: torch.Tensor,
attention_mask: torch.Tensor = None,
kv_cache: tuple[torch.Tensor, torch.Tensor] = None,
):
batch_size, seq_len, _ = x.shape

# 线性投影
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)

# ==========================================
# TODO 1: Reshape xq, xk, xv 以适配多头
# 提示:
# 1. reshape(batch, seq_len, num_heads, head_dim)
# 2. transpose(1, 2) 将 heads 维移到第 2 维
# 3. xk, xv 使用 num_kv_heads
# ==========================================
# xq = ???
# xk = ???
# xv = ???

# ==========================================
# TODO 2: 处理 KV Cache
# 提示:
# 1. 如果有 cache,在 dim=2 (seq_len 维) 拼接历史 K/V
# 2. 先拼接,再 repeat_kv
# ==========================================
if kv_cache is not None:
k_cache, v_cache = kv_cache
# xk = ???
# xv = ???

new_kv_cache = (xk, xv)

# 扩充 KV 头数
xk = repeat_kv(xk, self.num_queries_per_kv)
xv = repeat_kv(xv, self.num_queries_per_kv)

# ==========================================
# TODO 3: 计算注意力分数
# 提示:
# 1. scores = xq @ xk.transpose(-2, -1) / sqrt(head_dim)
# 2. 添加 attention_mask(广播)
# 3. probs = F.softmax(scores, dim=-1)
# 4. output = probs @ xv
# ==========================================
# scores = ???
if attention_mask is not None:
scores = scores + attention_mask

# probs = ???
# output = ???

# ==========================================
# TODO 4: 恢复形状
# 提示:
# 1. transpose(1, 2) 恢复 [B, S, H, D]
# 2. reshape(batch, seq_len, -1) 合并多头
# ==========================================
# output = ???

return self.o_proj(output), new_kv_cache

3.2 测试代码

def test_mha_mqa_gqa():
try:
batch_size, seq_len, hidden_dim, num_heads = 2, 16, 128, 4

# 测试 1:MHA
print("Testing MHA (Multi-Head Attention)…")
mha = GroupedQueryAttention(hidden_dim, num_heads, num_kv_heads=num_heads)
x = torch.randn(batch_size, seq_len, hidden_dim)
out, _ = mha(x)
assert out.shape == (batch_size, seq_len, hidden_dim)

# 测试 2:GQA
print("Testing GQA (Grouped-Query Attention)…")
gqa = GroupedQueryAttention(hidden_dim, num_heads, num_kv_heads=2)
out, _ = gqa(x)
assert out.shape == (batch_size, seq_len, hidden_dim)

# 测试 3:KV Cache
print("Testing KV Cache…")
prefill_len = 5
x_prefill = torch.randn(batch_size, prefill_len, hidden_dim)
_, kv_cache = mha(x_prefill)

x_decode = torch.randn(batch_size, 1, hidden_dim)
out_decode, new_kv_cache = mha(x_decode, kv_cache=kv_cache)
assert new_kv_cache[0].shape[2] == prefill_len + 1

print("\\n✅ All Tests Passed! Attention 算子实现通过测试。")

except NotImplementedError:
print("请先完成 TODO 部分的代码!")
raise

test_mha_mqa_gqa()


4. 参考代码与解析

4.1 完整实现

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) –> torch.Tensor:
"""将 KV 头复制 n_rep 次以匹配 Query 头数"""
batch, num_kv_heads, slen, head_dim = hidden_states.shape
if n_rep == 1:
return hidden_states
hidden_states = hidden_states.unsqueeze(2).expand(–1, –1, n_rep, –1, –1)
return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)

class GroupedQueryAttention(nn.Module):
def __init__(self, hidden_dim: int, num_heads: int, num_kv_heads: int = None):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads

assert num_heads % self.num_kv_heads == 0
assert hidden_dim % num_heads == 0, "hidden_dim 必须能被 num_heads 整除"

self.num_queries_per_kv = self.num_heads // self.num_kv_heads
self.head_dim = hidden_dim // num_heads

self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(num_heads * self.head_dim, hidden_dim, bias=False)

def forward(
self,
x: torch.Tensor,
attention_mask: torch.Tensor = None,
kv_cache: tuple[torch.Tensor, torch.Tensor] = None,
):
batch_size, seq_len, _ = x.shape
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)

# TODO 1: Reshape 为多头形式
xq = xq.reshape(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
xk = xk.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
xv = xv.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

# TODO 2: 处理 KV Cache
if kv_cache is not None:
k_cache, v_cache = kv_cache
xk = torch.cat([k_cache, xk], dim=2)
xv = torch.cat([v_cache, xv], dim=2)

new_kv_cache = (xk, xv)

# 扩充 KV 头数
xk = repeat_kv(xk, self.num_queries_per_kv)
xv = repeat_kv(xv, self.num_queries_per_kv)

# TODO 3: Scaled Dot-Product Attention
scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)

if attention_mask is not None:
scores = scores + attention_mask

probs = F.softmax(scores, dim=–1)
output = torch.matmul(probs, xv)

# TODO 4: 恢复形状
output = output.transpose(1, 2).reshape(batch_size, seq_len, –1)

return self.o_proj(output), new_kv_cache

5.实现理解

类型Q Head 数K/V Head 数一个 KV Head 给几个 Q HeadQ shape原始 K/V shape核心特点
MHA (H) (G=H) (H/G=1) [B,H,S,D] [B,H,S,D] 每个 Q Head 有自己的 K/V Head
GQA (H) (1<G<H) (H/G) [B,H,S,D] [B,G,S,D] 一组 Q Heads 共享一个 K/V Head
MQA (H) (G=1) (H) [B,H,S,D] [B,1,S,D] 所有 Q Heads 共享同一个 K/V Head

在这里插入图片描述

在这里插入图片描述 在这里插入图片描述

 

赞(0)
未经允许不得转载:171主机测评 » datawhale--llm-algo-leetcode4️⃣
分享到: 更多 (0)

评论 抢沙发

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址