Compare Qwen cache snapshots at matching stages

tools/issue109-model-reference.py

@@ -1,4 +1,8 @@
-"""Compare actual MTP verify/repair/commit rows and warm continuation to native."""
+"""Strict real MTP oracle: compare the same snapshot stage as the native bank.
+
+#114: snapshot_cache on a lazy trunk adds zero, creating a false signed-zero
+mismatch. Preserve lazy trunk bits and deep-clone only the committed MTP head.
+"""
 import argparse,hashlib,json,os,subprocess,sys
 from pathlib import Path
 checkout,root,native,out=map(Path,sys.argv[1:5])
@@ -17,8 +21,8 @@ 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
-from mtplx.cache_state import snapshot_cache
+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)
@@ -39,9 +43,11 @@ def wrap(fn):
         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(cache).states[35][0].astype(mx.float32))
+                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'
@@ -85,18 +91,16 @@ for phase in ['cold','warm']:
     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):
+    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(cache).states]
-    arrays=dict(trunk=states(final.final_trunk_cache),mtp=states(final.final_committed_mtp_cache))
+            return dict(shape=list(a.shape),sha256=bits)
+        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])):
@@ -108,18 +112,17 @@ for phase in ['cold','warm']:
                 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 in ([], [['trunk',35,0]]), zero_sign_differences
-    # Retain raw evidence of the prefill signed-zero mismatch; never hide it in
-    # the bit-exact verify/logit checks or substitute a numerical tolerance.
-    key=np.asarray(snapshot_cache(final.final_trunk_cache).states[35][0].astype(mx.float32))
+                assert a['shape']==b['shape'] and a['sha256']==b['sha256'],(phase,family,layer,slot,a,b)
+    # 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 np.argwhere(different).tolist() in ([], [[0,1,0,2]])
-    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)
+    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,raw_bits_exact=True,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,

tools/issue113-model-reference.py

@@ -1,132 +1,5 @@
-"""Strict default/off head-chunk oracle for real verify rows and matching bank snapshots."""
-import argparse,hashlib,json,os,subprocess,sys
+"""Use the shared strict real-request oracle for the head-chunk on/off controls."""
 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
+import runpy
 
-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)
+runpy.run_path(str(Path(__file__).with_name("issue109-model-reference.py")), run_name="__main__")

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)