410 lines
18 KiB
Python
410 lines
18 KiB
Python
"""真实 80B MTP 自投机基准 + 与非投机贪婪逐 token 一致性校验。
|
||
|
||
环境变量:K / MAXTOK / PROMPT / QN_CONFIG / MTP_OUT(其余主模型路径见 validate_mtp)。
|
||
"""
|
||
import json
|
||
import os
|
||
import statistics
|
||
import time
|
||
|
||
import mlx.core as mx
|
||
from mlx_lm.models.qwen3_next import ModelArgs
|
||
|
||
from mlx_streaming.core.moe import native_moe
|
||
from mlx_streaming.core.mem import snapshot, reset_peak
|
||
from mlx_streaming.mtp.drafter import MTPDrafter
|
||
from mlx_streaming.mtp.generate import forward_with_hidden, mtp_generate, prefill_chunked
|
||
from mlx_streaming.mtp.qwen3_next_mtp import load_mtp
|
||
from mlx_streaming.model_builder import build_streaming_model
|
||
|
||
from mlx_streaming import config as _cfg
|
||
QN_CONFIG = _cfg.qn_config()
|
||
MTP_OUT = _cfg.mtp_out()
|
||
PROMPT = os.environ.get("PROMPT", "用三句话解释什么是混合专家模型。")
|
||
MAXTOK = int(os.environ.get("MAXTOK", "96"))
|
||
K = int(os.environ.get("K", "3"))
|
||
PIN_HOT = int(os.environ.get("PIN_HOT", "0"))
|
||
PIN_CAL_TOK = int(os.environ.get("PIN_CAL_TOK", "32"))
|
||
# 稳态测速:warmup 跑满 MAXTOK(热 baseline+spec 两条路径的 Metal kernel 与常驻专家池/预取),
|
||
# 再各重复 REPEAT 次取中位数,避免冷启动编译/补池污染绝对 tok/s。WARMUP_TOK=0 关 warmup。
|
||
WARMUP_TOK = int(os.environ.get("WARMUP_TOK", str(MAXTOK)))
|
||
REPEAT = int(os.environ.get("REPEAT", "3"))
|
||
|
||
|
||
def _spec_once(model, drafter, tok, n, k):
|
||
ids, stats = mtp_generate(model, drafter, tok,
|
||
mx.array([tok.encode(PROMPT)]),
|
||
n, K=k, ids_mode=True, profile=True)
|
||
tps = round(stats["tokens"] / stats["wall_s"], 2)
|
||
return ids, stats, tps
|
||
|
||
|
||
def _baseline_greedy(model, tok, prompt, n):
|
||
cache = model.make_cache()
|
||
ids = mx.array([tok.encode(prompt)])
|
||
t0 = time.perf_counter()
|
||
# prefill 分块:把整段 prompt 的激活峰值压到与 decode 同稳态(见 config.prefill_chunk)。
|
||
logits, _ = prefill_chunked(model, ids, cache)
|
||
_dump_margin = bool(os.environ.get("DUMP_MARGIN"))
|
||
out = []
|
||
for _ in range(n):
|
||
lg = logits[:, -1, :]
|
||
nxt = int(mx.argmax(lg))
|
||
out.append(nxt)
|
||
if _dump_margin:
|
||
# 诊断:每步 top-2 logit 与差值,用于判定路径间发散是「FP 近平局」还是「真错」。
|
||
v = lg.reshape(-1)
|
||
top2 = mx.argpartition(-v, 2)[:2]
|
||
t2 = [int(x) for x in top2.tolist()]
|
||
vv = {t: float(v[t]) for t in t2}
|
||
st = sorted(t2, key=lambda t: -vv[t])
|
||
print(f"MARGIN step={len(out)-1} top1={st[0]}({vv[st[0]]:.5f}) "
|
||
f"top2={st[1]}({vv[st[1]]:.5f}) gap={vv[st[0]]-vv[st[1]]:.6f}", flush=True)
|
||
cur = mx.array([[nxt]])
|
||
mx.eval(cur)
|
||
if len(out) >= n:
|
||
break
|
||
logits, _ = forward_with_hidden(model, cur, cache)
|
||
return out, round(n / (time.perf_counter() - t0), 2)
|
||
|
||
|
||
def main():
|
||
reset_peak()
|
||
model, tok, store = build_streaming_model()
|
||
with open(QN_CONFIG) as f:
|
||
args = ModelArgs.from_dict(json.load(f))
|
||
mtp = load_mtp(args, MTP_OUT, quantize=True)
|
||
mtp.embed_tokens = model.model.embed_tokens # 共享主模型 embedding
|
||
drafter = MTPDrafter(mtp, model.lm_head)
|
||
|
||
# warmup:同时跑 baseline + spec 两条路径,编译 Metal kernel(含 multistate/batch verify)
|
||
# 并把专家常驻池/预取热起来,确保后续测的是稳态而非冷启动。默认 warmup=MAXTOK 跑满全长。
|
||
if WARMUP_TOK > 0:
|
||
_baseline_greedy(model, tok, PROMPT, WARMUP_TOK)
|
||
_spec_once(model, drafter, tok, WARMUP_TOK, K)
|
||
|
||
# ---- baseline 稳态:重复 REPEAT 次取中位数;最后一次清零统计供命中率口径 ----
|
||
base = None
|
||
base_tps_runs = []
|
||
for r in range(REPEAT):
|
||
if r == REPEAT - 1:
|
||
store.reset_stats()
|
||
base, bt = _baseline_greedy(model, tok, PROMPT, MAXTOK)
|
||
base_tps_runs.append(bt)
|
||
base_tps = statistics.median(base_tps_runs)
|
||
base_miss, base_hit = store.misses, store.hits
|
||
base_prefetch_loads = store._resident.prefetch_loads
|
||
base_prefetch_hits = store._resident.prefetch_hits
|
||
|
||
# 可选:baseline 之后、spec 之前校准每层热专家，并预取/钉入 resident pool。
|
||
# 这样 disk_load_ratio 的分母仍是未 pin baseline，能直接衡量 pin 是否压低 spec miss。
|
||
if PIN_HOT > 0:
|
||
store.record = True
|
||
_baseline_greedy(model, tok, PROMPT, PIN_CAL_TOK)
|
||
for li in store.recorded_layers():
|
||
store.pin(li, store.hot(li, PIN_HOT))
|
||
store.record = False
|
||
store.reset_stats()
|
||
|
||
# 诊断:开 handler 触发时刻探针(仅覆盖 spec 阶段,enable 会清零旧日志)。
|
||
_hprof = bool(os.environ.get("STAGING_HPROF"))
|
||
if _hprof:
|
||
from mlx_streaming import native_moe_ext as _Nhp
|
||
_Nhp.staging_hprof_enable(True)
|
||
|
||
# ---- spec 稳态:重复 REPEAT 次取中位数;最后一次清零统计 + reset_peak + 保留其 stats ----
|
||
ids, stats, spec_tps_runs = None, None, []
|
||
for r in range(REPEAT):
|
||
if r == REPEAT - 1:
|
||
store.reset_stats()
|
||
reset_peak()
|
||
from mlx_streaming.core.profiling import tprof_reset, union_reset
|
||
tprof_reset() # 探针只统计最终测量轮(与命中率/内存口径一致)
|
||
union_reset() # 并集专家数也只统计最终轮
|
||
ids, stats, st = _spec_once(model, drafter, tok, MAXTOK, K)
|
||
spec_tps_runs.append(st)
|
||
spec_tps = statistics.median(spec_tps_runs)
|
||
spec_miss, spec_hit = store.misses, store.hits
|
||
spec_prefetch_loads = store._resident.prefetch_loads
|
||
spec_prefetch_hits = store._resident.prefetch_hits
|
||
after = snapshot()
|
||
proj = stats.get("proj_no_replay_tps", 0.0)
|
||
result = {
|
||
"K": K,
|
||
"max_tokens": MAXTOK,
|
||
"warmup_tok": WARMUP_TOK,
|
||
"repeat": REPEAT,
|
||
"exact_match": ids == base,
|
||
"n_mismatch": sum(1 for a, b in zip(ids, base) if a != b),
|
||
"avg_accept_len": stats["avg_accept_len"],
|
||
"steps": stats["steps"],
|
||
"verify_mode": stats.get("verify_mode"),
|
||
"direct_commits": stats.get("direct_commits"),
|
||
"fallback_replays": stats.get("fallback_replays"),
|
||
"replayed_tokens": stats.get("replayed_tokens"),
|
||
"spec_tok_per_s": spec_tps,
|
||
"baseline_tok_per_s": base_tps,
|
||
"speedup": round(spec_tps / max(base_tps, 1e-6), 2),
|
||
"spec_tps_runs": spec_tps_runs,
|
||
"baseline_tps_runs": base_tps_runs,
|
||
"spec_tps_minmax": [min(spec_tps_runs), max(spec_tps_runs)],
|
||
# 分段计时与「重放免费」投机上限
|
||
"t_draft_s": stats.get("t_draft_s"),
|
||
"t_snap_s": stats.get("t_snap_s"),
|
||
"t_verify_s": stats.get("t_verify_s"),
|
||
"t_commit_s": stats.get("t_commit_s"),
|
||
"t_replay_s": stats.get("t_replay_s"),
|
||
"t_sync_s": stats.get("t_sync_s"),
|
||
"t_finalize_s": stats.get("t_finalize_s"),
|
||
"proj_no_replay_tps": proj,
|
||
"proj_no_replay_speedup": round(proj / max(base_tps, 1e-6), 2),
|
||
"baseline_disk_loads": base_miss,
|
||
"baseline_prefetch_loads": base_prefetch_loads,
|
||
"baseline_prefetch_hits": base_prefetch_hits,
|
||
"baseline_hit_rate": round(base_hit / max(base_hit + base_miss, 1), 3),
|
||
"spec_disk_loads": spec_miss,
|
||
"spec_prefetch_loads": spec_prefetch_loads,
|
||
"spec_prefetch_hits": spec_prefetch_hits,
|
||
"spec_hit_rate": round(spec_hit / max(spec_hit + spec_miss, 1), 3),
|
||
"disk_load_ratio": round(spec_miss / max(base_miss, 1), 2),
|
||
# 双源 acquire 分路计数:n_miss==0 走全 GPU 快路径;任一路由 miss 则整层落 host 慢路径
|
||
# (.tolist 全批同步 + demand 读盘)。fallback 占比高 → 即使 hit 高,慢路径仍按"层"频繁触发。
|
||
"gpu_fastpath": getattr(store._resident, "gpu_fastpath", None),
|
||
"gpu_fallback": getattr(store._resident, "gpu_fallback", None),
|
||
"pin_hot": PIN_HOT,
|
||
"pin_cal_tok": PIN_CAL_TOK,
|
||
"pinned_experts": store.pinned_count(),
|
||
"expert_slots": store.capacity,
|
||
"mlx_active_gb": round(after.mlx_active_bytes / 1e9, 2),
|
||
"mlx_peak_gb": round(after.mlx_peak_bytes / 1e9, 2),
|
||
"rss_gb": round(after.rss_bytes / 1e9, 2),
|
||
# 内存分块(清缓冲后的真实常驻):A 权重 / B 专家池 / C staging / D MTP / F 激活
|
||
"mem_breakdown": _mem_breakdown(model, store, mtp),
|
||
"prefill_chunk": _cfg.prefill_chunk(),
|
||
"native_stage_cache": native_moe.stage_cache_stats(),
|
||
"bg_stats": (store._bg.stats() if getattr(store, "_bg", None) is not None else None),
|
||
"window_prof": _window_prof(),
|
||
"predict_recall": _predict_recall(),
|
||
"miss_attrib": _miss_attrib(),
|
||
"prefetch_tprof": _prefetch_tprof(stats.get("wall_s")),
|
||
"union_experts": _union_prof(),
|
||
}
|
||
# 噪声地板测量口径:DUMP_IDS=1 时把 baseline greedy 与 spec 的完整 token 序列打进日志,
|
||
# 供跨进程 run-to-run 逐位对比(默认关闭,不污染常规输出)。
|
||
if os.environ.get("DUMP_IDS"):
|
||
print("DUMP_BASE_IDS " + json.dumps(list(base)))
|
||
print("DUMP_SPEC_IDS " + json.dumps(list(ids)))
|
||
# 字节真值校验自证口径:开 STG_VERIFY 时,把两处校验器的累计计数
|
||
# (ok/bad/calls)打进日志。关键:让「0 BAD」可判真伪——若 calls==0 说明本配置根本
|
||
# 没触发该校验器(如 STG_VERIFY 在 zerocopy_dual 路径不接线),此时 0 BAD 是空结论。
|
||
if os.environ.get("STG_VERIFY"):
|
||
from mlx_streaming.core.cache import resident_pool as _rp_mod
|
||
from mlx_streaming.core.cache import virtual_pool as _vp_mod
|
||
_vsum = {
|
||
"STG_VERIFY.resident(verify_acquire_bytes)": dict(_rp_mod._stg_verify_state),
|
||
"STG_VERIFY.virtual(_verify_native_bytes)": dict(_vp_mod._stg_verify_state),
|
||
}
|
||
print("VERIFY_SUMMARY " + json.dumps(_vsum, ensure_ascii=False))
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
if _hprof:
|
||
# dump (gen, layer, t_fire) 原始日志,供离线分析回调触发时刻分布。
|
||
_flat = _Nhp.staging_hprof_get() # 扁平 [gen,layer,t, ...]
|
||
_path = os.environ.get("STAGING_HPROF_OUT", "/tmp/ab/hprof.jsonl")
|
||
os.makedirs(os.path.dirname(_path), exist_ok=True)
|
||
_n = 0
|
||
with open(_path, "w") as _f:
|
||
for i in range(0, len(_flat), 3):
|
||
_f.write(json.dumps([int(_flat[i]), int(_flat[i + 1]), float(_flat[i + 2])]) + "\n")
|
||
_n += 1
|
||
print(f"[hprof] wrote {_n} handler records to {_path}")
|
||
|
||
|
||
def _tree_nbytes(obj):
|
||
from mlx.utils import tree_flatten
|
||
return sum(v.nbytes for _, v in tree_flatten(obj) if isinstance(v, mx.array))
|
||
|
||
|
||
def _mem_breakdown(model, store, mtp):
|
||
"""把 decode 稳态内存拆成各块(清掉 MLX 可回收缓冲后量真实常驻)。
|
||
|
||
A 常驻非专家权重 / B 专家常驻池 / C staging 预取 / D MTP drafter(扣共享 embedding)/
|
||
F 激活+临时(= active底 − A−B−C−D)。返回 GiB 字典。
|
||
"""
|
||
GIB = 1024 ** 3
|
||
A = _tree_nbytes(model.parameters())
|
||
rp = store._resident
|
||
B = sum(v.nbytes for pool in rp._pools.values() for v in pool.values())
|
||
C = 0
|
||
stg = getattr(store, "_staging", None)
|
||
if stg is not None:
|
||
from mlx.utils import tree_flatten
|
||
C = sum(v.nbytes for _, v in tree_flatten(getattr(stg, "__dict__", {}))
|
||
if isinstance(v, mx.array))
|
||
# D:MTP drafter 全部权重减去与主模型共享的 embedding(run 里 mtp.embed_tokens = 主模型的)
|
||
D = _tree_nbytes(mtp.parameters())
|
||
emb = getattr(getattr(mtp, "embed_tokens", None), "weight", None)
|
||
shared = 0
|
||
if isinstance(emb, mx.array):
|
||
# 共享 embedding 在 A 已计入;若 mtp 的 embed 与主模型同一对象则从 D 扣除避免重复
|
||
shared = _tree_nbytes(mtp.embed_tokens) if hasattr(mtp, "embed_tokens") else 0
|
||
D = max(0, D - shared)
|
||
mx.clear_cache()
|
||
active = mx.get_active_memory()
|
||
peak = mx.get_peak_memory()
|
||
F = max(0, active - (A + B + C + D))
|
||
pool_rows = sum(rp.allocated_slots(l) for l in rp._pools)
|
||
return {
|
||
"A_weights_gib": round(A / GIB, 3),
|
||
"B_expert_pool_gib": round(B / GIB, 3),
|
||
"B_pool_rows": pool_rows,
|
||
"C_staging_gib": round(C / GIB, 3),
|
||
"D_mtp_drafter_gib": round(D / GIB, 3),
|
||
"F_activation_temp_gib": round(F / GIB, 3),
|
||
"active_live_gib": round(active / GIB, 3),
|
||
"peak_gib": round(peak / GIB, 3),
|
||
}
|
||
|
||
|
||
def _union_prof():
|
||
"""按 seq 分桶汇报每层路由专家并集大小(avg/max/p99/分层);seq=K 桶即 MTP verify 的专家并集。
|
||
|
||
需 UNION_PROF=1 才有数据。返回 {seq: {avg_union, max_union, p99_union, per_layer...}};
|
||
verify 桶 = seq==K(verify_in=[x, d_1..d_{K-1}] 恰 K 个 token,见 mtp/generate.py)。
|
||
U_max 即该桶各层并集的全局最大值,是池 cap 下限的正确性依据(cap 必须 ≥ U_max 才保证不溢出)。
|
||
"""
|
||
from mlx_streaming.core.profiling import UNION_PROF as U, UNION_SAMPLES as S
|
||
if not U:
|
||
return None
|
||
|
||
def _p(vals, q):
|
||
# 最近秩法分位数(vals 已升序):q∈[0,1]。样本少时直接给保守上界。
|
||
if not vals:
|
||
return None
|
||
idx = max(0, min(len(vals) - 1, int(round(q * (len(vals) - 1)))))
|
||
return vals[idx]
|
||
|
||
out = {}
|
||
for seq in sorted(U):
|
||
s, n = U[seq]
|
||
rec = {"avg_union": round(s / max(1, n), 2), "n_layer_calls": n}
|
||
# 由分层原始样本汇总出该 seq 桶的 U_max / p99（cap 下限的正确性依据）。
|
||
layer_map = S.get(seq)
|
||
if layer_map:
|
||
flat = sorted(v for lst in layer_map.values() for v in lst)
|
||
rec["max_union"] = flat[-1]
|
||
rec["p99_union"] = _p(flat, 0.99)
|
||
rec["p50_union"] = _p(flat, 0.50)
|
||
rec["n_samples"] = len(flat)
|
||
# 分层分布：每层的 max / avg / 样本数（按层号排序）。
|
||
per_layer = {}
|
||
for li in sorted(layer_map):
|
||
lst = layer_map[li]
|
||
per_layer[li] = {
|
||
"max": max(lst),
|
||
"avg": round(sum(lst) / len(lst), 2),
|
||
"n": len(lst),
|
||
}
|
||
rec["per_layer"] = per_layer
|
||
rec["max_layer_idx"] = max(per_layer, key=lambda li: per_layer[li]["max"])
|
||
out[f"seq{seq}"] = rec
|
||
# 便捷:verify 桶 = seq==K(verify_in=[x, d_1..d_{K-1}] 恰 K 个 token)
|
||
verify_seq = K
|
||
verify = out.get(f"seq{verify_seq}")
|
||
return {"by_seq": out, "verify_seq": verify_seq,
|
||
"verify_avg_union": (verify or {}).get("avg_union"),
|
||
"verify_max_union": (verify or {}).get("max_union"),
|
||
"verify_p99_union": (verify or {}).get("p99_union")}
|
||
|
||
|
||
def _miss_attrib():
|
||
from mlx_streaming.core.profiling import MISS_ATTRIB as M
|
||
if not M["n"]:
|
||
return None
|
||
routed = max(1, M["routed"])
|
||
miss = M["miss_A_predicted"] + M["miss_B_unpredicted"]
|
||
out = {
|
||
"routed": M["routed"],
|
||
"hit_rate": round(M["resident_hit"] / routed, 4),
|
||
"miss_rate": round(miss / routed, 4),
|
||
# A：预测到却没进池（budget/时序/驱逐）；B：没预测到（召回缺口）
|
||
"miss_A_predicted": M["miss_A_predicted"],
|
||
"miss_B_unpredicted": M["miss_B_unpredicted"],
|
||
"A_share_of_miss": round(M["miss_A_predicted"] / max(1, miss), 4),
|
||
"B_share_of_miss": round(M["miss_B_unpredicted"] / max(1, miss), 4),
|
||
}
|
||
if M["dec_n"]:
|
||
# decode/verify 热路径专用（与上面 prefill-主导的全量分开看）：
|
||
dr = max(1, M["dec_routed"])
|
||
dmiss = M["dec_miss_A"] + M["dec_miss_B"]
|
||
out["decode"] = {
|
||
"routed": M["dec_routed"],
|
||
"hit_rate": round(M["dec_resident_hit"] / dr, 4),
|
||
"miss_rate": round(dmiss / dr, 4),
|
||
"miss_A_predicted": M["dec_miss_A"],
|
||
"miss_B_unpredicted": M["dec_miss_B"],
|
||
"A_share_of_miss": round(M["dec_miss_A"] / max(1, dmiss), 4),
|
||
"B_share_of_miss": round(M["dec_miss_B"] / max(1, dmiss), 4),
|
||
# miss_A 细分:时序(pread 没完成) vs 驱逐(就绪过但 acquire 前不在池)
|
||
"miss_A_timing": M["dec_miss_A_timing"],
|
||
"miss_A_evicted": M["dec_miss_A_evicted"],
|
||
"A_timing_share": round(M["dec_miss_A_timing"] / max(1, M["dec_miss_A"]), 4),
|
||
"A_evicted_share": round(M["dec_miss_A_evicted"] / max(1, M["dec_miss_A"]), 4),
|
||
}
|
||
return out
|
||
|
||
|
||
def _predict_recall():
|
||
from mlx_streaming.core.profiling import PREDICT_RECALL_PROF as P
|
||
if not P["n"]:
|
||
return None
|
||
return {"recall": round(P["hit"] / max(1, P["routed"]), 4),
|
||
"avg_routed": round(P["routed"] / P["n"], 2), "n": P["n"]}
|
||
|
||
|
||
def _window_prof():
|
||
from mlx_streaming.core.profiling import WINDOW_PROF
|
||
if not WINDOW_PROF["n"]:
|
||
return None
|
||
return {"avg_ms": round(WINDOW_PROF["sum_s"] / WINDOW_PROF["n"] * 1000, 4),
|
||
"n": WINDOW_PROF["n"]}
|
||
|
||
|
||
def _prefetch_tprof(wall_s=None):
|
||
"""预取 host 墙钟探针汇总(PREFETCH_TPROF=1 才有数据)。
|
||
|
||
汇报每段:总秒数、占最终测量轮 wall 的百分比、单次调用均值(ms)。
|
||
注意:这是"主线程不可重叠的 host 时间";gate matmul / pool scatter 的 GPU 执行落在
|
||
前向末尾统一 eval,不在此口径,需用消融(NATIVE_NO_SUBMIT/NATIVE_NO_PROMOTE)差值测。
|
||
"""
|
||
from mlx_streaming.core.profiling import PREFETCH_TPROF as T
|
||
if not (T["predict_n"] or T["promote_n"] or T["submit_n"]):
|
||
return None
|
||
|
||
def _seg(s, n):
|
||
d = {"total_s": round(T[s], 4)}
|
||
if wall_s:
|
||
d["pct_wall"] = round(T[s] / wall_s * 100, 2)
|
||
if n:
|
||
d["avg_ms"] = round(T[s] / n * 1000, 4)
|
||
return d
|
||
|
||
host_total = T["predict_s"] + T["submit_s"] + T["promote_s"]
|
||
out = {
|
||
"predict": _seg("predict_s", T["predict_n"]),
|
||
"submit": _seg("submit_s", T["submit_n"]),
|
||
"promote": _seg("promote_s", T["promote_n"]),
|
||
# promote 内部细分(take 锁读 / route 成员同步 / place 切片+scatter 入图)
|
||
"promote_take": _seg("take_s", T["promote_n"]),
|
||
"promote_route": _seg("route_s", T["promote_n"]),
|
||
"promote_place": _seg("place_s", T["promote_n"]),
|
||
"host_total_s": round(host_total, 4),
|
||
"host_total_pct_wall": (round(host_total / wall_s * 100, 2) if wall_s else None),
|
||
"counts": {"predict_n": T["predict_n"], "submit_n": T["submit_n"],
|
||
"promote_n": T["promote_n"], "place_experts": T["place_experts"]},
|
||
}
|
||
return out
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|