"""Strict default/off head-chunk oracle for real verify rows and matching bank snapshots.""" import argparse,hashlib,json,os,subprocess,sys from pathlib import Path checkout,root,native,out=map(Path,sys.argv[1:5]) assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112' sys.path.insert(0,str(checkout)) from mtplx.server.openai import _server_runtime_env_overrides os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root),generation_mode='mtp'),None)) os.environ['MTPLX_SUSTAINED_PREFILL']='1' os.environ['MTPLX_PREFILL_CHUNK_SIZE']='2048' os.environ['MTPLX_CONTEXT_COPY']='0' import mlx.core as mx import numpy as np from mlx_lm.utils import load_model from mtplx.models import qwen4_exp as qwen from mtplx.runtime import _load_tokenizer_resilient,MTPLXRuntime from mtplx.mtp_patch import MTPContract from mtplx.generation import generate_mtpk from mtplx.sampling import SamplerConfig from mtplx.session_bank import SessionBank,_lazy_snapshot_enabled from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid from mtplx.attention_context import current_attention_phase out.mkdir(parents=True,exist_ok=True) print(json.dumps(dict(case='environment',overrides={k:v for k,v in os.environ.items() if k.startswith('MTPLX_')})),flush=True) tokenizer=_load_tokenizer_resilient(root,json.loads((root/'config.json').read_text())) model,_=load_model(root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen.Model,qwen.ModelArgs)) model.post_weight_load(root) assert model.attach_mtp(root) rt=MTPLXRuntime(model,tokenizer,root,True,MTPContract()) phase='cold'; index=0; nested=0 def wrap(fn): def call(input_ids,*args,**kwargs): global index,nested outer=nested==0 nested+=1 try: result=fn(input_ids,*args,**kwargs) finally: nested-=1 if outer and current_attention_phase() == 'prefill' and isinstance(result,tuple) and result[0] is not None and result[1] is not None: cache=kwargs.get('cache',args[0] if args else None) if cache is not None: key=np.asarray((snapshot_cache_lazy_hybrid(cache) if _lazy_snapshot_enabled() else snapshot_cache(cache)).states[35][0].astype(mx.float32)) (out/f'{phase}-prefill-key35.f32').write_bytes(key.tobytes()) for part,a in zip(['logits','hidden'],result[:2]): # The suffix owner retains only the final hidden row (generation.py:4689). if part=='hidden':a=a[:,-1:,:] data=np.asarray(a.astype(mx.float32)).tobytes() (out/f'{phase}-prefill-{part}.f32').write_bytes(data) path=native/f'{phase}-prefill-{part}.f32' print(json.dumps(dict(case='prefill_compare',phase=phase,part=part,shape=list(a.shape),exact=path.exists() and path.read_bytes()==data)),flush=True) if outer and current_attention_phase()=='decode_verify': expected=json.loads((native/phase/f'{index}.json').read_text()) ids=np.asarray(input_ids).reshape(-1).tolist() assert expected['verify'] and ids==expected['ids'], (phase,index,ids,expected['ids']) receipt=dict(case='verify_forward',phase=phase,index=index,ids=ids) for part,a in zip(['logits','hidden'],result[:2]): data=np.asarray(a.astype(mx.float32)).tobytes() (out/f'{phase}-{index}-{part}.f32').write_bytes(data) ref=(native/phase/f'{index}-{part}.f32').read_bytes() receipt[part+'_exact']=data==ref receipt[part+'_sha256']=hashlib.sha256(data).hexdigest() if data!=ref: x=np.frombuffer(data,np.float32); y=np.frombuffer(ref,np.float32) receipt[part+'_different']=int(np.count_nonzero(x!=y)) receipt[part+'_max_abs']=float(np.max(np.abs(x-y))) print(json.dumps(receipt),flush=True) assert receipt['logits_exact'] and receipt['hidden_exact'],receipt index+=1 return result return call rt.forward_ar=wrap(rt.forward_ar) rt.forward_ar_capture=wrap(rt.forward_ar_capture) bank=SessionBank() for phase in ['cold','warm']: index=0 expected=json.loads((native/f'{phase}.json').read_text()) ids=expected['prompt'] result=generate_mtpk(rt,ids,max_tokens=12,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),speculative_depth=3, seed=1,stop_token_ids=set(tokenizer.eos_token_ids),mtp_history_policy='committed',verify_strategy='batched',capture_final_state=True, session_bank=bank,session_id='issue113',store_prefix_snapshot=False, token_callback=lambda ids:print(json.dumps(dict(case='tokens',phase=phase,ids=ids)),flush=True)) assert result.tokens==expected['tokens'],(result.tokens,expected['tokens']) assert result.stats.cached_tokens==expected['cached'],(result.stats.cached_tokens,expected['cached']) assert index==len(list((native/phase).glob('*.json'))) final=result.final_state assert final.safe_to_commit for part,a in [('logits',final.final_logits),('hidden',final.final_hidden)]: data=np.asarray(a.astype(mx.float32)).tobytes() assert data==(native/f'{phase}-{part}.f32').read_bytes(),(phase,part) def states(cache,lazy=False): def array(a): if a is None: return None values=np.asarray(a.astype(mx.float32)).copy() bits=hashlib.sha256(values.tobytes()).hexdigest() values[values==0]=0 return dict(shape=list(a.shape),sha256=bits,value_sha256=hashlib.sha256(values.tobytes()).hexdigest()) return [[array(a) for a in state] if state is not None else None for state in (snapshot_cache_lazy_hybrid(cache) if lazy else snapshot_cache(cache)).states] arrays=dict(trunk=states(final.final_trunk_cache,_lazy_snapshot_enabled()),mtp=states(final.final_committed_mtp_cache)) (out/f'{phase}-arrays.json').write_text(json.dumps(arrays)) native_arrays=json.loads((native/f'{phase}-arrays.json').read_text()) zero_sign_differences=[] for family in arrays: assert len(arrays[family])==len(native_arrays[family]) for layer,(actual,want) in enumerate(zip(arrays[family],native_arrays[family])): if actual is None or want is None: assert actual==want continue assert len(actual)==len(want) for slot,(a,b) in enumerate(zip(actual,want)): if a is None or b is None: assert a==b continue assert a['shape']==b['shape'] and a['value_sha256']==b['value_sha256'],(phase,family,layer,slot,a,b) if a['sha256']!=b['sha256']: zero_sign_differences.append([family,layer,slot]) assert zero_sign_differences==[], zero_sign_differences # Compare the same lazy trunk snapshot used by the server/native bank. # An extra deep clone would add zero and change negative-zero bits. key=np.asarray((snapshot_cache_lazy_hybrid(final.final_trunk_cache) if _lazy_snapshot_enabled() else snapshot_cache(final.final_trunk_cache)).states[35][0].astype(mx.float32)) (out/f'{phase}-key35.f32').write_bytes(key.tobytes()) want=np.fromfile(native/f'{phase}-key35.f32',np.float32).reshape(key.shape) different=key.view(np.uint32)!=want.view(np.uint32) assert np.all(key[different]==0) and np.all(want[different]==0) assert not np.any(different) for part in ["key35","logits","hidden"]: assert (out/f"{phase}-prefill-{part}.f32").read_bytes()==(native/f"{phase}-prefill-{part}.f32").read_bytes(),(phase,part,"prefill differs") print(json.dumps(dict(case='cache_arrays',phase=phase,values_exact=True,zero_sign_differences=zero_sign_differences,key35_bit_coordinates=np.argwhere(different).tolist())),flush=True) head=snapshot_cache(final.final_committed_mtp_cache) key=ids+result.tokens assert bank.put(runtime=rt,token_ids=key,cache=final.final_trunk_cache,logits=final.final_logits,hidden=final.final_hidden, hidden_variant='post_norm',session_id='issue113',mtp_history_policy='committed',mtp_history_snapshot=head, snapshot_epoch=len(key),mtp_snapshot_epoch=len(key)) is not None print(json.dumps(dict(case='turn',phase=phase,tokens=result.tokens,cached=result.stats.cached_tokens,verify_calls=index,final_exact=True)),flush=True)