RadarAI· MarkTechPost（RSS）·· 3 小时前AI 评分58
用 llama.cpp 运行 NVIDIA Dory 1.0 4B 的逐环实测教程
NVIDIA Dory 1.0 4B 用 llama.cpp 运行的逐环实测教程
AI 导读
MarkTechPost 发布教程，讲解如何用 llama.cpp 运行 NVIDIA Nemotron Labs 的 Dory 1.0 4B 循环推理模型，该模型中间层构成共享块，每个 token 可重复应用 1 至 4 次而不增加参数。
正文 · AI 翻译
MarkTechPost· Sana Hassan·2026-11-08 03:22· 2 小时前AI 评分59
如何用 llama.cpp 运行 NVIDIA Dory 4B：逐循环指南
AI 导读
MarkTechPost 发布一篇教程，讲解如何用 llama.cpp 运行 NVIDIA Nemotron Labs 的 Dory 1.0 4B 循环推理模型。
正文 · 原文
在本教程中，我们使用 Dory 1.0 4B，这是 NVIDIA Nemotron Labs 推出的一个小型推理模型，其中间层构成一个共享块，对每个 token 可应用一次、两次、三次或四次。每多一次传递（即循环），就会为每个 token 消耗更多计算量，却不增加任何参数，因此一个检查点可以在四种不同的计算预算下运行，无需重新训练。我们构建 NVIDIA 模型卡所指定的 llama.cpp 分支，下载 3.7 GB 的 GGUF，并直接从文件中读取架构，以预测每个循环在参数和缓存内存上的开销。随后我们将这些预测与运行中的服务器进行核对，测量每个循环次数下的解码速度，测试更廉价的末循环缓存是否会改变模型的预测，并运行配对的 GSM8K 评估，观察随着循环次数增加，准确率、答案长度和每道已解题目的耗时如何变化。
import os
import re
import sys
import json
import time
import math
import random
import signal
import shutil
import hashlib
import platform
import statistics
import traceback
import subprocess
import collections
import urllib.request
from pathlib import Path
RESULTS = {} # step -> one-line result, for the summary table
CLAIMS = [] # (step, claim, held)
CURRENT = ["1"]
def banner(title):
print("\n" + "=" * 78)
print(title)
print("=" * 78)
def section(name):
def wrap(fn):
def run(*a, **kw):
CURRENT[0] = name.split(".")[0]
banner(name)
t0 = time.time()
try:
out = fn(*a, **kw)
RESULTS[name] = out if isinstance(out, str) else "ok"
return out
except Exception as e:
RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
traceback.print_exc(limit=3)
return None
finally:
print(f"\n[{time.time() - t0:.1f} s]")
return run
return wrap
def claim(text, held):
CLAIMS.append((CURRENT[0], text, bool(held)))
return held
banner("1. Build the pinned Dory llama.cpp fork and record the environment")
FORK = "https://github.com/brianchmielNV/llama.cpp.git"
FORK_COMMIT = "98a74c97ea09ac72f6ceda1559808fc0a2f447ce" # the commit NVIDIA's model card pins; ggml-org PR #30306 is still a draft
GGUF_REPO = "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-GGUF"
GGUF_REV = "df53755cb8741843af8e019e088819fbef17e29b"
GGUF_FILE = "NVIDIA-Nemotron-Labs-Dory1.0-4B-Q5_K_M.gguf" # the card's commands call it dory-q5_k_m.gguf
GGUF_SIZE = 3_673_288_864
GGUF_SHA256 = "96080be49ffd67b916018b7bdef8c484762576fb52780e10ab98ac4978a660e0"
GSM8K_REV = "740312add88f781978c0658806c59bc2815b9866"
WORK = Path(os.environ.get("DORY_WORK") or ("/content/dory" if os.path.isdir("/content") else "dory_work")).resolve()
WORK.mkdir(parents=True, exist_ok=True)
(WORK / "logs").mkdir(exist_ok=True)
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "gguf==0.19.0", "huggingface_hub", "pandas", "pyarrow"], check=True)
SYSTEM = platform.system()
GPU = ""
if shutil.which("nvidia-smi"):
q = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,compute_cap", "--format=csv,noheader"],
capture_output=True, text=True)
GPU = q.stdout.strip().splitlines()[0] if q.returncode == 0 and q.stdout.strip() else ""
BACKEND = "metal" if SYSTEM == "Darwin" else ("cuda" if GPU else "cpu")
SRC = WORK / "llama.cpp-dory"
SERVER_BIN = SRC / "build" / "bin" / "llama-server"
free_gb = shutil.disk_usage(WORK).free / 1e9
print(f"{SYSTEM} {platform.machine()}, Python {platform.python_version()}, {os.cpu_count()} CPU cores, "
f"backend {BACKEND}{' (' + GPU + ')' if GPU else ''}, {free_gb:.1f} GB free disk")
if not SERVER_BIN.exists():
t0 = time.time()
if not (SRC / ".git").exists():
subprocess.run(["git", "init", "-q", str(SRC)], check=True)
subprocess.run(["git", "-C", str(SRC), "remote", "add", "origin", FORK], check=True)
subprocess.run(["git", "-C", str(SRC), "fetch", "-q", "--depth", "1", "origin", FORK_COMMIT], check=True)
subprocess.run(["git", "-C", str(SRC), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True)
flags = ["-DCMAKE_BUILD_TYPE=Release", "-DLLAMA_BUILD_SERVER=ON", "-DLLAMA_BUILD_TESTS=OFF",
"-DLLAMA_BUILD_EXAMPLES=OFF", "-DLLAMA_BUILD_APP=OFF", "-DLLAMA_BUILD_UI=OFF",
"-DLLAMA_USE_PREBUILT_UI=ON", "-DLLAMA_OPENSSL=OFF"]
if BACKEND == "metal": # the card's Apple-silicon flags
flags += ["-DCMAKE_OSX_ARCHITECTURES=arm64", "-DGGML_METAL=ON", "-DGGML_METAL_EMBED_LIBRARY=ON",
"-DGGML_BLAS=ON", "-DGGML_BLAS_VENDOR=Apple"]
elif BACKEND == "cuda": # kernels for this GPU only (a T4 is 7.5); flash-attention kernels are unused with --flash-attn off
cap = GPU.split(",")[-1].strip().replace(".", "")
flags += ["-DGGML_CUDA=ON", f"-DCMAKE_CUDA_ARCHITECTURES={cap}", "-DGGML_CUDA_FA=OFF", "-DGGML_NATIVE=OFF"]
env = dict(os.environ, HF_UI_VERSION="b11551")
with open(WORK / "build.log", "w") as log:
subprocess.run(["cmake", "-S", str(SRC), "-B", str(SRC / "build")] + flags, check=True, env=env, stdout=log, stderr=log)
subprocess.run(["cmake", "--build", str(SRC / "build"), "--parallel", str(os.cpu_count()), "--target", "llama-server"],
check=True, env=env, stdout=log, stderr=log)
print(f"built llama-server in {time.time() - t0:.0f} s")
head = subprocess.run(["git", "-C", str(SRC), "rev-parse", "HEAD"], capture_output=True, text=True).stdout.strip()
assert head == FORK_COMMIT, f"checked out {head}, expected {FORK_COMMIT}"
version = subprocess.run([str(SERVER_BIN), "--version"], capture_output=True, text=True)
print(f"fork commit {head[:12]}; {(version.stdout + version.stderr).strip().splitlines()[-1]}")
RESULTS["1. Build the pinned Dory llama.cpp fork and record the environment"] = f"fork {head[:12]}, backend {BACKEND}"
我们从 NVIDIA 模型卡所指定的 Dory 分支的确切提交构建 llama-server，因为 llama.cpp 中的 Dory 支持仍是一个开放的草稿拉取请求，没有任何已发布的二进制文件能加载该模型。该单元格会检测平台，在 Mac 上传递模型卡中的 Metal 标志，或在 Colab 上仅为所连接的 GPU 编译 CUDA 内核，跳过模型卡中 –flash-attn off 从未使用的 flash-attention 内核，从而使 T4 构建保持可控。随后它会检查检出的提交是否为所指定的提交，以便后续每个数字都能追溯到代码的一个已知版本。
@section("2. Download the Q5_K_M GGUF at a pinned revision and check its SHA-256")
def download_step():
global GGUF
from huggingface_hub import hf_hub_download
t0 = time.time()
GGUF = Path(hf_hub_download(GGUF_REPO, GGUF_FILE, revision=GGUF_REV, local_dir=str(WORK)))
h = hashlib.sha256()
with open(GGUF, "rb") as f:
for block in iter(lambda: f.read(1 << 24), b""):
h.update(block)
size, digest = GGUF.stat().st_size, h.hexdigest()
print(f"{GGUF.name}: {size:,} bytes, sha256 {digest[:16]}... ({time.time() - t0:.0f} s)")
assert size == GGUF_SIZE and digest == GGUF_SHA256, "the file does not match the card's size and checksum"
claim("the GGUF matches the size and SHA-256 published on the card", True)
return f"{size / 1e9:.2f} GB, checksum matches the card"
download_step()
Q5_K_M GGUF 来自 Hugging Face 的一个固定仓库修订版本，我们在任何其他操作触及它之前先对其进行哈希校验。模型卡中的命令将该文件称为 dory-q5_k_m.gguf，但仓库将其命名为 NVIDIA-Nemotron-Labs-Dory1.0-4B-Q5_K_M.gguf，因此逐字复制模型卡中的命令会失败；将真实名称、大小和 SHA-256 一并固定，可以同时捕获错误文件和下载损坏。3.67 GB 的文件与模型卡公布的校验和完全匹配，因此后续的每次测量都使用 NVIDIA 所发布的相同字节。
@section("3. Read the architecture from the GGUF: three layer groups and a default of four loops")
def architecture_step():
global META, GROUPS, PATTERN
import gguf
reader = gguf.GGUFReader(str(GGUF))
META = {k: f.contents() for k, f in reader.fields.items() if k.startswith("dory.")}
for k in ("dory.block_count", "dory.input_layer_count", "dory.recurrent_layer_count", "dory.output_layer_count",
"dory.num_loops", "dory.recurrent_kv_cache_mode", "dory.attention.sliding_window", "dory.context_length"):
print(f"{k:34s} {META[k]}")
PATTERN = META["dory.layer_pattern"].split("-")[:-1] # one symbol per attention layer: S sliding window, * full
n_in, n_rec = META["dory.input_layer_count"], META["dory.recurrent_layer_count"]
print(f"\nlayer entries alternate attention and MLP: {len(PATTERN)} attention layers "
f"({PATTERN.count('S')} sliding-window, {PATTERN.count('*')} full attention)")
GROUPS = collections.defaultdict(lambda: {"params": 0, "bytes": 0})
dtypes = collections.Counter()
for t in reader.tensors:
m = re.match(r"blk\.(\d+)\.", t.name)
if m:
i = int(m.group(1))
g = "input block" if i < n_in else "recurrent block" if i < n_in + n_rec else "output block"
else:
g = "embedding" if t.name.startswith("token_embd") else "LM head"
GROUPS[g]["params"] += int(math.prod(int(x) for x in t.shape))
GROUPS[g]["bytes"] += int(t.n_bytes)
dtypes[t.tensor_type.name] += 1
print(f"\n{'group':16s} {'params':>9s} {'bytes':>9s}")
for g in ("embedding", "input block", "recurrent block", "output block", "LM head"):
print(f"{g:16s} {GROUPS[g]['params'] / 1e9:7.3f} B {GROUPS[g]['bytes'] / 1e9:6.3f} GB")
print(f"tensor types: {dict(dtypes)}")
claim("the GGUF stores 16 input, 40 recurrent and 16 output layer entries (8, 20 and 8 blocks)",
(n_in, n_rec, META["dory.output_layer_count"]) == (16, 40, 16))
claim("the GGUF's built-in loop count is 4, not the 2 the card recommends", META["dory.num_loops"] == 4)
total = sum(v["params"] for v in GROUPS.values())
return (f"{total / 1e9:.2f} B parameters, {GROUPS['recurrent block']['params'] / 1e9:.2f} B of them in the shared "
f"recurrent block; built-in num_loops = {META['dory.num_loops']}")
architecture_step()
GGUF 头信息在加载模型之前就告诉我们模型是如何连接的。Dory 的 72 个层条目分为 16 个输入组、40 个共享循环组和 16 个输出组，注意力和 MLP 交替出现，每三个滑动窗口层对应一个全注意力层。按组对张量大小求和显示，循环块占模型 39.8 亿参数中的 18.4 亿。头信息还包含一个陷阱：其内置循环次数为 4，而模型卡建议为 2，因此未加覆盖启动的服务器会静默地以最昂贵的设置运行。
LOOPS = (1, 2, 3, 4)
CTX, UBATCH = 4096, 512
def predicted_kv(loops, mode, ctx=CTX, ubatch=UBATCH):
"""KV cache under the fork's allocation rule: per_loop keeps one cache per loop for every recurrent attention layer,
last_loop keeps one in all. Full-attention layers hold ctx cells; sliding-window layers hold window + ubatch, padded to 256."""
n_in, n_rec = META["dory.input_layer_count"] // 2, META["dory.recurrent_layer_count"] // 2
rec, fixed = PATTERN[n_in:n_in + n_rec], PATTERN[:n_in] + PATTERN[n_in + n_rec:]
copies = loops if mode == "per_loop" else 1
full = fixed.count("*") + copies * rec.count("*")
swa = fixed.count("S") + copies * rec.count("S")
swa_cells = min(ctx, -(-(META["dory.attention.sliding_window"] + ubatch) // 256) * 256)
per_cell = 2 * META["dory.attention.head_count_kv"] * META["dory.attention.key_length"] * 2 # K and V, 2 bytes each
return full, swa, (full * ctx + swa * swa_cells) * per_cell / 2**20
@section("4. Predict what each extra loop costs: parameters per token and KV-cache size")
def cost_model_step():
global PER_TOKEN
once = ("input block", "output block", "LM head") # the embedding is a one-row lookup, so it costs no matrix work
PER_TOKEN = {}
print(f"{'loops':>5s} {'params per token':>17s} {'vs 1 loop':>10s} {'weight bytes read':>18s} {'vs 1 loop':>10s}")
for n in LOOPS:
p = sum(GROUPS[g]["params"] for g in once) + n * GROUPS["recurrent block"]["params"]
b = sum(GROUPS[g]["bytes"] for g in once) + n * GROUPS["recurrent block"]["bytes"]
PER_TOKEN[n] = (p, b)
print(f"{n:5d} {p / 1e9:15.2f} B {p / PER_TOKEN[1][0]:9.2f}x {b / 1e9:16.2f} GB {b / PER_TOKEN[1][1]:9.2f}x")
print(f"\npredicted KV cache at --ctx-size {CTX}, f16 keys and values, micro-batch {UBATCH}:")
print(f"{'loops':>5s} {'per_loop: full + SWA layers':>28s} {'MiB':>7s} {'last_loop: full + SWA layers':>29s} {'MiB':>7s}")
for n in LOOPS:
a, b = predicted_kv(n, "per_loop"), predicted_kv(n, "last_loop")
print(f"{n:5d} {a[0]:>21d} + {a[1]:<4d} {a[2]:7.1f} {b[0]:>22d} + {b[1]:<4d} {b[2]:7.1f}")
full_ctx = META["dory.context_length"]
print(f"\nat the full {full_ctx:,}-token context: per_loop {predicted_kv(2, 'per_loop', full_ctx)[2] / 1024:.1f} GiB at 2 loops "
f"and {predicted_kv(4, 'per_loop', full_ctx)[2] / 1024:.1f} GiB at 4; last_loop "
f"{predicted_kv(4, 'last_loop', full_ctx)[2] / 1024:.1f} GiB at any loop count")
return (f"4 loops apply {PER_TOKEN[4][0] / PER_TOKEN[1][0]:.2f}x the parameters per token of 1 loop; predicted KV at 4 loops "
f"{predicted_kv(4, 'per_loop')[2]:.0f} MiB per_loop vs {predicted_kv(4, 'last_loop')[2]:.0f} MiB last_loop")
cost_model_step()
根据这些组大小，我们可以在运行任何内容之前预测一个循环的开销。一个 token 会经过输入块、每个循环一次的循环块、输出块和 LM head，因此 4 个循环应用的参数量是 1 个循环的 2.51 倍，读取的权重字节数是 2.28 倍，这对内存受限的解码器来说是更好的参考。KV 缓存遵循 fork 的分配规则：per_loop 会为每个循环注意力层的每个循环保留单独的缓存，在 4,096 token 上下文下从 1 个循环的 306 MiB 增长到 4 个循环的 816 MiB，而 last_loop 保持在 306 MiB。
PORT = 8090
KV_LINE = re.compile(r"llama_kv_cache: size =\s*([\d.]+) MiB \(\s*(\d+) cells,\s*(\d+) layers")
class DoryServer:
"""llama-server with one loop count and KV mode. Both are read when the model loads, so each setting is a fresh start."""
def __init__(self, loops=None, mode="per_loop", ngl=999, ubatch=UBATCH):
self.loops, self.mode, self.ngl, self.ubatch = loops, mode, ngl, ubatch
self.name = f"{'default' if loops is None else loops}-{mode}-ngl{ngl}-ub{ubatch}"
def __enter__(self):
cmd = [str(SERVER_BIN), "-m", str(GGUF), "--host", "127.0.0.1", "--port", str(PORT), "-ngl", str(self.ngl),
"--ctx-size", str(CTX), "--parallel", "1", "--batch-size", str(self.ubatch), "--ubatch-size", str(self.ubatch),
"--cache-type-k", "f16", "--cache-type-v", "f16", "--flash-attn", "off", "--no-context-shift",
"--jinja", "--reasoning", "on", "--temp", "1.0", "--top-p", "0.95", "--top-k", "0", "--min-p", "0",
"--repeat-penalty", "1"]
if self.loops is not None:
cmd += ["--override-kv", f"dory.num_loops=int:{self.loops},dory.recurrent_kv_cache_mode=str:{self.mode}"]
self.log = WORK / "logs" / f"server-{self.name}.log"
with open(self.log, "w") as f:
self.proc = subprocess.Popen(cmd, stdout=f, stderr=subprocess.STDOUT, start_new_session=True)
t0 = time.time()
while True:
try:
with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2) as r:
if r.status == 200:
break
except Exception:
pass
if self.proc.poll() is not None or time.time() - t0 > 600:
self.stop()
raise RuntimeError(f"llama-server did not start; see {self.log}")
time.sleep(0.5)
self.load_s = time.time() - t0
text = self.log.read_text(errors="ignore")
sizes = KV_LINE.findall(text) # the full-attention cache first, then the sliding-window cache
self.kv_mib = sum(float(s[0]) for s in sizes)
self.kv_layers = tuple(int(s[2]) for s in sizes)
m = re.search(r"offloaded (\d+)/(\d+) layers", text)
self.offloaded = f"{m.group(1)}/{m.group(2)}" if m else "n/a"
return self
def __exit__(self, *exc):
self.stop()
def stop(self):
if self.proc.poll() is None:
os.killpg(self.proc.pid, signal.SIGTERM)
try:
self.proc.wait(60)
except subprocess.TimeoutExpired:
os.killpg(self.proc.pid, signal.SIGKILL)
self.proc.wait()
def post(self, path, payload, timeout=1800):
req = urllib.request.Request(f"http://127.0.0.1:{PORT}{path}", data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=timeout) as r:
return json.loads(r.read())
def split_reasoning(message):
"""The server returns the thinking in reasoning_content when it can parse it; otherwise it stays inside <think> tags."""
content, reasoning = message.get("content") or "", message.get("reasoning_content") or ""
m = re.search(r"<think>(.*?)</think>", content, re.S)
if m:
reasoning, content = reasoning + m.group(1), content.replace(m.group(0), "")
return reasoning.strip(), content.strip()
@section("5. Start the server, prove the loop override took effect, and catch the silent default")
def override_step():
rows = []
for loops, mode in ((2, "per_loop"), (2, "last_loop"), (None, "per_loop")):
with DoryServer(loops, mode) as s:
want = predicted_kv(META["dory.num_loops"] if loops is None else loops, mode)
rows.append((s.name, s.kv_layers, s.kv_mib, want, s.offloaded, s.load_s))
if loops == 2 and mode == "per_loop":
r = s.post("/v1/chat/completions", {"messages": [{"role": "user", "content": "What is 17 multiplied by 23?"}],
"max_tokens": 1024, "seed": 1, "cache_prompt": False})
reasoning, answer = split_reasoning(r["choices"][0]["message"])
print(f"2 loops, a first request: {r['usage']['completion_tokens']} tokens, "
f"{len(reasoning.split())} words of reasoning, answer: {answer[:120]!r}")
print(f" fields in the message: {sorted(r['choices'][0]['message'])}; "
f"timings: {sorted(r.get('timings', {}))[:6]}...\n")
print(f"{'server':28s} {'KV layers (full, SWA)':>22s} {'KV MiB':>8s} {'predicted':>16s} {'offloaded':>10s} {'load s':>7s}")
for name, layers, mib, want, off, load in rows:
print(f"{name:28s} {str(layers):>22s} {mib:8.1f} {str(want[:2]) + f' {want[2]:.1f}':>16s} {off:>10s} {load:7.1f}")
ok = all(layers == want[:2] for _, layers, _, want, _, _ in rows)
claim("the --override-kv loop count and cache mode are applied (KV layer counts match the prediction)", ok)
default = rows[-1]
claim("without an override the server runs the GGUF's 4 loops (the most expensive setting)",
default[1] == predicted_kv(4, "per_loop")[:2])
return f"override applied: {ok}; no override -> {default[1]} KV layers, {default[2]:.0f} MiB, i.e. 4 loops"
override_step()
循环计数和缓存模式在模型加载时读取，因此每个设置都需要自己的服务器进程；DoryServer 辅助程序会为每个配置启动一个进程，等待健康检查，并从服务器日志中解析 KV 缓存行。我们将解析出的层数与预测进行比较，以证明覆盖确实生效，而不是信任标志。我们还会启动一个不带任何覆盖的服务器作为阴性对照，并发送一个聊天请求，以查看服务器如何将模型的推理与答案分开。
CONFIGS = [(1, "per_loop"), (2, "per_loop"), (2, "last_loop"), (3, "per_loop"), (3, "last_loop"), (4, "per_loop"), (4, "last_loop")]
SPEED_PROMPT = ("Explain, step by step, how a recurrent block that is applied several times per token changes the amount of "
"computation a language model spends on each token, and what that means for memory and latency.")
@section("6. Measure KV memory and decode speed for every loop count and cache mode")
def sweep_step():
global SPEED
SPEED = {}
print(f"{'loops':>5s} {'mode':10s} {'KV MiB':>8s} {'predicted':>9s} {'prefill tok/s':>14s} {'decode tok/s':>13s} "
f"{'decode vs 1 loop':>17s} {'params ratio':>13s}")
for loops, mode in CONFIGS:
with DoryServer(loops, mode) as s:
runs = [s.post("/completion", {"prompt": SPEED_PROMPT, "n_predict": 192, "temperature": 0, "ignore_eos": True,
"cache_prompt": False, "seed": 1})["timings"] for _ in range(3)]
dec = statistics.median(t["predicted_per_second"] for t in runs)
pre = statistics.median(t["prompt_per_second"] for t in runs)
SPEED[(loops, mode)] = {"kv_mib": s.kv_mib, "kv_layers": s.kv_layers, "decode": dec, "prefill": pre,
"decode_runs": [round(t["predicted_per_second"], 2) for t in runs]}
base = SPEED[(1, "per_loop")]["decode"]
print(f"{loops:5d} {mode:10s} {s.kv_mib:8.1f} {predicted_kv(loops, mode)[2]:9.1f} {pre:14.1f} {dec:13.1f} "
f"{base / dec:16.2f}x {PER_TOKEN[loops][0] / PER_TOKEN[1][0]:12.2f}x")
flat = len({round(SPEED[(n, 'last_loop')]['kv_mib'], 1) for n in (2, 3, 4)} | {round(SPEED[(1, 'per_loop')]['kv_mib'], 1)}) == 1
grows = SPEED[(1, "per_loop")]["kv_mib"] < SPEED[(2, "per_loop")]["kv_mib"] < SPEED[(3, "per_loop")]["kv_mib"] < SPEED[(4, "per_loop")]["kv_mib"]
claim("last_loop keeps the KV cache at the one-loop size for every loop count", flat)
claim("per_loop KV memory grows with every added loop", grows)
slow = SPEED[(1, "per_loop")]["decode"] / SPEED[(4, "per_loop")]["decode"]
json.dump({f"{k[0]}-{k[1]}": v for k, v in SPEED.items()}, open(WORK / "speed.json", "w"), indent=1)
return (f"KV {SPEED[(1, 'per_loop')]['kv_mib']:.0f} MiB at 1 loop; 4 loops: {SPEED[(4, 'per_loop')]['kv_mib']:.0f} MiB per_loop, "
f"{SPEED[(4, 'last_loop')]['kv_mib']:.0f} MiB last_loop; decode {slow:.2f}x slower at 4 loops than at 1")
sweep_step()
辅助程序就位后，我们为七种设置各启动一个服务器，即两种缓存模式下的 1 到 4 个循环，并将测得的 KV 内存记录在预测旁边。解码速度来自服务器自身在固定提示上的计时，贪婪生成 192 个 token 三次并取中位数，因此行与行之间唯一变化的是循环设置。将单循环速度除以每行的速度得到测得的减速，我们将其与预测比率并列，以查看解码行为是计算受限还是内存受限。
def top_probs(response):
"""Next-token candidates as {token: logprob}, from either of the server's two probability formats."""
first = response["completion_probabilities"][0]
if "top_logprobs" in first:
return {c["token"]: c["logprob"] for c in first["top_logprobs"]}
return {c["tok_str"]: math.log(max(c["prob"], 1e-30)) for c in first["probs"]}
GATE_PROMPTS = ["The capital of France is", "def fibonacci(n):\n ", "Water boils at a temperature of",
"In 1969, the first humans to walk on the Moon were", "The derivative of x squared is"]
@section("7. Correctness gates: the same seed gives the same answer, and the GPU agrees with the CPU")
def gates_step():
payload = {"messages": [{"role": "user", "content": "Name three prime numbers greater than 50."}],
"max_tokens": 512, "seed": 7, "cache_prompt": False}
probe = lambda s, p: top_probs(s.post("/completion", {"prompt": p, "n_predict": 1, "n_probs": 10, "temperature": 0,
"cache_prompt": False}))
with DoryServer(2, "per_loop") as s:
a, b = (s.post("/v1/chat/completions", payload) for _ in range(2))
same = a["choices"][0]["message"] == b["choices"][0]["message"]
accel = [probe(s, p) for p in GATE_PROMPTS]
print(f"two sampled chat answers with seed 7 at temperature 1.0: identical = {same}")
with DoryServer(2, "per_loop", ngl=0) as s:
cpu = [probe(s, p) for p in GATE_PROMPTS]
agree, gaps = 0, []
for p, x, y in zip(GATE_PROMPTS, accel, cpu):
tx, ty = max(x, key=x.get), max(y, key=y.get)
agree += tx == ty
gaps += [abs(x[t] - y[t]) for t in set(x) & set(y)]
print(f" {p[:40]!r:44s} {BACKEND}: {tx!r:12s} CPU: {ty!r:12s}")
print(f"{BACKEND} vs CPU: top-1 token agrees on {agree} of {len(GATE_PROMPTS)} prompts; largest log-probability gap "
f"over shared top-10 tokens {max(gaps):.3f}")
claim("a fixed seed reproduces a sampled answer exactly", same)
claim(f"the {BACKEND} build agrees with the CPU build on the top-1 token", agree >= len(GATE_PROMPTS) - 1)
return f"seeded answers identical: {same}; {BACKEND} vs CPU top-1 agreement {agree}/{len(GATE_PROMPTS)}"
gates_step()
只有当测量可重复且加速构建与参考 CPU 路径计算相同的模型时，设置之间的比较才有意义。在显卡采样温度下，使用相同种子的两个聊天请求必须返回相同的消息，否则逐问题比较会将模型行为与采样噪声混在一起。然后我们启动一个仅 CPU 的服务器，不卸载任何层，并在五个提示上将最可能的下一个 token 与 GPU 或 Metal 服务器进行比较，因为 Dory 支持尚未在每个后端上游审查。
def load_gsm8k():
import pandas as pd
from huggingface_hub import hf_hub_download
path = hf_hub_download("openai/gsm8k", "main/test-00000-of-00001.parquet", repo_type="dataset", revision=GSM8K_REV,
local_dir=str(WORK / "gsm8k"))
df = pd.read_parquet(path)
df["gold"] = df.answer.str.extract(r"####\s*([-\d,\.]+)")[0].str.replace(",", "").astype(float)
return df
def tv_distance(p, q):
"""Total-variation distance between two top-k next-token distributions, each renormalised over its own top k."""
zp, zq = sum(math.exp(v) for v in p.values()), sum(math.exp(v) for v in q.values())
return 0.5 * sum(abs(math.exp(p.get(t, -1e9)) / zp - math.exp(q.get(t, -1e9)) / zq) for t in set(p) | set(q))
@section("8. What an extra loop changes: next-token distributions on 50 prompts")
def mechanism_step():
global GSM8K
GSM8K = load_gsm8k()
prompts = list(GSM8K.sample(50, random_state=11).question)
configs = [(n, m, ub) for ub in (UBATCH, 16) for n in (1, 2, 4) for m in ("per_loop", "last_loop")]
dist, lengths = {}, []
for n, m, ub in configs:
with DoryServer(n, m, ubatch=ub) as s:
got = [s.post("/completion", {"prompt": p, "n_predict": 1, "n_probs": 20, "temperature": 0, "cache_prompt": False})
for p in prompts]
dist[(n, m, ub)] = [top_probs(r) for r in got]
lengths = lengths or [r["timings"]["prompt_n"] for r in got]
print(f"50 GSM8K questions as raw prompts, {min(lengths)} to {max(lengths)} tokens: one chunk at a {UBATCH}-token "
f"micro-batch, {math.ceil(min(lengths) / 16)} to {math.ceil(max(lengths) / 16)} chunks at 16\n")
print(f"{'comparison':52s} {'same top-1':>10s} {'mean TV':>8s} {'max |dlogp|':>12s}")
def compare(label, a, b):
x, y = dist[a], dist[b]
same = sum(max(p, key=p.get) == max(q, key=q.get) for p, q in zip(x, y))
tv = statistics.mean(tv_distance(p, q) for p, q in zip(x, y))
mx = max(abs(p[t] - q[t]) for p, q in zip(x, y) for t in set(p) & set(q))
print(f"{label:52s} {same:>7d}/50 {tv:8.4f} {mx:12.5f}")
return same, tv, mx
out = {}
out["modes_1"] = compare("1 loop: per_loop vs last_loop", (1, "per_loop", UBATCH), (1, "last_loop", UBATCH))
for n in (2, 4):
out[f"modes_{n}"] = compare(f"{n} loops: per_loop vs last_loop, one {UBATCH}-token chunk",
(n, "per_loop", UBATCH), (n, "last_loop", UBATCH))
out[f"chunks_{n}"] = compare(f"{n} loops: per_loop vs last_loop, 16-token chunks", (n, "per_loop", 16), (n, "last_loop", 16))
print()
out["1v2"] = compare("1 loop vs 2 loops (per_loop)", (1, "per_loop", UBATCH), (2, "per_loop", UBATCH))
out["2v4"] = compare("2 loops vs 4 loops (per_loop)", (2, "per_loop", UBATCH), (4, "per_loop", UBATCH))
out["1v4"] = compare("1 loop vs 4 loops (per_loop)", (1, "per_loop", UBATCH), (4, "per_loop", UBATCH))
claim("with the prompt in one chunk, last_loop gives the same next-token distribution as per_loop",
all(out[f"modes_{n}"][2] == 0 for n in (1, 2, 4)))
claim("split into small chunks, last_loop and per_loop disagree at 2 and 4 loops",
all(out[f"chunks_{n}"][1] > 0 for n in (2, 4)))
return (f"top-1 agreement 1 vs 4 loops {out['1v4'][0]}/50 (mean TV {out['1v4'][1]:.3f}); last_loop = per_loop in one chunk: "
f"{all(out[f'modes_{n}'][2] == 0 for n in (1, 2, 4))}")
mechanism_step()
下一个 token 分布几乎不花成本地展示了循环在模型内部改变了什么。对于用作原始提示的 50 个 GSM8K 问题，我们在每个设置下请求前 20 个下一个 token 候选，并通过 top-1 一致性和总变差距离进行比较。fork 的代码暗示了一个明确的预测：last_loop 在每个循环上覆盖循环层的缓存，因此只有当提示的较早块仅保留其最终循环键时，它才可能与 per_loop 不同，而当整个提示适合一个微批次时，这种情况永远不会发生。用 16 token 的微批次重复探测正好测试了这一点。
EVAL_CONFIGS = [(1, "per_loop"), (2, "per_loop"), (2, "last_loop"), (4, "per_loop"), (4, "last_loop")]
N_QUESTIONS = 30
MAX_TOKENS = 2048
INSTRUCTION = "\n\nSolve the problem. End your reply with a last line of the form 'Final answer: <number>'."
def grade(text, gold):
m = re.findall(r"Final answer:\s*\$?\s*(-?[\d,]*\.?\d+)", text)
nums = m or re.findall(r"-?[\d,]*\.?\d+", text)
if not nums:
return False
try:
return abs(float(nums[-1].replace(",", "")) - gold) < 1e-6
except ValueError:
return False
@section("9. Paired GSM8K evaluation across loop counts and cache modes")
def eval_step():
global EVAL, SAMPLE
SAMPLE = GSM8K.sample(N_QUESTIONS, random_state=20261106).reset_index()
EVAL = {}
for loops, mode in EVAL_CONFIGS:
path = WORK / f"gsm8k-{loops}-{mode}.jsonl"
done = {json.loads(l)["i"]: json.loads(l) for l in open(path)} if path.exists() else {}
todo = [i for i in range(len(SAMPLE)) if i not in done]
if todo:
with DoryServer(loops, mode) as s, open(path, "a") as f:
for i in todo:
row = SAMPLE.iloc[i]
t0 = time.time()
r = s.post("/v1/chat/completions", {"messages": [{"role": "user", "content": row.question + INSTRUCTION}],
"max_tokens": MAX_TOKENS, "seed": 1000 + i, "cache_prompt": False})
reasoning, answer = split_reasoning(r["choices"][0]["message"])
rec = {"i": i, "correct": grade(answer, row.gold), "tokens": r["usage"]["completion_tokens"],
"finish": r["choices"][0]["finish_reason"], "wall_s": round(time.time() - t0, 2),
"decode_tps": round(r["timings"]["predicted_per_second"], 2), "answer": answer[-200:]}
f.write(json.dumps(rec) + "\n")
f.flush()
done[i] = rec
EVAL[(loops, mode)] = [done[i] for i in range(len(SAMPLE))]
rows = EVAL[(loops, mode)]
acc = sum(r["correct"] for r in rows)
print(f"{loops} loops {mode:9s}: {acc:2d}/{len(rows)} correct, median {statistics.median(r['tokens'] for r in rows):6.0f} "
f"tokens, {sum(r['finish'] == 'length' for r in rows)} cut off at {MAX_TOKENS}, "
f"median {statistics.median(r['wall_s'] for r in rows):5.1f} s per question")
return "; ".join(f"{n}/{m[:4]} {sum(r['correct'] for r in EVAL[(n, m)])}/{N_QUESTIONS}" for n, m in EVAL_CONFIGS)
eval_step()
配对 GSM8K 评估在五种设置下对同一批带种子的抽样问题提问，开启推理，使用显卡的采样设置，并为每个问题固定种子，因此每个问题在每次循环计数下都处于相同条件。对于每条回复，我们记录其相对于数据集最终答案的正确性、生成的 token 数（包括推理）、回复是否触及 2,048 token 上限，以及实际耗时。结果到达时按设置分别写入一个 JSONL 文件，因此中断的 Colab 会话可以从停止处继续。
def mcnemar(a, b):
"""Exact two-sided McNemar p-value from paired booleans."""
only_a = sum(x and not y for x, y in zip(a, b))
only_b = sum(y and not x for x, y in zip(a, b))
n, k = only_a + only_b, min(only_a, only_b)
p = 1.0 if n == 0 else min(1.0, 2 * sum(math.comb(n, i) for i in range(k + 1)) / 2**n)
return only_a, only_b, p
def wilson(k, n, z=1.96):
c = (k + z * z / 2) / (n + z * z)
h = z * math.sqrt(k * (n - k) / n + z * z / 4) / (n + z * z)
return 100 * (c - h), 100 * (c + h)
@section("10. Statistics: accuracy, tokens and time per solved question")
def stats_step():
global STATS
rng = random.Random(7)
ref = EVAL[(1, "per_loop")]
print(f"{'setting':18s} {'accuracy':>9s} {'95% CI':>13s} {'median tokens':>14s} {'tokens vs 1 loop [95% CI]':>27s} "
f"{'s per solved':>13s}")
STATS = {}
for key in EVAL_CONFIGS:
rows = EVAL[key]
k, n = sum(r["correct"] for r in rows), len(rows)
ratios = [r["tokens"] / q["tokens"] for r, q in zip(rows, ref)]
boots = sorted(statistics.median(rng.choices(ratios, k=len(ratios))) for _ in range(2000))
per_solved = sum(r["wall_s"] for r in rows) / max(k, 1)
STATS[key] = {"acc": k / n, "ratio": statistics.median(ratios), "ratio_ci": (boots[50], boots[1949]), "per_solved": per_solved,
"tokens": statistics.median(r["tokens"] for r in rows)}
lo, hi = wilson(k, n)
print(f"{key[0]} loops {key[1]:9s} {100 * k / n:8.1f}% {lo:5.1f}-{hi:5.1f}% {STATS[key]['tokens']:14.0f} "
f"{STATS[key]['ratio']:12.2f} [{boots[50]:.2f}, {boots[1949]:.2f}] {per_solved:12.1f}")
print("\npaired accuracy contrasts (McNemar exact; 'only A' = questions only the first setting solved)")
pairs = [((2, "per_loop"), (1, "per_loop")), ((4, "per_loop"), (1, "per_loop")), ((4, "per_loop"), (2, "per_loop")),
((2, "last_loop"), (2, "per_loop")), ((4, "last_loop"), (4, "per_loop"))]
for a, b in pairs:
oa, ob, p = mcnemar([r["correct"] for r in EVAL[a]], [r["correct"] for r in EVAL[b]])
print(f" {a[0]} {a[1]:9s} vs {b[0]} {b[1]:9s}: only A {oa:2d}, only B {ob:2d}, p = {p:.3f}")
STATS[(a, b)] = (oa, ob, p)
fewer = STATS[(4, "per_loop")]["ratio_ci"][1] < 1
claim("on GSM8K, 4 loops need fewer generated tokens than 1 loop (95% CI of the median ratio below 1)", fewer)
return (f"accuracy {', '.join(f'{k[0]}/{k[1][:4]} {100 * STATS[k]['acc']:.0f}%' for k in EVAL_CONFIGS)}; "
f"tokens at 4 loops {STATS[(4, 'per_loop')]['ratio']:.2f}x those at 1 loop")
stats_step()
三十个配对问题无法确定微小的准确率差异，因此统计单元格的设计目的是说明该样本能显示什么、不能显示什么。准确率附带 Wilson 区间，每个对比对仅由一个设置解决的问题使用精确 McNemar 检验。Token 数按相对于一次循环的逐问题比率中位数进行比较，并附带 bootstrap 区间，以此检验显卡关于更多循环需要更少生成 token 的说法。每个已解决问句的耗时将速度、长度和准确率合并为用户实际付出的代价。
@section("11. The card's claims against this run")
def claims_step():
s, sp = STATS, SPEED
rows = [
("The GGUF runs 2 loops by default", f"built-in num_loops = {META['dory.num_loops']}", META["dory.num_loops"] == 2),
("Last-loop caching keeps the KV footprint at the one-loop size",
f"{sp[(1, 'per_loop')]['kv_mib']:.0f} MiB at 1 loop, {sp[(4, 'last_loop')]['kv_mib']:.0f} at 4 last_loop",
round(sp[(1, "per_loop")]["kv_mib"]) == round(sp[(4, "last_loop")]["kv_mib"])),
("Per-loop KV memory grows with the loop count",
" / ".join(f"{sp[(n, 'per_loop')]['kv_mib']:.0f}" for n in LOOPS) + " MiB",
sp[(4, "per_loop")]["kv_mib"] > sp[(1, "per_loop")]["kv_mib"]),
("Loop count trades compute per token without a model swap",
f"decode {sp[(1, 'per_loop')]['decode']:.1f} -> {sp[(4, 'per_loop')]['decode']:.1f} tok/s, same file",
sp[(4, "per_loop")]["decode"] < sp[(1, "per_loop")]["decode"]),
("More loops solve tasks with fewer generated tokens",
f"GSM8K median ratio at 4 loops {s[(4, 'per_loop')]['ratio']:.2f} [{s[(4, 'per_loop')]['ratio_ci'][0]:.2f}, "
f"{s[(4, 'per_loop')]['ratio_ci'][1]:.2f}]", s[(4, "per_loop")]["ratio_ci"][1] < 1),
("Last-loop accuracy stays close to per-loop",
f"4 loops: {100 * s[(4, 'last_loop')]['acc']:.0f}% vs {100 * s[(4, 'per_loop')]['acc']:.0f}%, "
f"McNemar p = {s[((4, 'last_loop'), (4, 'per_loop'))][2]:.2f}", s[((4, "last_loop"), (4, "per_loop"))][2] > 0.05),
]
print(f"{'claim (model card)':60s} {'this run':48s} verdict")
for text, found, ok in rows:
print(f"{text:60s} {found:48s} {'holds' if ok else 'does not hold'}")
claim(text, ok)
print(f"\nnot testable here: the benchmark table (SWE-bench, Terminal-Bench, GPQA ...), the 256K context, NVFP4 pre-training;"
f"\nwith {N_QUESTIONS} paired questions only large accuracy differences are resolvable")
return f"{sum(bool(ok) for *_, ok in rows)} of {len(rows)} card claims hold in this run"
claims_step()
我们不是重述模型卡，而是对照本次运行来检验其说法。每一行将模型卡中的一句话与检验它的测量配对：缓存大小声明和循环默认值对应 token 效率与末次循环缓存的准确率，结论来自保存的结果而非手动输入。笔记本无法触及的说法，例如基准表、256K 上下文和 NVFP4 预训练，被列为未测试，而不是悄悄丢弃。
banner("SUMMARY")
for name, res in RESULTS.items():
print(f" {name:<100s} {res}")
print("\nchecks decided in this run:")
for step, text, held in CLAIMS:
print(f" [{'held' if held else 'did not hold'}] step {step}: {text}")
failed = [k for k, v in RESULTS.items() if str(v).startswith("SKIPPED")]
print(f"\n{sum(h for *_, h in CLAIMS)} of {len(CLAIMS)} checks held; steps failed or skipped: {failed or 'none'}")
print(f"llama.cpp fork {FORK_COMMIT[:12]}, GGUF @ {GGUF_REV[:7]}, backend {BACKEND}{' ' + GPU if GPU else ''}, "
f"Python {platform.python_version()}")
摘要会打印每一步的一行结果以及本次运行做出的每项检查，标记哪些通过、哪些失败，同时给出 fork 提交、GGUF 修订版和后端，因此该笔记本同时充当确切验证了什么以及在何种硬件上验证的记录。
总之，Dory 将测试时计算转变为加载时设置：一个 3.7 GB 文件可作为四种不同模型运行，从其 20 块循环核心的单次通过到四次，而仅 GGUF 头部就能预测每种设置在参数和缓存内存上的开销。笔记本对照服务器检查这些预测，测量速度，用代码给出的精确预测检验末次循环缓存，并在配对 GSM8K 样本上评估准确率和答案长度。无论测量结果如何，有两条实用经验成立：始终显式设置循环计数，因为 GGUF 默认为四；并固定 fork 提交，因为 llama.cpp 中的 Dory 支持仍是未完成的草案。
在此查看完整代码。另外，欢迎在 Twitter 上关注我们，别忘了加入我们的 150k+ ML SubReddit 并订阅我们的新闻通讯。等等！你在 Telegram 上吗？现在你也可以在 Telegram 上加入我们。
需要与我们合作推广你的 GitHub 仓库、Hugging Face 页面、产品发布或网络研讨会等？联系我们
文章如何使用 llama.cpp 运行 NVIDIA Dory 4B：逐循环指南首次出现在MarkTechPost上。
来源：MarkTechPost · marktechpost.com
来源：RadarAI · aihot.news