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