Day 04 我們把 The answer to 2 + 3 is 送進 Qwen3 0.6B Base，最後得到 [1, 9, 151936] 的 logits。意思是 9 個 token 各自都有 151,936 個原始分數，只是模型真正要生成下一個 token 時，關心的是最後一個位置。那 151,936 個分數並不會自己排列成一句話，程式還得從中選出一個 token，把它接回原來的文字，然後再算一次。

假如仍沿用昨天的 greedy decoding，選擇規則只會有一行：

next_token = logits[:, -1, :].argmax(
    dim=-1,
    keepdim=True,
)

-1 取出序列的最後一個位置，argmax 則找出分數最高的 token ID。在昨天的例子裡，這一步選中了代表 5 的 token。接著我們把它接到輸入後面，原本的 9 個 token 變成 10 個，再讓模型預測第 11 個。如此反覆下去，才會得到畫面上看似連續出現的一段文字。

generated = input_ids

for _ in range(max_new_tokens):
    with torch.inference_mode():
        logits = model(
            input_ids=generated,
            use_cache=False,
        ).logits

    next_token = logits[:, -1, :].argmax(
        dim=-1,
        keepdim=True,
    )
    generated = torch.cat(
        [generated, next_token],
        dim=1,
    )

這就是 autoregressive generation，模型不會一次輸出完整答案，而是用目前所有內容預測下一個 token，再把自己的輸出變成下一輪輸入。即使換成能寫出數千個 token 推理軌跡的 reasoning model，最底下仍然是同一套逐 token 生成。逐 token 只描述解碼方式，模型會否推理，仍仰賴它經過什麼樣的訓練，以及它在生成途中如何使用計算量。

然而，上面的程式雖然直觀，卻有種笨牛精神：生成第一個 token 時，它讀過完整 prompt；第二次又把完整 prompt 和剛生成的 token 從頭讀一遍；到了第 64 次，前面已經算過 63 次的部分依然全部重算。句子愈長，模型揹回起點的行李也就愈多。

會發生這件事，是因為下一個 token 確實需要參考前文。Transformer 的 attention 會替每個位置產生 Query、Key 和 Value。若暫時用比較口語的說法，Query 表示現在這個位置想尋找什麼，Key 用來和它比對，Value 則是比對之後真正匯入的內容；Causal attention 又限制了每個位置只能看見自己與先前的 token，不能偷看尚未生成的未來。

當新的 token 接到句尾時，舊 token 的 Key 和 Value 並沒有跟著改變。未來的 token 還會需要它們，舊的 Query 卻不必再拿出來使用，因為新的 attention 輸出會由新 token 的 Query 重新查詢所有 Key。KV cache 所做的事情就是把各層已算過的 Key 和 Value 留下來，免得每次生成都假裝自己第一次見到整段文字。Sebastian Raschka 的 KV cache 說明也特別指出，cache 保存的是每一層的 Key 與 Value，而不是把模型的答案預先存起來。

有了 cache 之後，生成會分成兩個節奏。第一次仍要處理完整 prompt，這一段通常叫做 prefill；從第二次開始進入 decode，每一輪只送入剛生成的那個 token，同時把前一輪留下的 past_key_values 傳回模型。

generated = input_ids
next_input = input_ids
past_key_values = None

for _ in range(max_new_tokens):
    with torch.inference_mode():
        output = model(
            input_ids=next_input,
            past_key_values=past_key_values,
            use_cache=True,
        )

    next_token = output.logits[:, -1, :].argmax(
        dim=-1,
        keepdim=True,
    )
    past_key_values = output.past_key_values
    generated = torch.cat(
        [generated, next_token],
        dim=1,
    )
    next_input = next_token

next_input = next_token 是整段程式的轉折點。第一輪的輸入仍是完整 prompt，往後的模型輸入卻只剩 [1, 1]，因為其餘內容已經以 Key 和 Value 的形式留在 cache 裡。Hugging Face 的 Transformers cache 文件也採用相同概念，預設的 Dynamic Cache 會隨著生成逐步增長。

為了讓這個差異在小模型上也看得出來，我把 The quick brown fox jumps over the lazy dog. 重複 192 次，最後再接回昨天的加法句，組成一段 1,929-token prompt。這並非能啟動複雜推理的提示詞，僅是一段夠長、每次又能保持一樣的 context。模型接著生成 64 個 token，實際記錄到的輸入形狀如下：

step 1:  model input [1, 1929], cache length 1929
step 2:  model input [1, 1],    cache length 1930
step 64: model input [1, 1],    cache length 1992

第 64 步的 cache 長度是 1,992，因為在生成第 64 個 token 以前，原本的 1,929-token prompt 已經多了 63 個新 token。我特別留下這道簡單加法，因為它很適合用來抓 cache 的錯位；只要位置少一格或多一格，模型讀到的上下文就可能不再是原來那一段。

在比較速度以前，我先確認三種寫法產生相同結果：Transformers 內建的 generate()、每次重算完整 prefix 的迴圈，以及上面保留 KV cache 的迴圈。三條路徑生成的 64 個 token 完全相同。這裡再次使用 greedy decoding，理由仍然是排除抽樣造成的差異；如果連生成內容都不同，接下來的快慢比較也就失去共同基準。

接著才是計時。我使用同一張 RTX 3090 Ti，以 BF16、batch size 1 執行，各跑三次取中位數，結果如下：

tokens identical: true

cache-free: 12,269.5 ms
KV cache:    4,430.0 ms
speedup:     2.77x

peak allocated VRAM:
2,507,531,264 -> 2,082,331,136 bytes

在這段 1,929-token context 上，KV cache 將生成時間縮短到原來的大約 36%。短 prompt 尤其容易被其他固定成本淹沒，因此我沒有把先前零星測過、卻沒有完整保留下來的短序列數字也塞進文章裡湊熱鬧。

記憶體結果也值得一提，KV cache 本身當然會占用空間，而且 context 愈長，保存的 Key 和 Value 通常也愈多；它不是一個讓記憶體憑空消失的技巧。今天量到的 peak allocated VRAM 反而下降，是因為無 cache 的簡單寫法每一步都產生整段序列的 logits，而 cache 版本在 decode 階段只產生新位置的 logits。這是我們本日這份實作與工作量共同完成的結果，並非開啟 KV cache 總是比較省 VRAM。

同樣地，cache 也沒有讓 attention 變成毫無限額，新 token 的 Query 仍要和保留下來的 Key 互動，模型依舊得讀取逐漸變長的前文。它省掉的是舊 token 一次又一次重做的投影與中間計算，而不是消除長上下文本身的成本。

昨天那 151,936 個 logits，現在已經變成完整的過程：選出一個 token、接回序列、更新 attention 狀態，再進入下一輪。今天故意把每次選擇固定成 argmax，才有辦法乾淨地確認兩條生成路徑沒有走散；但現代模型通常不會永遠只拿分數最高的那一個，如我 Day 02 所說，greedy 的作法已經過時了。明天我們把隨機性放回來，看看 temperature、top-p 與 seed 如何讓同一組權重說出不同的答案，也看看所謂「創意想法」和「胡說八道」之間間隔多遠。