Match Qwen head-chunked attention execution

tools/issue113-model-reference.py

@@ -0,0 +1,132 @@
+"""Strict default/off head-chunk oracle for real verify rows and matching bank snapshots."""
+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,_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
+
+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)

tools/issue113-reference.py

@@ -0,0 +1,51 @@
+"""Pinned head-chunk plan, actual SDPA subcall shapes and exact BF16 outputs."""
+import hashlib,io,json,os,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.models import qwen4_exp as qwen
+assert mx.__version__=='0.32.2'
+plans=[dict(rows=r,heads=h,kv=kv,plan=qwen._verify_sdpa_head_chunk_plan(r,h,kv)) for r in [-1,0,1,2,3,4,8,9,16,31,32,33,64] for h in [0,1,2,12,24,33,36,64] for kv in [0,1,2,3,5]]
+def pattern(shape,salt):
+ n=int(np.prod(shape));v=((np.arange(n,dtype=np.uint32)*17+salt*13)%257).astype(np.int32)-128
+ return mx.array((v.astype(np.float32)/128).reshape(shape)).astype(mx.bfloat16)
+geometries=[(1,24,2,r,65) for r in [1,2,3,4,8,9,16,32,33]]+[(2,24,2,4,17),(1,64,1,1,17),(1,14,2,5,17),(1,12,2,5,17),(1,1,1,4,17),(1,24,2,4,2),(1,24,2,4,1024),(1,24,2,4,4097),(1,5,2,4,17)]
+original=mx.fast.scaled_dot_product_attention
+calls=[]
+def traced(q,k,v,**kwargs):
+ calls.append(dict(q=list(q.shape),k=list(k.shape),v=list(v.shape),mask=None if kwargs['mask'] is None else list(kwargs['mask'].shape)))
+ return original(q,k,v,**kwargs)
+mx.fast.scaled_dot_product_attention=traced
+cases=[]
+for batch,heads,kv,rows,total in geometries:
+ for mode in ['none','bool','add','heads','local','two_dim']:
+  for enabled in [False,True]:
+   os.environ['MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK']='1' if enabled else '0'
+   calls.clear();case=dict(batch=batch,heads=heads,kv=kv,rows=rows,total=total,mode=mode,enabled=enabled)
+   q=pattern((batch,rows,heads,256),93).transpose(0,2,1,3)
+   k=pattern((batch,kv,total+7,256),94)[:,:,2:total+2]
+   v=pattern((batch,kv,total+7,256),95)[:,:,2:total+2]
+   mask=None
+   if mode!='none':
+    mh=heads if mode=='heads' else (4 if mode=='local' else 1)
+    mask=np.array([j<=total-rows+i and (j+h)%3!=1 and i!=1 for h in range(mh) for i in range(rows) for j in range(total)],dtype=np.bool_).reshape(1,mh,rows,total)
+    mask=mx.array(mask)
+    if mode=='add':mask=mx.where(mask,mx.array(0,dtype=mx.bfloat16),mx.array(-10,dtype=mx.bfloat16))
+    if mode=='two_dim':mask=mask.reshape(rows,total)
+   try:
+    y=qwen._verify_sdpa(q,k,v,scale=.0625,mask=mask)
+    dot=io.StringIO();mx.export_to_dot(dot,output=y)
+    case['fused_primitives']=dot.getvalue().count('ScaledDotProductAttention')
+    mx.eval(y)
+    case['sha256']=hashlib.sha256(np.asarray(y.view(mx.uint16)).tobytes()).hexdigest()
+    case['shape']=list(y.shape)
+   except ValueError as e:case['error']=str(e)
+   case['calls']=list(calls);cases.append(case)
+   del q,k,v,mask
+   mx.clear_cache()
+assert len(cases)==216
+out.write_text(json.dumps(dict(revision='9882703f3105363ddc37eca9f97aa09a1d387112',mlx=mx.__version__,plans=plans,cases=cases),separators=(',',':'))+'\n')
+print(json.dumps(dict(plans=len(plans),cases=len(cases),errors=sum('error' in c for c in cases),output=str(out))))