テキスト生成と KV キャッシュ

モジュール 01 第 3 課では DeepSeek で温度の実験をしましたが、そのときモデルはブラックボックスでした。今のモデルは自分で学習させたものなので、開けて各ステップで実際に何をしているのか見られます。

python generate.py

モデルが出すのは確率だけ

モデルは文章を読み終えると、最後の位置で語彙表の 6289 文字それぞれにスコアを付け、softmax で確率に変えます。

== 1. 读完一段文字，模型给下一个字的概率（前 8 名）
  「白日依山尽，黄河入海」→ 无 6.9%  间 4.5%  多 4.3%  深 3.7%  边 3.0%  空 2.6%  难 2.3%  遥 2.1%
  「床前明月光，疑是地上」→ 天 3.2%  花 2.9%  楼 2.4%  清 1.8%  看 1.7%  游 1.7%  人 1.6%  枝 1.3%
  「春风」→ 。 14.2%  ， 5.1%  吹 3.3%  满 2.9%  雨 1.9%  起 1.9%  尽 1.4%  落 1.3%

「黄河入海」（黄河は海に入って）の後は、原詩では「流」（流れる）ですが、「流」は上位 8 位に入っていません。「疑是地上」（地上の……かと疑う）の後の「霜」（霜）も入っていません。前の課で確かめたとおり、このモデルは暗唱しません。この二つの有名な詩を覚えておらず、学んだ規則に沿って「妥当な」文字を出しているだけです。

ついでに面白い発見を一つ。『全唐詩』に収められている「静夜思」は「床前看月光，疑是地上霜。举头望山月，低头思故乡」で、教科書の「床前明月光……举头望明月」とは違います。私たちになじみのある版は、明・清の時代の選集に由来するというのが一般的な見方です。ですから、たとえモデルがこの詩を暗唱していたとしても、暗唱するのは「看月光」のほうなのです。

「春风」（春風）の後で最も確率が高いのは句点（14.2%）です。学習データで数えてみると、「春风」は合計 797 回現れ、そのうち 111 回はすぐ後に句点が続いていて 13.9%、モデルが出した確率とほぼ同じでした。モデルが学んだのは、こうした文字の唐詩での現れ方の規則なのです。

文章を生成するとは、同じことを繰り返すことです。確率を受け取り、文字を一つ選び、後ろにつなげ、次の文字の確率を計算する。どう「選ぶ」かが、何が書かれるかを決めます。

温度

モジュール 01 で温度を扱いました。スコアを温度で割ってから softmax を取ります。温度が 1 より小さいと高いスコアと低いスコアの差が広がり、モデルはより「保守的」になります。1 より大きいと差が縮まり、モデルはより「大胆」になります。

logits = logits[:, -1, :] / temperature
...
next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)

同じ乱数シードで、三つの温度でそれぞれ 3 首書かせます。

  温度 0.3：
    白云无限意，白发不如何。何事无人在，无人不可知。
    一望东山路，千行万里愁。夜深人不见，秋色月无声。
    江南春水上，江上月明时。独有清风起，何人别有期。
  温度 1.0：
    白露起春光，知君得舞酲。霜时飏先没，桐死势悠扬。
    白发摩牛放，真关万趾兴。夜寒休绕郡，秋杀杜娘宫。玉镜收新镜，金尊坐晚麕。何言吊鱼道，双锡空余年。
    戍楼映江岸，淼危棹度斜。欲分凝炯娩，归盖乱精冰。
  温度 1.5：
    淡露依春似映兹，襞舞凋年雪路揉筝衬情。死势悠扬双觉处，摩牛放将溉噭胞。……

温度 0.3 が最もなめらかで、「一望东山路，千行万里愁。夜深人不见，秋色月无声」（ひとたび東の山道を望めば、千行に万里の愁い。夜は深く人は見えず、秋の色、月は声もない）は本物と見まがうほどです。代償は、使う字がありふれていて、繰り返しも多いことです。「无人」（人はいない）が 2 回、「江」（川）が 2 回出てきて、たくさん読むとどれも似たように感じられます。

温度 1.0 は、モデルが出した元の確率どおりに字を選ぶことです。使う字は豊かになりますが、「麕」「酲」「炯娩」といった珍しい字が出始め、意味もつながりにくくなります。

温度 1.5 では完全に崩れます。形式が壊れ（「襞舞凋年雪路揉筝衬情」は一句が 9 文字）、珍しい字だらけです。温度が高いと確率の分布が平らになり、6000 以上ある文字のうち、もともと確率がきわめて低かったものも、合計するとかなりの確率を持っていくからです。一度おかしな字が選ばれると、その後の文脈がおかしくなり、誤りがどんどん積み重なっていきます。

top-k：上位の字の中からだけ選ぶ

温度が高いときの問題は「ロングテール」です。確率の低い何千もの字が、合計するとあまりにも大きな割合を占めてしまいます。top-k のやり方は単純です。確率が最も高い k 個の字だけを残し、残りはすべて捨て、その k 個の中から確率に従って選びます。

if top_k is not None:
    kth = torch.topk(logits, top_k).values[:, -1:]
    logits = logits.masked_fill(logits < kth, float("-inf"))

同じく温度 1.5 に、top-k=20 を加えます。

== 3. 温度 1.5，但只从概率最高的 20 个字里选（top-k=20）
    白头白马在西陵，自是西林下路尘。今日不辞青史士，可堪回首白苹州。
    万国风尘里，千金雨色深。不知归未到，应似到人中。
    东北山川尽不愁，三千春尽一枝愁。春寒未入千门色，雨尽空闻九夜寒。……

形式が戻り、珍しい字も消えましたが、温度 0.3 より変化があります。高い温度が「大胆さ」を、top-k が「度を越さないこと」を受け持ち、二つを一緒に使うのが最も効果的です。

もう一つよく使われる方法が top-p（核サンプリング）です。個数を固定するのではなく、最も高いものから順に足していき、累積の確率が p（たとえば 0.9）に達したところで止めます。確率が集中しているときは数文字だけが残り、分散しているときはもっと多く残ります。モジュール 01 で API を使ったときに、このパラメータを見ました。

続きの生成と折句

生成の書き出しは空でなくてもかまいません。前半の句を与えれば、モデルはその続きを書きます。

== 4. 给出开头，让它续写
    床前明月光，已有四邻情。知君得舞袖，何奈在山期。
    大漠孤烟直，高楼楚水清。雁雕摩露下，鸿劒鼓歌中。夜雨初分郡，秋风又过巴。此才应不遇，空与钓鱼舟。
    人工智能好，言蠹不见人。徒思一岁月，独在岘生春。

「大漠孤烟直」（果てしない砂漠に一筋の煙がまっすぐ立ちのぼる）は、8 句そろった完全な五言律詩に続けられました。「人工智能」（人工知能）という唐の時代にはなかった言葉からも、同じように続けられます。この 4 文字はそれぞれ唐詩に出てきたことがあるからです。

さらに進んで、生成の途中に手を加えることもできます。折句（藏头诗）は、各句の最初の文字が指定されたものでなければなりません。まず指定された文字を入れ、それからモデルに 1 文字ずつその句を書き終えさせます。書いている間は句読点と改行を出させず、最後に私たちが句読点を付けます。

@torch.no_grad()
def acrostic(heads, n_char=7, temperature=0.8):
    ids = encode("\n")
    for i, head in enumerate(heads):
        ids = torch.cat([ids, encode(head)], dim=1)  # 先把指定的字放进去
        for _ in range(n_char - 1):  # 再让模型一个字一个字地写完这一句
            logits = model(ids)[0][0, -1] / temperature
            logits[BANNED] = float("-inf")
            next_id = torch.multinomial(F.softmax(logits, dim=-1), 1)
            ids = torch.cat([ids, next_id[None]], dim=1)
        ids = torch.cat([ids, encode("，。"[i % 2])], dim=1)  # 标点由我们加，保证格式
    return decode(ids[0].tolist()).strip()

== 5. 藏头诗：每句开头的字由我们指定，其余由模型写
    春色欲依春似春，眠花先得舞花年。不知诗句先酬情，觉后还随楚水人。
    学得重雕出，无情忘却回。止行空向兴，境静更闻闻。

各句の書き出しをつなげると、「春眠不觉」（春の眠りに気づかず）と「学无止境」（学びに終わりはない）になります。

このテクニックには正式な名前があります。制約付きデコーディングです。文字を選ぶ前に、要件を満たさない選択肢のスコアを負の無限大にするのです。モジュール 02 第 4 課で DeepSeek に JSON を出力させたとき、サーバー側が出力の形式が正しいことを保証していたのも、同じ原理です。各ステップで、JSON を不正にしてしまうトークンをすべて除外していたのです。

重複した計算

今度は生成の効率を見ます。generate メソッドに戻ると、最も単純な書き方は、各ステップで文章全体をモデルに送り込むものです。

logits, _ = self(idx)  # 不用缓存：每一步都把整段序列从头算一遍

50 番目の文字を書くとき、モデルは前の 49 文字すべてを処理しなければなりません。51 番目の文字を書くときには、この 49 文字をもう一度処理し直し、さらに 50 番目の文字も処理します。

しかしよく考えてみると、前の 49 文字の計算結果は変わっていません。因果マスクがあるので、各位置は前しか見ず、後ろに新しい文字が加わっても前の位置の計算にはまったく影響しません。本当に新しいのは、最後の位置だけです。

最後の位置でアテンションを計算するには何が必要でしょうか。自分自身の q と、前のすべての位置の k と v です。ですから各層の k と v を保存しておき、各ステップでは新しい文字の q、k、v だけを計算し、新しい k と v を保存したものの後ろにつなげれば十分です。これが KV キャッシュです。

実装

gpt.py のアテンションでは、キャッシュのために数行増えるだけです。

if cache is not None:  # KV 缓存：把之前算过的 K、V 接在前面，不用重算
    if "k" in cache:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v

注意すべき細部が一つあります。因果マスクです。キャッシュを使わないとき、q と k は同じ長さで、マスクは正方形の下三角です。キャッシュを使うとき、q は新しく来た 1 文字だけなのに、k にはすべての文字があり、しかも新しい文字は最後にあって、すべての k を見られます。ですからマスクをずらす必要があります。

total = k.size(2)
mask = torch.ones(T, total, dtype=torch.bool, device=x.device).tril(diagonal=total - T)

位置の埋め込みも合わせて調整する必要があります。新しい文字は 0 番目の位置ではなく、start 番目の位置にあります。これが forward の start 引数の役割です。

生成のとき、最初のステップではプロンプト全体を送り込んで、すべての位置の k と v をキャッシュに保存します。その後は各ステップで最新の 1 文字だけを送ります。

if use_cache:
    # 第一步把整段提示词送进去；之后每步只送最新的一个词元，前面的 K、V 从缓存里取
    logits, _ = self(idx[:, start:], caches=caches, start=start)
    start = idx.size(1)

測ってみる

== 6. KV 缓存：结果一样吗？快了多少？
  同样的随机种子，写 120 个字：用缓存和不用缓存的结果完全一样
  写  32 个字：不用缓存     23 毫秒，用缓存    12 毫秒，快了 1.8 倍
  写  64 个字：不用缓存     52 毫秒，用缓存    26 毫秒，快了 2.0 倍
  写 120 个字：不用缓存    130 毫秒，用缓存    49 毫秒，快了 2.6 倍
  这个模型每个词元的缓存：2 × 4 层 × 128 维 × 4 字节 = 4 KB，128 个词元共 512 KB

まず正しさを確かめます。同じ乱数シードなら、キャッシュを使っても使わなくても、書かれる 120 文字はまったく同じです。キャッシュは重複した計算を省くだけで、結果は変わりません。

次に速度です。書く文字が多いほど、多く省けます。32 文字なら 1.8 倍、120 文字なら 2.6 倍速くなりました。キャッシュを使わないと、n 番目の文字を書くのに n 文字を処理する必要があり、合計の計算量は長さの二乗で増えます。キャッシュを使えば、各ステップで 1 文字を処理するだけです。

私たちのモデルでは、高速化は 2 倍あまりで、それほど驚くほどではありません。モデルが小さすぎ、文章が短すぎるからです。各ステップの本当の計算は 1 ミリ秒にも満たず、Python 自体が関数を呼んだりテンソルをつなげたりするオーバーヘッドが大きな割合を占めていて、その部分はキャッシュでは省けません。LLM では各ステップの計算量がずっと大きく、文章もたびたび数千、数万トークンになるので、KV キャッシュが省く量は桁違いの差になります。

キャッシュの代償：メモリ

KV キャッシュはメモリを時間と引き換えにしています。トークンごとに、各層で k と v を一つずつ保存する必要があります。私たちのモデルならトークンあたり 4 KB だけで、128 トークンで 512 KB、取るに足りません。

LLM では話が違います。数十層あり、各層が数千次元のモデルでは、トークンあたりのキャッシュが数百 KB から 1 MB 以上になることもあります。十数万トークンの長い会話なら、キャッシュだけで数十 GB の GPU メモリを占め、モデルそのものより大きくなることもよくあります。多くのユーザーに同時にサービスを提供するとき、これが最大の GPU メモリの消費になります。

ですから前の課の最後で触れた改良の多くは、KV キャッシュを狙ったものです。複数のヘッドで k と v を共有する（グループ化クエリアテンション）、あるいは k と v をより小さなベクトルに圧縮してから保存する（DeepSeek のマルチヘッド潜在アテンション、MLA）。第 10 モジュールで vLLM を扱うときには、サーバーがこうしたキャッシュをどう管理しているかも見ます。

モジュール 06 第 4 課で扱った「キャッシュにヒットした入力は安い」は、別のレベルのキャッシュです。サービス提供者はよく使われる前置き（たとえば固定の system プロンプト）の KV キャッシュを保存しておき、次のリクエストが同じ前置きを使うときは、計算し直さずにそのまま取り出すので、ずっと安くできるのです。

練習問題

1. top-p サンプリングを実装してください。確率を大きい順に並べ、累積の確率が p に達するまでの字を残します。温度 1.5、p=0.9 で何首か書き、top-k=20 と比べてください。
2. 「決まった形式」の制約付きデコーディングを書いてください。七言絶句だけを生成させ、8、16、24 番目の文字は読点か句点だけ、ほかの位置には句読点を出させないようにします。
3. gpt.py の block_size を 512 に変えて学習し直し（あるいはランダムに初期化するだけでも）、500 文字を書くとき、キャッシュを使う場合と使わない場合でそれぞれどれくらいかかるか測ってください。

確認テスト

1. 温度と top-k はそれぞれ何を変えますか？なぜよく一緒に使うのですか？

温度は softmax の前にスコア全体を拡大・縮小し、確率の分布の「尖り具合」を変えます。温度が低いとスコアの高い少数の字に集中し、高いとより均等になります。top-k は確率が最も高い k 個の字だけを残し、残りを捨てます。温度を上げれば多様性が生まれますが、確率の低い大量の字が割合を持っていってしまうので、top-k でそのロングテールを取り除けば、変化がありつつ度を越さないようにできます。

2. KV キャッシュが生成の結果を変えないのはなぜですか？

因果マスクがあるので、各位置の計算は自分と前の位置だけに依存し、後ろに新しい文字を加えても前の位置の k と v は変わりません。キャッシュはこうした変わらない結果を保存して使い回し、重複した計算を省くだけなので、数学的にはまったく同等です。

3. KV キャッシュの代償は何ですか？なぜ LLM ではこの代償が重要なのですか？

代償はメモリです。トークンごとに、各層で k と v を一つずつ保存する必要があります。LLM は層が多く次元も大きいので、トークンあたりのキャッシュが大きく、長いコンテキストや多くのユーザーが同時に使う場合には、キャッシュの占める GPU メモリがモデルそのものを超えることもあります。そのため多くのモデルの構造上の改良（グループ化クエリアテンション、マルチヘッド潜在アテンションなど）は、これを小さくするためのものなのです。