59 lines
4.1 KiB
Python
59 lines
4.1 KiB
Python
"""Pinned full-model cold/warm prefill, including explicit hidden-state capture."""
|
|
import argparse,hashlib,json,os,subprocess,sys
|
|
from pathlib import Path
|
|
parser=argparse.ArgumentParser(description=__doc__)
|
|
for name in ['checkout','root','native','out']:parser.add_argument(name,type=Path)
|
|
parser.add_argument('--hidden',action='store_true')
|
|
a=parser.parse_args()
|
|
assert subprocess.check_output(['git','-C',str(a.checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
|
|
sys.path.insert(0,str(a.checkout))
|
|
from mtplx.server.openai import _server_runtime_env_overrides
|
|
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(a.root),generation_mode='mtp' if a.hidden else 'ar'),None))
|
|
os.environ['MTPLX_SUSTAINED_PREFILL']='1'
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mlx_lm.utils import load_model
|
|
from mtplx.models import qwen4_exp
|
|
from mtplx.runtime import MTPLXRuntime,_load_tokenizer_resilient
|
|
from mtplx.sampling import SamplerConfig
|
|
from mtplx.generation import generate_ar,restore_or_prefill_prompt_state
|
|
from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid
|
|
from mtplx.session_bank import SessionBank,_lazy_snapshot_enabled
|
|
from mtplx.kernels import moe_sorted_gather_nax as kernel
|
|
assert mx.__version__=='0.32.2'
|
|
a.out.mkdir(parents=True,exist_ok=True)
|
|
model,config=load_model(a.root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen4_exp.Model,qwen4_exp.ModelArgs))
|
|
model.post_weight_load(a.root)
|
|
tokenizer=_load_tokenizer_resilient(a.root,config)
|
|
from mtplx.mtp_patch import MTPContract
|
|
if a.hidden: assert model.attach_mtp(a.root)
|
|
rt=MTPLXRuntime(model,tokenizer,a.root,a.hidden,MTPContract() if a.hidden else None)
|
|
def fingerprint(v):
|
|
values=np.asarray(v.astype(mx.float32)).astype('<f4')
|
|
assert np.isfinite(values).all()
|
|
return dict(shape=list(v.shape),sha256=hashlib.sha256(values.tobytes()).hexdigest())
|
|
for chunk in [2048,4096]:
|
|
os.environ['MTPLX_PREFILL_CHUNK_SIZE']=str(chunk)
|
|
bank=SessionBank()
|
|
for phase in ['cold','warm']:
|
|
name=f'{chunk}-{phase}' if a.hidden else f'{chunk}-0-{phase}'
|
|
expected=json.loads((a.native/(name+'.json')).read_text())
|
|
prompt=expected['state']['ids'] if a.hidden else expected['state']['ids'][:expected['metrics']['prompt']]
|
|
if a.hidden:
|
|
result=restore_or_prefill_prompt_state(rt,prompt,mtp_history_policy='cycle',session_bank=bank,session_id='issue112',capture_hidden=True,store_prefix_snapshot=False,prefill_callback=lambda event:print(json.dumps(dict(case='progress',chunk=chunk,phase=phase,event=event),default=str),flush=True))
|
|
cache,logits,hidden,ids,cached=result.trunk_cache,result.logits,result.hidden,prompt,result.cached_tokens
|
|
else:
|
|
result=generate_ar(rt,prompt,max_tokens=1,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),seed=1,stop_token_ids=set(),capture_final_state=True,session_bank=bank,session_id='issue112',store_prefix_snapshot=False,token_callback=lambda ids:print(json.dumps(dict(case='tokens',chunk=chunk,phase=phase,ids=ids)),flush=True))
|
|
final=result.final_state
|
|
assert final.safe_to_commit
|
|
cache,logits,hidden,ids,cached=final.final_trunk_cache,final.final_logits,final.final_hidden,prompt+result.tokens,result.stats.cached_tokens
|
|
states=[[None if v is None else fingerprint(v) for v in state] if state is not None else None for state in (snapshot_cache_lazy_hybrid(cache) if _lazy_snapshot_enabled() else snapshot_cache(cache)).states]
|
|
record=dict(state=dict(ids=ids,logits=fingerprint(logits),states=states),hidden=None if hidden is None else fingerprint(hidden))
|
|
(a.out/(name+'.json')).write_text(json.dumps(record))
|
|
assert record['state']==expected['state'],(chunk,phase,'state differs')
|
|
assert record['hidden']==expected['hidden'],(chunk,phase,'hidden differs')
|
|
assert cached==(expected['cached'] if a.hidden else expected['metrics']['cached'])
|
|
print(json.dumps(dict(case='model_oracle',chunk=chunk,phase=phase,hidden=a.hidden,exact=True,cached=cached,kernel=kernel.stats())),flush=True)
|
|
if phase=='cold':
|
|
assert bank.put(runtime=rt,token_ids=ids,cache=cache,logits=logits,hidden=hidden,hidden_variant='post_norm',session_id='issue112',mtp_history_policy='cycle',snapshot_epoch=len(ids)) is not None
|