# KV Cache Size Calculator

> Source: <https://gist.github.com/wszqkzqk/2702ba95fd20e9176b366ce1575a556c>
> Published: 2026-08-30 07:11:28+00:00

| #!/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() |