File size: 8,789 Bytes
7218492 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """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
# The trapezoidal key scale at token t also reads dt/trap at t+1.
# A state is committed only after this look-ahead token is present.
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
|