507 lines
26 KiB
Python
507 lines
26 KiB
Python
"""Qwen3-Next MTP 自投机贪婪解码循环。
|
||
|
||
见 spec/plan: docs/superpowers/{specs,plans}/2026-06-07-qwen3next-mtp-self-speculation*。
|
||
|
||
核心:每步主模型前向出 hidden -> MTP 自回归抽 K 草稿 -> 主模型并行验证 ->
|
||
接受最长命中前缀 -> cache 快照/恢复/重放回滚(统一处理 ArraysCache 与 KVCache),
|
||
保证输出与非投机贪婪逐 token 等价。
|
||
|
||
cache 校验/提交机制见 mtp/kv_cache.py；drafter 接口见 mtp/drafter.py。
|
||
"""
|
||
import time
|
||
|
||
import mlx.core as mx
|
||
from mlx_lm.models.base import create_attention_mask, create_ssm_mask
|
||
|
||
from mlx_streaming import config
|
||
from mlx_streaming.mtp.kv_cache import (
|
||
enable_qwen3next_speculative_checkpoints, begin_speculative_checkpoints,
|
||
_snapshot, _restore, commit_verified_prefix, commit_verified_snapshot,
|
||
tile_caches, commit_tree_row)
|
||
|
||
|
||
def _batch_direct_commit_guaranteed(model) -> bool:
|
||
"""判断该模型的 batch verify 是否「必定」能直接提交(commit_verified_prefix 恒成功)。
|
||
|
||
commit 成功的充要条件是每个 cache 都满足:可裁剪(KVCache)或已捕获 per-token checkpoint。
|
||
- 非线性层 → KVCache,可裁剪,必成功;
|
||
- 线性层 → 必须是被 patch 过的 Qwen3NextGatedDeltaNet,verify 前向才会写 `_spec_checkpoints`。
|
||
全部满足才返回 True——此时每步的 `snap_m` 回退快照永远用不上,可安全省略以压低在途峰值。
|
||
通用/玩具递归层(未 patch)返回 False,保留 snap_m 走安全 replay 兜底(与原行为逐 token 等价)。
|
||
"""
|
||
import mlx_streaming.mtp.kv_cache as _kv
|
||
if not getattr(_kv, "_QWEN3NEXT_CHECKPOINTS_PATCHED", False):
|
||
return False
|
||
try:
|
||
from mlx_lm.models.qwen3_next import Qwen3NextGatedDeltaNet
|
||
layers = model.model.layers
|
||
except (ImportError, AttributeError):
|
||
return False
|
||
for l in layers:
|
||
if getattr(l, "is_linear", False) and not isinstance(
|
||
getattr(l, "linear_attn", None), Qwen3NextGatedDeltaNet):
|
||
return False
|
||
return True
|
||
|
||
|
||
def forward_with_hidden(model, ids, cache, compute_logits: bool = True):
|
||
"""跑主模型层循环 + 最终 norm,返回 (logits(1,L,V), hidden(1,L,H))。
|
||
|
||
logits 恒由 final-norm 后的 H 经 lm_head 得到(保证与贪婪一致)。
|
||
返回给 MTP drafter 的 hidden:**默认用 final-norm 之前的 h**——这与 MTP 训练/验证
|
||
(`capture_prenorm_hidden`)消费的输入一致;喂 norm 后的 hidden 会双重归一化、显著压低
|
||
草稿接受率(不影响正确性,因 verify 走主模型)。`MTP_HIDDEN=post_norm` 可切回旧行为对照。
|
||
|
||
compute_logits=False:跳过 lm_head(2048×151936 的大投影),只更新 cache 并返回 hidden。
|
||
供分块 prefill 的中间块用(中间块不需要 logits,只需 cache 因果累积)。
|
||
"""
|
||
inner = model.model
|
||
h = inner.embed_tokens(ids)
|
||
layers = inner.layers
|
||
has_full = any(not l.is_linear for l in layers)
|
||
has_linear = any(l.is_linear for l in layers)
|
||
fa_idx = next((i for i, l in enumerate(layers) if not l.is_linear), 0)
|
||
ssm_idx = next((i for i, l in enumerate(layers) if l.is_linear), 0)
|
||
fa_mask = create_attention_mask(h, cache[fa_idx]) if has_full else None
|
||
ssm_mask = create_ssm_mask(h, cache[ssm_idx]) if has_linear else None
|
||
for layer, c in zip(layers, cache):
|
||
mask = ssm_mask if layer.is_linear else fa_mask
|
||
h = layer(h, mask=mask, cache=c)
|
||
H = inner.norm(h)
|
||
hidden = H if config.mtp_hidden() == "post_norm" else h
|
||
logits = model.lm_head(H) if compute_logits else None
|
||
return logits, hidden
|
||
|
||
|
||
def prefill_chunked(model, ids, cache, chunk: "int | None" = None):
|
||
"""分块 prefill:把 ids 按 chunk 切片逐块喂入,增量更新 cache,返回最后一块的 (logits, hidden)。
|
||
|
||
整段 prefill 一次前向的激活峰值 ∝ prompt 长度(每 MoE 层瞬时物化大量唯一专家 + 长序列激活)。
|
||
分块后每块只算 chunk 个 token,逐块 eval 释放上一块瞬时图,峰值压回 ∝chunk,与 decode 同稳态。
|
||
|
||
数值等价:全注意力层逐块因果累积与整段完全等价;线性层(gated-delta)分块在 chunk 边界处有
|
||
重结合的末位浮点差异(与 decode 本身逐 token 同源,可忽略)。中间块跳过 lm_head 省大投影。
|
||
chunk<=0 或 prompt 不超过一块时,退化为整段 prefill。
|
||
"""
|
||
if chunk is None:
|
||
chunk = config.prefill_chunk()
|
||
L = ids.shape[1]
|
||
if chunk <= 0 or L <= chunk:
|
||
return forward_with_hidden(model, ids, cache)
|
||
last_start = ((L - 1) // chunk) * chunk
|
||
for s in range(0, last_start, chunk):
|
||
_, h = forward_with_hidden(model, ids[:, s:s + chunk], cache, compute_logits=False)
|
||
mx.eval(h) # 物化本块、写 cache、释放上一块瞬时图
|
||
return forward_with_hidden(model, ids[:, last_start:], cache)
|
||
|
||
|
||
def forward_with_hidden_stepwise(model, ids, cache, capture_snapshots: bool = False):
|
||
"""逐 token 走主模型解码路径,可在每个 token 后保存 cache 快照。
|
||
|
||
这是 exact direct-commit 的保守路径:verify 与 baseline 都走 N=1 cache 更新,
|
||
接受前缀时直接恢复到对应 token 后的快照,不再 replay accepted prefix。
|
||
"""
|
||
logits_parts, hidden_parts, snaps = [], [], []
|
||
for i in range(ids.shape[1]):
|
||
logits, H = forward_with_hidden(model, ids[:, i:i + 1], cache)
|
||
mx.eval(logits, H)
|
||
logits_parts.append(logits)
|
||
hidden_parts.append(H)
|
||
if capture_snapshots and i < ids.shape[1] - 1:
|
||
snaps.append(_snapshot(cache))
|
||
logits = mx.concatenate(logits_parts, axis=1)
|
||
H = mx.concatenate(hidden_parts, axis=1)
|
||
if capture_snapshots:
|
||
return logits, H, snaps
|
||
return logits, H
|
||
|
||
|
||
# ----------------------------------------------------------------- 接受判定
|
||
def accept_prefix(drafts, preds):
|
||
"""drafts: MTP 抽的 K 个草稿; preds: 主模型对应位置真实下一 token(len==K)。
|
||
|
||
返回命中草稿数 matched(最长使 preds[j]==drafts[j] 的前缀长度)。
|
||
发射 token 由调用方组装为 drafts[:matched] + [重放末位 bonus],统一处理纠正/全命中,
|
||
保证贪婪下逐 token 与非投机等价。
|
||
"""
|
||
matched = 0
|
||
for d, p in zip(drafts, preds):
|
||
if int(p) == int(d):
|
||
matched += 1
|
||
else:
|
||
break
|
||
return matched
|
||
|
||
|
||
# --------------------------------------------------------- 完整树形验证(batch-of-paths)
|
||
def tree_verify_step(model, drafter, x, H_last, x_ids, mtp_cache, main_cache, K, P, snap_d):
|
||
"""一次 batched 前向并行验证 P 条候选路径,提交接受最长的赢家路径。
|
||
|
||
返回 (new_tokens, x_new, rH, accepted_in, matched)。
|
||
|
||
原理:每条路径 `[x, d_1..d_{K-1}]` 是普通线性序列,拍到 batch 维后线性层/全注意力层都走成熟
|
||
batch 前向(逐 row 与 batch=1 等价,见 benchmarks/test_layerwise.py)。batch=P 的计算量加宽了
|
||
每层预取窗口(hit_rate),多路径又提升接受长度(accept_len)。验证后按赢家 row 提取 checkpoint
|
||
/trim 提交回 batch=1 主 cache,后续解码逐 token 从赢家路径续,与非投机贪婪等价。
|
||
"""
|
||
paths = drafter.draft_paths(H_last, x_ids, mtp_cache, K, P)
|
||
verify_in = mx.array([[x] + p[: K - 1] for p in paths]) # (P, K)
|
||
|
||
snap_m = _snapshot(main_cache) # tile 前的 batch=1 快照(回退用)
|
||
tile_caches(main_cache, len(paths))
|
||
begin_speculative_checkpoints(main_cache)
|
||
vlogits, vH = forward_with_hidden(model, verify_in, main_cache)
|
||
mx.eval(vlogits, vH)
|
||
|
||
best_w, matched, preds = 0, -1, None
|
||
for wi, p in enumerate(paths):
|
||
preds_i = [int(t) for t in mx.argmax(vlogits[wi], axis=-1)]
|
||
m = accept_prefix(p, preds_i)
|
||
if m > matched: # 平局取更靠前的路径(top-1 优先)
|
||
best_w, matched, preds = wi, m, preds_i
|
||
drafts = paths[best_w]
|
||
accepted_len = min(matched + 1, K)
|
||
|
||
committed = commit_tree_row(main_cache, verified_len=K,
|
||
accepted_len=accepted_len, row=best_w)
|
||
if snap_d is not None:
|
||
_restore(mtp_cache, snap_d)
|
||
|
||
if not committed:
|
||
# 理论上 checkpoint 齐全不会走到;稳妥回退:恢复 batch=1 主 cache 后重放赢家 accepted prefix。
|
||
_restore(main_cache, snap_m)
|
||
accepted_in = mx.array([[x] + drafts[:matched]])
|
||
rlogits, rH = forward_with_hidden(model, accepted_in, main_cache)
|
||
mx.eval(rlogits)
|
||
bonus = int(mx.argmax(rlogits[:, -1, :]))
|
||
return drafts[:matched] + [bonus], bonus, rH, accepted_in, matched
|
||
|
||
rH = vH[best_w:best_w + 1, :accepted_len, :]
|
||
accepted_in = verify_in[best_w:best_w + 1, :accepted_len]
|
||
if matched == K:
|
||
return drafts[:K], drafts[-1], rH, accepted_in, matched
|
||
return drafts[:matched] + [preds[matched]], preds[matched], rH, accepted_in, matched
|
||
|
||
|
||
# ----------------------------------------------------------------- 主循环
|
||
def mtp_generate(model, drafter, tok, prompt, max_tokens, K=3, ids_mode=False,
|
||
profile=False, on_tokens=None, main_cache=None, cached_len=0):
|
||
"""贪婪 MTP 自投机。
|
||
|
||
drafter 需提供 draft(H_last(1,1,H), x_ids(1,1), mtp_cache, K) -> list[int](长度 K);
|
||
可选 make_cache()->list 与 sync(rH, replay_in, mtp_cache)。
|
||
ids_mode=True 时 prompt 已是 ids(1,L) 且返回 token id 列表(测试用)。
|
||
on_tokens:可选流式钩子,每步用「本步真正写入 produced 的 token id 列表」回调一次
|
||
(prefill 的首 token 作为第一次回调);返回 True 表示请求尽快停止生成。
|
||
|
||
跨轮复用(main_cache/cached_len):传入已持有前 cached_len 个 prompt token 的 main_cache,
|
||
则只 prefill `prompt[:, cached_len:]`(等价于一次多 token decode,复用历史 KV/递归态,不重算)。
|
||
调用方须保证 prompt[:cached_len] 与 cache 中已有 token 完全一致(否则结果错误)。
|
||
不变式:返回时 main_cache 恰好持有 `prompt + produced[:-1]`(produced[-1] 为 pending,未入 cache),
|
||
调用方可据此拼下轮的 cached prefix。main_cache=None 时内部新建(默认,整段 prefill)。
|
||
"""
|
||
enable_qwen3next_speculative_checkpoints()
|
||
if main_cache is None:
|
||
main_cache = model.make_cache()
|
||
cached_len = 0
|
||
mtp_cache = drafter.make_cache() if hasattr(drafter, "make_cache") else None
|
||
ids = prompt if ids_mode else mx.array([tok.encode(prompt)])
|
||
if cached_len < 0 or cached_len >= ids.shape[1]:
|
||
raise ValueError(
|
||
f"cached_len={cached_len} 必须落在 [0, prompt_len={ids.shape[1]}) 内,至少留 1 个 token 供 prefill")
|
||
|
||
# prefill(分块):只喂尚未在 cache 中的后缀 ids[:, cached_len:](cached_len==0 即整段);
|
||
# 得到第 1 个 pending token x 与其 hidden;分块把激活峰值压到与 decode 同稳态。
|
||
logits, H = prefill_chunked(model, ids[:, cached_len:], main_cache)
|
||
x = int(mx.argmax(logits[:, -1, :]))
|
||
H_last = H[:, -1:, :]
|
||
produced = [x]
|
||
_stop_from_initial = on_tokens is not None and on_tokens([x])
|
||
mx.eval(H_last)
|
||
t0 = time.perf_counter()
|
||
n_steps = 0
|
||
# profile 分段累加(秒):把 sync/finalize 拆开,避免隐藏在 replay 名下。
|
||
t_draft = t_verify = t_replay = t_snap = 0.0
|
||
t_commit = t_sync = t_finalize = 0.0
|
||
direct_commits = fallback_replays = replayed_tokens = 0
|
||
# matched 直方图:accept_hist[j] = 恰好命中 j 个草稿(j∈0..K)的步数。用于实测每位置接受率,
|
||
# 不做几何分布假设。emitted/step = min(matched+1,K)，故上限为 K(本实现 verify 只喂 K-1 草稿)。
|
||
accept_hist = [0] * (max(K, config.depth_max() if config.adaptive_depth() else K) + 1)
|
||
# top-k 覆盖探针:每位置统计模型真实 token 落在 MTP top-1/2/3 的步数(树形救回上界)。
|
||
topk_probe = config.accept_topk() if profile else 0
|
||
tk_n = [0] * K
|
||
tk_cover1 = [0] * K
|
||
tk_cover2 = [0] * K
|
||
tk_cover3 = [0] * K
|
||
verify_mode = config.mtp_verify_mode()
|
||
tree_mode = config.tree_top2()
|
||
tree_p1_mode = tree_mode and config.tree_top2_p1() # 第2 位(pos1)救回,附加在 tree_top2 上
|
||
tree_verify_mode = config.tree_verify() # 完整树形验证(batch-of-paths),优先级最高
|
||
# 置信度门控动态深度:仅在 plain 单链路径生效(与 tree/tree_verify/step 互斥)。
|
||
adaptive_mode = config.adaptive_depth() and not tree_mode and not tree_verify_mode \
|
||
and verify_mode != "step"
|
||
adaptive_rescue_mode = adaptive_mode and config.adaptive_rescue() # 深链叠加 pos0 救回
|
||
conf_tau_v = config.conf_tau()
|
||
depth_max_v = max(1, config.depth_max())
|
||
tree_P = max(1, config.tree_branches())
|
||
tree_rescues = 0 # 第1 位 top-2 成功救回(B 链首命中)的步数
|
||
tree_rescues_p1 = 0 # 第2 位 top-2 成功救回(C 链第2 token 命中)的步数
|
||
# 纯 batch 直接提交路径 + 模型保证 commit 恒成功时,跳过每步一次「全 cache 深拷贝 + eval」
|
||
# (snap_m 只用于 tree 救回 / step / replay 回退,这些路径都不满足下方条件)。省 ~72MiB 在途
|
||
# 峰值和一次同步栅栏,数值完全不变(bit-exact,只改内存调度)。模型结构恒定,循环外算一次即可。
|
||
_skip_snap = (verify_mode != "step") and (not tree_mode) and (not tree_verify_mode) \
|
||
and (not adaptive_rescue_mode) and _batch_direct_commit_guaranteed(model)
|
||
|
||
while len(produced) < max_tokens and not _stop_from_initial:
|
||
x0 = x
|
||
x_ids = mx.array([[x]])
|
||
|
||
snap_d = _snapshot(mtp_cache) if mtp_cache else None
|
||
prev_H_last = H_last
|
||
|
||
# 完整树形验证:一次 batched 前向验证 P 条路径,提交赢家。自带 step 尾处理后 continue。
|
||
if tree_verify_mode and verify_mode != "step":
|
||
if profile:
|
||
_tic = time.perf_counter()
|
||
new_tokens, x, rH, accepted_in, matched = tree_verify_step(
|
||
model, drafter, x, H_last, x_ids, mtp_cache, main_cache, K, tree_P, snap_d)
|
||
accept_hist[matched] += 1
|
||
direct_commits += 1
|
||
if profile:
|
||
t_verify += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
_stop = False
|
||
_n_before = len(produced)
|
||
for t in new_tokens:
|
||
produced.append(t)
|
||
if len(produced) >= max_tokens:
|
||
break
|
||
# 只把「本步真正写入 produced」的 token 交给回调:多 token 步命中 max_tokens
|
||
# 上限时,截断掉的尾巴不应上报给流式消费者(避免超报被丢弃的 token)。
|
||
if on_tokens is not None and on_tokens(produced[_n_before:]):
|
||
_stop = True
|
||
H_last = rH[:, -1:, :]
|
||
if mtp_cache is not None and hasattr(drafter, "sync"):
|
||
drafter.sync(prev_H_last, rH, accepted_in, mtp_cache)
|
||
if profile:
|
||
t_sync += time.perf_counter() - _tic
|
||
n_steps += 1
|
||
mx.eval(x, H_last)
|
||
if _stop:
|
||
break
|
||
continue
|
||
|
||
if profile:
|
||
_tic = time.perf_counter()
|
||
tree_b = None
|
||
tree_c = None
|
||
step_k = K # 本步验证深度(动态深度时可变)
|
||
if tree_mode and verify_mode != "step":
|
||
# chainA(全 top-1)、chainB(第1 位 top-2 → pos0 救回)、chainC(第2 位 top-2 → pos1 救回)
|
||
drafts, tree_b, tree_c = drafter.draft_tree(
|
||
H_last, x_ids, mtp_cache, K, pos1=tree_p1_mode)
|
||
draft_cands = None
|
||
elif adaptive_mode:
|
||
# 置信度门控:返回可变长度(1..depth_max)草稿链,本步深度随置信度自适应。
|
||
if adaptive_rescue_mode:
|
||
# 合并路径:变长 chainA + pos0 分支 chainB(深链才有),供下方 pos0 救回。
|
||
drafts, tree_b = drafter.draft_adaptive_tree(
|
||
H_last, x_ids, mtp_cache, depth_max_v, conf_tau_v)
|
||
else:
|
||
drafts = drafter.draft_adaptive(H_last, x_ids, mtp_cache, depth_max_v, conf_tau_v)
|
||
draft_cands = None
|
||
step_k = len(drafts)
|
||
elif topk_probe > 0:
|
||
drafts, draft_cands = drafter.draft(H_last, x_ids, mtp_cache, K, topk=topk_probe)
|
||
else:
|
||
drafts, draft_cands = drafter.draft(H_last, x_ids, mtp_cache, K), None # 长度 K
|
||
verify_in = mx.array([[x] + drafts[: step_k - 1]]) # [x, d_1..d_{step_k-1}]
|
||
if profile:
|
||
mx.eval(verify_in)
|
||
t_draft += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
|
||
snap_m = None if _skip_snap else _snapshot(main_cache)
|
||
if profile:
|
||
t_snap += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
|
||
verify_snaps = None
|
||
if verify_mode == "step":
|
||
vlogits, vH, verify_snaps = forward_with_hidden_stepwise(
|
||
model, verify_in, main_cache, capture_snapshots=True)
|
||
else:
|
||
begin_speculative_checkpoints(main_cache)
|
||
vlogits, vH = forward_with_hidden(model, verify_in, main_cache)
|
||
mx.eval(vlogits, vH)
|
||
preds = [int(t) for t in mx.argmax(vlogits[0], axis=-1)] # 长度 K
|
||
matched = accept_prefix(drafts, preds)
|
||
# 最小树救回:A 链首草稿被拒(matched==0)且 B 候选=模型真实 token(=preds[0])时,
|
||
# 恢复 cache、改验 B 链。preds[0] 只依赖 [x],两链一致;故 d1b==preds[0] 即 B 首命中。
|
||
# lossless(2026-07-05 定论,见 reports/tree-top2-rescue-2026-07-05.md):曾疑此救回非
|
||
# bit-lossless,实为 demand_dual 的一个 bug——inds(argpartition 切片)是非连续视图,C++
|
||
# 按连续读 data() 致 seq≥2 的 token≥1 装错专家,spec/救回全体失真。修 demand_dual 用
|
||
# contiguous() 物化后,seq=K verify 与 seq=1 解码逐位等价,spec 与本救回均 bit-lossless。
|
||
if tree_b is not None and matched == 0 and tree_b[0] == preds[0]:
|
||
_restore(main_cache, snap_m)
|
||
begin_speculative_checkpoints(main_cache)
|
||
# 重建 verify_in 为 B 链:后续 accepted_in=verify_in[:, :accepted_len] 才会取到 chainB,
|
||
# 否则残留 chainA 的错误首草稿会经 sync 污染 MTP cache、拉低后续草稿质量(低估救回收益)。
|
||
verify_in = mx.array([[x] + tree_b[:step_k - 1]])
|
||
vlogits, vH = forward_with_hidden(model, verify_in, main_cache)
|
||
mx.eval(vlogits, vH)
|
||
preds = [int(t) for t in mx.argmax(vlogits[0], axis=-1)]
|
||
drafts = tree_b
|
||
matched = accept_prefix(drafts, preds)
|
||
tree_rescues += 1
|
||
# 第2 位(pos1)救回:A 链首草稿命中(matched==1)但第2 草稿被拒,且 chainC 的第2 token
|
||
# =模型真值 preds[1] 时,改验 chainC=[x, d1a, d2c, ...]。preds[1] 只依赖 [x, d1a],A/C
|
||
# 前缀一致故 preds[0]/preds[1] 不变、重验后 matched≥2;d3c 可能续命中把接受长度顶到 3。
|
||
# 与 pos0 救回同构(纯重验、只提交模型 argmax token → bit-lossless,并集不放大)。
|
||
elif tree_c is not None and matched == 1 and len(tree_c) > 1 and tree_c[1] == preds[1]:
|
||
_restore(main_cache, snap_m)
|
||
begin_speculative_checkpoints(main_cache)
|
||
verify_in = mx.array([[x] + tree_c[:step_k - 1]])
|
||
vlogits, vH = forward_with_hidden(model, verify_in, main_cache)
|
||
mx.eval(vlogits, vH)
|
||
preds = [int(t) for t in mx.argmax(vlogits[0], axis=-1)]
|
||
drafts = tree_c
|
||
matched = accept_prefix(drafts, preds)
|
||
tree_rescues_p1 += 1
|
||
accept_hist[matched] += 1
|
||
if draft_cands is not None:
|
||
# preds[i]=模型在位置 i 的真实下一 token;查它是否在 MTP 该位置 top-1/2/3 候选里。
|
||
for i in range(K):
|
||
p = preds[i]
|
||
c = draft_cands[i]
|
||
tk_n[i] += 1
|
||
if p == c[0]:
|
||
tk_cover1[i] += 1
|
||
if p in c[:2]:
|
||
tk_cover2[i] += 1
|
||
if p in c[:3]:
|
||
tk_cover3[i] += 1
|
||
if profile:
|
||
t_verify += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
|
||
accepted_len = min(matched + 1, step_k)
|
||
|
||
if verify_snaps is not None:
|
||
# step 模式：逐 token 解码路径产生的快照可精确 direct commit。
|
||
committed = commit_verified_snapshot(main_cache, verify_snaps,
|
||
accepted_len, verified_len=K)
|
||
else:
|
||
# batch 模式：一次并行验证后按接受长度直接提交，零 replay。
|
||
# - 可裁剪 cache（KVCache）：trim 掉 rejected 后缀即精确。
|
||
# - 递归状态 cache（Qwen3-Next gated-delta 的 conv/ssm）：verify 前向里用
|
||
# multistate kernel 捕获的 per-token checkpoint 与 baseline kernel 逐 bit 等价
|
||
# （见 mtp/kv_cache.py 与 tests/test_gated_delta_multistate.py），可精确直提交。
|
||
# commit_verified_prefix 一次处理混合 cache：KV 走 trim、Arrays 走 checkpoint 直提交。
|
||
# 若任一递归 cache 缺 checkpoint（异常情况）则返回 False，自动回退到下方 replay。
|
||
committed = commit_verified_prefix(main_cache, verified_len=step_k,
|
||
accepted_len=accepted_len)
|
||
if snap_d is not None:
|
||
_restore(mtp_cache, snap_d)
|
||
|
||
accepted_in = verify_in[:, :accepted_len]
|
||
if profile:
|
||
t_commit += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
if committed:
|
||
direct_commits += 1
|
||
# 直接使用验证前向产生的 hidden 同步 MTP cache,避免主模型 replay。
|
||
rH = vH[:, :accepted_len, :]
|
||
if matched == step_k:
|
||
new_tokens = drafts[:step_k]
|
||
x = drafts[-1]
|
||
else:
|
||
new_tokens = drafts[:matched] + [preds[matched]]
|
||
x = preds[matched]
|
||
else:
|
||
fallback_replays += 1
|
||
if snap_m is None:
|
||
# 理论不可达:batch 直接提交路径 checkpoint 齐全,commit 必成功。若真走到这里,
|
||
# 说明前提被破坏(如未捕获 checkpoint),此时无快照可回退,直接抛错定位而非静默错算。
|
||
raise RuntimeError("batch verify commit failed but snapshot was skipped")
|
||
# fallback:回滚到验证前,重放 accepted prefix,保证 ArraysCache 正确。
|
||
_restore(main_cache, snap_m)
|
||
accepted_in = mx.array([[x0] + drafts[:matched]])
|
||
replayed_tokens += accepted_in.shape[1]
|
||
rlogits, rH = forward_with_hidden(model, accepted_in, main_cache)
|
||
mx.eval(rlogits)
|
||
bonus = int(mx.argmax(rlogits[:, -1, :]))
|
||
new_tokens = drafts[:matched] + [bonus]
|
||
x = bonus
|
||
if profile:
|
||
t_replay += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
|
||
_n_before = len(produced)
|
||
for t in new_tokens:
|
||
produced.append(t)
|
||
if len(produced) >= max_tokens:
|
||
break
|
||
# 只把「本步真正写入 produced」的 token 交给回调:多 token 步命中 max_tokens
|
||
# 上限时,截断掉的尾巴不应上报给流式消费者(避免超报被丢弃的 token)。
|
||
_stop = on_tokens is not None and on_tokens(produced[_n_before:])
|
||
|
||
H_last = rH[:, -1:, :]
|
||
if mtp_cache is not None and hasattr(drafter, "sync"):
|
||
drafter.sync(prev_H_last, rH, accepted_in, mtp_cache)
|
||
if profile:
|
||
t_sync += time.perf_counter() - _tic
|
||
_tic = time.perf_counter()
|
||
n_steps += 1
|
||
mx.eval(x, H_last)
|
||
if profile:
|
||
t_finalize += time.perf_counter() - _tic
|
||
if _stop:
|
||
break
|
||
|
||
produced = produced[:max_tokens]
|
||
wall = time.perf_counter() - t0
|
||
# cache 实际驻留 token 数(以可裁剪 cache 的 offset 为准),供跨轮复用精确对账:
|
||
# 末步多 token 跨 max_tokens 会 over-commit(cache 领先于 produced),据此识别并禁用复用。
|
||
resident_tokens = None
|
||
for c in main_cache:
|
||
if getattr(c, "is_trimmable", None) and c.is_trimmable() and hasattr(c, "offset"):
|
||
resident_tokens = int(c.offset)
|
||
break
|
||
stats = {
|
||
"steps": n_steps,
|
||
"tokens": len(produced),
|
||
"resident_tokens": resident_tokens,
|
||
"avg_accept_len": round(len(produced) / max(n_steps, 1), 3),
|
||
"wall_s": round(wall, 3),
|
||
"verify_mode": verify_mode,
|
||
"direct_commits": direct_commits,
|
||
"fallback_replays": fallback_replays,
|
||
"replayed_tokens": replayed_tokens,
|
||
"accept_hist": accept_hist, # [恰好命中0,1,...,K 个草稿的步数]
|
||
"tree_rescues": tree_rescues, # 最小树第1 位成功救回步数(tree_top2 开时)
|
||
"tree_rescues_p1": tree_rescues_p1, # 最小树第2 位成功救回步数(tree_top2 开时)
|
||
}
|
||
if topk_probe > 0:
|
||
stats["topk_probe"] = {
|
||
"topk": topk_probe, "n": tk_n,
|
||
"cover_top1": tk_cover1, "cover_top2": tk_cover2, "cover_top3": tk_cover3,
|
||
}
|
||
if profile:
|
||
seg = t_draft + t_snap + t_verify + t_commit + t_replay + t_sync + t_finalize
|
||
stats.update({
|
||
"t_draft_s": round(t_draft, 3),
|
||
"t_snap_s": round(t_snap, 3),
|
||
"t_verify_s": round(t_verify, 3),
|
||
"t_commit_s": round(t_commit, 3),
|
||
"t_replay_s": round(t_replay, 3),
|
||
"t_sync_s": round(t_sync, 3),
|
||
"t_finalize_s": round(t_finalize, 3),
|
||
# 「重放免费」上限:草稿+验证(不含快照,因 per-token 检查点设计无需全 cache 快照)
|
||
"proj_no_replay_tps": round(len(produced) / max(t_draft + t_verify, 1e-6), 2),
|
||
"measured_seg_tps": round(len(produced) / max(seg, 1e-6), 2),
|
||
})
|
||
if ids_mode:
|
||
return produced, stats
|
||
return tok.decode(produced), stats
|