KV Cache Size Calculator

A developer released a Python tool that calculates and plots KV cache size versus context length from HuggingFace config.json files. The tool supports standard MHA/GQA, MLA, hybrid architectures, and sliding window with per-layer compression, and can load configs from local files, URLs, or HuggingFace/ModelScope repo IDs.

| /usr/bin/env python3 | | | """根据 HuggingFace config.json 计算并绘制 KV cache 大小随上下文长度的变化。 | | | 支持的注意力结构: | | | - 标准 MHA / GQA: KV = 2 L n kv heads head dim dtype bytes | | | - MLA kv lora rank, 如 DeepSeek-V3 / Kimi / GLM : | | | 每层只缓存压缩的 c kv kv lora rank 和 rope key qk rope head dim , | | | V 由 c kv 重构不单独占用缓存: | | | KV = L kv lora rank + qk rope head dim kv bytes | | | - 混合架构 linear attention + 周期性 full attention : | | | 只有 full attention 层的缓存随上下文线性增长; | | | linear attention KDA / gated delta net 等 是常数大小的递归状态 | | | delta 规则状态 + 短卷积状态, 默认按 fp32 估算并叠加, --no-linear-state 可关闭 。 | | | - 滑窗 + 按层压缩 KV DeepSeek V4 的 CSA/HCA, compress ratios : | | | sliding kv = active layers sliding window head dim kv bytes 常数 | | | compressed kv = Σ {ratio 0} floor tokens-1 /ratio head dim kv bytes | | | indexer kv = ratio4 layers floor tokens/4 index head dim fp4 bytes | | | 生产环境默认 KV 为 FP8 1 B , indexer 为 FP4 0.5 B 。 | | | 用法: | | | python3 kv cache calc.py configs/Kimi-K3.json configs/DeepSeek-V4-Flash.json | | | python3 kv cache calc.py moonshotai/Kimi-K3 HuggingFace repo id | | | python3 kv cache calc.py https://.../config.json 直接 URL | | | python3 kv cache calc.py cfg.json --check "Kimi-K3:1048576:27.41851" | | | """ | | | import argparse | | | import json | | | import math | | | import os | | | import sys | | | import urllib.request | | | from collections import Counter | | | DTYPE BYTES = { | | | "float32": 4, "float": 4, "fp32": 4, | | | "float16": 2, "fp16": 2, "half": 2, | | | "bfloat16": 2, "bf16": 2, | | | "float8 e4m3fn": 1, "float8 e5m2": 1, "fp8": 1, "int8": 1, | | | "float4": 0.5, "fp4": 0.5, "nvfp4": 0.5, "mxfp4": 0.5, | | | } | | | 层类型名字中包含这些关键字的层, 其 KV cache 随上下文线性增长 | | | FULL ATTN TYPES = "full attention", "deepseek sparse attention" | | | def dtype bytes name: str, default: float = 2 - float: | | | return DTYPE BYTES.get str name .lower .replace "torch.", "" , default | | | MS ORG ALIASES = { ModelScope 上与 HuggingFace 组织名不同的映射 | | | "zai-org": "ZhipuAI", 智谱 | | | } | | | def load config src: str, hub: str = "auto" - dict: | | | """从本地文件、URL 或模型仓库 id 加载 config.json. | | | hub: "huggingface" | "modelscope" | "auto" 先 HF 后 MS | | | """ | | | if os.path.exists src : | | | with open src as f: | | | return json.load f | | | if src.startswith "http" : | | | with urllib.request.urlopen src as r: | | | return json.load r | | | urls = | | | if hub in "huggingface", "auto" : | | | urls.append f"https://huggingface.co/{src}/resolve/main/config.json" | | | if hub in "modelscope", "auto" : | | | org, , name = src.partition "/" | | | alias = MS ORG ALIASES.get org, org | | | if alias = org: 别名优先, 原名留作回退 | | | urls.append f"https://modelscope.cn/models/{alias}/{name}/resolve/master/config.json" | | | urls.append f"https://modelscope.cn/models/{src}/resolve/master/config.json" | | | last err = None | | | for url in urls: | | | try: | | | print f"正在下载 {url} ...", file=sys.stderr | | | with urllib.request.urlopen url as r: | | | return json.load r | | | except Exception as e: 换下一个 hub/别名重试 | | | last err = e | | | raise RuntimeError f"无法获取 {src} 的 config.json: {last err}" | | | def model name src: str, cfg: dict - str: | | | if os.path.exists src : | | | base = os.path.splitext os.path.basename src 0 | | | else: | | | base = src.rstrip "/" .split "/" -1 | | | if base == "config.json": | | | base = src.split "/" -3 | | | arch = cfg.get "architectures" or "unknown" 0 | | | return f"{base} {arch} " | | | def full attn layers text: dict, cfg: dict - list: | | | """确定哪些层是 full attention KV 随上下文增长 , 返回层索引列表""" | | | 1 显式的 layer types 列表 如 Qwen flash-next | | | if "layer types" in text: | | | return i for i, t in enumerate text "layer types" | | | if any k in t for k in FULL ATTN TYPES | | | 2 linear attn config.full attn layers 如 Kimi K3, GLM | | | lac = text.get "linear attn config" or cfg.get "linear attn config" or {} | | | if "full attn layers" in lac: | | | return list lac "full attn layers" | | | 3 普通稠密模型: 所有层 | | | return list range text.get "num hidden layers", 0 | | | def linear state bytes text: dict, n linear: int, state bytes: float - float: | | | """估算每个 linear attention 层的常数递归状态 可选叠加项 . | | | 包含两部分, 均为常数大小 不随上下文增长 : | | | - delta 规则状态 num heads, d k, d v , 生产实现通常为 fp32 | | | - 短卷积状态 kernel size x 各投影维度 | | | """ | | | if n linear <= 0: | | | return 0 | | | if "linear num value heads" in text: Qwen 风格 gated deltanet: K/V head 分列 | | | dstate = text "linear num value heads" text "linear key head dim" | | | text "linear value head dim" | | | kernel = int text.get "linear conv kernel dim", 4 | | | conv = kernel text "linear num key heads" text "linear key head dim" | | | + text "linear num value heads" text "linear value head dim" | | | return dstate + conv state bytes n linear | | | lac = text.get "linear attn config" or {} | | | if lac: Kimi / GLM 风格 KDA | | | h, d = lac.get "num heads", 0 , lac.get "head dim", 0 | | | kernel = int lac.get "short conv kernel size", 4 | | | return h d d + kernel h d state bytes n linear | | | return 0 信息不足, 不估算 | | | def make info src, cfg, , dtype, db, n layers, n full, n linear, n mtp, | | | per token, const, max pos, kind, layer desc=None, | | | parts=None, extra desc=None : | | | """构造统一的模型信息; parts 为 名称, fn tokens - bytes , 用于分项展示""" | | | def fn tokens, pt=per token, c=const : | | | return pt tokens + c | | | if parts is None: | | | parts = "KV cache", fn | | | total fn = lambda t, p=parts: sum f t for , f in p if len parts 1 else parts 0 1 | | | return { | | | "name": model name src, cfg , "kind": kind, "dtype": dtype, "db": db, | | | "n layers": n layers, "n full": n full, "n linear": n linear, | | | "n mtp": n mtp, "per token bytes": per token, "const bytes": const, | | | "per layer bytes": per token // max n full, 1 , | | | "max pos": max pos, "layer desc": layer desc, "parts": parts, | | | "fn": total fn, "extra desc": extra desc, | | | } | | | def analyze compress text, src, cfg, dtype, kv db, indexer db : | | | """DeepSeek V4 风格: 每层 = 滑窗 + 按 compress ratio 压缩的 latent KV; | | | 压缩比最小的层带 indexer 稀疏打分器 , 另有一份压缩缓存.""" | | | ratios = int r for r in text "compress ratios" | | | n layers = text.get "num hidden layers", len ratios | | | main = ratios :n layers | | | hd = text "head dim" | | | win = int text.get "sliding window" or 0 | | | idx hd = int text.get "index head dim" or 0 | | | nonzero = r for r in main if r 0 | | | 官方实现中 indexer 挂在 compress ratio==4 的层上; 泛化为压缩比最小的层 | | | idx ratio = min nonzero if idx hd and nonzero else None | | | n idx = sum 1 for r in main if idx ratio and r == idx ratio | | | active = len main 滑窗缓存存在于每个主层 含 ratio=0 的纯滑窗层 | | | const kv = active win hd kv db | | | def kv fn t : | | | 只统计已完成的压缩组 进行中的组在滑窗/压缩状态里 | | | return sum max t - 1, 0 // r for r in nonzero hd kv db + const kv | | | def idx fn t : | | | return n idx t // idx ratio idx hd indexer db if idx ratio else 0 | | | parts = "KV cache 滑窗+压缩 ", kv fn | | | if idx hd: | | | parts.append f"Indexer cache FP4, {n idx}层x1/{idx ratio} ", idx fn | | | dist = ", ".join f"{c}层x1/{r}" for r, c in sorted Counter nonzero .items | | | kind = f"压缩KV稀疏注意力 latent head dim={hd}, window={win}" | | | + f", indexer={idx hd}维x1/{idx ratio}" if idx ratio else "" | | | + f", KV精度={kv db}B, indexer精度={indexer db}B " | | | layer desc = f"共 {n layers} 层 = {dist}, 滑窗 {win}" | | | + "; MTP 层仅滑窗, 不计入" if len ratios n layers else "" | | | return make info src, cfg, dtype=dtype, db=kv db, n layers=n layers, | | | n full=len nonzero , n linear=active - len nonzero , | | | n mtp=len ratios - n layers if len ratios n layers else 0, | | | per token=0, const=const kv, per token 由 fn 精确给出 | | | max pos=text.get "max position embeddings" , | | | kind=kind, layer desc=layer desc, parts=parts | | | def default kv dtype text: dict, cfg: dict, model dtype: str - str: | | | """按官方配置推断 KV cache 精度: | | | - fp8 量化且未排除 attention 输入投影 k/v -- fp8 如 MiMo/DSv4 | | | - attention 被显式排除量化 如 GLM 的 attn mha、Kimi 的 self attn -- 模型精度 bf16 | | | - 无量化配置 -- 模型精度 | | | """ | | | q = text.get "quantization config" or cfg.get "quantization config" or {} | | | if q.get "quant method" == "fp8": | | | ign = str x for x in q.get "ignored layers" or q.get "modules to not convert" | | | or q.get "ignore" or | | | attn excluded = any "attn" in s and "o proj" not in s for s in ign | | | if not attn excluded: | | | return "fp8" | | | return str model dtype .lower .replace "torch.", "" | | | def analyze src: str, cfg: dict, dtype override: str | None, | | | kv dtype override: str | None, indexer dtype: str, | | | include linear state: bool, state dtype: str - dict: | | | text = cfg.get "text config", cfg | | | dtype = dtype override or text.get "dtype" or text.get "torch dtype" or "bfloat16" | | | db = dtype bytes dtype | | | if text.get "compress ratios" : DeepSeek V4 风格 CSA/HCA | | | kv d = kv dtype override or "fp8" 生产默认 FP8 attention cache | | | return analyze compress text, src, cfg, kv d, dtype bytes kv d, 1 , | | | dtype bytes indexer dtype, 0.5 | | | kv d = kv dtype override or default kv dtype text, cfg, dtype | | | db = dtype bytes kv d | | | MiMo V2 风格: hybrid layer pattern 逐层标注, 0=full/global attention, | | | 1=滑窗 SWA 层; 滑窗层只保留常数大小的窗口缓存 如 128 个 token . | | | 官方口径: MiMo-V2.5-Pro 为 10 GA + 60 SWA 6:1 , 长上下文缓存省 ~7 倍. | | | if text.get "hybrid layer pattern" : | | | pattern = int x for x in text "hybrid layer pattern" | | | n full = pattern.count 0 | | | n swa = len pattern - n full | | | kv heads = text.get "num key value heads", text.get "num attention heads", 0 | | | head dim = text.get "head dim" or text.get "hidden size", 0 // max text.get "num attention heads", 1 , 1 | | | v head dim = text.get "v head dim", head dim 部分模型 QK/V head dim 不同 | | | win = int text.get "sliding window" or text.get "sliding window size" or 0 | | | per entry = head dim + v head dim kv heads db | | | const = n swa win per entry | | | per token = n full per entry | | | kind = f"GQA+滑窗混合 n kv heads={kv heads}, qk head dim={head dim}, " | | | f"v head dim={v head dim}, full层={n full}, 滑窗层={n swa}x窗{win} " | | | layer desc = f"共 {len pattern } 层 = {n full} 层 full attention + " | | | f"{n swa} 层滑窗 常数 {human const } " | | | return make info src, cfg, dtype=kv d, db=db, n layers=len pattern , | | | n full=n full, n linear=n swa, n mtp=0, | | | per token=per token, const=const, | | | max pos=text.get "max position embeddings" , kind=kind, | | | layer desc=layer desc | | | full idx = full attn layers text, cfg | | | n layers = text.get "num hidden layers", len full idx | | | n full = len full idx | | | MTP / nextn 层: 若其注意力是 full attention, 同样占用随上下文增长的缓存 | | | mtp = text.get "num nextn predict layers" or text.get "mtp num hidden layers" or 0 | | | mtp cfg = text.get "mtp" or {} | | | if mtp and mtp cfg.get "layer types" and not any | | | k in t for k in FULL ATTN TYPES for t in mtp cfg "layer types" : | | | mtp = 0 | | | if text.get "kv lora rank" : MLA: 缓存 c kv + k pe, V 不单独存 | | | mla rank = text "kv lora rank" | | | rope dim = text.get "qk rope head dim", 0 | | | per layer = mla rank + rope dim db | | | kind = f"MLA kv lora rank={mla rank}, qk rope head dim={rope dim} " | | | else: 标准 MHA / GQA | | | kv heads = text.get "num key value heads", text.get "num attention heads", 0 | | | head dim = text.get "head dim" | | | if not head dim: | | | head dim = text.get "qk nope head dim", 0 or 0 + \ | | | text.get "qk rope head dim", 0 or 0 | | | if not head dim: | | | head dim = text.get "hidden size", 0 // max text.get "num attention heads", 1 , 1 | | | per layer = 2 kv heads head dim db K 和 V 各一份 | | | kind = f"GQA n kv heads={kv heads}, head dim={head dim} " | | | per token = per layer n full + mtp | | | n linear = n layers - n full | | | const = linear state bytes text, n linear, dtype bytes state dtype \ | | | if include linear state else 0 | | | DSA 稀疏层 deepseek sparse attention 的 indexer 也有一份随上下文增长的 | | | 压缩缓存 每 index kpool 个 token 存 index head dim 维 , 如 GLM / DSv3.2 系 | | | parts = None | | | idx hd, kpool = text.get "index head dim" , text.get "index kpool" | | | n sparse = sum 1 for t in text.get "layer types", if "sparse" in t if idx hd else 0 | | | if n sparse and idx hd and kpool: | | | def idx fn t, n=n sparse, r=int kpool , d=int idx hd , b=db : | | | return n t // r d b | | | parts = "KV cache", lambda t, pt=per token, c=const: pt t + c , | | | f"Indexer cache {n sparse}层x1/{kpool} ", idx fn | | | kind += f" + DSA稀疏 indexer {idx hd}维x1/{kpool} " | | | return make info src, cfg, dtype=kv d, db=db, n layers=n layers, n full=n full, | | | n linear=n linear, n mtp=mtp, per token=per token, const=const, | | | max pos=text.get "max position embeddings" , kind=kind, | | | parts=parts | | | def human n: float - str: | | | for unit in "B", "KiB", "MiB", "GiB", "TiB" : | | | if abs n < 1024 or unit == "TiB": | | | return f"{n:.0f} B" if unit == "B" else f"{n:.4f} {unit}" | | | n /= 1024 | | | def kv bytes info: dict, ctx: int, batch: int = 1 - float: | | | return info "fn" ctx batch | | | def print report info: dict, batch: int, checks: list - None: | | | print "=" 78 | | | print f"模型: {info 'name' }" | | | print f"注意力类型: {info 'kind' }, dtype={info 'dtype' }" | | | if info.get "layer desc" : | | | print f"层数: {info 'layer desc' }" | | | else: | | | print f"层数: 共 {info 'n layers' } 层 = {info 'n full' } 层 full attention" | | | + f" + {info 'n mtp' } 层 MTP" if info "n mtp" else "" | | | + f" + {info 'n linear' } 层 linear attention 常数状态, 未计入 " | | | if info "n linear" else "" | | | max ctx = info "max pos" or 131072 | | | per tok = info "fn" max ctx / max ctx | | | if info "per token bytes" : | | | line = f"每层缓存: {human info 'per layer bytes' } / token" | | | if info "const bytes" : | | | line += f"\n每 token 缓存: {human info 'per token bytes' } + 常数 {human info 'const bytes' }" | | | else: | | | line += f"\n每 token 缓存: {human info 'per token bytes' }" | | | print line | | | else: | | | print f"每 token 缓存: {human per tok } 按最大上下文折算 " | | | if len info "parts" 1: | | | for label, f in info "parts" : | | | print f" - {label:<28} @ {max ctx:,}: {human f max ctx }" | | | print f" {'合计':<30} @ {max ctx:,}: {human info 'fn' max ctx }" | | | print f"最大上下文: {max ctx:,} tokens" | | | print "-" 78 | | | print f"{'上下文长度': 14} | {'KV cache / 序列': 18}" + | | | f" | {'x batch='+str batch : 18}" if batch 1 else "" | | | points = {1024, 4096, 8192, 16384, 32768, 65536, 131072, | | | 262144, 524288, 1048576, max ctx} | | | points = sorted p for p in points if p <= max ctx | | | for ctx in points: | | | b = kv bytes info, ctx, batch | | | line = f"{ctx: 14,} | {human b : 18}" | | | if batch 1: | | | line += f" | {human b batch : 18}" | | | print line | | | print "-" 78 | | | for c in checks: | | | name, ctx, expected = c | | | if name.lower not in info "name" .lower : | | | continue | | | got = kv bytes info, ctx, batch / 2 30 | | | diff = got - expected | | | print f"核对 {name} @ {ctx:,} : 计算 {got:.5f} GiB vs 参考 {expected:.5f} GiB" | | | f" 差 {diff:+.5f} GiB, {diff/expected 100:+.3f}% " | | | print | | | def short name info: dict - str: | | | return info "name" .split " " 0 | | | def fmt per tok b: float - str: | | | return f"{b/1024:.2f}".rstrip "0" .rstrip "." + " KiB/tok" if b = 1024 \ | | | else f"{b:.0f} B/tok" | | | def plot infos: list, batch: int, out: str, max ctx cap: int | None : | | | import matplotlib | | | matplotlib.use "Agg" | | | import matplotlib.pyplot as plt | | | fig, ax = plt.subplots figsize= 10, 7 | | | markers = "o", "s", "^", "D", "v", "p", "h", "X" | | | endpoints = axvline x, 最终 y, 标签 用于错开标注 | | | for i, info in enumerate infos : | | | max ctx = min info "max pos" or 131072, max ctx cap or math.inf | | | lo, hi = 10, int math.log2 max ctx | | | xs = 2 k for k in range lo, hi + 1 | | | ys = kv bytes info, x, batch / 2 30 for x in xs | | | per tok = info "fn" max ctx / max ctx | | | ax.plot xs, ys, marker=markers i % len markers , ms=4, lw=1.5, | | | label=f"{short name info } {fmt per tok per tok } · KV {info 'dtype' }" | | | endpoints.append xs -1 , ys -1 , f"{ys -1 :.2f}" | | | if info "max pos" and info "max pos" <= max ctx cap or math.inf : | | | ax.axvline info "max pos" , ls=":", lw=0.8, alpha=0.5 | | | 端点标注按 y 排序后做简单的错位, 避免重叠 | | | endpoints.sort key=lambda e: e 1 | | | placed = | | | for x, y, text in endpoints: | | | for py in placed: 对数轴上保持约 25% 的间距, placed 已按升序 | | | if abs math.log10 y - math.log10 py < 0.22: | | | y = py 1.25 | | | placed.append y | | | ax.annotate text, x, y , textcoords="offset points", xytext= -6, 0 , | | | ha="right", va="center", fontsize=8, alpha=0.8, | | | bbox=dict fc="white", ec="none", alpha=0.5, pad=0.6 | | | ax.set xscale "log", base=2 | | | ax.set yscale "log" | | | ax.set xlabel "Context length tokens " | | | ax.set ylabel "KV cache size GiB " | | | ax.grid True, which="both", ls="--", alpha=0.3 | | | ax.legend fontsize=8.5, ncol=3, loc="lower center", | | | bbox to anchor= 0.5, 1.01 , frameon=False | | | fig.suptitle "KV cache vs context length" + f" batch={batch} " if batch 1 else "" , | | | y=0.975 | | | fig.tight layout rect= 0, 0, 1, 0.92 | | | fig.savefig out, dpi=300 | | | svg = os.path.splitext out 0 + ".svg" | | | fig.savefig svg | | | print f"图已保存到 {out} 300 dpi 和 {svg}" | | | def main : | | | ap = argparse.ArgumentParser description="根据 config.json 计算/绘制 KV cache 随上下文长度的变化" | | | ap.add argument "configs", nargs="+", help="config.json 路径 / 模型仓库 id / URL" | | | ap.add argument "--hub", choices= "huggingface", "modelscope", "auto" , | | | default="auto", help="模型仓库来源 默认 auto: HF 失败回退 ModelScope " | | | ap.add argument "--dtype", help="覆盖模型 dtype 如 bf16/fp16 ", default=None | | | ap.add argument "--kv-dtype", help="覆盖 KV cache 精度 如 bf16/fp8/fp4/int8 ", default=None | | | ap.add argument "--indexer-dtype", help="indexer 缓存精度 默认 fp4 ", default="fp4" | | | ap.add argument "--batch", type=int, default=1, help="batch 大小 默认 1 " | | | ap.add argument "--max-context", type=int, default=None, help="绘制曲线的最大上下文长度" | | | ap.add argument "--no-linear-state", action="store false", dest="linear state", | | | help="不计入 linear attention 层的常数递归状态估算" | | | ap.set defaults linear state=True | | | ap.add argument "--state-dtype", default="float32", help="递归状态 dtype 默认 float32 " | | | ap.add argument "--check", action="append", default= , | | | help='核对数据, 格式 "名字:上下文长度:GiB数值", 可多次指定' | | | ap.add argument "-o", "--out", default="kv cache.png", help="输出图片路径" | | | args = ap.parse args | | | checks = | | | for c in args.check: | | | name, ctx, val = c.rsplit ":", 2 | | | checks.append name, int ctx , float val | | | infos = | | | for src in args.configs: | | | info = analyze src, load config src, args.hub , args.dtype, args.kv dtype, | | | args.indexer dtype, args.linear state, args.state dtype | | | infos.append info | | | print report info, args.batch, checks | | | if args.out: | | | plot infos, args.batch, args.out, args.max context | | | if name == " main ": | | | main |