Match Qwen compiled verification gate math

tools/issue109-gate-reference.py

@@ -0,0 +1,36 @@
+"""Exact attention gate laws for all finite BF16/F16 values and dense F32 inputs."""
+import hashlib,json,subprocess,sys
+from pathlib import Path
+checkout,out=map(Path,sys.argv[1:3])
+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.attention_math import attention_gate
+from mtplx.attention_context import attention_phase
+from mtplx.models.qwen4_exp import _verify_inject
+out.mkdir(parents=True,exist_ok=True)
+for name,dtype in [('bf16',mx.bfloat16),('f16',mx.float16),('f32',mx.float32)]:
+    if name=='bf16':
+        values=(np.arange(65536,dtype=np.uint32)<<16).view(np.float32)
+    elif name=='f16':
+        values=np.arange(65536,dtype=np.uint16).view(np.float16).astype(np.float32)
+    else:
+        values=np.concatenate([np.linspace(-30,30,32768,dtype=np.float32),np.random.default_rng(5).normal(0,4,32768).astype(np.float32)])
+    values=values[np.isfinite(values)]
+    gate=mx.array(values).astype(dtype)
+    output=mx.array(((np.arange(values.size,dtype=np.int64)*17)%97-48).astype(np.float32)/32).astype(dtype)
+    mx.eval(gate,output)
+    for part,a in [('gate',gate),('output',output)]:
+        (out/f'{name}-{part}.f32').write_bytes(np.asarray(a.astype(mx.float32)).tobytes())
+    for phase in ['prefill','ar_decode','decode_verify']:
+        with attention_phase(phase): result=attention_gate(output,gate)
+        raw=np.asarray(result.astype(mx.float32)).tobytes()
+        (out/f'{name}-{phase}.f32').write_bytes(raw)
+        print(json.dumps(dict(case='gate',dtype=name,phase=phase,count=values.size,sha256=hashlib.sha256(raw).hexdigest())),flush=True)
+        if name=='bf16':
+            inject=_verify_inject(gate,4) if phase=='decode_verify' else 2.0*mx.sigmoid(gate/4)
+            raw=np.asarray(inject.astype(mx.float32)).tobytes()
+            (out/f'inject-{phase}.f32').write_bytes(raw)
+            print(json.dumps(dict(case='inject',dtype=name,phase=phase,count=values.size,sha256=hashlib.sha256(raw).hexdigest())),flush=True)
+

tools/issue109-model-reference.py

@@ -0,0 +1,128 @@
+"""Compare actual MTP verify/repair/commit rows and warm continuation to native."""
+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
+from mtplx.cache_state import snapshot_cache
+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(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]):
+                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='issue109',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):
+        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))
+    (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 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))
+    (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)
+    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='issue109',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)