39 lines
1.9 KiB
Python
39 lines
1.9 KiB
Python
"""Installed gate/up row-map oracle, using the identical native fixture inputs."""
|
|
import hashlib,json,subprocess,sys
|
|
from pathlib import Path
|
|
checkout,root,native=map(Path,sys.argv[1:4])
|
|
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'
|
|
index=json.loads((root/'model.safetensors.index.json').read_text())['weight_map']
|
|
base='language_model.model.layers.0.mlp.switch_mlp.'
|
|
weights=mx.load(str(root/index[base+'gate_proj.weight']))
|
|
def parts(proj):return [weights[base+proj+'.'+part] for part in ['weight','scales','biases']]
|
|
gate=parts('gate_proj');up=parts('up_proj')
|
|
gu=[mx.concatenate(a,axis=1) for a in zip(gate,up)]
|
|
mx.eval(*gu,*gate,*up)
|
|
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":"installed_gate_up"' not in line:continue
|
|
case=json.loads(line[line.index('{'):]);count+=1
|
|
tokens,rows=case['tokens'],case['rows']
|
|
v=((np.arange(tokens*2560,dtype=np.uint32)*17+7*13)%257).astype(np.int32)-128
|
|
x=mx.array((v.astype(np.float32)/128).reshape(tokens,1,2560)).astype(mx.bfloat16)
|
|
row_map=mx.arange(rows,dtype=mx.uint32)%tokens
|
|
ids=mx.array(np.arange(rows,dtype=np.uint32)*512//rows)
|
|
g,u=(gate,up) if case['separate'] else (gu,None)
|
|
kw=dict(group_size=64,bits=4,swiglu=True,up=u)
|
|
y=kernel.gather_rows_qmm(x,row_map,*g,ids,**kw)
|
|
stock=kernel._stock(x,row_map,*g,ids,**kw)
|
|
assert y is not None and digest(y)==digest(stock)==case['sha256'],case
|
|
print(json.dumps(dict(case='installed_oracle',native=case,exact=True,kernel=kernel.stats())),flush=True)
|
|
del x,row_map,ids,y,stock
|
|
mx.clear_cache()
|
|
assert count==12,count
|