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,