KV Cache - 読んだトークンの計算結果を取っておく仕組みとコンテキスト長

"""KV キャッシュの実験で共通に使う道具（Python 標準ライブラリだけ）。

- Ollama の HTTP API を呼ぶ
- /api/ps を記録する
- サーバーのログ（systemd の journal）から、読み込み時の "kv cache" などの行を拾う
- 長さを変えられる詰め物の文章を作る
"""
import json
import os
import subprocess
import time
import urllib.request

HOST = os.environ.get("OLLAMA_HOST", "127.0.0.1:11435")
# ログを読むコマンド。著者の環境では Ollama を systemd のユーザーサービス local-llm で動かしている
LOG_UNIT = os.environ.get("KVLAB_LOG_UNIT", "local-llm")
MODELS_DIR = os.environ.get("OLLAMA_MODELS", os.path.expanduser("~/.ollama/models"))


def api(path, body=None, timeout=900):
    data = None if body is None else json.dumps(body).encode()
    req = urllib.request.Request(f"http://{HOST}{path}", data=data,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def ps():
    """/api/ps を短い形にする（名前・size・size_vram・context_length）。"""
    out = []
    for m in api("/api/ps")["models"]:
        out.append({"name": m["name"], "size": m["size"], "size_vram": m["size_vram"],
                    "context_length": m.get("context_length")})
    return out


def ps_line(rows):
    return "; ".join(f'{r["name"]} size={r["size"]/1e9:.2f}GB vram={r["size_vram"]/1e9:.2f}GB '
                     f'ctx={r["context_length"]}' for r in rows) or "(none)"


def blob_of(model):
    """モデル名から重みの blob（sha256-...）を引く。ログの "starting runner" 行と照合するため。"""
    name, tag = model.split(":")
    path = os.path.join(MODELS_DIR, "manifests/registry.ollama.ai/library", name, tag)
    with open(path) as f:
        man = json.load(f)
    for layer in man["layers"]:
        if layer["mediaType"] == "application/vnd.ollama.image.model":
            return layer["digest"].replace(":", "-")
    raise RuntimeError(model)


def load_log(model, since_epoch):
    """since_epoch 以降のログから、model を読み込んだ時の行（KvSize と device.go の内訳）を返す。"""
    try:
        txt = subprocess.run(["journalctl", "--user", "-u", LOG_UNIT, "--since", f"@{since_epoch:.0f}",
                              "--no-pager", "-o", "cat"], capture_output=True, text=True,
                             timeout=30).stdout
    except Exception as e:  # ログが読めない環境では空で返す
        return {"error": str(e)}
    blob = blob_of(model)
    lines = txt.splitlines()
    found = None
    for i, ln in enumerate(lines):
        if "starting runner" in ln and blob in ln:
            found = i
    if found is None:
        return {}
    info = {}
    for ln in lines[found:found + 80]:
        if "KvSize:" in ln and "Operation:commit" in ln:
            info["KvSize"] = int(ln.split("KvSize:")[1].split()[0])
            info["KvCacheType"] = ln.split("KvCacheType:")[1].split(" ")[0] or "(empty)"
        if "source=device.go" in ln:
            msg = ln.split('msg="')[1].split('"')[0]
            dev = ln.split("device=")[1].split()[0] if "device=" in ln else "-"
            size = ln.split('size="')[1].split('"')[0]
            info[f"{msg} {dev}"] = size
            if msg == "total memory":
                break
    return info


def unload(model):
    try:
        api("/api/generate", {"model": model, "keep_alive": 0})
    except Exception:
        pass
    # 実際に消えるまで待つ（他の作業の都合で長引いた時は記録して先へ進む）
    for _ in range(60):
        if all(r["name"] != model for r in ps()):
            return True
        time.sleep(1)
    return False


def filler(n_lines, seed=0):
    """数えやすい詰め物。1 行はおよそ 25 トークン（qwen3 で prompt_eval_count を見て調整）。"""
    rows = []
    for i in range(n_lines):
        k = (i * 7919 + seed * 104729) % 100000
        rows.append(f"Record {i + 1}: sensor S{k % 97:02d} reported {k % 1000} units "
                    f"at hour {k % 24} on day {k % 31 + 1}.")
    return "\n".join(rows)


def generate(model, prompt, num_ctx=None, num_predict=64, keep_alive="10m"):
    opts = {"temperature": 0, "seed": 1, "num_predict": num_predict}
    if num_ctx is not None:
        opts["num_ctx"] = num_ctx
    body = {"model": model, "prompt": prompt, "stream": False, "think": False,
            "options": opts, "keep_alive": keep_alive}
    t0 = time.time()
    r = api("/api/generate", body)
    r["wall"] = time.time() - t0
    r["thinking_chars"] = len(r.get("thinking") or "")
    return r

"""実験 1: num_ctx を変えて読み込み、/api/ps の size と、ログの KV キャッシュの大きさを記録する。

同じ num_ctx で短いプロンプトと長いプロンプトを送り、size が入力の長さで変わるかも見る。
"""
import json
import sys
import time

from kvlab import filler, generate, load_log, ps, ps_line, unload

CASES = [
    ("qwen3:1.7b-q8_0", None), ("qwen3:1.7b-q8_0", 2048), ("qwen3:1.7b-q8_0", 4096),
    ("qwen3:1.7b-q8_0", 8192), ("qwen3:1.7b-q8_0", 16384), ("qwen3:1.7b-q8_0", 32768),
    ("qwen3:1.7b-q8_0", 40960),
    # 重みの形式だけを変える（KV キャッシュは同じ大きさになるはず）
    ("qwen3:1.7b-fp16", 4096), ("qwen3:1.7b-fp16", 32768),
    ("qwen3:1.7b-q4_K_M", None), ("qwen3:1.7b-q4_K_M", 4096),
]

results = []
for model, ctx in CASES:
    print(f"## {model} num_ctx={ctx}", flush=True)
    print("  ps before:", ps_line(ps()), flush=True)
    gone = unload(model)
    t0 = time.time() - 2
    r = generate(model, "Say hi.", num_ctx=ctx, num_predict=1, keep_alive="5m")
    time.sleep(1)
    after = ps()
    mine = [x for x in after if x["name"] == model]
    log = load_log(model, t0)
    row = {"model": model, "num_ctx": ctx, "unloaded_first": gone,
           "ps": mine[0] if mine else None, "log": log,
           "load_duration_s": r["load_duration"] / 1e9, "thinking_chars": r["thinking_chars"]}
    # 同じ num_ctx のまま長い入力（約 1,800 トークン）を送り、size が変わるかを見る
    r2 = generate(model, filler(70) + "\nHow many records are listed above? Answer with a number.",
                  num_ctx=ctx, num_predict=8, keep_alive="5m")
    mine2 = [x for x in ps() if x["name"] == model]
    row["long_prompt_eval_count"] = r2["prompt_eval_count"]
    row["ps_after_long"] = mine2[0] if mine2 else None
    row["long_load_duration_s"] = r2["load_duration"] / 1e9
    print("  ps after :", ps_line(after), flush=True)
    print("  log      :", json.dumps(log, ensure_ascii=False), flush=True)
    print(f"  long prompt {r2['prompt_eval_count']} tok -> size "
          f"{(mine2[0]['size'] / 1e9) if mine2 else float('nan'):.3f} GB", flush=True)
    results.append(row)
    unload(model)

json.dump(results, open(sys.argv[1] if len(sys.argv) > 1 else "memory.json", "w"),
          indent=1, ensure_ascii=False)

model              num_ctx  loaded        ps size     after 1.8k log kv cache calc kv cache
qwen3:1.7b-q8_0       None   40960  7,070,831,488  7,070,831,488      4.4 GiB     4.375 GiB
qwen3:1.7b-q8_0       2048    2048  2,517,914,496  2,517,914,496    224.0 MiB     224.0 MiB
qwen3:1.7b-q8_0       4096    4096  2,754,892,672  2,754,892,672    448.0 MiB     448.0 MiB
qwen3:1.7b-q8_0       8192    8192  3,237,237,632  3,237,237,632    896.0 MiB     896.0 MiB
qwen3:1.7b-q8_0      16384   16384  4,189,344,640  4,189,344,640      1.8 GiB     1.750 GiB
qwen3:1.7b-q8_0      32768   32768  6,106,141,568  6,106,141,568      3.5 GiB     3.500 GiB
qwen3:1.7b-q8_0      40960   40960  7,070,831,488  7,070,831,488      4.4 GiB     4.375 GiB
qwen3:1.7b-fp16       4096    4096  4,605,426,560  4,605,426,560    448.0 MiB     448.0 MiB
qwen3:1.7b-fp16      32768   32768  7,956,675,456  7,956,675,456      3.5 GiB     3.500 GiB
qwen3:1.7b-q4_K_M     None   40960  6,210,022,272  6,210,022,272      4.4 GiB     4.375 GiB
qwen3:1.7b-q4_K_M     4096    4096  1,894,083,456  1,894,083,456    448.0 MiB     448.0 MiB

"""実験 2: num_ctx を 32768 に固定し、入力の長さだけを変えて、
最初のトークンまでの時間（prompt_eval_duration）と生成の速さを測る。

毎回、先頭に違う札（nonce）を置き、前回の KV キャッシュを使い回させない。
3 巡して、長さごとに中央値と最小・最大を出す。
"""
import json
import statistics
import sys
import time
import uuid

from kvlab import filler, generate, ps, ps_line

MODEL = "qwen3:1.7b-q8_0"
NUM_CTX = 32768
LINES = [19, 39, 79, 159, 319, 639, 1170]  # およそ 512〜30,000 トークン
ROUNDS = 3
TAIL = "\nIgnore the records. Write the integers from 1 to 300 separated by commas."

generate(MODEL, "Say hi.", num_ctx=NUM_CTX, num_predict=1)  # 読み込みを先に済ませる
rows = []
for rnd in range(ROUNDS):
    for n in LINES:
        before = ps()
        prompt = f"Session {uuid.uuid4().hex}\n" + filler(n, seed=rnd) + TAIL
        r = generate(MODEL, prompt, num_ctx=NUM_CTX, num_predict=64)
        after = ps()
        row = {"round": rnd + 1, "lines": n,
               "prompt_eval_count": r["prompt_eval_count"],
               "prompt_eval_s": r["prompt_eval_duration"] / 1e9,
               "eval_count": r["eval_count"], "eval_s": r["eval_duration"] / 1e9,
               "load_s": r["load_duration"] / 1e9, "thinking_chars": r["thinking_chars"],
               "ps_before": before, "ps_after": after, "time": time.strftime("%H:%M:%S")}
        rows.append(row)
        print(f"round {rnd + 1} prompt {row['prompt_eval_count']:6d} tok "
              f"prefill {row['prompt_eval_s']:7.3f} s ({row['prompt_eval_count'] / row['prompt_eval_s']:7.1f} tok/s) "
              f"| gen {row['eval_count']} tok {row['eval_count'] / row['eval_s']:5.1f} tok/s "
              f"| load {row['load_s']:.2f} s | thinking {row['thinking_chars']} chars", flush=True)
        print("    ps after:", ps_line(after), flush=True)

print("\n# median (min - max) over rounds")
for n in LINES:
    rs = [x for x in rows if x["lines"] == n]
    pe = [x["prompt_eval_s"] for x in rs]
    gen = [x["eval_count"] / x["eval_s"] for x in rs]
    print(f"prompt {statistics.median(x['prompt_eval_count'] for x in rs):6.0f} tok  "
          f"prefill {statistics.median(pe):7.3f} s ({min(pe):.3f}-{max(pe):.3f})  "
          f"gen {statistics.median(gen):5.1f} tok/s ({min(gen):.1f}-{max(gen):.1f})")

json.dump(rows, open(sys.argv[1] if len(sys.argv) > 1 else "prefill.json", "w"), indent=1)

# median (min - max) over rounds
prompt    542 tok  prefill   0.122 s (0.120-0.122)  gen  92.0 tok/s (91.9-93.2)
prompt   1047 tok  prefill   0.236 s (0.235-0.276)  gen  88.4 tok/s (81.8-91.2)
prompt   2046 tok  prefill   0.435 s (0.435-0.450)  gen  85.2 tok/s (84.8-85.6)
prompt   4129 tok  prefill   1.048 s (1.035-1.067)  gen  76.6 tok/s (75.6-77.1)
prompt   8305 tok  prefill   2.647 s (2.632-2.666)  gen  66.1 tok/s (65.6-66.5)
prompt  16692 tok  prefill   7.581 s (7.550-7.905)  gen  51.9 tok/s (47.6-52.2)
prompt  30759 tok  prefill  21.030 s (20.626-21.253)  gen  38.8 tok/s (36.2-39.0)

"""実験 3: 同じプレフィックスを送り直した時に、KV キャッシュが使い回されるかを見る。

Ollama 0.22.1 の prompt_eval_count は、使い回した分も含めた入力の全トークン数を返す。
使い回せたかどうかは、prompt_eval_duration（キャッシュに無かった分を計算した時間）で見る。
"""
import json
import sys
import uuid

from kvlab import filler, generate, ps, ps_line

MODEL = "qwen3:1.7b-q8_0"
NUM_CTX = 32768
Q1 = "\nWhich sensor appears in Record 5? Answer briefly."
Q2 = "\nWhich sensor appears in Record 9? Answer briefly."

rows = []
generate(MODEL, "Say hi.", num_ctx=NUM_CTX, num_predict=1)
for rnd in range(3):
    head = f"Session {uuid.uuid4().hex}\n"
    lines = filler(320, seed=10 + rnd).split("\n")
    base = head + "\n".join(lines)
    mid = len(lines) // 2
    changed_mid = head + "\n".join(lines[:mid] + ["Record X: this line was edited."] + lines[mid + 1:])
    cases = [
        ("1 first send", base + Q1),
        ("2 same prompt again", base + Q1),
        ("3 same prefix, new question", base + Q2),
        ("4 one line edited in the middle", changed_mid + Q1),
        ("5 one char added at the head", "X" + base + Q1),
    ]
    for label, prompt in cases:
        r = generate(MODEL, prompt, num_ctx=NUM_CTX, num_predict=16)
        row = {"round": rnd + 1, "case": label, "prompt_eval_count": r["prompt_eval_count"],
               "prompt_eval_s": r["prompt_eval_duration"] / 1e9, "load_s": r["load_duration"] / 1e9,
               "thinking_chars": r["thinking_chars"], "response": r["response"], "ps": ps()}
        rows.append(row)
        print(f"round {rnd + 1} {label:34s} prompt {row['prompt_eval_count']:5d} tok "
              f"prompt_eval {row['prompt_eval_s']:6.3f} s | load {row['load_s']:.2f} s "
              f"| thinking {row['thinking_chars']} chars | {r['response'].strip()[:40]!r}", flush=True)
    print("    ps:", ps_line(ps()), flush=True)

json.dump(rows, open(sys.argv[1] if len(sys.argv) > 1 else "prefix.json", "w"), indent=1,
          ensure_ascii=False)

"""実験 4: 入力が num_ctx を超えた時に何が起きるか。

先頭に合言葉を置き、約 3,000 トークンの詰め物の後で合言葉を尋ねる。
num_ctx 2048（足りない）と 8192（足りる）で比べる。
"""
import json
import subprocess
import sys

from kvlab import LOG_UNIT, filler, generate, ps, ps_line

MODEL = "qwen3:1.7b-q8_0"
PROMPT = ("The secret code is BLUE-FALCON-42. Remember it.\n" + filler(120, seed=7) +
          "\nWhat is the secret code given at the very beginning? Answer with the code only.")

rows = []
for ctx in (2048, 8192):
    print(f"## num_ctx={ctx}  ps before: {ps_line(ps())}", flush=True)
    # 送る直前のログの位置（cursor）を覚え、その後に出た行だけを見る
    cur = subprocess.run(["journalctl", "--user", "-u", LOG_UNIT, "-n", "0", "--show-cursor"],
                         capture_output=True, text=True).stdout.split("cursor: ")[-1].strip()
    try:
        r = generate(MODEL, PROMPT, num_ctx=ctx, num_predict=32)
        row = {"num_ctx": ctx, "prompt_eval_count": r["prompt_eval_count"],
               "done_reason": r.get("done_reason"), "thinking_chars": r["thinking_chars"],
               "response": r["response"]}
    except Exception as e:  # エラーで返る版もあるので、その時は本文を記録する
        body = e.read().decode() if hasattr(e, "read") else str(e)
        row = {"num_ctx": ctx, "error": body}
    # サーバーのログに切り詰めの警告が出たかを見る
    log = subprocess.run(["journalctl", "--user", "-u", LOG_UNIT, f"--after-cursor={cur}",
                          "--no-pager", "-o", "cat"], capture_output=True, text=True).stdout
    row["server_log"] = [ln.split("source=")[1] for ln in log.splitlines()
                         if "truncating input prompt" in ln]
    rows.append(row)
    print(json.dumps(row, ensure_ascii=False), flush=True)
    print(f"   ps after: {ps_line(ps())}", flush=True)

json.dump(rows, open(sys.argv[1] if len(sys.argv) > 1 else "truncate.json", "w"), indent=1,
          ensure_ascii=False)