Compare Qwen cache snapshots at matching stages

tools/issue114-snapshot-reference.py

@@ -0,0 +1,38 @@
+"""Isolate the extra oracle snapshot that changed the reported key's zero sign."""
+import hashlib,json,subprocess,sys
+from pathlib import Path
+checkout,native,old_oracle,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))
+import mlx.core as mx
+import numpy as np
+from mtplx.models.qwen4_exp import QSACache
+from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid
+assert mx.__version__=='0.32.2'
+def data(a):return np.asarray(a.view(mx.uint16 if a.dtype==mx.bfloat16 else mx.uint32)).tobytes()
+def cache(a):
+ c=QSACache();c.state=(a,a,None,None);return c
+cases=[]
+for dtype in [mx.bfloat16,mx.float32]:
+ bits=np.array([0x80000000,0,0xbf800000,0x3f800000,0,0x80000000,0x40000000,0xc0000000],dtype=np.uint32)
+ a=mx.array(bits.view(np.float32).reshape(1,2,1,4)).astype(dtype)
+ c=cache(a);raw=c.state[0];lazy=snapshot_cache_lazy_hybrid([c]);deep=snapshot_cache([c])
+ row=dict(dtype=str(dtype),shape=list(a.shape),raw=data(raw).hex(),lazy=data(lazy.states[0][0]).hex(),deep=data(deep.states[0][0]).hex())
+ assert row['raw']==row['lazy'] and row['raw']!=row['deep']
+ cases.append(row)
+traces=[]
+for phase in ['cold-prefill','cold','warm']:
+ f=native/(phase+'-key35.f32');before=f.read_bytes();shape=(1,2,len(before)//(2*256*4),256)
+ a=mx.array(np.frombuffer(before,dtype='<f4').copy().reshape(shape)).astype(mx.bfloat16)
+ def f32(x):return np.asarray(x.astype(mx.float32)).astype('<f4').tobytes()
+ assert f32(a)==before,'input reconstruction must retain every bit'
+ c=cache(a);raw=f32(c.state[0]);lazy=f32(snapshot_cache_lazy_hybrid([c]).states[0][0]);deep=f32(snapshot_cache([c]).states[0][0])
+ original=(old_oracle/(phase+'-key35.f32')).read_bytes()
+ assert raw==lazy==before
+ assert deep==original,'extra clone must reproduce the original mismatch exactly'
+ x=np.frombuffer(raw,np.uint32).reshape(shape);y=np.frombuffer(deep,np.uint32).reshape(shape)
+ coords=np.argwhere(x!=y).tolist()
+ for at in coords:assert x[tuple(at)]==0x80000000 and y[tuple(at)]==0
+ traces.append(dict(phase=phase,shape=shape,coordinates=coords,raw_sha256=hashlib.sha256(raw).hexdigest(),lazy_sha256=hashlib.sha256(lazy).hexdigest(),deep_sha256=hashlib.sha256(deep).hexdigest(),original_oracle_sha256=hashlib.sha256(original).hexdigest()))
+out.write_text(json.dumps(dict(revision='9882703f3105363ddc37eca9f97aa09a1d387112',mlx='0.32.2',cases=cases,real_key_traces=traces),indent=2)+'\n')
+print(json.dumps(dict(cases=len(cases),real_key_traces=traces)),flush=True)