| """Heterogeneous LFM3 decoder stack with exact source-operation transport. |
| |
| Machine cells use the compiled circuit's original sparse ReLU/head operations. |
| No SiLU approximation, normalization, widening or numerical conversion is added. |
| Language cache indices remain those of the unmodified language layers. |
| """ |
| from collections import OrderedDict |
| import inspect,weakref |
| import torch |
| from torch import nn |
| from .computer.circuit import CompiledComputer,CircuitStage |
|
|
| class MachineDecoderLayer(nn.Module): |
| decoder_role='machine' |
| def language_identity(self,hidden_states):return hidden_states |
|
|
| class ComputerControlDecoderLayer(MachineDecoderLayer): |
| decoder_role='computer_control' |
| def __init__(self,computer): |
| super().__init__();self.spec=computer.spec |
| self.embedding=computer.embedding;self.output=computer.output;self.programs=computer.programs |
| def forward(self,hidden_states=None,*,phase='language',args=(),kwargs=None): |
| if phase=='language':return hidden_states |
| port=self._port();kwargs=kwargs or {} |
| if phase=='boot':return CompiledComputer.boot(port,*args,**kwargs) |
| if phase=='embed':return CompiledComputer.embed(port,*args,**kwargs) |
| raise ValueError('Unknown control phase: '+phase) |
|
|
| class OrderedSparseDownProjection(nn.Module): |
| """Exact w2 adapter: retain coefficient multiplication and scatter order.""" |
| def __init__(self,data):super().__init__();self.data=data |
| def forward(self,values,template): |
| values=values*self.data.coefficient |
| return torch.zeros_like(template).index_add_(1,self.data.output,values) |
|
|
| class ExactGatedMLP(nn.Module): |
| """LFM gated-MLP decomposition with exact ReLU/sparse projection adapters.""" |
| def __init__(self,operation): |
| super().__init__();self.w1=operation.b;self.w3=operation.a |
| self.w2=OrderedSparseDownProjection(operation.data);self.activation=nn.ReLU() |
| def forward(self,hidden_states,active): |
| values=self.w3(hidden_states)*self.activation(self.w1(hidden_states)) |
| return self.w2(values,hidden_states)*active[:,None] |
|
|
| class ExactRecurrentAttention(nn.Module): |
| """Exact memory attention adapter; persistent writes are source operations.""" |
| def __init__(self,operation): |
| super().__init__();self.spec=operation.spec;self.data=operation.data |
| for name in operation.spec['maps']:self.add_module(name,getattr(operation,name)) |
| def forward(self,x,state,active): |
| d=self.data;memory,history=state;flat=torch.cat((x,history),dim=1) |
| B,V=x.shape[0],d.khead.numel() |
| q=self.q(flat).relu().reshape(B,-1,128)[:,d.khead] |
| k=self.k(flat).relu().reshape(B,-1,128)[:,d.khead] |
| v=self.v(flat).relu().reshape(B,V,128) |
| z=self.z(flat).relu().reshape(B,V,128) |
| a,beta=self.a(flat),self.b(flat)+d.bconst |
| factor=torch.where(a==0,1.,torch.where(a==110,.5,0.)) |
| updated=memory*factor[:,:,None,None];row=k.argmax(-1) |
| write=(k.amax(-1)>0)&(beta==1) |
| selector=torch.arange(128,device=x.device)[None,None,:]==row[:,:,None] |
| updated=torch.where((write[:,:,None]&selector)[:,:,:,None],v[:,:,None,:],updated) |
| memory.copy_(torch.where(active[:,None,None,None],updated,memory)) |
| read=((q>0)[:,:,:,None]*memory).sum(2) |
| reference=read[:,:,:2].square().sum(-1,keepdim=True) |
| denominator=(reference/128.).sqrt() |
| normalized=torch.where(reference>0,read/torch.where(denominator>0,denominator,1.),0.) |
| delta=(normalized*z*d.oweight).reshape(B,-1) |
| return torch.zeros_like(x).index_add_(1,d.output.flatten(),delta)*active[:,None] |
| def commit_history(self,x,state,active): |
| memory,history=state |
| new_history=torch.cat((x,history[:,:2*x.shape[1]]),dim=1) |
| history.copy_(torch.where(active[:,None],new_history,history)) |
|
|
| class ComputerExecutionDecoderLayer(MachineDecoderLayer): |
| decoder_role='computer_execution' |
| def __init__(self,operations): |
| super().__init__();self.stage_specs=[operation.spec for operation in operations];self.triton_widths={} |
| self.operator_norm=nn.Identity();self.ffn_norm=nn.Identity();self.use_sandwich_norm=False |
| head=next((op for op in operations if op.kind=='head'),None);ffn=next((op for op in operations if op.kind=='ffn'),None) |
| self.self_attn=ExactRecurrentAttention(head) if head is not None else None |
| self.feed_forward=ExactGatedMLP(ffn) if ffn is not None else None |
| def forward(self,hidden_states,state=None,active=None,backend='torch',scratch=None,*,phase='language'): |
| if active is None:return hidden_states |
| if backend=='triton': |
| for spec in self.stage_specs: |
| view=StageAdapter(self,spec) |
| if view.kind not in self.triton_widths:raise RuntimeError("Call circuit.prepare('triton') before CUDA capture.") |
| view.widths=self.triton_widths[view.kind] |
| hidden_states,_=CircuitStage.triton(view,hidden_states,state if view.kind=='head' else None,active,scratch) |
| return hidden_states,state |
| |
| |
| if self.self_attn is not None: |
| residual=hidden_states |
| attention_output=self.self_attn(self.operator_norm(hidden_states),state,active) |
| hidden_states=residual+attention_output |
| self.self_attn.commit_history(residual,state,active) |
| if self.feed_forward is not None: |
| residual=hidden_states |
| ffn_output=self.feed_forward(self.ffn_norm(hidden_states),active) |
| hidden_states=residual+ffn_output |
| return hidden_states,state |
|
|
| class StageAdapter: |
| """Parameter-free view for stable logical state names and program boot.""" |
| def __init__(self,cell,spec):self.cell=cell;self.spec=spec;self.kind=spec['kind'];self.widths=None |
| @property |
| def data(self):return self.cell.feed_forward.w2.data if self.kind=='ffn' else self.cell.self_attn.data |
| def __getattr__(self,name): |
| if self.kind=='ffn' and name in ('a','b'):return getattr(self.cell.feed_forward,'w3' if name=='a' else 'w1') |
| if self.kind=='head':return getattr(self.cell.self_attn,name) |
| raise AttributeError(name) |
| def logical_modules(self): |
| for name in self.spec['maps']:yield name,getattr(self,name) |
| yield 'data',self.data |
| def prepare_triton(self): |
| CircuitStage.prepare_triton(self);self.cell.triton_widths[self.kind]=self.widths |
|
|
| class MemoryDecoderLayer(MachineDecoderLayer): |
| decoder_role='memory' |
| def __init__(self,ram,files):super().__init__();self.ram=ram;self.files=files |
| def forward(self,hidden_states=None,*,phase='language',component='ram',method='forward',args=(),kwargs=None): |
| if phase=='language':return hidden_states |
| module=getattr(self,component) |
| return module(*args,**(kwargs or {})) if method=='forward' else getattr(module,method)(*args,**(kwargs or {})) |
|
|
| class ServiceDecoderLayer(MachineDecoderLayer): |
| def __init__(self,module,role):super().__init__();self.payload=module;self.decoder_role=role |
| def forward(self,hidden_states=None,*,phase='language',method='forward',args=(),kwargs=None): |
| if phase=='language':return hidden_states |
| return self.payload(*args,**(kwargs or {})) if method=='forward' else getattr(self.payload,method)(*args,**(kwargs or {})) |
|
|
| class DecoderModelMixin: |
| def forward(self,*args,execution_phase=None,**kwargs): |
| if execution_phase is None or execution_phase=='language':return super().forward(*args,**kwargs) |
| if args:raise ValueError('Machine execution uses explicit typed arguments') |
| if execution_phase=='computer': |
| state=kwargs.pop('computer_state') |
| if kwargs:raise ValueError('Unexpected computer execution fields') |
| |
| |
| port=self._computer_port |
| active=(state.current!=port.spec['halt'])&state.valid |
| x=port.embed(state.current);scratch=torch.zeros_like(x) if port.backend=='triton' else None |
| hi=0 |
| for cell in self.layers[1:1+self.computer_block_count]: |
| head=state.heads[hi] if cell.self_attn is not None else None |
| x,_=cell(x,head,active,port.backend,scratch) |
| hi+=cell.self_attn is not None |
| out=port.output |
| family=(x[:,out.family].round()*out.family_weight).sum(1).long() |
| payload=(x[:,out.payload].round()*out.payload_weight).sum(1).long() |
| f=family.clamp(0,port.embedding.start.numel()-1) |
| valid=((family>=0)&(family<port.embedding.start.numel())&(payload>=0)&(payload<port.embedding.size[f])) |
| state.valid.copy_(state.valid&(valid|~active));predicted=port.embedding.start[f]+payload |
| state.current.copy_(torch.where(active,predicted,state.current));return state.current |
| index=kwargs.pop('decoder_index');method=kwargs.pop('decoder_method','forward') |
| call_args=kwargs.pop('decoder_args',());call_kwargs=kwargs.pop('decoder_kwargs',{}) |
| component=kwargs.pop('decoder_component',None) |
| if kwargs:raise ValueError('Unexpected decoder execution fields') |
| layer=self.layers[index] |
| if isinstance(layer,ComputerControlDecoderLayer):return layer(phase=execution_phase,args=call_args,kwargs=call_kwargs) |
| if isinstance(layer,MemoryDecoderLayer):return layer(phase=execution_phase,component=component,method=method,args=call_args,kwargs=call_kwargs) |
| return layer(phase=execution_phase,method=method,args=call_args,kwargs=call_kwargs) |
|
|
| class StageView: |
| def __init__(self,owner,specs,mapping):self.owner=weakref.ref(owner);self.specs=specs;self.mapping=mapping |
| def __len__(self):return len(self.specs) |
| def __iter__(self):return iter(tuple(self[i] for i in range(len(self)))) |
| def __getitem__(self,index): |
| if isinstance(index,slice):return list(self)[index] |
| return StageAdapter(self.owner().layers[self.mapping[index]],self.specs[index]) |
|
|
| class ComputerPort: |
| """Parameter-free compatibility API; execution always enters decoder.forward.""" |
| def __init__(self,owner,spec,mapping,backend): |
| self.owner=weakref.ref(owner);self.spec=spec;self.backend=backend;self.stages=StageView(owner,spec['stages'],mapping) |
| @property |
| def embedding(self):return self.owner().layers[0].embedding |
| @property |
| def output(self):return self.owner().layers[0].output |
| @property |
| def programs(self):return self.owner().layers[0].programs |
| @property |
| def device(self):return self.embedding.base.device |
| def __call__(self,state):return self.owner()(execution_phase='computer',computer_state=state) |
| def embed(self,tokens):return self.owner()(execution_phase='embed',decoder_index=0,decoder_args=(tokens,)) |
| def boot(self,*args,**kwargs):return self.owner()(execution_phase='boot',decoder_index=0,decoder_args=args,decoder_kwargs=kwargs) |
| def prepare(self,backend='torch'):return CompiledComputer.prepare(self,backend) |
| def _logical_modules(self): |
| yield 'embedding',self.embedding;yield 'output',self.output |
| for i,stage in enumerate(self.stages): |
| for name,module in stage.logical_modules():yield 'stages.'+str(i)+'.'+name,module |
| yield 'programs',self.programs |
| def state_dict(self,destination=None,prefix='',keep_vars=False): |
| result=OrderedDict() if destination is None else destination |
| for name,module in self._logical_modules():module.state_dict(destination=result,prefix=prefix+name+'.',keep_vars=keep_vars) |
| return result |
| def named_parameters(self,prefix='',recurse=True,remove_duplicate=True): |
| seen=set() |
| for name,module in self._logical_modules(): |
| for key,value in module.named_parameters(recurse=recurse,remove_duplicate=remove_duplicate): |
| if remove_duplicate and id(value) in seen:continue |
| seen.add(id(value));yield ((prefix+'.') if prefix else '')+name+'.'+key,value |
| def parameters(self,recurse=True):return (value for _,value in self.named_parameters(recurse=recurse)) |
|
|
| class ComponentPort: |
| def __init__(self,owner,index,component=None):self.owner=weakref.ref(owner);self.index=index;self.component=component |
| def _module(self): |
| layer=self.owner().layers[self.index] |
| return getattr(layer,self.component) if self.component else layer.payload |
| def _call(self,method,args,kwargs): |
| return self.owner()(execution_phase='service',decoder_index=self.index,decoder_component=self.component,decoder_method=method,decoder_args=args,decoder_kwargs=kwargs) |
| def __call__(self,*args,**kwargs):return self._call('forward',args,kwargs) |
| def __getattr__(self,name): |
| value=getattr(self._module(),name) |
| if callable(value) and not isinstance(value,nn.Module) and name not in {'state_dict','named_parameters','parameters','named_buffers','buffers','named_modules','modules','children','named_children','get_submodule'}: |
| return lambda *args,**kwargs:self._call(name,args,kwargs) |
| return value |
|
|
| def install_decoder_stack(model): |
| backbone=model.language_model.model |
| if getattr(backbone,'computer_decoder_count',0):raise ValueError('Decoder stack already installed') |
| original=model.computer;stage_count=len(original.stages) |
| groups=[];i=0 |
| while i<stage_count: |
| group=[i] |
| if original.stages[i].kind=='head' and i+1<stage_count and original.stages[i+1].kind=='ffn':group.append(i+1) |
| groups.append(group);i+=len(group) |
| mapping=[0]*stage_count |
| for index,group in enumerate(groups,1): |
| for source in group:mapping[source]=index |
| prefix=[ComputerControlDecoderLayer(original)]+[ComputerExecutionDecoderLayer([original.stages[j] for j in group]) for group in groups] |
| indices={'memory':len(prefix)};prefix.append(MemoryDecoderLayer(model.ram,model.files)) |
| if model.ram_bridge is not None:indices['bridge']=len(prefix);prefix.append(ServiceDecoderLayer(model.ram_bridge,'bridge')) |
| indices['analytics']=len(prefix);prefix.append(ServiceDecoderLayer(model.analytics,'analytics')) |
| language_layers=list(backbone.layers);backbone.layers=nn.ModuleList(prefix+language_layers) |
| backbone.computer_decoder_count=len(prefix);backbone.computer_block_count=len(groups);backbone.language_decoder_count=len(language_layers) |
| backbone.__class__=type('LFMputerDecoderModel',(DecoderModelMixin,type(backbone)),{}) |
| computer=ComputerPort(backbone,original.spec,mapping,original.backend) |
| object.__setattr__(backbone,'_computer_port',computer);object.__setattr__(prefix[0],'_port',weakref.ref(computer)) |
| for name in ['computer','ram','files','analytics']: |
| model._modules.pop(name) |
| if model.ram_bridge is not None:model._modules.pop('ram_bridge') |
| object.__setattr__(model,'computer',computer) |
| object.__setattr__(model,'ram',ComponentPort(backbone,indices['memory'],'ram')) |
| object.__setattr__(model,'files',ComponentPort(backbone,indices['memory'],'files')) |
| object.__setattr__(model,'analytics',ComponentPort(backbone,indices['analytics'])) |
| if 'bridge' in indices:object.__setattr__(model,'ram_bridge',ComponentPort(backbone,indices['bridge'])) |
| model.decoder_layout={'computer_stage_count':stage_count,'computer_block_count':len(groups),'stage_groups':groups,'stage_to_decoder':mapping,'computer_stage_kinds':[s.kind for s in original.stages],'machine_layer_count':len(prefix),'language_layer_count':len(language_layers),'total_layer_count':len(backbone.layers),'memory_index':indices['memory'],'bridge_index':indices.get('bridge'),'analytics_index':indices['analytics'],'roles':[layer.decoder_role for layer in prefix]+['language']*len(language_layers),'language_cache_indices':'unchanged; machine state is explicit'} |
| model.decoder_layout.update(machine_hidden_size=original.spec['dims'],language_hidden_size=backbone.config.hidden_size) |
| if hasattr(model,'lora_names'): |
| model.lora_names=[logical_to_decoder_key('language_model.'+name,model.decoder_layout)[len('language_model.'):] for name in model.lora_names] |
| model._runners={};return model |
|
|
| def logical_to_decoder_key(key,layout): |
| base='language_model.model.layers.';p=layout['machine_layer_count'] |
| if key.startswith(base): |
| i,tail=key[len(base):].split('.',1);return base+str(int(i)+p)+'.'+tail |
| if key.startswith('computer.stages.'): |
| i,tail=key[len('computer.stages.'):].split('.',1);kind=layout['computer_stage_kinds'][int(i)] |
| if kind=='ffn': |
| name,rest=tail.split('.',1);tail={'a':'feed_forward.w3','b':'feed_forward.w1','data':'feed_forward.w2.data'}[name]+'.'+rest |
| else:tail='self_attn.'+tail |
| return base+str(layout['stage_to_decoder'][int(i)])+'.'+tail |
| for name in ['embedding','output','programs']: |
| prefix='computer.'+name+'.' |
| if key.startswith(prefix):return base+'0.'+name+'.'+key[len(prefix):] |
| for prefix,index,child in [('ram.',layout['memory_index'],'ram.'),('files.',layout['memory_index'],'files.'),('ram_bridge.',layout.get('bridge_index'),'payload.'),('analytics.',layout['analytics_index'],'payload.')]: |
| if key.startswith(prefix): |
| if index is None:raise ValueError('Bridge tensor without a bridge layer') |
| return base+str(index)+'.'+child+key[len(prefix):] |
| return key |
|
|
| def decoder_to_logical_key(key,layout): |
| base='language_model.model.layers.' |
| if not key.startswith(base):return key |
| i,tail=key[len(base):].split('.',1);i=int(i);p=layout['machine_layer_count'] |
| if i>=p:return base+str(i-p)+'.'+tail |
| if i==0:return 'computer.'+tail |
| if 1<=i<=layout['computer_block_count']: |
| kind='ffn' if tail.startswith('feed_forward.') else 'head' |
| source=next(j for j,n in enumerate(layout['stage_to_decoder']) if n==i and layout['computer_stage_kinds'][j]==kind) |
| if kind=='ffn': |
| for before,after in [('feed_forward.w3.','a.'),('feed_forward.w1.','b.'),('feed_forward.w2.data.','data.')]: |
| if tail.startswith(before):return 'computer.stages.'+str(source)+'.'+after+tail[len(before):] |
| raise ValueError('Unmapped machine FFN tensor: '+key) |
| assert tail.startswith('self_attn.');return 'computer.stages.'+str(source)+'.'+tail[len('self_attn.'):] |
| for index,child,prefix in [(layout['memory_index'],'ram.','ram.'),(layout['memory_index'],'files.','files.'),(layout.get('bridge_index'),'payload.','ram_bridge.'),(layout['analytics_index'],'payload.','analytics.')]: |
| if i==index and tail.startswith(child):return prefix+tail[len(child):] |
| raise ValueError('Unmapped decoder tensor: '+key) |
|
|
