| """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() |
|
|
