| """Experimental V4.1 MLX text runtime, not an oMLX plugin. |
| |
| Formula reference: DeepSeek-AI's MIT-licensed V4.1 inference/model.py, |
| inference/engram.py and inference/kernel.py. This implementation uses MLX |
| operations and optional resident weights. Vision is unsupported; native MTP |
| and serial greedy verification are provided separately by dspark.py. |
| """ |
| import argparse |
| from collections import OrderedDict |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import struct |
| import time |
|
|
| import mlx.core as mx |
| import numpy as np |
|
|
|
|
| def rms(x, weight, eps): |
| f=x.astype(mx.float32) |
| return (f*mx.rsqrt(mx.mean(f*f,axis=-1,keepdims=True)+eps)*weight).astype(x.dtype) |
|
|
|
|
| def hc_mix(x, weight, scale, base, hc, iterations, eps, norm_eps): |
| f=x.reshape(*x.shape[:-2],-1).astype(mx.float32) |
| m=(f@weight.T)*mx.rsqrt(mx.mean(f*f,axis=-1,keepdims=True)+norm_eps) |
| pre=mx.sigmoid(m[...,:hc]*scale[0]+base[:hc])+eps |
| post=2*mx.sigmoid(m[...,hc:2*hc]*scale[1]+base[hc:2*hc]) |
| c=(m[...,2*hc:]*scale[2]+base[2*hc:]).reshape(*f.shape[:-1],hc,hc) |
| c=mx.softmax(c,axis=-1)+eps |
| c=c/(mx.sum(c,axis=-2,keepdims=True)+eps) |
| for _ in range(iterations-1): |
| c=c/(mx.sum(c,axis=-1,keepdims=True)+eps) |
| c=c/(mx.sum(c,axis=-2,keepdims=True)+eps) |
| return pre,post,c |
|
|
|
|
| def hc_pre(x, pre): |
| return mx.sum(x.astype(mx.float32)*pre[...,None],axis=-2).astype(x.dtype) |
|
|
|
|
| def hc_post(x, residual, post, comb): |
| return (post[...,None]*x[...,None,:]+mx.einsum('...ij,...id->...jd',comb,residual.astype(mx.float32))).astype(x.dtype) |
|
|
|
|
| def rotary(x, position, dim, base, original=0, factor=1, beta_fast=32, beta_slow=1, inverse=False): |
| freq=1/(base**(mx.arange(0,dim,2,dtype=mx.float32)/dim)) |
| if original: |
| def corrected(r):return dim*math.log(original/(r*2*math.pi))/(2*math.log(base)) |
| low=max(math.floor(corrected(beta_fast)),0);high=min(math.ceil(corrected(beta_slow)),dim-1) |
| smooth=1-mx.clip((mx.arange(dim//2)-low)/max(high-low,1e-3),0,1) |
| freq=freq/factor*(1-smooth)+freq*smooth |
| angle=freq*position*(-1 if inverse else 1) |
| tail=x[...,-dim:].astype(mx.float32).reshape(*x.shape[:-1],dim//2,2) |
| a,b=tail[...,0],tail[...,1] |
| out=mx.stack((a*mx.cos(angle)-b*mx.sin(angle),a*mx.sin(angle)+b*mx.cos(angle)),axis=-1) |
| return mx.concatenate((x[...,:-dim],out.reshape(*x.shape[:-1],dim).astype(x.dtype)),axis=-1) |
|
|
|
|
| def round_fp8(x): |
| a=mx.abs(x);step=2**(mx.maximum(mx.floor(mx.log2(mx.maximum(a,2**-20))),-6)-3) |
| return (mx.sign(x)*mx.minimum(mx.round(a/step)*step,448)).astype(mx.float32) |
|
|
|
|
| def round_fp4(x): |
| |
| levels=mx.array([0,1,2,4,.5,1.5,3,6],dtype=mx.float32) |
| chosen=mx.argmin(mx.abs(mx.abs(x)[...,None]-levels),axis=-1) |
| return mx.sign(x)*levels[chosen] |
|
|
|
|
| def quantize_cache(x, bits, group, scale_format='e8m0'): |
| shape=x.shape;v=x.astype(mx.float32).reshape(*shape[:-1],-1,group) |
| maximum=mx.max(mx.abs(v),axis=-1,keepdims=True) |
| if bits==8: |
| scale=2**mx.ceil(mx.log2(mx.maximum(maximum,1e-4)/448)) |
| out=round_fp8(v/scale)*scale |
| else: |
| if scale_format=='e4m3':scale=round_fp8(mx.maximum(maximum,6*2**-9)/6) |
| else:scale=2**mx.ceil(mx.log2(mx.maximum(maximum,6*2**-126)/6)) |
| out=round_fp4(v/scale)*scale |
| return out.reshape(shape).astype(x.dtype) |
|
|
|
|
| def sparse_attention(q, kv, sink): |
| scores=(q.astype(mx.float32)@kv.astype(mx.float32).T)*(q.shape[-1]**-.5) |
| p=mx.softmax(mx.concatenate((scores,sink[:,None]),axis=-1),axis=-1) |
| return (p[:,:-1]@kv.astype(mx.float32)).astype(q.dtype) |
|
|
|
|
| compiled_hc_mix=mx.compile(hc_mix) |
| compiled_rms=mx.compile(rms) |
|
|
|
|
| class Weights: |
| """Read only requested arrays/embedding rows, bounded LRU on the GPU.""" |
| def __init__(self,path,cache_bytes=512*1024**2): |
| self.path=Path(path);self.config=json.loads((self.path/'config.json').read_text()) |
| self.entries={};self.cache=OrderedDict();self.used=0;self.limit=cache_bytes |
| self.resident={};self.resident_bytes=0 |
| self.disk_bytes_read=0;self.disk_read_calls=0;self.disk_bytes_by_key={} |
| self.engram_cache_rows=0;self.engram_rows=OrderedDict() |
| self.engram_cache_hits=0;self.engram_cache_misses=0 |
| index=json.loads((self.path/'model.safetensors.index.json').read_text())['weight_map'] |
| for name in sorted(set(index.values())): |
| p=self.path/name |
| with p.open('rb') as f: |
| n=struct.unpack('<Q',f.read(8))[0];h=json.loads(f.read(n)) |
| for key,s in h.items(): |
| if key=='__metadata__':continue |
| self.entries[key]=dict(s,path=p,offset=8+n+s['data_offsets'][0]) |
| assert set(index)==set(self.entries) |
| self.q=self.config['quantization'] |
|
|
| @staticmethod |
| def backbone_key(key): |
| return '.engram.embed.' not in key and not key.startswith(('mtp.','vision.','aligner.','image_')) |
|
|
| def preload_backbone(self,budget_bytes=None,progress=False): |
| import psutil |
| keys=[k for k in self.entries if self.backbone_key(k)] |
| total=sum(self.entries[k]['data_offsets'][1]-self.entries[k]['data_offsets'][0] for k in keys) |
| reserve=32*1024**3 |
| device=mx.device_info() |
| safe=min(psutil.virtual_memory().available-reserve, |
| device['max_recommended_working_set_size']-mx.get_active_memory()-reserve) |
| budget=safe if budget_bytes is None else min(safe,budget_bytes) |
| if total>budget:raise MemoryError(f'Backbone needs {total} bytes; safe budget is {budget}') |
| if self.resident:raise RuntimeError('Backbone already preloaded') |
| start=time.monotonic();last=start |
| if progress:print(json.dumps({'phase':'resident_load','bytes':total,'tensors':len(keys),'safe_budget_bytes':budget}),flush=True) |
| for i,key in enumerate(keys): |
| value=self.read(key,_pin=True) |
| self.resident[key]=value;self.resident_bytes+=value.nbytes |
| if progress and time.monotonic()-last>=15: |
| print(json.dumps({'phase':'resident_load','loaded_bytes':self.resident_bytes,'completed':i+1,'total':len(keys),'seconds':time.monotonic()-start}),flush=True) |
| last=time.monotonic() |
| mx.clear_cache() |
| if progress:print(json.dumps({'phase':'resident_ready','resident_bytes':self.resident_bytes,'seconds':time.monotonic()-start}),flush=True) |
|
|
| def record_read(self,key,count): |
| self.disk_bytes_read+=count;self.disk_read_calls+=1 |
| self.disk_bytes_by_key[key]=self.disk_bytes_by_key.get(key,0)+count |
|
|
| def read(self,key,rows=None,_pin=False): |
| if key in self.resident: |
| value=self.resident[key] |
| if rows is None:return value |
| rows=np.asarray(rows,dtype=np.int64).ravel() |
| assert np.all((rows>=0)&(rows<value.shape[0])) |
| return value[mx.array(rows.astype(np.int32))] |
| if rows is None and key in self.cache: |
| self.cache.move_to_end(key);return self.cache[key] |
| s=self.entries[key];dt={'U32':'<u4','U8':'u1','BF16':'<u2','F32':'<f4'}[s['dtype']] |
| width=np.dtype(dt).itemsize;shape=s['shape'] |
| if rows is None: |
| count=math.prod(shape)*width |
| if count>self.limit*2 and not _pin:raise MemoryError('Whole-tensor read blocked; use row lookup for '+key) |
| with s['path'].open('rb') as f:f.seek(s['offset']);blob=f.read(count) |
| assert len(blob)==count |
| self.record_read(key,count) |
| a=np.frombuffer(blob,dtype=dt).reshape(shape).copy() |
| else: |
| rows=np.asarray(rows,dtype=np.int64).ravel();stride=math.prod(shape[1:])*width |
| assert np.all((rows>=0)&(rows<shape[0])) |
| a=np.empty((len(rows),*shape[1:]),dtype=dt) |
| with s['path'].open('rb') as f: |
| for i,row in enumerate(rows): |
| f.seek(s['offset']+int(row)*stride);blob=f.read(stride);assert len(blob)==stride |
| self.record_read(key,stride) |
| a[i]=np.frombuffer(blob,dtype=dt).reshape(shape[1:]) |
| out=mx.array(a) |
| if s['dtype']=='BF16':out=out.view(mx.bfloat16) |
| mx.eval(out) |
| if rows is None and not _pin and out.nbytes<=self.limit: |
| while self.cache and self.used+out.nbytes>self.limit: |
| _,old=self.cache.popitem(last=False);self.used-=old.nbytes |
| self.cache[key]=out;self.used+=out.nbytes |
| return out |
|
|
| def quant_config(self,base): |
| cfg=self.q.get(base,self.q) |
| return {k:cfg[k] for k in ('bits','group_size','mode')} |
|
|
| def linear(self,base,x): |
| if base+'.scales' in self.entries: |
| cfg=self.quant_config(base) |
| bias=self.read(base+'.biases') if base+'.biases' in self.entries else None |
| y=mx.quantized_matmul(x,self.read(base+'.weight'),self.read(base+'.scales'),bias,transpose=True,**cfg) |
| else: |
| w=self.read(base+'.weight');y=x.astype(w.dtype)@w.T |
| if base+'.bias' in self.entries:y=y+self.read(base+'.bias') |
| return y.astype(x.dtype) |
|
|
| def embedding(self,base,ids): |
| if '.engram.embed' in base and self.engram_cache_rows>0: |
| original=np.asarray(ids);flat=[int(i) for i in original.ravel()] |
| missing=list(dict.fromkeys(i for i in flat if (base,i) not in self.engram_rows)) |
| resolved={} |
| for i in flat: |
| if (base,i) in self.engram_rows: |
| resolved[i]=self.engram_rows[base,i];self.engram_rows.move_to_end((base,i)) |
| self.engram_cache_hits+=1 |
| else:self.engram_cache_misses+=1 |
| if missing: |
| values=self._embedding_uncached(base,missing);mx.eval(values) |
| for n,i in enumerate(missing): |
| resolved[i]=values[n];self.engram_rows[base,i]=values[n] |
| while len(self.engram_rows)>self.engram_cache_rows:self.engram_rows.popitem(last=False) |
| result=mx.stack([resolved[i] for i in flat]) |
| return result.reshape(*original.shape,result.shape[-1]) |
| return self._embedding_uncached(base,ids) |
|
|
| def _embedding_uncached(self,base,ids): |
| ids=np.asarray(ids);w=self.read(base+'.weight',ids) |
| if base+'.scales' in self.entries: |
| s=self.read(base+'.scales',ids);b=self.read(base+'.biases',ids) if base+'.biases' in self.entries else None |
| w=mx.dequantize(w,s,b,**self.quant_config(base)) |
| return w.reshape(*ids.shape,w.shape[-1]) |
|
|
| def grouped(self,base,x,groups): |
| shape=self.entries[base+'.weight']['shape'];rows=shape[0]//groups |
| out=[] |
| for g in range(groups): |
| ids=np.arange(g*rows,(g+1)*rows);w=self.read(base+'.weight',ids) |
| if base+'.scales' in self.entries: |
| s=self.read(base+'.scales',ids);b=self.read(base+'.biases',ids) |
| y=mx.quantized_matmul(x[g:g+1],w,s,b,transpose=True,**self.quant_config(base)) |
| else:y=x[g:g+1]@w.T |
| out.append(y) |
| return mx.concatenate(out,axis=-1) |
|
|
|
|
| def is_prime(n): |
| if n<2:return False |
| if n%2==0:return n==2 |
| return all(n%d for d in range(3,math.isqrt(n)+1,2)) |
|
|
|
|
| class EngramHash: |
| def __init__(self,c,tokenizer): |
| from tokenizers import Regex,normalizers |
| normalizer=normalizers.Sequence([normalizers.NFKC(),normalizers.NFD(),normalizers.StripAccents(),normalizers.Lowercase(),normalizers.Replace(Regex(r'[ \t\r\n]+'),' '),normalizers.Replace(Regex(r'^ $'),'\ue000'),normalizers.Strip(),normalizers.Replace('\ue000',' ')]) |
| keys={};self.token_map=[] |
| for token in range(c['vocab_size']): |
| text=tokenizer.decode([token],skip_special_tokens=False) |
| key=tokenizer.id_to_token(token) if '\ufffd' in text else (normalizer.normalize_str(text) or text) |
| if key not in keys:keys[key]=len(keys) |
| self.token_map.append(keys[key]) |
| assert len(keys)==c['engram_compressed_vocab_size'], 'Engram tokenizer map mismatch' |
| self.pad=self.token_map[c['engram_pad_token_id']];self.history=[];self.tables={} |
| seen=set();order=c['engram_max_ngram_size'] |
| for layer,expected in zip(c['engram_layer_ids'],c['engram_num_embeddings']): |
| primes=[] |
| for _ in range(order-1): |
| current=c['engram_vocab_size']-1 |
| for _ in range(c['engram_n_heads']): |
| current+=1 |
| while current in seen or not is_prime(current):current+=1 |
| seen.add(current);primes.append(current) |
| assert sum(primes)==expected |
| rng=np.random.default_rng(10007*layer) |
| bound=max(1,((2**63-1)//len(keys))//2) |
| mult=rng.integers(0,bound,size=order,dtype=np.int64)*2+1 |
| self.tables[layer]=(primes,np.cumsum([0,*primes[:-1]]),mult) |
| self.order=order;self.heads=c['engram_n_heads'] |
|
|
| def push(self,token):self.history.append(self.token_map[token]) |
|
|
| def ids(self,layer): |
| primes,offsets,mult=self.tables[layer] |
| tokens=[self.history[-1-i] if i<len(self.history) else self.pad for i in range(self.order)] |
| rolling=int(tokens[0])*int(mult[0]);ids=[] |
| for i in range(1,self.order): |
| rolling^=int(tokens[i])*int(mult[i]) |
| for head in range(self.heads): |
| j=(i-1)*self.heads+head;ids.append(rolling%primes[j]+int(offsets[j])) |
| return ids |
|
|
|
|
| class TextRuntime: |
| def __init__(self,path,max_tokens=256,resident_backbone=False,execution_mode='reference'): |
| from tokenizers import Tokenizer |
| self.w=Weights(path);self.c=self.w.config['text_config'];self.max_tokens=max_tokens |
| assert self.w.config['model_type']=='deepseek_v41' |
| self.tokenizer=Tokenizer.from_file(str(Path(path)/'tokenizer.json')) |
| self.hash=EngramHash(self.c,self.tokenizer) |
| self.position=0;self.states={};self.shared={} |
| if execution_mode not in ('reference','deferred','compiled'):raise ValueError('Unknown execution mode') |
| self.execution_mode=execution_mode |
| if resident_backbone:self.w.preload_backbone(progress=True) |
|
|
| def norm(self,name,x): |
| fn=compiled_rms if getattr(self,'execution_mode','reference')=='compiled' else rms |
| return fn(x,self.w.read(name+'.weight'),self.c['rms_norm_eps']) |
|
|
| def rotate(self,x,pos,ratio,inverse=False): |
| c=self.c;r=c.get('rope_scaling',{}) |
| return rotary(x,pos,c['qk_rope_head_dim'],c['compress_rope_theta'] if ratio else c['rope_theta'],r.get('original_max_position_embeddings',0) if ratio else 0,r.get('factor',1),r.get('beta_fast',32),r.get('beta_slow',1),inverse) |
|
|
| def mixes(self,base,kind,x): |
| c=self.c |
| fn=compiled_hc_mix if getattr(self,'execution_mode','reference')=='compiled' else hc_mix |
| return fn(x,self.w.read(base+'.hc_'+kind+'_fn'),self.w.read(base+'.hc_'+kind+'_scale'),self.w.read(base+'.hc_'+kind+'_base'),c['hc_mult'],c['hc_sinkhorn_iters'],c['hc_eps'],c['rms_norm_eps']) |
|
|
| def engram(self,base,h,layer): |
| ids=self.hash.ids(layer) |
| values=self.w.embedding(base+'.embed',ids).reshape(1,-1) |
| kv=self.w.linear(base+'.wkv',values).astype(mx.float32) |
| c=self.c;dim=c['hidden_size'];hc=c['hc_mult'];eps=c['rms_norm_eps'] |
| key=kv[:,:hc*dim].reshape(1,hc,dim);value=kv[:,hc*dim:] |
| x=h.astype(mx.float32);weight=self.w.read(base+'.q_weight').astype(mx.float32)*self.w.read(base+'.k_weight').astype(mx.float32) |
| std=mx.rsqrt(mx.mean(x*x,axis=-1)+eps)*mx.rsqrt(mx.mean(key*key,axis=-1)+eps) |
| dot=mx.sum(x*key*weight,axis=-1)*std*dim**-.5 |
| gate=mx.sigmoid(mx.where(dot<0,-1,1)*mx.sqrt(mx.maximum(mx.abs(dot),1e-6))) |
| return (x+gate[...,None]*value[:,None,:]).astype(h.dtype) |
|
|
| def expert(self,base,x,routing=None): |
| a=self.w.linear(base+'.w1',x).astype(mx.float32);b=self.w.linear(base+'.w3',x).astype(mx.float32) |
| limit=self.c['swiglu_limit'] |
| if limit>0:a=mx.minimum(a,limit);b=mx.clip(b,-limit,limit) |
| y=a*mx.sigmoid(a)*b |
| if routing is not None:y=y*routing |
| return self.w.linear(base+'.w2',y.astype(x.dtype)) |
|
|
| def moe(self,base,x): |
| c=self.c |
| logits=x.astype(mx.float32)@self.w.read(base+'.gate.weight').astype(mx.float32).T |
| scores=mx.sqrt(mx.logaddexp(logits,mx.zeros_like(logits))) |
| k=c['num_experts_per_tok'] |
| picks=mx.argsort(scores+self.w.read(base+'.gate.bias'),axis=-1)[:,-k:] |
| selected=mx.take_along_axis(scores,picks,axis=-1) |
| if c['norm_topk_prob'] and k>1:selected=selected/(mx.sum(selected,axis=-1,keepdims=True)+1e-20) |
| selected=selected*c['routed_scaling_factor'];mx.eval(picks,selected) |
| result=mx.zeros_like(x).astype(mx.float32) |
| for index,weight in zip(picks.tolist()[0],selected.tolist()[0]): |
| result=result+self.expert(base+f'.experts.{index}',x,weight).astype(mx.float32) |
| if getattr(self,'execution_mode','reference')=='reference':mx.eval(result) |
| return (result+self.expert(base+'.shared_experts',x).astype(mx.float32)).astype(x.dtype) |
|
|
| def attention(self,base,x,layer): |
| c=self.c;t=self.position;ratio=c['compress_ratios'][layer] |
| st=self.states.setdefault(layer,{'window':[],'pool':[],'score':[],'compressed':[],'index':[]}) |
| qr=self.norm(base+'.q_norm',self.w.linear(base+'.wq_a',x)) |
| q=self.w.linear(base+'.wq_b',qr).reshape(c['num_attention_heads'],c['head_dim']) |
| q=self.rotate(q,t,ratio) |
| kv=self.norm(base+'.kv_norm',self.w.linear(base+'.wkv',x)) |
| kv=quantize_cache(self.rotate(kv,t,ratio),8,32) |
| st['window']=(st['window']+[kv])[-c['sliding_window']:] |
| selected=mx.concatenate(st['window'],axis=0) |
| latent=None |
| if ratio: |
| if layer in c['kv_source_layer_ids']: |
| pool=self.w.linear(base+'.compressor.wkv',x.astype(mx.float32) if ratio>1 else x) |
| if ratio>1: |
| st['pool'].append(pool);st['score'].append(self.w.linear(base+'.compressor.wgate',x.astype(mx.float32))) |
| if (t+1)%ratio==0: |
| a=mx.concatenate(st['pool'],axis=0);b=mx.concatenate(st['score'],axis=0) |
| latent=mx.sum(a*mx.softmax(b,axis=0),axis=0,keepdims=True).astype(x.dtype) |
| st['pool']=[];st['score']=[] |
| else:latent=pool |
| if latent is not None:latent=self.norm(base+'.compressor.norm',latent) |
| self.shared['compressed']=st['compressed'];self.shared['index']=st['index'] |
| if layer in c['index_source_layer_ids']: |
| if layer in c['kv_source_layer_ids'] and latent is not None: |
| ik=self.norm(base+'.indexer.k_norm',self.w.linear(base+'.indexer.wk',latent)) |
| st['index'].append(quantize_cache(self.rotate(ik,t+1-ratio,ratio),4,32)) |
| n=(t+1)//ratio |
| if n: |
| iq=self.w.linear(base+'.indexer.wq_b',qr).reshape(c['index_n_heads'],c['index_head_dim']) |
| iq=quantize_cache(self.rotate(iq,t,ratio),4,32) |
| keys=mx.concatenate(self.shared['index'],axis=0) |
| importance=self.w.linear(base+'.indexer.weights_proj',x).reshape(-1)*c['index_head_dim']**-.5*c['index_n_heads']**-.5 |
| scores=mx.sum(mx.maximum(iq.astype(mx.float32)@keys.astype(mx.float32).T,0)*importance[:,None],axis=0) |
| if layer==c['candidate_source_layer_id']: |
| block=c['candidate_block_size'];padded=mx.pad(scores,(0,(-n)%block),constant_values=-float('inf')) |
| block_scores=mx.max(padded.reshape(-1,block),axis=-1) |
| block_scores=mx.where(mx.arange(len(block_scores))==(n-1)//block,float('inf'),block_scores) |
| picks=mx.argsort(block_scores)[-min(c['candidate_topk_blocks'],len(block_scores)):] |
| self.shared['candidates']=mx.any(mx.arange(n)[:,None]//block==picks[None,:],axis=-1) |
| elif layer>c['candidate_source_layer_id']: |
| scores=mx.where(self.shared['candidates'],scores,-float('inf')) |
| self.shared['topk']=mx.sort(mx.argsort(scores)[-min(c['index_topk'],n):]) |
| else:self.shared['topk']=mx.array([],dtype=mx.int32) |
| if latent is not None: |
| st['compressed'].append(quantize_cache(self.rotate(latent,t+1-ratio,ratio),4,16,'e4m3')) |
| if self.shared['compressed']: |
| compressed=mx.concatenate(self.shared['compressed'],axis=0) |
| selected=mx.concatenate((selected,compressed[self.shared['topk']]),axis=0) |
| out=sparse_attention(q,selected,self.w.read(base+'.attn_sink')) |
| out=self.rotate(out,t,ratio,inverse=True).reshape(c['o_groups'],-1) |
| out=self.w.grouped(base+'.wo_a',out,c['o_groups']) |
| return self.w.linear(base+'.wo_b',out) |
|
|
| def step(self,token,progress=False): |
| if self.position>=self.max_tokens:raise ValueError('Reference context limit reached') |
| if token==self.w.config.get('image_token_id'):raise NotImplementedError('Vision is not validated') |
| c=self.c;self.hash.push(int(token));self.shared={} |
| x=self.w.embedding('embed',[token]);h=mx.repeat(x[:,None,:],c['hc_mult'],axis=1) |
| pre=mx.array([[1.]+[0.]*(c['hc_mult']-1)],dtype=mx.float32) |
| targets=[] |
| for layer in range(c['num_hidden_layers']): |
| base=f'layers.{layer}' |
| if layer in c['engram_layer_ids']:h=self.engram(base+'.engram',h,layer) |
| if layer in c['dspark_target_layer_ids']:targets.append(mx.mean(h,axis=1)) |
| ap,post,comb=self.mixes(base,'attn',h) |
| a=self.attention(base+'.attn',self.norm(base+'.attn_norm',hc_pre(h,pre)),layer) |
| h=hc_post(a,h,post,comb) |
| pre,post,comb=self.mixes(base,'ffn',h) |
| h=hc_post(self.moe(base+'.ffn',self.norm(base+'.ffn_norm',hc_pre(h,ap))),h,post,comb) |
| mx.eval(h,pre) |
| if progress:print(json.dumps({'position':self.position,'layer':layer,'finite':bool(mx.all(mx.isfinite(h)))}),flush=True) |
| if not bool(mx.all(mx.isfinite(h))):raise FloatingPointError('Nonfinite hidden state') |
| logits=self.w.linear('head',self.norm('norm',hc_pre(h,pre)).astype(mx.float32)) |
| mx.eval(logits);self.position+=1 |
| return logits, mx.concatenate(targets,axis=-1) |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser();p.add_argument('--model',required=True);p.add_argument('--token-id',type=int);p.add_argument('--load-only',action='store_true') |
| p.add_argument('--resident-backbone',action='store_true') |
| p.add_argument('--execution-mode',choices=('reference','deferred','compiled'),default='compiled') |
| args=p.parse_args();start=time.monotonic() |
| model=TextRuntime(args.model,resident_backbone=args.resident_backbone,execution_mode=args.execution_mode) |
| print(json.dumps({'adapter':'isolated-experimental-v41-text','indexed_tensors':len(model.w.entries),'mtp_tensors':sum(k.startswith('mtp.') for k in model.w.entries),'engram_map_verified':True,'omlx_integrated':False,'mtp_generation_tested':False}),flush=True) |
| if not args.load_only: |
| if args.token_id is None:raise ValueError('Explicit token-id required for one-token diagnostic') |
| logits,_=model.step(args.token_id,progress=True) |
| print(json.dumps({'next_token_id':int(mx.argmax(logits)),'seconds':time.monotonic()-start,'peak_MLX_bytes':mx.get_peak_memory(),'quality_validated':False}),flush=True) |
|
|
|
|
| if __name__=='__main__':main() |
|
|
