本文記錄了我學(xué)習(xí) AI Infra以 nano-vllm 為切入點(diǎn)過(guò)程中圍繞 LLM 推理機(jī)制產(chǎn)生的幾個(gè)核心問題及思考。起因在學(xué)習(xí) nano-vllm 源碼時(shí)我遇到了一個(gè)看似簡(jiǎn)單但細(xì)想很深的問題KV Cache 為什么只緩存 K 和 V不緩存 Q這個(gè)問題把我拉回到了 Attention 機(jī)制本身進(jìn)而引出了對(duì) Decoder-Only 架構(gòu)逐層數(shù)據(jù)流的完整梳理。問題一KV Cache 為什么只存 K 和 V回到 Attention 公式Attention(Qt,K,V)softmax(Qt?[K1,…,Kt]Td)?[V1,…,Vt]\text{Attention}(Q_t, K, V) \text{softmax}\left(\frac{Q_t \cdot [K_1, \dots, K_t]^T}{\sqrtcj3uvhbzvbv}\right) \cdot [V_1, \dots, V_t]Attention(Qt?,K,V)softmax(d?Qt??[K1?,…,Kt?]T?)?[V1?,…,Vt?]關(guān)鍵不對(duì)稱性矩陣形狀含義QtQ_tQt?[1, d]只有當(dāng)前 1 個(gè)token 的 QueryK1..tK_{1..t}K1..t?[t, d]所有歷史token 的 KeyV1..tV_{1..t}V1..t?[t, d]所有歷史token 的 Value用生成過(guò)程推演假設(shè)已生成 “我 喜歡 吃”要生成第 4 個(gè) token生成第 4 個(gè) token Q? x? · W_q ← 只需要當(dāng)前 1 個(gè) Q Attention softmax(Q? · [K?,K?,K?,K?]? / √d) · [V?,V?,V?,V?] 生成第 5 個(gè) token Q? x? · W_q ← Q? 再也不會(huì)被用到 Attention softmax(Q? · [K?,K?,K?,K?,K?]? / √d) · [V?,V?,V?,V?,V?] ↑ K?~K? 又要用 ↑ V?~V? 又要用結(jié)論Q 是一次性提問者每步只需當(dāng)前一個(gè)用完即棄無(wú)跨步復(fù)用需求K 和 V 是被反復(fù)查詢的知識(shí)庫(kù)每一步生成都要回顧全部歷史必須緩存一個(gè)類比把 Attention 想象成圖書館查資料Q 你手里的一次性提問紙條問完就扔K 每本書封面上的索引標(biāo)簽?zāi)阈枰袝臉?biāo)簽V 書里的實(shí)際內(nèi)容找到后你要讀所有相關(guān)書圖書館KV Cache存的是書KV不是你的紙條Q。不緩存的代價(jià)無(wú) Cache暴力重算有 KV Cache每步計(jì)算量O(t)O(1)生成 n 個(gè) tokenO(n2)O(n)代價(jià)無(wú)額外顯存顯存占用KV Cache 本質(zhì)是用空間換時(shí)間。問題二多層 Decoder 中每層的輸入輸出到底是什么全局?jǐn)?shù)據(jù)流Token IDs: [我, 喜, 歡, 吃] │ ▼ Embedding: [4, 4096] ← 第一層的輸入 │ ▼ ┌─ Layer 1 ─┐ in: [4, 4096] out: [4, 4096] ├─ Layer 2 ─┤ in: [4, 4096] out: [4, 4096] ├─ Layer 3 ─┤ in: [4, 4096] out: [4, 4096] │ ... │ └─ Layer N ─┘ in: [4, 4096] out: [4, 4096] │ ▼ RMSNorm LM Head: [4, 4096] → [4, vocab_size] │ ▼ 取最后一個(gè)位置 → softmax → 下一個(gè) token核心事實(shí)每一層的輸入輸出形狀完全相同都是[seq_len, hidden_dim]。層與層之間變的不是形狀是語(yǔ)義的深度。第一層的輸入Embedding 層的輸出。每個(gè) token 被映射為一個(gè) 4096 維向量此時(shí)只包含該 token 自身的詞義尚未與任何其他 token 交互。每層內(nèi)部做了什么輸入 x: [seq_len, 4096] │ ├──→ RMSNorm → Self-Attention → (殘差) │ │ │ ▼ ├──→ RMSNorm → FFN (SwiGLU) → (殘差) │ ▼ 輸出: [seq_len, 4096]兩個(gè)子模塊各司其職模塊作用類比Self-Attentiontoken 之間交流橫向融合上下文開會(huì)討論FFN/MLP每個(gè) token 獨(dú)立思考縱向深化表示會(huì)后獨(dú)立思考Prefill vs DecodePrefill處理 promptDecode逐 token 生成每層輸入[n, 4096][1, 4096]Q[n, 4096][1, 4096]K, V[n, 4096]新算并存入 cache[1, 4096]新算并追加到 cache參與計(jì)算的 K, V[n, 4096][t, 4096]全部歷史從 cache 讀問題三我的核心感悟逐層語(yǔ)義完善把上面兩個(gè)問題串起來(lái)我得到的理解是每個(gè) token 的向量從孤立的詞義出發(fā)每經(jīng)過(guò)一層通過(guò) Attention 吸收上下文、通過(guò) FFN 深化推理維度不變但語(yǔ)義逐層豐富最終變成一個(gè)濃縮了全文信息、足以預(yù)測(cè)下一個(gè)詞的表示。具體地Embedding 后 吃 只知道自己是吃 Layer 1 后 吃 知道前面是喜歡 Layer 5 后 吃 知道這是主語(yǔ)喜歡吃的結(jié)構(gòu) Layer 20 后 吃 知道語(yǔ)境大概率要接一個(gè)食物名詞 Layer N 后 吃 的向量已經(jīng)準(zhǔn)備好去預(yù)測(cè)蘋果/米飯/火鍋...三個(gè)讓這個(gè)過(guò)程成立的關(guān)鍵設(shè)計(jì)Causal Mask每個(gè) token 只能看到自己 前文保證自回歸生成的合法性殘差連接每層不是推倒重來(lái)而是在上一層基礎(chǔ)上疊加一個(gè)增量修正維度恒定4096 維貫穿始終讓層與層之間可以無(wú)限堆疊這就像畫畫第一層打輪廓后面每層加一點(diǎn)細(xì)節(jié)和色彩最終成為一幅完整的畫。