Port Qwen row-map SwiGLU expert gathers

tools/issue112-reference.py

@@ -0,0 +1,45 @@
+"""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