File size: 9,944 Bytes
802f1a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Numerical gates for the isolated V4.1 reference implementation."""
import unittest
import numpy as np
import mlx.core as mx
from runtime import hc_mix, hc_post, rotary, sparse_attention, quantize_cache, TextRuntime, Weights
import hashlib
import json
import tempfile
from pathlib import Path
from unittest.mock import patch
class Residency(unittest.TestCase):
def test_engram_row_cache_hits_and_eviction(self):
with tempfile.TemporaryDirectory() as tmp:
path=Path(tmp);self.fixture(path);w=Weights(path)
w.engram_cache_rows=2
a=w.embedding('layers.1.engram.embed',[1,2]);before=w.disk_bytes_read
b=w.embedding('layers.1.engram.embed',[2,1])
np.testing.assert_array_equal(a[::-1],b)
self.assertEqual(w.disk_bytes_read,before,'Warm Engram rows still read disk')
w.embedding('layers.1.engram.embed',[3])
self.assertEqual(len(w.engram_rows),2)
before=w.disk_bytes_read;w.embedding('layers.1.engram.embed',[2])
self.assertGreater(w.disk_bytes_read,before)
def fixture(self, path):
arrays={name:mx.arange(32,dtype=mx.float32).reshape(8,4) for name in
('embed.weight','layers.1.engram.embed.weight','mtp.0.weight','vision.weight')}
mx.save_safetensors(str(path/'model.safetensors'),arrays)
(path/'config.json').write_text(json.dumps({'quantization':{}}))
(path/'model.safetensors.index.json').write_text(json.dumps({'weight_map':{k:'model.safetensors' for k in arrays}}))
def test_resident_rows_and_engram_disk_reads(self):
with tempfile.TemporaryDirectory() as tmp:
path=Path(tmp);self.fixture(path)
streamed=Weights(path);expected=streamed.read('embed.weight',[6,1,6])
w=Weights(path);w.preload_backbone(budget_bytes=1024)
self.assertEqual(set(w.resident),{'embed.weight'})
self.assertEqual(w.resident_bytes,128)
before=w.disk_bytes_read
np.testing.assert_array_equal(w.read('embed.weight',[6,1,6]),expected)
np.testing.assert_array_equal(w.read('embed.weight'),mx.arange(32).reshape(8,4))
self.assertEqual(w.disk_bytes_read,before)
w.read('layers.1.engram.embed.weight',[1,2])
self.assertEqual(w.disk_bytes_read-before,32)
def test_budget_refusal_before_read(self):
with tempfile.TemporaryDirectory() as tmp:
path=Path(tmp);self.fixture(path);w=Weights(path)
with self.assertRaises(MemoryError):w.preload_backbone(budget_bytes=127)
self.assertEqual(w.disk_bytes_read,0)
self.assertFalse(w.resident)
class Kernels(unittest.TestCase):
def test_deferred_moe_limits_sync_and_preserves_result(self):
m=TextRuntime.__new__(TextRuntime)
m.c={'num_experts_per_tok':3,'norm_topk_prob':True,'routed_scaling_factor':1.5}
class Store:
def read(self,key):
return mx.ones((4,8)) if key.endswith('.weight') else mx.arange(4,dtype=mx.float32)
m.w=Store();m.expert=lambda base,x,routing=None:x*(1 if routing is None else routing)
x=mx.arange(8,dtype=mx.float32)[None,:]
m.execution_mode='reference';expected=m.moe('ffn',x);mx.eval(expected)
m.execution_mode='deferred'
with patch('runtime.mx.eval',wraps=mx.eval) as evaluate:
actual=m.moe('ffn',x)
self.assertLessEqual(evaluate.call_count,1,'Per-expert GPU waits returned')
np.testing.assert_array_equal(actual,expected)
def test_compiled_mixes_matches_reference(self):
import runtime
x=mx.arange(32,dtype=mx.float32).reshape(1,4,8)*.01
w=mx.ones((24,32))*.1;scale=mx.ones((3,));base=mx.zeros((24,))
a=runtime.hc_mix(x,w,scale,base,4,20,1e-6,1e-20)
b=runtime.compiled_hc_mix(x,w,scale,base,4,20,1e-6,1e-20)
for left,right in zip(a,b):np.testing.assert_allclose(left,right,rtol=1e-4,atol=1e-6)
def test_hc_against_numpy(self):
rng=np.random.default_rng(42)
x=rng.normal(size=(1,4,8)).astype('float32')
w=rng.normal(size=(24,32)).astype('float32')*.1
base=rng.normal(size=24).astype('float32')*.1
scale=np.array([.5,.2,.7],dtype='float32')
pre,post,comb=hc_mix(mx.array(x),mx.array(w),mx.array(scale),mx.array(base),4,20,1e-6,1e-20)
m=x.reshape(1,-1)@w.T/np.sqrt(np.mean(x*x,axis=(1,2),keepdims=False)[:,None]+1e-20)
expected_pre=1/(1+np.exp(-(m[:,:4]*scale[0]+base[:4])))+1e-6
expected_post=2/(1+np.exp(-(m[:,4:8]*scale[1]+base[4:8])))
c=(m[:,8:]*scale[2]+base[8:]).reshape(1,4,4)
c=np.exp(c-c.max(-1,keepdims=True));c=c/c.sum(-1,keepdims=True)+1e-6
c=c/(c.sum(-2,keepdims=True)+1e-6)
for _ in range(19):
c=c/(c.sum(-1,keepdims=True)+1e-6);c=c/(c.sum(-2,keepdims=True)+1e-6)
np.testing.assert_allclose(pre,expected_pre,rtol=2e-5,atol=1e-6)
np.testing.assert_allclose(post,expected_post,rtol=2e-5,atol=1e-6)
np.testing.assert_allclose(comb,c,rtol=2e-5,atol=1e-6)
y=rng.normal(size=(1,8)).astype('float32')
expected=y[:,None,:]*expected_post[:,:,None]+np.einsum('bij,bid->bjd',c,x)
np.testing.assert_allclose(hc_post(mx.array(y),mx.array(x),post,comb),expected,rtol=2e-5,atol=1e-6)
def test_attention_sink(self):
q=np.array([[1.,0.],[0.,1.]],dtype='float32')
kv=np.array([[2.,0.],[0.,3.]],dtype='float32')
sink=np.array([1.,2.],dtype='float32')
scores=q@kv.T/np.sqrt(2)
z=np.concatenate([scores,sink[:,None]],axis=-1)
p=np.exp(z-z.max(-1,keepdims=True));p/=p.sum(-1,keepdims=True)
expected=p[:,:2]@kv
np.testing.assert_allclose(sparse_attention(mx.array(q),mx.array(kv),mx.array(sink)),expected,rtol=1e-5)
def test_rotary_inverse(self):
x=mx.array(np.random.default_rng(3).normal(size=(2,16)).astype('float32'))
z=rotary(rotary(x,17,8,10000),17,8,10000,inverse=True)
np.testing.assert_allclose(z,x,rtol=1e-5,atol=1e-6)
def test_cache_quantization(self):
x=mx.array(np.arange(-64,64,dtype='float32').reshape(2,64)/13)
for bits,group,scale in ((8,32,'e8m0'),(4,32,'e8m0'),(4,16,'e4m3')):
y=quantize_cache(x,bits,group,scale)
self.assertEqual(y.shape,x.shape)
self.assertTrue(bool(mx.all(mx.isfinite(y))))
np.testing.assert_allclose(quantize_cache(y,bits,group,scale),y,rtol=1e-5,atol=1e-6)
def test_tiny_model_shared_cache_and_engram(self):
def make():
m=TextRuntime.__new__(TextRuntime)
c=dict(hidden_size=64,hc_mult=4,rms_norm_eps=1e-20,hc_sinkhorn_iters=20,hc_eps=1e-6,num_attention_heads=2,head_dim=64,qk_rope_head_dim=32,o_groups=2,sliding_window=4,compress_ratios=[0,2,1,1],compress_rope_theta=160000,rope_theta=10000,kv_source_layer_ids=[1,2],index_source_layer_ids=[1,2,3],index_n_heads=2,index_head_dim=64,index_topk=2,candidate_source_layer_id=2,candidate_block_size=2,candidate_topk_blocks=1,swiglu_limit=10,num_experts_per_tok=2,norm_topk_prob=True,routed_scaling_factor=1.5,num_hidden_layers=4,engram_layer_ids=[1],dspark_target_layer_ids=[2,3])
class Hash:
def __init__(self):self.history=[]
def push(self,t):self.history.append(t)
def ids(self,layer):return [self.history[-1],self.history[-2] if len(self.history)>1 else 0]
class Store:
config={}
def random(self,key,shape):
seed=int.from_bytes(hashlib.sha256(key.encode()).digest()[:8],'little')
return mx.array(np.random.default_rng(seed).normal(0,.04,size=shape).astype('float32'))
def read(self,key):
if key.endswith('_fn'):shape=(24,256)
elif key.endswith('_base'):shape=(24,)
elif key.endswith('_scale'):shape=(3,)
elif key.endswith('attn_sink'):shape=(2,)
elif key.endswith('gate.weight'):shape=(4,64)
elif key.endswith('gate.bias'):shape=(4,)
elif key.endswith(('q_weight','k_weight')):return mx.ones((4,64))
elif key.endswith('.weight'):return mx.ones((64,))
else:raise AssertionError(key)
return self.random(key,shape)
def linear(self,base,x):
out=64
if base.endswith('engram.wkv'):out=320
elif base.endswith('indexer.wq_b'):out=128
elif base.endswith('weights_proj'):out=2
elif base.endswith('.wq_b'):out=128
elif base=='head':out=32
return x@self.random(base,(out,x.shape[-1])).T
def embedding(self,base,ids):return self.random(base,(64,64))[mx.array(ids)]
def grouped(self,base,x,groups):return mx.concatenate([x[g:g+1]@self.random(base+str(g),(32,x.shape[-1])).T for g in range(groups)],axis=-1)
m.w=Store();m.c=c;m.hash=Hash();m.max_tokens=16;m.position=0;m.states={};m.shared={}
return m
self.tiny_factory=make
a,b,cmodel=make(),make(),make()
b.execution_mode='deferred';cmodel.execution_mode='compiled'
for token in [1,2,3,4,5,6]:
x,h=a.step(token);y,_=b.step(token);z,_=cmodel.step(token)
self.assertEqual(x.shape,(1,32));self.assertEqual(h.shape,(1,128))
self.assertTrue(bool(mx.all(mx.isfinite(x))))
np.testing.assert_allclose(x,y,rtol=1e-6,atol=1e-6)
np.testing.assert_allclose(x,z,rtol=1e-4,atol=1e-5)
self.assertEqual(len(a.states[1]['compressed']),3)
self.assertEqual(len(a.states[2]['compressed']),6)
self.assertEqual(len(a.states[3]['window']),4)
if __name__=='__main__':unittest.main()
|