39 lines
2.5 KiB
Python
39 lines
2.5 KiB
Python
"""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)
|