52 lines
3.0 KiB
Python
52 lines
3.0 KiB
Python
"""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))))
|