56 lines
2.9 KiB
Python
56 lines
2.9 KiB
Python
"""Pinned stock guard, installed projection geometry, and fused-oracle controls."""
|
|
import hashlib,json,subprocess,sys
|
|
from pathlib import Path
|
|
checkout,root=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 mlx.nn as nn
|
|
import numpy as np
|
|
from mtplx import moe_sorted_gather as msg
|
|
from mtplx.kernels import moe_sorted_gather_nax as kernel
|
|
assert mx.__version__=='0.32.2' and msg.guard_active()
|
|
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']]
|
|
gu=[mx.concatenate(a,axis=1) for a in zip(parts('gate_proj'),parts('up_proj'))]
|
|
down=parts('down_proj')
|
|
mx.eval(*gu,*down)
|
|
def pattern(rows,k):
|
|
v=((np.arange(rows*k,dtype=np.uint32)*17+7*13)%257).astype(np.int32)-128
|
|
return mx.array((v.astype(np.float32)/128).reshape(rows,1,k)).astype(mx.bfloat16)
|
|
def digest(a): return hashlib.sha256(np.asarray(a.view(mx.uint16)).tobytes()).hexdigest()
|
|
for rows in [20480,32760,32770,34040,40950,40960]:
|
|
for name,layer,k,n in [('gate_up',gu,2560,1280),('down',down,640,2560)]:
|
|
x=pattern(rows,k)
|
|
ids=mx.array(np.arange(rows,dtype=np.uint32)*512//rows)
|
|
for left in ([False,True] if rows==34040 else [False]):
|
|
lhs=mx.arange(rows-1,-1,-1,dtype=mx.int32) if left else None
|
|
before=msg.stats()
|
|
y=msg.gather_qmm(x,*layer,lhs_indices=lhs,rhs_indices=ids,group_size=64,bits=4,sorted_indices=True)
|
|
mx.eval(y)
|
|
assert y.shape==(rows,1,n)
|
|
print(json.dumps(dict(case='projection',projection=name,rows=rows,k=k,n=n,left=left,sha256=digest(y),pad=msg.pad_rows(rows),guard_before=before,guard_after=msg.stats())),flush=True)
|
|
del y
|
|
del x,ids
|
|
mx.clear_cache()
|
|
for tokens in [2048,3277,4095,4096]:
|
|
rows=tokens*10
|
|
x=pattern(tokens,2560)
|
|
row_map=mx.arange(rows,dtype=mx.uint32)%tokens
|
|
ids=mx.array(np.arange(rows,dtype=np.uint32)*512//rows)
|
|
projection=msg.gather_qmm(x[row_map],*gu,rhs_indices=ids,group_size=64,bits=4,sorted_indices=True)
|
|
gate,up=mx.split(projection,2,axis=-1)
|
|
stock=nn.silu(gate)*up
|
|
before=kernel.stats()
|
|
fused=msg.swiglu_rows(x,row_map,*gu,ids,group_size=64,bits=4)
|
|
mx.eval(stock,fused)
|
|
assert digest(stock)==digest(fused)
|
|
after=kernel.stats()
|
|
assert after['calls']>before['calls'],(before,after)
|
|
assert after['canaries']>0 and after['canary_failures']==0,after
|
|
print(json.dumps(dict(case='optimized_oracle_control',tokens=tokens,rows=rows,exact=True,sha256=digest(stock),kernel_before=before,kernel_after=after)),flush=True)
|
|
del x,row_map,ids,projection,gate,up,stock,fused
|
|
mx.clear_cache()
|