46 lines
2.2 KiB
Python
46 lines
2.2 KiB
Python
"""Pinned row-map kernel, guarded stock chain and native synthetic bit receipts."""
|
|
import hashlib,json,subprocess,sys
|
|
from pathlib import Path
|
|
checkout,native=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.kernels import moe_sorted_gather_nax as kernel
|
|
assert mx.__version__=='0.32.2'
|
|
def pattern(count,salt,scale=1):
|
|
a=((np.arange(count,dtype=np.uint32)*17+salt*13)%257).astype(np.int32)-128
|
|
return mx.array(a.astype(np.float32)/128*scale).astype(mx.bfloat16)
|
|
def pack(e,n,k,bits,group,salt,dtype):
|
|
words=np.arange(e*n*k//32*bits,dtype=np.uint32)*np.uint32(2654435761)+np.uint32(12345)
|
|
return (mx.array(words.reshape(e,n,k*bits//32)),pattern(e*n*k//group,salt,1/64).reshape(e,n,k//group).astype(dtype),pattern(e*n*k//group,salt+1,1/64).reshape(e,n,k//group).astype(dtype))
|
|
def digest(a):
|
|
mx.eval(a)
|
|
return hashlib.sha256(np.asarray(a.view(mx.uint16)).tobytes()).hexdigest()
|
|
count=0
|
|
for line in native.read_text().splitlines():
|
|
if '"case":"synthetic"' not in line:continue
|
|
case=json.loads(line[line.index('{'):])
|
|
count+=1
|
|
dtype={'BF16':mx.bfloat16,'F16':mx.float16}[case['dtype']]
|
|
bits,group,rows=case['bits'],case['group'],case['rows']
|
|
gate=pack(128,256,256,bits,group,7,dtype)
|
|
up=pack(128,256,256,bits,group,9,dtype) if case['separate'] else None
|
|
x=pattern(512*256,8).reshape(512,1,256).astype(dtype)
|
|
r=np.arange(rows,dtype=np.uint32)
|
|
ids=mx.array((r*125//rows)//2*2+1)
|
|
row_map=mx.array((r*13)%512)
|
|
assert kernel.applies(x,row_map,*gate,ids,group_size=group,bits=bits,mode='affine',up=up)
|
|
kwargs=dict(group_size=group,bits=bits,swiglu=True,up=up)
|
|
before=kernel.stats()
|
|
y=kernel.gather_rows_qmm(x,row_map,*gate,ids,**kwargs)
|
|
assert y is not None,kernel.stats()
|
|
stock=kernel._stock(x,row_map,*gate,ids,**kwargs)
|
|
actual=digest(y)
|
|
assert actual==digest(stock)==case['sha256'],(case,actual,digest(stock))
|
|
print(json.dumps(dict(case='synthetic_oracle',native=case,exact=True,kernel_before=before,kernel_after=kernel.stats())),flush=True)
|
|
del x,row_map,ids,gate,up,y,stock
|
|
mx.clear_cache()
|
|
|
|
assert count==72,count
|