| """Request-owned guarded hybrid cache for the fresh current-recipe graph."""
|
| from contextlib import contextmanager
|
| import torch
|
| from transformers.cache_utils import Cache
|
| from .runtime import execution_identity
|
|
|
|
|
| def first_changed(a,b,dim=0):
|
| a=a.movedim(dim,0).contiguous();b=b.movedim(dim,0).contiguous()
|
| common=min(a.shape[0],b.shape[0])
|
| if not common:return 0
|
| av=a[:common].reshape(common,-1).view(torch.uint8)
|
| bv=b[:common].reshape(common,-1).view(torch.uint8)
|
| changed=(av!=bv).any(-1).nonzero().flatten()
|
| return int(changed[0]) if changed.numel() else common
|
|
|
|
|
| def fingerprint(core):
|
| return tuple((n,p.data_ptr(),p._version,str(p.device),str(p.dtype),tuple(p.shape))
|
| for n,p in list(core.named_parameters())+list(core.named_buffers()))
|
|
|
|
|
| def execution_fingerprint(core):
|
| tensor=next(core.parameters(),None)
|
| if tensor is None:tensor=next(core.buffers(),None)
|
| return execution_identity(tensor.device if tensor is not None else 'cpu')
|
|
|
|
|
| class CachedLinears:
|
| def __init__(self,base,cache):self.base=base;self.cache=cache
|
| def __getattr__(self,name):return getattr(self.base,name)
|
| def __call__(self,x,weight,bias=None,*,name,seed=0,bridge=False):
|
| rows=x.reshape(-1,x.shape[-1])
|
| key=(name,seed,bridge,weight.data_ptr(),weight._version,tuple(weight.shape),tuple(weight.stride()),
|
| x.dtype,str(x.device),None if bias is None else (bias.data_ptr(),bias._version))
|
| previous=self.cache.linears.get(key)
|
| start=0
|
| if previous is not None:
|
| old_input,old_output=previous
|
| changed=first_changed(rows,old_input)
|
| if changed==rows.shape[0]:
|
| value=old_output[:len(rows)]
|
| self.cache.stats['linear_hits']+=1
|
| self.cache.stats['reused_rows']+=len(rows)
|
| self.cache.linears[key]=(rows.detach().clone(),value.detach().clone())
|
| return value.reshape(*x.shape[:-1],weight.shape[0])
|
| start=(changed//128)*128
|
| self.cache.stats['reused_rows']+=start
|
| if changed<min(len(rows),len(old_input)):self.cache.stats['row_refreshes']+=1
|
| tail=self.base(rows[start:],weight,bias,name=name,seed=seed,bridge=bridge,activation_row_offset=start)
|
| value=tail if not start else torch.cat((previous[1][:start],tail),0)
|
| self.cache.stats['computed_rows']+=len(rows)-start
|
| self.cache.linears[key]=(rows.detach().clone(),value.detach().clone())
|
| return value.reshape(*x.shape[:-1],weight.shape[0])
|
|
|
|
|
| class FidelHybridCache(Cache):
|
| """Projection histories and attention KV; Mamba scan replay is explicit."""
|
| def __init__(self,owner,core,prompt_length,context_limit=8192,mode='guarded'):
|
| super().__init__(layers=[])
|
| self.mode=mode;self.scans={}
|
| self.owner=owner;self.weight_fingerprint=fingerprint(core)
|
| self.execution_fingerprint=execution_fingerprint(core)
|
| self.prompt_length=prompt_length;self.context_limit=context_limit
|
| self.position=0;self.closed=False;self.token_ids=None
|
| self.linears={};self.attention={}
|
| self.stats={'linear_hits':0,'reused_rows':0,'computed_rows':0,'row_refreshes':0,
|
| 'attention_rows_reused':0,'attention_rows_computed':0,'decode_calls':0,
|
| 'scan_chunks_reused':0,'scan_chunks_computed':0,'scan_prefix_refreshes':0}
|
| def get_seq_length(self,layer_idx=0):return self.position
|
| def get_max_cache_shape(self,layer_idx=0):return self.context_limit
|
| def get_usable_length(self,new_seq_length,layer_idx=0):return self.position
|
| @property
|
| def is_compileable(self):return False
|
| def reorder_cache(self,beam_idx):
|
| if beam_idx.numel()!=1 or int(beam_idx[0])!=0:raise ValueError('Beam expansion is unsupported')
|
| def reset(self):
|
| self.linears.clear();self.attention.clear();self.scans.clear();self.token_ids=None
|
| self.position=0;self.closed=True
|
| def scan_for(self,index,live):
|
| def execute(q,k,v,**kwargs):
|
| from .mamba_scan import scan,ScanResult
|
|
|
|
|
| boundary=((live-1)//16)*16
|
| prior=self.scans.get(index);start=0;state=None
|
| tensors={'q':(q,1),'k':(k,1),'v':(v,1),'z':(kwargs['z'],1),
|
| 'dt':(kwargs['dt'],2),'trap':(kwargs['trap'],2)}
|
| for name,value in kwargs['prepared'].items():
|
| if name=='angle_state':continue
|
| tensors['prepared/'+name]=(value,3 if name=='segsum' else (1 if name=='angles' else 2))
|
| if prior is not None and prior['boundary']:
|
| old_boundary=prior['boundary'];valid=old_boundary<=boundary
|
| for name,(value,axis) in tensors.items():
|
| if name=='prepared/segsum':
|
| valid=valid and first_changed(value,prior['inputs'][name],2)>=old_boundary//16
|
| else:
|
| needed=old_boundary+(1 if name in ('dt','trap') else 0)
|
| valid=valid and first_changed(value,prior['inputs'][name],axis)>=needed
|
| if not valid:break
|
| if valid:start=old_boundary;state=prior['state']
|
| else:self.stats['scan_prefix_refreshes']+=1
|
| result=scan(q,k,v,**kwargs,start=start,initial_state=state,checkpoint_at=boundary)
|
| output=result.output if not start else torch.cat((prior['output'][:,:start],result.output),1)
|
| self.stats['scan_chunks_reused']+=start//16
|
| self.stats['scan_chunks_computed']+=(q.shape[1]-start)//16
|
| self.scans[index]={'boundary':boundary,'state':result.checkpoint,'output':output.detach().clone(),
|
| 'inputs':{name:value.detach().clone() for name,(value,axis) in tensors.items()}}
|
| return ScanResult(output,result.state,result.last_key,result.trace,result.checkpoint,boundary)
|
| return execute
|
|
|
| def attend(self,q,k,v):
|
| prior=self.attention.get(7);start=0;outputs=[]
|
| if prior is not None:
|
| start=min(first_changed(q,prior['q'],2),first_changed(k,prior['k'],2),first_changed(v,prior['v'],2))
|
| start=min(start,q.shape[2])
|
| if start:outputs.append(prior['output'][:,:,:start])
|
| scale=q.shape[-1]**-.5
|
| for pos in range(start,q.shape[2]):
|
| from .operators import causal_attention_row
|
| outputs.append(causal_attention_row(q[:,:,pos:pos+1],k[:,:,:pos+1],v[:,:,:pos+1])[2])
|
| out=torch.cat(outputs,2)
|
| self.stats['attention_rows_reused']+=start
|
| self.stats['attention_rows_computed']+=q.shape[2]-start
|
| self.attention[7]={'q':q.detach().clone(),'k':k.detach().clone(),'v':v.detach().clone(),
|
| 'output':out.detach().clone()}
|
| return out
|
|
|
|
|
| @contextmanager
|
| def activated(core,cache):
|
| engine=core.engine;prior=getattr(core,'_active_cache',None)
|
| object.__setattr__(core,'engine',CachedLinears(engine,cache))
|
| object.__setattr__(core,'_active_cache',cache)
|
| try:yield
|
| finally:
|
| object.__setattr__(core,'engine',engine)
|
| object.__setattr__(core,'_active_cache',prior)
|
|
|
|
|
| def cached_forward(core,new_ids,prompt_length,cache,owner,mode='guarded'):
|
| if cache is None:cache=FidelHybridCache(owner,core,prompt_length,mode=mode)
|
| if not isinstance(cache,FidelHybridCache) or cache.owner!=owner or cache.closed or cache.mode!=mode:
|
| raise ValueError('Use a live cache owned by this model')
|
| if cache.execution_fingerprint!=execution_fingerprint(core):
|
| cache.reset();raise ValueError('Execution backend changed; start a fresh request')
|
| if cache.prompt_length!=prompt_length or cache.weight_fingerprint!=fingerprint(core):
|
| cache.reset();raise ValueError('Prompt boundary or model tensors changed; start a fresh request')
|
| if new_ids.ndim!=2 or new_ids.shape[0]!=1 or new_ids.shape[1]<1:
|
| raise ValueError('Provide one or more new tokens for one request')
|
| full=new_ids if cache.token_ids is None else torch.cat((cache.token_ids,new_ids),1)
|
| if full.shape[1]>cache.context_limit:raise ValueError('Context limit exceeded')
|
| try:
|
| with activated(core,cache):logits=core.logits(full,prompt_length)[-1:].unsqueeze(0)
|
| cache.token_ids=full.detach().clone();cache.position=full.shape[1]
|
| cache.stats['decode_calls']+=1
|
| return logits,cache
|
| except BaseException:
|
| cache.reset();raise
|
|
|
